Modification of Pipeline-2

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

Author: Amit Raj

QEfficient/diffusers/models/pytorch_transforms.py

@@ -35,6 +35,7 @@
     QEffFluxAttnProcessor,
     QEffFluxSingleTransformerBlock,
     QEffFluxTransformer2DModel,
+    QEffFluxTransformer2DModelOF,
     QEffFluxTransformerBlock,
 )
 
@@ -83,7 +84,7 @@ def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
 
 
 class OnnxFunctionTransform(ModuleMappingTransform):
-    _module_mapping = {FluxTransformer2DModel: QEffFluxTransformer2DModel}
+    _module_mapping = {QEffFluxTransformer2DModel, QEffFluxTransformer2DModelOF}
 
     @classmethod
     def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:

QEfficient/diffusers/models/transformers/transformer_flux.py

@@ -22,7 +22,6 @@
 )
 
 from QEfficient.diffusers.models.normalization import (
-    QEffAdaLayerNormContinuous,
     QEffAdaLayerNormZero,
     QEffAdaLayerNormZeroSingle,
 )
@@ -253,58 +252,6 @@ def forward(
 
 
 class QEffFluxTransformer2DModel(FluxTransformer2DModel):
-    def __init__(
-        self,
-        patch_size: int = 1,
-        in_channels: int = 64,
-        out_channels: Optional[int] = None,
-        num_layers: int = 19,
-        num_single_layers: int = 38,
-        attention_head_dim: int = 128,
-        num_attention_heads: int = 24,
-        joint_attention_dim: int = 4096,
-        pooled_projection_dim: int = 768,
-        guidance_embeds: bool = False,
-        axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
-    ):
-        super().__init__(
-            patch_size=patch_size,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            num_layers=num_layers,
-            num_single_layers=num_single_layers,
-            attention_head_dim=attention_head_dim,
-            num_attention_heads=num_attention_heads,
-            joint_attention_dim=joint_attention_dim,
-            pooled_projection_dim=pooled_projection_dim,
-            guidance_embeds=guidance_embeds,
-            axes_dims_rope=axes_dims_rope,
-        )
-
-        self.transformer_blocks = nn.ModuleList(
-            [
-                QEffFluxTransformerBlock(
-                    dim=self.inner_dim,
-                    num_attention_heads=num_attention_heads,
-                    attention_head_dim=attention_head_dim,
-                )
-                for _ in range(num_layers)
-            ]
-        )
-
-        self.single_transformer_blocks = nn.ModuleList(
-            [
-                QEffFluxSingleTransformerBlock(
-                    dim=self.inner_dim,
-                    num_attention_heads=num_attention_heads,
-                    attention_head_dim=attention_head_dim,
-                )
-                for _ in range(num_single_layers)
-            ]
-        )
-
-        self.norm_out = QEffAdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
-
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -448,3 +395,31 @@ def forward(
             return (output,)
 
         return Transformer2DModelOutput(sample=output)
+
+
+class QEffFluxTransformer2DModelOF(QEffFluxTransformer2DModel):
+    def __qeff_init__(self):
+        self.transformer_blocks = nn.ModuleList()
+        self._block_classes = set()
+
+        for _ in range(self.config.num_layers):
+            BlockClass = QEffFluxTransformerBlock
+            block = BlockClass(
+                dim=self.inner_dim,
+                num_attention_heads=self.config.num_attention_heads,
+                attention_head_dim=self.config.attention_head_dim,
+            )
+            self.transformer_blocks.append(block)
+            self._block_classes.add(BlockClass)
+
+        self.single_transformer_blocks = nn.ModuleList()
+
+        for _ in range(self.config.num_single_layers):
+            SingleBlockClass = QEffFluxSingleTransformerBlock
+            single_block = SingleBlockClass(
+                dim=self.inner_dim,
+                num_attention_heads=self.config.num_attention_heads,
+                attention_head_dim=self.config.attention_head_dim,
+            )
+            self.single_transformer_blocks.append(single_block)
+            self._block_classes.add(SingleBlockClass)

QEfficient/diffusers/pipelines/flux/config/default_flux_compile_config.json renamed to QEfficient/diffusers/pipelines/flux/flux_config.json

@@ -1,5 +1,5 @@
 {
-  "description": "Example compilation configuration for Flux pipeline",
+  "description": "Default configuration for Flux pipeline",
   "model_type": "flux",
 
   "modules": 
@@ -52,9 +52,7 @@
                                         {
                                           "batch_size": 1,
                                           "seq_len": 256,
-                                          "steps": 1,
-                                          "num_layers": 1,
-                                          "num_single_layers": 1
+                                          "steps": 1
                                         },
                       "compilation": 
                                         {

QEfficient/diffusers/pipelines/flux/pipeline_flux.py

@@ -19,7 +19,6 @@
 
 from QEfficient.diffusers.pipelines.config_manager import config_manager, set_module_device_ids
 from QEfficient.diffusers.pipelines.pipeline_utils import (
-    QEffClipTextEncoder,
     QEffFluxTransformerModel,
     QEffTextEncoder,
     QEffVAE,
@@ -38,11 +37,13 @@ class QEFFFluxPipeline(FluxPipeline):
     """
 
     def __init__(self, model, use_onnx_function, *args, **kwargs):
-        self.text_encoder = QEffClipTextEncoder(model.text_encoder)
+        self.text_encoder = QEffTextEncoder(model.text_encoder)
         self.text_encoder_2 = QEffTextEncoder(model.text_encoder_2)
         self.transformer = QEffFluxTransformerModel(model.transformer, use_onnx_function=use_onnx_function)
         self.vae_decode = QEffVAE(model, "decoder")
         self.use_onnx_function = use_onnx_function
+
+        # Add all modules of FluxPipeline
         self.has_module = [
             ("text_encoder", self.text_encoder),
             ("text_encoder_2", self.text_encoder_2),
@@ -78,6 +79,10 @@ def __init__(self, model, use_onnx_function, *args, **kwargs):
         self.latent_width = self.width // self.vae_scale_factor
         self.cl = (self.latent_height * self.latent_width) // 4
 
+        self.text_encoder_2.model.config.max_position_embeddings = (
+            self.text_encoder.model.config.max_position_embeddings
+        )
+
     @classmethod
     def from_pretrained(
         cls,
@@ -140,6 +145,15 @@ def export(self, export_dir: Optional[str] = None) -> str:
                 export_kwargs=export_kwargs,
             )
 
+    def get_default_config_path():
+        """
+        Returns the default configuration file path for Flux pipeline.
+
+        Returns:
+            str: Path to the default flux_config.json file.
+        """
+        return os.path.join(os.path.dirname(__file__), "flux_config.json")
+
     def compile(
         self,
         compile_config: Optional[str] = None,
@@ -193,7 +207,6 @@ def _get_t5_prompt_embeds(
         num_images_per_prompt: int = 1,
         max_sequence_length: int = 512,
         device_ids: Optional[List[int]] = None,
-        dtype: Optional[torch.dtype] = None,
     ):
         """
         Get T5 prompt embeddings for the given prompt(s).
@@ -203,7 +216,6 @@ def _get_t5_prompt_embeds(
             num_images_per_prompt (int, defaults to 1): Number of images to generate per prompt.
             max_sequence_length (int, defaults to 256): Maximum sequence length for tokenization.
             device ids (Optional[torch.device], optional): The device to place tensors on QAIC device ids.
-            dtype (Optional[torch.dtype], optional): The data type for tensors.
 
         Returns:
             torch.Tensor: The T5 prompt embeddings with shape (batch_size * num_images_per_prompt, seq_len, hidden_size).
@@ -245,12 +257,12 @@ def _get_t5_prompt_embeds(
         self.text_encoder_2.qpc_session.set_buffers(text_encoder_2_output)
 
         aic_text_input = {"input_ids": text_input_ids.numpy().astype(np.int64)}
-        prompt_embeds = torch.tensor(self.text_encoder_2.qpc_session.run(aic_text_input)["last_hidden_state"])
+        import time
 
-        # # # AIC Testing
-        # prompt_embeds_pytorch = self.text_encoder_2.model(text_input_ids, output_hidden_states=False)
-        # mad = torch.abs(prompt_embeds_pytorch["last_hidden_state"] - prompt_embeds).mean()
-        # print(">>>>>>>>>>>> MAD for text-encoder-2 - T5 => Pytorch vs AI 100:", mad)
+        start_time = time.time()
+        prompt_embeds = torch.tensor(self.text_encoder_2.qpc_session.run(aic_text_input)["last_hidden_state"])
+        end_time = time.time()
+        print(f"T5 Text encoder inference time: {end_time - start_time:.4f} seconds")
 
         _, seq_len, _ = prompt_embeds.shape
         # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
@@ -303,26 +315,21 @@ def _get_clip_prompt_embeds(
             self.text_encoder.qpc_session = QAICInferenceSession(str(self.text_encoder.qpc_path), device_ids=device_ids)
 
         text_encoder_output = {
-            "pooler_output": np.random.rand(batch_size, embed_dim).astype(np.int32),
             "last_hidden_state": np.random.rand(batch_size, self.tokenizer_max_length, embed_dim).astype(np.int32),
+            "pooler_output": np.random.rand(batch_size, embed_dim).astype(np.int32),
         }
 
         self.text_encoder.qpc_session.set_buffers(text_encoder_output)
 
         aic_text_input = {"input_ids": text_input_ids.numpy().astype(np.int64)}
-        aic_embeddings = self.text_encoder.qpc_session.run(aic_text_input)
-        # aic_text_encoder_emb = aic_embeddings["pooler_output"]
 
-        # # # # [TEMP] CHECK ACC # #
-        # prompt_embeds_pytorch = self.text_encoder.model(text_input_ids, output_hidden_states=False)
-        # pt_pooled_embed = prompt_embeds_pytorch["pooler_output"].detach().numpy()
-        # mad = np.max(np.abs(pt_pooled_embed - aic_text_encoder_emb))
-        # print(f">>>>>>>>>>>> CLIP text encoder pooled embed MAD: ", mad) ## 0.0043082903 ##TODO : Clean up
-        ### END CHECK ACC ###
+        import time
 
-        # Use pooled output of CLIPTextModel
+        start_time = time.time()
+        aic_embeddings = self.text_encoder.qpc_session.run(aic_text_input)
+        end_time = time.time()
+        print(f"CLIP Text encoder inference time: {end_time - start_time:.4f} seconds")
         prompt_embeds = torch.tensor(aic_embeddings["pooler_output"])
-        # prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
 
         # duplicate text embeddings for each generation per prompt, using mps friendly method
         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
@@ -491,23 +498,6 @@ def __call__(
             [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
             is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
             images.
-
-        Examples:
-            ```python
-            # Basic text-to-image generation
-            from QEfficient import QEFFFluxPipeline
-            pipeline = QEFFFluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell")
-            pipeline.compile(num_devices_text_encoder=1, num_devices_transformer=4, num_devices_vae_decoder=1)
-
-            generator = torch.manual_seed(42)
-            # NOTE: guidance_scale <=1 is not supported
-            image = pipeline("A cat holding a sign that says hello world",
-                guidance_scale=0.0,
-                num_inference_steps=4,
-                max_sequence_length=256,
-                generator=generator).images[0]
-            image.save("flux-schnell_aic.png")
-            ```
         """
         device = "cpu"
 
@@ -663,7 +653,7 @@ def __call__(
                 start_time = time.time()
                 outputs = self.transformer.qpc_session.run(inputs_aic)
                 end_time = time.time()
-                print(f"Time : {end_time - start_time:.2f} seconds")
+                print(f"Transformers inference time : {end_time - start_time:.2f} seconds")
 
                 noise_pred = torch.from_numpy(outputs["output"])
 
@@ -711,8 +701,10 @@ def __call__(
             self.vae_decode.qpc_session.set_buffers(output_buffer)
 
             inputs = {"latent_sample": latents.numpy()}
+            start_time = time.time()
             image = self.vae_decode.qpc_session.run(inputs)
-
+            end_time = time.time()
+            print(f"Decoder Text encoder inference time: {end_time - start_time:.4f} seconds")
             image_tensor = torch.from_numpy(image["sample"])
             image = self.image_processor.postprocess(image_tensor, output_type=output_type)