NeuronModel API modification done

optimum/commands/export/neuronx.py

@@ -224,7 +224,7 @@ def parse_args_neuronx(parser: "ArgumentParser"):
         type=float,
         help="Scaling factors for the IP-Adapters.",
     )
-    
+
     # Static Input Shapes
     input_group = parser.add_argument_group("Input shapes")
     doc_input = "that the Neuronx-cc compiler exported model will be able to take as input."

optimum/exporters/neuron/__main__.py

@@ -36,10 +36,10 @@
     DIFFUSION_MODEL_VAE_ENCODER_NAME,
     ENCODER_NAME,
     NEURON_FILE_NAME,
-    LoRAAdapterArguments,
-    IPAdapterArguments,
     ImageEncoderArguments,
     InputShapesArguments,
+    IPAdapterArguments,
+    LoRAAdapterArguments,
     is_neuron_available,
     is_neuronx_available,
     is_transformers_neuronx_available,
@@ -298,7 +298,7 @@ def infer_stable_diffusion_shapes_from_diffusers(
                 hidden_size=model.image_encoder.vision_model.embeddings.position_embedding.weight.shape[1],
                 projection_dim=getattr(model.image_encoder.config, "projection_dim", None),
             )
-    
+
     # Format with `InputShapesArguments`
     for sub_model_name in input_shapes.keys():
         input_shapes[sub_model_name] = InputShapesArguments(**input_shapes[sub_model_name])

optimum/exporters/neuron/base.py

@@ -23,7 +23,7 @@
 from optimum.utils import logging
 
 from ...exporters.base import ExportConfig
-from ...neuron.utils import InputShapesArguments, ImageEncoderArguments, is_neuron_available
+from ...neuron.utils import ImageEncoderArguments, InputShapesArguments, is_neuron_available
 
 
 if TYPE_CHECKING:
@@ -192,8 +192,8 @@ def __init__(
             "vae_scale_factor": input_shapes.vae_scale_factor,
             "encoder_hidden_size": input_shapes.encoder_hidden_size,
             "image_encoder_shapes": ImageEncoderArguments(
-                sequence_length=getattr(input_shapes.image_encoder_shapes, "sequence_length", None), 
-                hidden_size=getattr(input_shapes.image_encoder_shapes, "hidden_size", None), 
+                sequence_length=getattr(input_shapes.image_encoder_shapes, "sequence_length", None),
+                hidden_size=getattr(input_shapes.image_encoder_shapes, "hidden_size", None),
                 projection_dim=getattr(input_shapes.image_encoder_shapes, "projection_dim", None),
             ),
         }

optimum/exporters/neuron/convert.py

