Flux support with Custom config

Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>
Signed-off-by: Amit Raj <amitraj@qti.qualcommm.com>

Author: quic-amitraj

QEfficient/diffusers/models/normalization.py

@@ -4,10 +4,10 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 # ----------------------------------------------------------------------------
-import numbers
-from typing import Dict, Optional, Tuple
+from typing import Optional, Tuple
+
 import torch
-from diffusers.models.normalization import AdaLayerNormZero, AdaLayerNormZeroSingle,AdaLayerNormContinuous
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
 
 
 class QEffAdaLayerNormZero(AdaLayerNormZero):
@@ -21,13 +21,14 @@ def forward(
         scale_msa: Optional[torch.Tensor] = None,
         # emb: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        if self.emb is not None:
-            emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)
+        # if self.emb is not None:
+        #     emb = self.emb(timestep, class_labels, hidden_dtype=hidden_dtype)
         # emb = self.linear(self.silu(emb))
         # shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=1)
         x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
         return x
 
+
 class QEffAdaLayerNormZeroSingle(AdaLayerNormZeroSingle):
     def forward(
         self,
@@ -39,11 +40,12 @@ def forward(
         x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
         return x
 
+
 class QEffAdaLayerNormContinuous(AdaLayerNormContinuous):
     def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
         # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT)
         # emb = self.linear(self.silu(conditioning_embedding).to(x.dtype))
         emb = conditioning_embedding
         scale, shift = torch.chunk(emb, 2, dim=1)
         x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
-        return x
+        return x

QEfficient/diffusers/models/pytorch_transforms.py

@@ -5,12 +5,18 @@
 #
 # -----------------------------------------------------------------------------
 from typing import Tuple
-import torch
-from torch import nn
+
 from diffusers.models.attention import JointTransformerBlock
 from diffusers.models.attention_processor import Attention, JointAttnProcessor2_0
-from diffusers.models.normalization import RMSNorm, AdaLayerNormZero, AdaLayerNormZeroSingle, AdaLayerNormContinuous
-from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock, FluxTransformer2DModel, FluxAttnProcessor,FluxAttention
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle, RMSNorm
+from diffusers.models.transformers.transformer_flux import (
+    FluxAttention,
+    FluxAttnProcessor,
+    FluxSingleTransformerBlock,
+    FluxTransformer2DModel,
+    FluxTransformerBlock,
+)
+from torch import nn
 
 from QEfficient.base.pytorch_transforms import ModuleMappingTransform
 from QEfficient.customop.rms_norm import CustomRMSNormAIC
@@ -19,14 +25,25 @@
     QEffAttention,
     QEffJointAttnProcessor2_0,
 )
-from QEfficient.diffusers.models.transformers.transformer_flux import  QEffFluxSingleTransformerBlock, QEffFluxTransformerBlock, QEffFluxTransformer2DModel, QEffFluxAttnProcessor, QEffFluxAttention
-from QEfficient.diffusers.models.normalization import QEffAdaLayerNormZero, QEffAdaLayerNormZeroSingle, QEffAdaLayerNormContinuous
+from QEfficient.diffusers.models.normalization import (
+    QEffAdaLayerNormContinuous,
+    QEffAdaLayerNormZero,
+    QEffAdaLayerNormZeroSingle,
+)
+from QEfficient.diffusers.models.transformers.transformer_flux import (
+    QEffFluxAttention,
+    QEffFluxAttnProcessor,
+    QEffFluxSingleTransformerBlock,
+    QEffFluxTransformer2DModel,
+    QEffFluxTransformerBlock,
+)
+
 
 class CustomOpsTransform(ModuleMappingTransform):
     _module_mapping = {
         RMSNorm: CustomRMSNormAIC,
-        nn.RMSNorm: CustomRMSNormAIC #  for torch.nn.RMSNorm
-      }
+        nn.RMSNorm: CustomRMSNormAIC,  #  for torch.nn.RMSNorm
+    }
 
     @classmethod
     def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
@@ -42,21 +59,23 @@ class AttentionTransform(ModuleMappingTransform):
         FluxSingleTransformerBlock: QEffFluxSingleTransformerBlock,
         FluxTransformerBlock: QEffFluxTransformerBlock,
         FluxTransformer2DModel: QEffFluxTransformer2DModel,
