Fix-3

Signed-off-by: Amit Raj <[email protected]>

Author: quic-amitraj

QEfficient/diffusers/models/attention_processor.py

@@ -123,7 +123,6 @@ def __call__(
             hidden_states = torch.bmm(attention_probs, value)
         else:  # self-attention, use blocked attention
             # QKV done with block-attention (a la FlashAttentionV2)
-            print(f"{query.shape = }, {key.shape = }, {value.shape = }")
             query_block_size = self.query_block_size
             query_seq_len = query.size(-2)
             num_blocks = (query_seq_len + query_block_size - 1) // query_block_size

QEfficient/diffusers/models/pytorch_transforms.py

@@ -23,6 +23,11 @@
 class CustomOpsTransform(ModuleMappingTransform):
     _module_mapping = {RMSNorm: CustomRMSNormAIC}
 
+    @classmethod
+    def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
+        model, transformed = super().apply(model)
+        return model, transformed
+
 
 class AttentionTransform(ModuleMappingTransform):
     _module_mapping = {

QEfficient/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion3.py

@@ -14,6 +14,15 @@
 import numpy as np
 import torch
 import torch.nn as nn
+from transformers import (
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    AutoModelForSpeechSeq2Seq,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+    TextStreamer,
+)
 
 import QEfficient
 from QEfficient.base.modeling_qeff import QEFFBaseModel
@@ -49,15 +58,6 @@
 )
 from QEfficient.utils.cache import to_hashable
 from QEfficient.utils.logging_utils import logger
-from transformers import (
-    AutoModel,
-    AutoModelForCausalLM,
-    AutoModelForImageTextToText,
-    AutoModelForSpeechSeq2Seq,
-    PreTrainedTokenizer,
-    PreTrainedTokenizerFast,
-    TextStreamer,
-)
 
 
 class QEFFTransformersBase(QEFFBaseModel):

QEfficient/transformers/models/modeling_auto.py

@@ -0,0 +1,6 @@
+from QEfficient import QEFFStableDiffusion3Pipeline
+
+pipeline = QEFFStableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-large")
+pipeline.compile(num_devices_text_encoder=1, num_devices_transformer=4, num_devices_vae_decoder=1)
+image = pipeline("A girl laughing", num_inference_steps=1, guidance_scale=0.0).images[0]
+image.save("new_testing.png")

examples/diffusers/__init__.py

@@ -0,0 +1,69 @@
+{
+  "Compile": {
+
+    "text_encoder_compile_config": {
+      "num_devices": 1,
+      "num_coresr": 16,
+      "mxfp6_matmul": false,
+      "batch_size": 1,
+      "onnx_path": null,
+      "compile_dir": null
+    },
+
+    "text_encoder_compile_config_2": {
+      "num_devices": 1,
+      "num_cores": 16,
+      "mxfp6_matmul": false,
+      "batch_size": 1,
+      "onnx_path": null,
+      "compile_dir": null
+    },
+
+    "text_encoder_compile_config_3": {
+      "num_devices": 1,
+      "num_cores": 16,
+      "mxfp6_matmul": false,
+      "batch_size": 1,
+      "onnx_path": null,
+      "compile_dir": null
+    },
+
+    "transformer_config": {
+      "num_devices": 4,
+      "num_cores": 16,
+      "mxfp6_matmul": false,
+      "fp16": true,
+      "batch_size": 1,
+      "onnx_path": null,
+      "compile_dir": null
+    },
+
+    "vae_compile_config": {
+      "num_devices": 1,
+      "num_cores": 16,
+      "mxfp6_matmul": false,
+      "batch_size": 1,
+      "onnx_path": null,
+      "compile_dir": null
+    }
+  },
+  
+  
+  "generate":{
+    "text_encoder_generate_config":{
+        "device_ids":[0]
+    },
+    "text_encoder_generate_config_2":{
+        "device_ids":[1]
+    },
+    "text_encoder_generate_config_3":{
+        "device_ids":[2]
+    },
+    "transformer_generate_config":{
+        "device_ids":[4,5,6,7]
+    },
+    "vae_generate_config":{
+        "device_ids":[3]
+    }
+  }
+}