Comments addressed-2

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

Author: Amit Raj

QEfficient/base/modeling_qeff.py

@@ -135,16 +135,19 @@ def model_name(self) -> str:
             mname = mname[4:]
         return mname
 
-    # TODO: Make get_model_config method abstract
     @property
+    @abstractmethod
     def get_model_config(self) -> Dict:
         """
         Get the model configuration as a dictionary.
 
+        This is an abstract property that must be implemented by all subclasses.
+        Typically returns: self.model.config.__dict__
+
         Returns:
-            Dict: The configuration dictionary of the underlying HuggingFace model
+            Dict: The configuration dictionary of the underlying model
         """
-        return self.model.config.__dict__
+        pass
 
     @abstractmethod
     def export(self, export_dir: Optional[str] = None) -> Path:

QEfficient/diffusers/pipelines/pipeline_module.py

@@ -44,6 +44,16 @@ class QEffTextEncoder(QEFFBaseModel):
     _pytorch_transforms = [CustomOpsTransform, T5ModelTransform]
     _onnx_transforms = [FP16ClipTransform, SplitTensorsTransform]
 
+    @property
+    def get_model_config(self) -> Dict:
+        """
+        Get the model configuration as a dictionary.
+
+        Returns:
+            Dict: The configuration dictionary of the underlying text encoder model
+        """
+        return self.model.config.__dict__
+
     def __init__(self, model: nn.Module) -> None:
         """
         Initialize the text encoder wrapper.
@@ -143,6 +153,16 @@ class QEffUNet(QEFFBaseModel):
     _pytorch_transforms = [CustomOpsTransform]
     _onnx_transforms = [FP16ClipTransform, SplitTensorsTransform]
 
+    @property
+    def get_model_config(self) -> Dict:
+        """
+        Get the model configuration as a dictionary.
+
+        Returns:
+            Dict: The configuration dictionary of the underlying UNet model
+        """
+        return self.model.config.__dict__
+
     def __init__(self, model: nn.Module) -> None:
         """
         Initialize the UNet wrapper.
@@ -211,6 +231,16 @@ class QEffVAE(QEFFBaseModel):
     _pytorch_transforms = [CustomOpsTransform]
     _onnx_transforms = [FP16ClipTransform, SplitTensorsTransform]
 
+    @property
+    def get_model_config(self) -> Dict:
+        """
+        Get the model configuration as a dictionary.
+
+        Returns:
+            Dict: The configuration dictionary of the underlying VAE model
+        """
+        return self.model.config.__dict__
+
     def __init__(self, model: nn.Module, type: str) -> None:
         """
         Initialize the VAE wrapper.
@@ -314,6 +344,16 @@ class QEffFluxTransformerModel(QEFFBaseModel):
     _pytorch_transforms = [AttentionTransform, NormalizationTransform, CustomOpsTransform]
     _onnx_transforms = [FP16ClipTransform, SplitTensorsTransform]
 
+    @property
+    def get_model_config(self) -> Dict:
+        """
+        Get the model configuration as a dictionary.
+
+        Returns:
+            Dict: The configuration dictionary of the underlying Flux transformer model
+        """
+        return self.model.config.__dict__
+
     def __init__(self, model: nn.Module) -> None:
         """
         Initialize the Flux transformer wrapper.

QEfficient/transformers/models/modeling_auto.py

@@ -289,6 +289,18 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k
 
         return cls(model, pretrained_model_name_or_path=pretrained_model_name_or_path, pooling=pooling, **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 export(self, export_dir: Optional[str] = None, use_onnx_subfunctions: bool = False) -> str:
         """
         Export the model to ONNX format using ``torch.onnx.export``.
@@ -2059,6 +2071,18 @@ def cloud_ai_100_generate(
             ),
         )
 
+    @property
+    def get_model_config(self) -> dict:
+        """
+        Get the configuration dictionary of the underlying HuggingFace model.
+
+        Returns
+        -------
+        dict
+            The configuration dictionary.
+        """
+        return self.model.config.__dict__
+
 
 class QEFFAutoModelForImageTextToText:
     """
@@ -2429,6 +2453,18 @@ def from_pretrained(
             **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 export(self, export_dir: Optional[str] = None, use_onnx_subfunctions: bool = False, **kwargs) -> str:
         """
         Export the model to ONNX format using ``torch.onnx.export``.
@@ -3149,6 +3185,18 @@ def __init__(self, model: nn.Module, **kwargs):
         self.num_layers = model.config.num_hidden_layers
         self.hash_params["qeff_auto_class"] = self.__class__.__name__
 
+    @property
+    def get_model_config(self) -> dict:
+        """
+        Get the configuration dictionary of the underlying HuggingFace model.
+
+        Returns
+        -------
+        dict
+            The configuration dictionary.
+        """
+        return self.model.config.__dict__
+
     def export(self, export_dir: Optional[str] = None, use_onnx_subfunctions: bool = False) -> str:
         """
         Export the model to ONNX format using ``torch.onnx.export``.
@@ -3521,6 +3569,10 @@ def from_pretrained(cls, pretrained_model_name_or_path, pooling=None, *args, **k
 
         return cls(model, pretrained_model_name_or_path=pretrained_model_name_or_path, pooling=pooling, **kwargs)
 
+    @property
+    def get_model_config(self) -> dict:
+        return self.model.config.__dict__
+
     def export(self, export_dir: Optional[str] = None, use_onnx_subfunctions: bool = False) -> str:
         """
         Exports the model to ``ONNX`` format using ``torch.onnx.export``.