-        FluxAttention : QEffFluxAttention,
-        FluxAttnProcessor: QEffFluxAttnProcessor
+        FluxAttention: QEffFluxAttention,
+        FluxAttnProcessor: QEffFluxAttnProcessor,
     }
 
     @classmethod
     def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
         model, transformed = super().apply(model)
         return model, transformed
 
+
 class NormalizationTransform(ModuleMappingTransform):
     _module_mapping = {
         AdaLayerNormZero: QEffAdaLayerNormZero,
         AdaLayerNormZeroSingle: QEffAdaLayerNormZeroSingle,
         AdaLayerNormContinuous: QEffAdaLayerNormContinuous,
     }
+
     @classmethod
     def apply(cls, model: nn.Module) -> Tuple[nn.Module, bool]:
         model, transformed = super().apply(model)

QEfficient/diffusers/models/transformers/transformer_flux.py

@@ -4,23 +4,33 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 # ----------------------------------------------------------------------------
-import os
-from typing import Any, Callable, Dict, List, Tuple, Optional, Union
+from typing import Any, Dict, Optional, Tuple, Union
 from venv import logger
 
+import numpy as np
 import torch
 import torch.nn as nn
-import numpy as np
-
-from diffusers.models.transformers.transformer_flux import FluxAttention,FluxSingleTransformerBlock, FluxTransformerBlock, FluxTransformer2DModel, FluxPosEmbed, FluxAttnProcessor, _get_qkv_projections
-from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.transformers.transformer_flux import (
+    FluxAttention,
+    FluxAttnProcessor,
+    FluxSingleTransformerBlock,
+    FluxTransformer2DModel,
+    FluxTransformerBlock,
+    _get_qkv_projections,
+)
+
+from QEfficient.diffusers.models.normalization import (
+    QEffAdaLayerNormContinuous,
+    QEffAdaLayerNormZero,
+    QEffAdaLayerNormZeroSingle,
+)
 
-from QEfficient.diffusers.models.normalization import QEffAdaLayerNormZero, QEffAdaLayerNormZeroSingle, QEffAdaLayerNormContinuous
 
 def qeff_apply_rotary_emb(
-    x: torch.Tensor,
-    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]]) -> Tuple[torch.Tensor, torch.Tensor]:
+    x: torch.Tensor, freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]]
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
     to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
