Rebased with main and fixed some issues

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

Author: Amit Raj

QEfficient/base/modeling_qeff.py

@@ -288,7 +288,7 @@ def _export(
         self.onnx_path = onnx_path
         return onnx_path
 
-    # @dump_qconfig
+    @dump_qconfig
     def _compile(
         self,
         onnx_path: Optional[str] = None,

QEfficient/diffusers/pipelines/flux/pipeline_flux.py

@@ -601,7 +601,7 @@ def __call__(
                     continue
 
                 timestep = t.expand(latents.shape[0]).to(latents.dtype)
-                temb = self.transformer.model.time_text_embed(timestep, pooled_prompt_embeds)
+                temb = self.transformer.model.time_text_embed(timestep.cpu(), pooled_prompt_embeds.cpu())
 
                 adaln_emb = []
                 for i in range(len(self.transformer.model.transformer_blocks)):

QEfficient/diffusers/pipelines/pipeline_module.py

@@ -6,7 +6,6 @@
 # ----------------------------------------------------------------------------
 
 import copy
-import hashlib
 
 import torch
 import torch.nn as nn
@@ -23,7 +22,6 @@
     T5ModelTransform,
 )
 from QEfficient.utils import constants
-from QEfficient.utils.cache import to_hashable
 
 
 class QEffTextEncoder(QEFFBaseModel):
@@ -57,6 +55,18 @@ def get_onnx_config(self):
 
         return example_inputs, dynamic_axes, output_names
 
+    @property
+    def get_model_config(self) -> dict:
+        """
+        Get the model configuration as a dictionary.
+
+        Returns
+        -------
+        dict
+            The configuration dictionary of the underlying HuggingFace model.
+        """
+        return self.model.config.__dict__
+
     def export(
         self,
         inputs,
@@ -76,15 +86,6 @@ def export(
     def compile(self, specializations, **compiler_options):
         self._compile(specializations=specializations, **compiler_options)
 
-    @property
-    def model_hash(self) -> str:
-        # Compute the hash with: model_config, continuous_batching, transforms
-        mhash = hashlib.sha256()
-        mhash.update(to_hashable(self.model.config.to_diff_dict()))
-        mhash.update(to_hashable(self._transform_names()))
-        mhash = mhash.hexdigest()[:16]
-        return mhash
-
     @property
     def model_name(self) -> str:
         mname = self.model.__class__.__name__
@@ -125,18 +126,21 @@ def export(
             export_kwargs=export_kwargs,
         )
 
+    @property
+    def get_model_config(self) -> dict:
+        """
+        Get the model configuration as a dictionary.
+
+        Returns
+        -------
+        dict
+            The configuration dictionary of the underlying HuggingFace model.
+        """
+        return self.model.config.__dict__
+
     def compile(self, specializations, **compiler_options):
         self._compile(specializations=specializations, **compiler_options)
 
-    @property
-    def model_hash(self) -> str:
-        # Compute the hash with: model_config, continuous_batching, transforms
-        mhash = hashlib.sha256()
-        mhash.update(to_hashable(dict(self.model.config)))
-        mhash.update(to_hashable(self._transform_names()))
-        mhash = mhash.hexdigest()[:16]
-        return mhash
-
     @property
     def model_name(self) -> str:
         mname = self.model.__class__.__name__
@@ -197,14 +201,16 @@ def compile(self, specializations, **compiler_options):
         self._compile(specializations=specializations, **compiler_options)
 
     @property
-    def model_hash(self) -> str:
-        # Compute the hash with: model_config, continuous_batching, transforms
-        mhash = hashlib.sha256()
-        mhash.update(to_hashable(dict(self.model.config)))
-        mhash.update(to_hashable(self._transform_names()))
-        mhash.update(to_hashable(self.type))
-        mhash = mhash.hexdigest()[:16]
-        return mhash
+    def get_model_config(self) -> dict:
+        """
+        Get the model configuration as a dictionary.
+
+        Returns
+        -------
+        dict
+            The configuration dictionary of the underlying HuggingFace model.
+        """
+        return self.model.config.__dict__
 
     @property
     def model_name(self) -> str:
@@ -250,13 +256,16 @@ def compile(self, specializations, **compiler_options):
         self._compile(specializations=specializations, **compiler_options)
 
     @property
-    def model_hash(self) -> str:
-        # Compute the hash with: model_config, continuous_batching, transforms
-        mhash = hashlib.sha256()
-        mhash.update(to_hashable(self.model.config.to_diff_dict()))
-        mhash.update(to_hashable(self._transform_names()))
-        mhash = mhash.hexdigest()[:16]
-        return mhash
+    def get_model_config(self) -> dict:
+        """
+        Get the model configuration as a dictionary.
+
+        Returns
+        -------
+        dict
+            The configuration dictionary of the underlying HuggingFace model.
+        """
+        return self.model.config.__dict__
 
     @property
     def model_name(self) -> str:
@@ -282,7 +291,8 @@ def __init__(self, model: nn.modules, use_onnx_function):
         if use_onnx_function:
             self._pytorch_transforms.append(OnnxFunctionTransform)
             model, _ = OnnxFunctionTransform.apply(model)
-        self.model = model
+        # Ensure the model and all its submodules are on CPU to avoid meta device issues
+        self.model = model.to("cpu")
 
     def get_onnx_config(self, batch_size=1, seq_length=256, cl=4096):
         example_inputs = {
@@ -313,6 +323,18 @@ def get_onnx_config(self, batch_size=1, seq_length=256, cl=4096):
 
         return example_inputs, dynamic_axes, output_names
 
+    @property
+    def get_model_config(self) -> dict:
+        """
+        Get the model configuration as a dictionary.
+
+        Returns
+        -------
+        dict
+            The configuration dictionary of the underlying HuggingFace model.
+        """
+        return self.model.config.__dict__
+
     def export(
         self,
         inputs,
@@ -347,17 +369,6 @@ def get_specializations(self, batch_size: int, seq_len: int, cl: int):
     def compile(self, specializations, **compiler_options):
         self._compile(specializations=specializations, **compiler_options)
 
-    @property
-    def model_hash(self) -> str:
-        # Compute the hash with: model_config, continuous_batching, transforms
-        mhash = hashlib.sha256()
-        dict_model_config = dict(self.model.config)
-        dict_model_config.pop("_use_default_values", None)
-        mhash.update(to_hashable(dict_model_config))
-        mhash.update(to_hashable(self._transform_names()))
-        mhash = mhash.hexdigest()[:16]
-        return mhash
-
     @property
     def model_name(self) -> str:
         mname = self.model.__class__.__name__

QEfficient/diffusers/pipelines/pipeline_utils.py

@@ -50,6 +50,7 @@ def set_module_device_ids(cls):
     for module_name, module_obj in cls.modules.items():
         module_obj.device_ids = config_modules[module_name]["execute"]["device_ids"]
 
+
 @dataclass(frozen=True)
 class ModulePerf:
     module_name: str

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.to_diff_dict()
+    model_params["config"] = qeff_model.model.config
     model_params["peft_config"] = getattr(qeff_model.model, "active_peft_config", None)
     model_params["applied_transform_names"] = qeff_model._transform_names()
     return model_params

QEfficient/utils/hash_utils.py

@@ -15,6 +15,9 @@
 def json_serializable(obj):
     if isinstance(obj, set):
         return sorted(obj)
+    # Handle objects with to_dict() method (e.g., transformers config objects)
+    if hasattr(obj, "to_dict") and callable(getattr(obj, "to_dict")):
+        return obj.to_dict()
     raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable")