Code cleaning and removed redundant code-2

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

Author: Amit Raj

QEfficient/diffusers/README.md

@@ -68,8 +68,6 @@ pip install "QEfficient[diffusers]"
 ---
 
 ## 🎯 Supported Models
-
-### Stable Diffusion 3.x Series
 - ✅ [`black-forest-labs/FLUX.1-schnell`](https://huggingface.co/black-forest-labs/FLUX.1-schnell)
 
 ---

QEfficient/diffusers/models/transformers/__init__.py

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

QEfficient/diffusers/pipelines/flux/pipeline_flux.py

@@ -5,7 +5,6 @@
 #
 # ----------------------------------------------------------------------------
 
-import logging
 import os
 import time
 from typing import Any, Callable, Dict, List, Optional, Union
@@ -28,9 +27,7 @@
     set_module_device_ids,
 )
 from QEfficient.generation.cloud_infer import QAICInferenceSession
-
-# Initialize logger for this module
-logger = logging.getLogger(__name__)
+from QEfficient.utils.logging_utils import logger
 
 
 class QEFFFluxPipeline(FluxPipeline):
@@ -224,18 +221,18 @@ def compile(self, compile_config: Optional[str] = None) -> None:
         # Compile each module with its specific configuration
         for module_name, module_obj in self.modules.items():
             module_config = self.custom_config["modules"]
-            specializations = [module_config[module_name]["specializations"]]
+            specializations = module_config[module_name]["specializations"]
             compile_kwargs = module_config[module_name]["compilation"]
 
             # Set dynamic specialization values based on image dimensions
             if module_name == "transformer":
-                specializations[0]["cl"] = self.cl
+                specializations["cl"] = self.cl
             elif module_name == "vae_decoder":
-                specializations[0]["latent_height"] = self.latent_height
-                specializations[0]["latent_width"] = self.latent_width
+                specializations["latent_height"] = self.latent_height
+                specializations["latent_width"] = self.latent_width
 
             # Compile the module to QPC format
-            module_obj.compile(specializations=specializations, **compile_kwargs)
+            module_obj.compile(specializations=[specializations], **compile_kwargs)
 
     def _get_t5_prompt_embeds(
         self,

QEfficient/diffusers/pipelines/pipeline_module.py

@@ -6,7 +6,6 @@
 # ----------------------------------------------------------------------------
 
 import copy
-import logging
 from typing import Dict, List, Tuple
 
 import torch
@@ -25,9 +24,6 @@
 )
 from QEfficient.utils import constants
 
-# Initialize logger for this module
-logger = logging.getLogger(__name__)
-
 
 class QEffTextEncoder(QEFFBaseModel):
     """

QEfficient/utils/_utils.py

@@ -532,7 +532,7 @@ def create_model_params(qeff_model, **kwargs) -> Dict:
     """
     model_params = copy.deepcopy(kwargs)
     model_params = {k: v for k, v in model_params.items() if k in KWARGS_INCLUSION_LIST}
-    model_params["config"] = qeff_model.model.config
+    model_params["config"] = qeff_model.model.config.to_diff_dict()
     model_params["peft_config"] = getattr(qeff_model.model, "active_peft_config", None)
     model_params["applied_transform_names"] = qeff_model._transform_names()
     return model_params

examples/diffusers/flux/flux_1_schnell.py

@@ -4,20 +4,48 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 # -----------------------------------------------------------------------------
+
+"""
+FLUX.1-schnell Image Generation Example
+
+This example demonstrates how to use the QEFFFluxPipeline to generate images
+using the FLUX.1-schnell model from Black Forest Labs. FLUX.1-schnell is a
+fast, distilled version of the FLUX.1 text-to-image model optimized for
+speed with minimal quality loss.
+
+Key Features:
+- Fast inference with only 4 steps
+- High-quality image generation from text prompts
+- Optimized for Qualcomm Cloud AI 100 using ONNX runtime
+- Deterministic output using fixed random seed
+
+Output:
+- Generates an image based on the text prompt
+- Saves the image as 'cat_with_sign.png' in the current directory
+"""
+
 import torch
 
 from QEfficient import QEFFFluxPipeline
 
-pipeline = QEFFFluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", use_onnx_function=True)
+# Initialize the FLUX.1-schnell pipeline from pretrained weights
+# use_onnx_function=True enables ONNX-based optimizations for faster compilation
+pipeline = QEFFFluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", use_onnx_function=False)
 
+# Generate an image from a text prompt
 output = pipeline(
     prompt="A cat holding a sign that says hello world",
     guidance_scale=0.0,
     num_inference_steps=4,
     max_sequence_length=256,
     generator=torch.manual_seed(42),
 )
+
+# Extract the generated image from the output
 image = output.images[0]
+
+# Save the generated image to disk
 image.save("cat_with_sign.png")
 
+# Print the output object (contains perf info)
 print(output)

examples/diffusers/flux/flux_1_shnell_custom.py

@@ -37,17 +37,17 @@
 # Option 1: Basic initialization with custom image dimensions
 pipeline = QEFFFluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell",
-    height=256,
-    width=256,
+    height=512,
+    width=512,
 )
 
 # Option 2: Advanced initialization with custom modules
 # Uncomment and modify to use your own custom components:
 #
 # pipeline = QEFFFluxPipeline.from_pretrained(
 #     "black-forest-labs/FLUX.1-schnell",
-#     height=256,
-#     width=256,
+#     height=512,
+#     width=512,
 #     text_encoder=custom_text_encoder,      # Your custom CLIP text encoder
 #     transformer=custom_transformer,         # Your custom transformer model
 #     tokenizer=custom_tokenizer,             # Your custom tokenizer

scripts/Jenkinsfile

@@ -22,6 +22,7 @@ pipeline {
                    . preflight_qeff/bin/activate &&
                    pip install --upgrade pip setuptools &&
                    pip install .[test] &&
+                   pip install .[diffusers] &&
                    pip install junitparser pytest-xdist &&
                    pip install librosa==0.10.2 soundfile==0.13.1 && #packages needed to load example for whisper testing
                    pip install --extra-index-url https://download.pytorch.org/whl/cpu timm==1.0.14 torchvision==0.22.0+cpu einops==0.8.1 && #packages to load VLMs
@@ -34,7 +35,7 @@ pipeline {
            parallel {
                stage('Run Non-CLI Non-QAIC Tests') {
                    steps {
-                       timeout(time: 25, unit: 'MINUTES') {
+                       timeout(time: 100, unit: 'MINUTES') {
                            sh '''
                            sudo docker exec ${BUILD_TAG} bash -c "
                            cd /efficient-transformers &&