@@ -39,12 +49,13 @@ def qeff_apply_rotary_emb(
     cos = cos[None, :, None, :]
     sin = sin[None, :, None, :]
     cos, sin = cos.to(x.device), sin.to(x.device)
-    B, S, H , D  = x.shape
-    x_real, x_imag = x.reshape(B, -1, H, D//2, 2).unbind(-1)
+    B, S, H, D = x.shape
+    x_real, x_imag = x.reshape(B, -1, H, D // 2, 2).unbind(-1)
     x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
     out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
     return out
 
+
 class QEffFluxAttnProcessor(FluxAttnProcessor):
     _attention_backend = None
     _parallel_config = None
@@ -102,6 +113,7 @@ def __call__(
         else:
             return hidden_states
 
+
 class QEffFluxAttention(FluxAttention):
     def __qeff_init__(self):
         processor = QEffFluxAttnProcessor()
@@ -158,6 +170,7 @@ def forward(
         encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:]
         return encoder_hidden_states, hidden_states
 
+
 class QEffFluxTransformerBlock(FluxTransformerBlock):
     def __init__(
         self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6
@@ -178,7 +191,6 @@ def __init__(
             eps=eps,
         )
 
-
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -187,15 +199,12 @@ def forward(
         image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
         joint_attention_kwargs: Optional[Dict[str, Any]] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
-
         temb1 = tuple(torch.split(temb[:6], 1))
         temb2 = tuple(torch.split(temb[6:], 1))
         norm_hidden_states = self.norm1(hidden_states, shift_msa=temb1[0], scale_msa=temb1[1])
         gate_msa, shift_mlp, scale_mlp, gate_mlp = temb1[-4:]
 
-        norm_encoder_hidden_states = self.norm1_context(
-            encoder_hidden_states, shift_msa=temb2[0], scale_msa=temb2[1]
-        )
+        norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, shift_msa=temb2[0], scale_msa=temb2[1])
 
         c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = temb2[-4:]
 
@@ -242,6 +251,7 @@ def forward(
 
         return encoder_hidden_states, hidden_states
 
+
 class QEffFluxTransformer2DModel(FluxTransformer2DModel):
     def __init__(
         self,
@@ -257,7 +267,6 @@ def __init__(
         guidance_embeds: bool = False,
         axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
     ):
-
         super().__init__(
             patch_size=patch_size,
             in_channels=in_channels,
@@ -296,7 +305,6 @@ def __init__(
 
         self.norm_out = QEffAdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
 
-
     def forward(
         self,
         hidden_states: torch.Tensor,

QEfficient/diffusers/pipelines/config_manager.py

@@ -0,0 +1,37 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# ----------------------------------------------------------------------------
+
+import os
+from typing import Optional
+
+from QEfficient.utils._utils import load_json
+
+
+def config_manager(cls, config_source: Optional[str] = None):
+    """
+    JSON-based compilation configuration manager for diffusion pipelines.
+
+    Supports loading configuration from JSON files only. Automatically detects
+    model type and handles model-specific requirements.
+    Initialize the configuration manager.
+
+    Args:
+        config_source: Path to JSON configuration file. If None, uses default config.
+    """
+    if config_source is None:
+        config_source = cls.get_default_config_path()
+
+    if not isinstance(config_source, str):
+        raise ValueError("config_source must be a path to JSON configuration file")
+
+    # Direct use of load_json utility - no wrapper needed
+    if not os.path.exists(config_source):
+        raise FileNotFoundError(f"Configuration file not found: {config_source}")
+
+    cls._compile_config = load_json(config_source)
+
+   

QEfficient/diffusers/pipelines/flux/__init__.py

@@ -3,4 +3,4 @@
 # Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
 # SPDX-License-Identifier: BSD-3-Clause
 #
-# ----------------------------------------------------------------------------
+# ----------------------------------------------------------------------------

QEfficient/diffusers/pipelines/flux/config/default_flux_compile_config.json

@@ -0,0 +1,97 @@
+{
+  "description": "Example compilation configuration for Flux pipeline",
+  "model_type": "flux",
+
+  "modules": 
+  {
+    "text_encoder": 
+                    {
+                      "specializations":{
+                                          "batch_size": 1,
+                                          "seq_len": 77
+                                        },
+                      "compilation": 
+                                        {
+                                          "onnx_path": null,
+                                          "compile_dir": null,
+                                          "mdp_ts_num_devices": 1,
+                                          "mxfp6_matmul": false,
+                                          "convert_to_fp16": true,
+                                          "aic_num_cores": 16
+                                        },
+                      "execute":
+                                        {
+                                          "device_ids": null
+                                        }  
+
+                    },
+    "text_encoder_2": 
+                    {
+                      "specializations": 
+                                        {
+                                          "batch_size": 1,
+                                          "seq_len": 256
+                                        },
+                      "compilation": 
+                                        {
+                                          "onnx_path": null,
+                                          "compile_dir": null,
+                                          "mdp_ts_num_devices": 1,
+                                          "mxfp6_matmul": false,
+                                          "convert_to_fp16": true,
+                                          "aic_num_cores": 16
+                                        },
+                      "execute":
+                                        {
+                                          "device_ids": null
+                                        }
+                    },
+    "transformer": 
+                    {
+                      "specializations": 
+                                        {
+                                          "batch_size": 1,
+                                          "seq_len": 256,
+                                          "steps": 1,
+                                          "num_layers": 1,
+                                          "num_single_layers": 1
+                                        },
+                      "compilation": 
+                                        {
+                                          "onnx_path": null,
+                                          "compile_dir": null,
+                                          "mdp_ts_num_devices": 4,
+                                          "mxfp6_matmul": true,
+                                          "convert_to_fp16": true,
+                                          "aic_num_cores": 16,
+                                          "mos": 1,
+                                          "mdts-mos": 1
+                                        },
+                      "execute":
+                                        {
+                                          "device_ids": null
+                                        }
+                    },
+    "vae_decoder": 
+                    {
+                      "specializations": 
+                                        {
+                                          "batch_size": 1,
+                                          "channels": 16
+                                        },
+                      "compilation": 
+                                        {
+                                          "onnx_path": null,
+                                          "compile_dir": null,
+                                          "mdp_ts_num_devices": 1,
+                                          "mxfp6_matmul": false,
+                                          "convert_to_fp16": true,
+                                          "aic_num_cores": 16
+                                        },
+                      "execute":
+                                        {
+                                          "device_ids": null
+                                        }
+                    }
+  }
+}