@@ -348,12 +348,6 @@ def export_models(
     failed_models = []
     total_compilation_time = 0
     compile_configs = {}
-    models_and_neuron_configs.pop("text_encoder")
-    # models_and_neuron_configs.pop("text_encoder_2")
-    models_and_neuron_configs.pop("unet")
-    models_and_neuron_configs.pop("vae_encoder")
-    models_and_neuron_configs.pop("vae_decoder")
-    # models_and_neuron_configs.pop("image_encoder")
     for i, model_name in enumerate(models_and_neuron_configs.keys()):
         logger.info(f"***** Compiling {model_name} *****")
         submodel, sub_neuron_config = models_and_neuron_configs[model_name]

optimum/exporters/neuron/model_configs/traced_configs.py

@@ -632,7 +632,7 @@ def inputs(self) -> List[str]:
             if self.image_encoder_output_hidden_states:
                 common_inputs += ["image_enc_hidden_states"]
             else:
-                common_inputs += ["image_embeds"]      
+                common_inputs += ["image_embeds"]
 
         return common_inputs
 

optimum/exporters/neuron/model_wrappers.py

@@ -571,8 +571,8 @@ def forward(self, input_ids, attention_mask):
 
 class CLIPVisionModelNeuronWrapper(torch.nn.Module):
     def __init__(
-        self, 
-        model, 
+        self,
+        model,
         input_names: List[str],
         output_hidden_states: bool = True,
     ):
@@ -585,9 +585,9 @@ def forward(self, pixel_values):
         vision_outputs = self.model.vision_model(pixel_values=pixel_values, output_hidden_states=self.output_hidden_states)
         pooled_output = vision_outputs[1]
         image_embeds = self.model.visual_projection(pooled_output)
-        
+
         outputs = (image_embeds, vision_outputs.last_hidden_state)
-        
+
         if self.output_hidden_states:
             outputs += (vision_outputs.hidden_states, )
         return outputs

optimum/exporters/neuron/utils.py

@@ -18,12 +18,11 @@
 import os
 from collections import OrderedDict
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union, Any
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import torch
 
 from ...neuron.utils import (
-    LoRAAdapterArguments,
     DECODER_NAME,
     DIFFUSION_MODEL_CONTROLNET_NAME,
     DIFFUSION_MODEL_TEXT_ENCODER_2_NAME,
@@ -33,6 +32,7 @@
     DIFFUSION_MODEL_VAE_DECODER_NAME,
     DIFFUSION_MODEL_VAE_ENCODER_NAME,
     ENCODER_NAME,
+    LoRAAdapterArguments,
     get_attention_scores_sd,
     get_attention_scores_sdxl,
     neuron_scaled_dot_product_attention,
@@ -235,7 +235,7 @@ def get_diffusion_models_for_export(
         unet_neuron_config.with_ip_adapter = getattr(unet.config, "encoder_hid_dim_type", None) == "ip_image_proj"
 
         models_for_export[DIFFUSION_MODEL_UNET_NAME] = (unet, unet_neuron_config)
-        
+
 
     # Diffusion Transformer
     transformer = None
@@ -340,7 +340,7 @@ def get_diffusion_models_for_export(
             input_shapes=image_encoder_input_shapes,
         )
         models_for_export["image_encoder"] = (image_encoder, image_encoder_neuron_config)
-        models_for_export[DIFFUSION_MODEL_UNET_NAME][1].image_encoder_output_hidden_states = output_hidden_states        
+        models_for_export[DIFFUSION_MODEL_UNET_NAME][1].image_encoder_output_hidden_states = output_hidden_states
 
     return models_for_export
 

optimum/neuron/modeling_diffusion.py

@@ -22,6 +22,7 @@
 import shutil
 from abc import abstractmethod
 from collections import OrderedDict
+from dataclasses import asdict
 from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
@@ -52,6 +53,7 @@
     DIFFUSION_MODEL_VAE_ENCODER_NAME,
     NEURON_FILE_NAME,
     DiffusersPretrainedConfig,
+    NeuronArgumentParser,
     check_if_weights_replacable,
     is_neuronx_available,
     replace_weights,
@@ -877,16 +879,8 @@ def _export(
         dynamic_batch_size: bool = False,
         output_hidden_states: bool = False,
         data_parallel_mode: Optional[Literal["none", "unet", "transformer", "all"]] = None,
-        lora_model_ids: Optional[Union[str, List[str]]] = None,
-        lora_weight_names: Optional[Union[str, List[str]]] = None,
-        lora_adapter_names: Optional[Union[str, List[str]]] = None,
-        lora_scales: Optional[Union[float, List[float]]] = None,
         controlnet_ids: Optional[Union[str, List[str]]] = None,
-        ip_adapter_ids: Optional[Union[str, List[str]]] = None,
-        ip_adapter_subfolders: Optional[Union[str, List[str]]] = None,
-        ip_adapter_weight_names: Optional[Union[str, List[str]]] = None,
-        ip_adapter_scales: Optional[Union[float, List[float]]] = None,
-        **kwargs_shapes,
+        **kwargs,
     ) -> "NeuronDiffusionPipelineBase":
         """
         Args:
@@ -967,9 +961,13 @@ def _export(
                 The name of the weight file to load. If a list is passed, it should have the same length as `ip_adapter_subfolders`.
             ip_adapter_scales (`Optional[Union[float, List[float]]]`, defaults to `None`):
                 Scaling factors for the IP-Adapters.
-            kwargs_shapes (`Dict[str, int]`):
-                Shapes to use during inference. This argument allows to override the default shapes used during the export.
         """
+        # Parse kwargs to their dataclass
+        parser = NeuronArgumentParser(**kwargs)
+        lora_args = parser.lora_args
+        ip_adapter_args =  parser.ip_adapter_args
+        kwargs_shapes = asdict(parser.input_shapes)
+
         if task is None:
             if cls.task is not None:
                 task = cls.task
@@ -1024,17 +1022,11 @@ def _export(
                 local_files_only=local_files_only,
                 token=token,
                 submodels=submodels,
+                lora_args=lora_args,
+                ip_adapter_args=ip_adapter_args,
                 output_hidden_states=output_hidden_states,
-                lora_model_ids=lora_model_ids,
-                lora_weight_names=lora_weight_names,
-                lora_adapter_names=lora_adapter_names,
-                lora_scales=lora_scales,
                 torch_dtype=torch_dtype,
                 controlnet_ids=controlnet_ids,
-                ip_adapter_ids=ip_adapter_ids,
-                ip_adapter_subfolders=ip_adapter_subfolders,
-                ip_adapter_weight_names=ip_adapter_weight_names,
-                ip_adapter_scales=ip_adapter_scales,
                 **input_shapes_copy,
             )
 
@@ -1090,6 +1082,8 @@ def _export(
                 model_name_or_path=model_id,
                 output=save_dir_path,
                 compiler_kwargs=compiler_kwargs,
+                lora_args=lora_args,
+                ip_adapter_args=ip_adapter_args,
                 torch_dtype=torch_dtype,
                 task=task,
                 dynamic_batch_size=dynamic_batch_size,
@@ -1107,15 +1101,7 @@ def _export(
                 do_validation=False,
                 submodels={"unet": unet_id},
                 output_hidden_states=output_hidden_states,
-                lora_model_ids=lora_model_ids,
-                lora_weight_names=lora_weight_names,
-                lora_adapter_names=lora_adapter_names,
-                lora_scales=lora_scales,
                 controlnet_ids=controlnet_ids,
-                ip_adapter_ids=ip_adapter_ids,
-                ip_adapter_subfolders=ip_adapter_subfolders,
-                ip_adapter_weight_names=ip_adapter_weight_names,
-                ip_adapter_scales=ip_adapter_scales,
                 library_name=cls.library_name,
                 **input_shapes,
             )

optimum/neuron/modeling_traced.py

@@ -31,8 +31,8 @@
 from ..exporters.tasks import TasksManager
 from .modeling_base import NeuronModel
 from .utils import (
-    InputShapesArguments,
     NEURON_FILE_NAME,
+    InputShapesArguments,
     check_if_weights_replacable,
     is_neuron_available,
     replace_weights,

optimum/neuron/utils/__init__.py

@@ -24,7 +24,8 @@
         "IPAdapterArguments",
         "ImageEncoderArguments",
         "InputShapesArguments",
-        "convert_neuronx_compiler_args_to_neuron", 
+        "NeuronArgumentParser",
+        "convert_neuronx_compiler_args_to_neuron",
         "store_compilation_config"
     ],
     "constant": [
@@ -91,11 +92,12 @@
 
 if TYPE_CHECKING:
     from .argument_utils import (
-        LoRAAdapterArguments,
-        IPAdapterArguments,
         ImageEncoderArguments,
         InputShapesArguments,
-        convert_neuronx_compiler_args_to_neuron, 
+        IPAdapterArguments,
+        LoRAAdapterArguments,
+        NeuronArgumentParser,
+        convert_neuronx_compiler_args_to_neuron,
         store_compilation_config,
     )
     from .constant import (

optimum/neuron/utils/argument_utils.py

@@ -15,7 +15,7 @@
 """Utilities related to CLI arguments."""
 
 import os
-from dataclasses import dataclass, is_dataclass, asdict
+from dataclasses import asdict, dataclass, fields, is_dataclass
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
 from ...utils import logging
@@ -36,7 +36,7 @@ class LoRAAdapterArguments:
     weight_names: Optional[Union[str, List[str]]] = None
     adapter_names: Optional[Union[str, List[str]]] = None
     scales: Optional[Union[float, List[float]]] = None
-    
+
     def __post_init__(self):
         if isinstance(self.model_ids, str):
             self.model_ids = [self.model_ids,]
@@ -86,6 +86,44 @@ class InputShapesArguments:
     image_encoder_shapes: Optional[ImageEncoderArguments] = None
 
 
+class DataclassParser:
+    def __init__(self, **kwargs):
+        for name, cls in self.__class__.__annotations__.items():
+            if is_dataclass(cls):
+                parsed_kwargs = {k: v for k, v in kwargs.items() if k in {f.name for f in fields(cls)}}
+                setattr(self, f"{name}", cls(**parsed_kwargs))
+
+
+class NeuronArgumentParser(DataclassParser):
+    input_shapes: InputShapesArguments
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        for name, value in kwargs.items():
+            if value is not None:
+                setattr(self, name, value)
+
+    @property
+    def lora_args(self):
+        _lora_args = LoRAAdapterArguments(
+            model_ids=getattr(self, "lora_model_ids", None),
+            weight_names=getattr(self, "lora_weight_names", None),
+            adapter_names=getattr(self, "lora_adapter_names", None),
+            scales=getattr(self, "lora_scales", None),
+        )
+        return _lora_args
+
+    @property
+    def ip_adapter_args(self):
+        _ip_adapter_args = IPAdapterArguments(
+            model_id=getattr(self, "ip_adapter_id", None),
+            subfolder=getattr(self, "ip_adapter_subfolder", None),
+            weight_name=getattr(self, "ip_adapter_weight_name", None),
+            scale=getattr(self, "ip_adapter_scale", None),
+        )
+        return _ip_adapter_args
+
+
 def validate_arg(
     args,
     arg_name: str,

optimum/neuron/utils/input_generators.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """Dummy input generation classes."""
 
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
@@ -195,7 +195,7 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int
         if input_name == "image_enc_hidden_states":
             shape = [self.batch_size, 1, self.image_encoder_shapes.sequence_length, self.image_encoder_shapes.hidden_size]
             return self.random_float_tensor(shape, framework=framework, dtype=float_dtype)
-        elif input_name == "image_embeds":   
+        elif input_name == "image_embeds":
             shape = [self.batch_size, 1, self.image_encoder_shapes.projection_dim]
             return self.random_float_tensor(shape, framework=framework, dtype=float_dtype)
         elif input_name == "ip_adapter_masks":