# coding=utf-8 # Copyright 2023 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Entry point to the optimum.exporters.neuron command line.""" import argparse import inspect import os os.environ["TORCHDYNAMO_DISABLE"] = "1" # Always turn off torchdynamo as it's incompatible with neuron from argparse import ArgumentParser from dataclasses import fields from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union import torch from requests.exceptions import ConnectionError as RequestsConnectionError from transformers import AutoConfig, AutoTokenizer, PretrainedConfig from optimum.exporters.error_utils import AtolError, OutputMatchError, ShapeError from optimum.exporters.tasks import TasksManager from optimum.utils import is_diffusers_available, logging from optimum.utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors from ...neuron.models.auto_model import get_neuron_model_class, has_neuron_model_class from ...neuron.utils import ( DECODER_NAME, DIFFUSION_MODEL_CONTROLNET_NAME, DIFFUSION_MODEL_TEXT_ENCODER_2_NAME, DIFFUSION_MODEL_TEXT_ENCODER_NAME, DIFFUSION_MODEL_TRANSFORMER_NAME, DIFFUSION_MODEL_UNET_NAME, DIFFUSION_MODEL_VAE_DECODER_NAME, DIFFUSION_MODEL_VAE_ENCODER_NAME, ENCODER_NAME, NEURON_FILE_NAME, ImageEncoderArguments, InputShapesArguments, IPAdapterArguments, LoRAAdapterArguments, is_neuron_available, is_neuronx_available, map_torch_dtype, ) from ...neuron.utils.version_utils import ( check_compiler_compatibility_for_stable_diffusion, ) from .base import NeuronExportConfig from .convert import export_models, validate_models_outputs from .model_configs import * # noqa: F403 from .utils import ( build_stable_diffusion_components_mandatory_shapes, check_mandatory_input_shapes, get_diffusion_models_for_export, get_encoder_decoder_models_for_export, replace_stable_diffusion_submodels, ) if is_neuron_available(): from ...commands.export.neuron import parse_args_neuron NEURON_COMPILER = "Neuron" if is_neuronx_available(): from ...commands.export.neuronx import parse_args_neuronx as parse_args_neuron # noqa: F811 NEURON_COMPILER = "Neuronx" if is_diffusers_available(): from diffusers import StableDiffusionXLPipeline if TYPE_CHECKING: from transformers import PreTrainedModel if is_diffusers_available(): from diffusers import DiffusionPipeline, ModelMixin, StableDiffusionPipeline logger = logging.get_logger() logger.setLevel(logging.INFO) def infer_compiler_kwargs(args: argparse.Namespace) -> Dict[str, Any]: # infer compiler kwargs auto_cast = None if args.auto_cast == "none" else args.auto_cast auto_cast_type = None if auto_cast is None else args.auto_cast_type compiler_kwargs = {"auto_cast": auto_cast, "auto_cast_type": auto_cast_type} if hasattr(args, "disable_fast_relayout"): compiler_kwargs["disable_fast_relayout"] = getattr(args, "disable_fast_relayout") if hasattr(args, "disable_fallback"): compiler_kwargs["disable_fallback"] = getattr(args, "disable_fallback") return compiler_kwargs def infer_task(model_name_or_path: str) -> str: try: return TasksManager.infer_task_from_model(model_name_or_path) except KeyError as e: raise KeyError( "The task could not be automatically inferred. Please provide the argument --task with the task " f"from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}" ) except RequestsConnectionError as e: raise RequestsConnectionError( f"The task could not be automatically inferred as this is available only for models hosted on the Hugging Face Hub. Please provide the argument --task with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}" ) # This function is not applicable for diffusers / sentence transformers models def get_input_shapes(task: str, args: argparse.Namespace) -> Dict[str, int]: neuron_config_constructor = get_neuron_config_class(task, args.model) input_args = neuron_config_constructor.func.get_input_args_for_task(task) return {name: getattr(args, name) for name in input_args} def get_neuron_config_class(task: str, model_id: str) -> NeuronExportConfig: config = AutoConfig.from_pretrained(model_id) model_type = config.model_type.replace("_", "-") if config.is_encoder_decoder: model_type = model_type + "-encoder" neuron_config_constructor = TasksManager.get_exporter_config_constructor( model_type=model_type, exporter="neuron", task=task, library_name="transformers", ) return neuron_config_constructor def normalize_sentence_transformers_input_shapes(args: argparse.Namespace) -> Dict[str, int]: args = vars(args) if isinstance(args, argparse.Namespace) else args if "clip" in args.get("model", "").lower(): mandatory_axes = {"text_batch_size", "image_batch_size", "sequence_length", "num_channels", "width", "height"} else: mandatory_axes = {"batch_size", "sequence_length"} if not mandatory_axes.issubset(set(args.keys())): raise AttributeError( f"Shape of {mandatory_axes} are mandatory for neuron compilation, while {mandatory_axes.difference(args.keys())} are not given." ) mandatory_shapes = {name: args[name] for name in mandatory_axes} return mandatory_shapes def customize_optional_outputs(args: argparse.Namespace) -> Dict[str, bool]: """ Customize optional outputs of the traced model, eg. if `output_attentions=True`, the attentions tensors will be traced. """ possible_outputs = ["output_attentions", "output_hidden_states"] customized_outputs = {} for name in possible_outputs: customized_outputs[name] = getattr(args, name, False) return customized_outputs def parse_optlevel(args: argparse.Namespace) -> Dict[str, bool]: """ (NEURONX ONLY) Parse the level of optimization the compiler should perform. If not specified apply `O2`(the best balance between model performance and compile time). """ if is_neuronx_available(): if args.O1: optlevel = "1" elif args.O2: optlevel = "2" elif args.O3: optlevel = "3" else: optlevel = "2" else: optlevel = None return optlevel def normalize_stable_diffusion_input_shapes( args: argparse.Namespace, ) -> Dict[str, Dict[str, int]]: args = vars(args) if isinstance(args, argparse.Namespace) else args mandatory_axes = set(getattr(inspect.getfullargspec(build_stable_diffusion_components_mandatory_shapes), "args")) mandatory_axes = mandatory_axes - { "sequence_length", # `sequence_length` is optional, diffusers will pad it to the max if not provided. # remove number of channels. "unet_or_transformer_num_channels", "vae_encoder_num_channels", "vae_decoder_num_channels", "num_images_per_prompt", # default to 1 } if not mandatory_axes.issubset(set(args.keys())): raise AttributeError( f"Shape of {mandatory_axes} are mandatory for neuron compilation, while {mandatory_axes.difference(args.keys())} are not given." ) mandatory_shapes = {name: args[name] for name in mandatory_axes} mandatory_shapes["num_images_per_prompt"] = args.get("num_images_per_prompt", 1) or 1 mandatory_shapes["sequence_length"] = args.get("sequence_length", None) input_shapes = build_stable_diffusion_components_mandatory_shapes(**mandatory_shapes) return input_shapes def infer_stable_diffusion_shapes_from_diffusers( input_shapes: Dict[str, Dict[str, int]], model: Union["StableDiffusionPipeline", "StableDiffusionXLPipeline"], has_controlnets: bool, ): if model.tokenizer is not None: max_sequence_length = model.tokenizer.model_max_length elif hasattr(model, "tokenizer_2") and model.tokenizer_2 is not None: max_sequence_length = model.tokenizer_2.model_max_length else: raise AttributeError( f"Cannot infer max sequence_length from {type(model)} as there is no tokenizer as attribute." ) vae_encoder_num_channels = model.vae.config.in_channels vae_decoder_num_channels = model.vae.config.latent_channels vae_scale_factor = 2 ** (len(model.vae.config.block_out_channels) - 1) or 8 height = input_shapes["unet_or_transformer"]["height"] scaled_height = height // vae_scale_factor width = input_shapes["unet_or_transformer"]["width"] scaled_width = width // vae_scale_factor # Text encoders if input_shapes["text_encoder"].get("sequence_length") is None: input_shapes["text_encoder"].update({"sequence_length": max_sequence_length}) if hasattr(model, "text_encoder_2"): input_shapes["text_encoder_2"] = input_shapes["text_encoder"] # UNet or Transformer unet_or_transformer_name = "transformer" if hasattr(model, "transformer") else "unet" unet_or_transformer_num_channels = getattr(model, unet_or_transformer_name).config.in_channels input_shapes["unet_or_transformer"].update( { "num_channels": unet_or_transformer_num_channels, "height": scaled_height, "width": scaled_width, } ) if input_shapes["unet_or_transformer"].get("sequence_length") is None: input_shapes["unet_or_transformer"]["sequence_length"] = max_sequence_length input_shapes["unet_or_transformer"]["vae_scale_factor"] = vae_scale_factor input_shapes[unet_or_transformer_name] = input_shapes.pop("unet_or_transformer") if unet_or_transformer_name == "transformer": input_shapes[unet_or_transformer_name]["encoder_hidden_size"] = model.text_encoder.config.hidden_size # VAE input_shapes["vae_encoder"].update({"num_channels": vae_encoder_num_channels, "height": height, "width": width}) input_shapes["vae_decoder"].update( {"num_channels": vae_decoder_num_channels, "height": scaled_height, "width": scaled_width} ) # ControlNet if has_controlnets: encoder_hidden_size = model.text_encoder.config.hidden_size if hasattr(model, "text_encoder_2"): encoder_hidden_size += model.text_encoder_2.config.hidden_size input_shapes["controlnet"] = { "batch_size": input_shapes[unet_or_transformer_name]["batch_size"], "sequence_length": input_shapes[unet_or_transformer_name]["sequence_length"], "num_channels": unet_or_transformer_num_channels, "height": scaled_height, "width": scaled_width, "vae_scale_factor": vae_scale_factor, "encoder_hidden_size": encoder_hidden_size, } # Image encoder if getattr(model, "image_encoder", None): input_shapes["image_encoder"] = { "batch_size": input_shapes[unet_or_transformer_name]["batch_size"], "num_channels": model.image_encoder.config.num_channels, "width": model.image_encoder.config.image_size, "height": model.image_encoder.config.image_size, } # IP-Adapter: add image_embeds as input for unet/transformer # unet has `ip_adapter_image_embeds` with shape [batch_size, 1, (self.image_encoder.config.image_size//patch_size)**2+1, self.image_encoder.config.hidden_size] as input if getattr(model.unet.config, "encoder_hid_dim_type", None) == "ip_image_proj": input_shapes[unet_or_transformer_name]["image_encoder_shapes"] = ImageEncoderArguments( sequence_length=model.image_encoder.vision_model.embeddings.position_embedding.weight.shape[0], 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]) return input_shapes def get_submodels_and_neuron_configs( model: Union["PreTrainedModel", "DiffusionPipeline"], input_shapes: Dict[str, int], task: str, output: Path, library_name: str, tensor_parallel_size: int = 1, subfolder: str = "", trust_remote_code: bool = False, dynamic_batch_size: bool = False, model_name_or_path: Optional[Union[str, Path]] = None, submodels: Optional[Dict[str, Union[Path, str]]] = None, output_attentions: bool = False, output_hidden_states: bool = False, controlnet_ids: Optional[Union[str, List[str]]] = None, lora_args: Optional[LoRAAdapterArguments] = None, ): is_encoder_decoder = ( getattr(model.config, "is_encoder_decoder", False) if isinstance(model.config, PretrainedConfig) else False ) if library_name == "diffusers": # TODO: Enable optional outputs for Stable Diffusion if output_attentions: raise ValueError(f"`output_attentions`is not supported by the {task} task yet.") models_and_neuron_configs, output_model_names = _get_submodels_and_neuron_configs_for_stable_diffusion( model=model, input_shapes=input_shapes, output=output, dynamic_batch_size=dynamic_batch_size, submodels=submodels, output_hidden_states=output_hidden_states, controlnet_ids=controlnet_ids, lora_args=lora_args, ) elif is_encoder_decoder: optional_outputs = {"output_attentions": output_attentions, "output_hidden_states": output_hidden_states} preprocessors = maybe_load_preprocessors( src_name_or_path=model_name_or_path, subfolder=subfolder, trust_remote_code=trust_remote_code, ) models_and_neuron_configs, output_model_names = _get_submodels_and_neuron_configs_for_encoder_decoder( model=model, input_shapes=input_shapes, tensor_parallel_size=tensor_parallel_size, task=task, output=output, dynamic_batch_size=dynamic_batch_size, model_name_or_path=model_name_or_path, preprocessors=preprocessors, **optional_outputs, ) else: # TODO: Enable optional outputs for encoders if output_attentions or output_hidden_states: raise ValueError( f"`output_attentions` and `output_hidden_states` are not supported by the {task} task yet." ) neuron_config_constructor = TasksManager.get_exporter_config_constructor( model=model, exporter="neuron", task=task, library_name=library_name, ) input_shapes = check_mandatory_input_shapes(neuron_config_constructor, task, input_shapes) input_shapes = InputShapesArguments(**input_shapes) neuron_config = neuron_config_constructor( model.config, dynamic_batch_size=dynamic_batch_size, input_shapes=input_shapes ) model_name = getattr(model, "name_or_path", None) or model_name_or_path model_name = model_name.split("/")[-1] if model_name else model.config.model_type output_model_names = {model_name: "model.neuron"} models_and_neuron_configs = {model_name: (model, neuron_config)} maybe_save_preprocessors(model_name_or_path, output, src_subfolder=subfolder) return models_and_neuron_configs, output_model_names def _get_submodels_and_neuron_configs_for_stable_diffusion( model: Union["PreTrainedModel", "DiffusionPipeline"], input_shapes: Dict[str, int], output: Path, dynamic_batch_size: bool = False, submodels: Optional[Dict[str, Union[Path, str]]] = None, output_hidden_states: bool = False, controlnet_ids: Optional[Union[str, List[str]]] = None, lora_args: Optional[LoRAAdapterArguments] = None, ): check_compiler_compatibility_for_stable_diffusion() model = replace_stable_diffusion_submodels(model, submodels) if is_neuron_available(): raise RuntimeError( "Stable diffusion export is not supported by neuron-cc on inf1, please use neuronx-cc on either inf2/trn1 instead." ) input_shapes = infer_stable_diffusion_shapes_from_diffusers( input_shapes=input_shapes, model=model, has_controlnets=controlnet_ids is not None, ) # Saving the model config and preprocessor as this is needed sometimes. model.scheduler.save_pretrained(output.joinpath("scheduler")) if getattr(model, "tokenizer", None) is not None: model.tokenizer.save_pretrained(output.joinpath("tokenizer")) if getattr(model, "tokenizer_2", None) is not None: model.tokenizer_2.save_pretrained(output.joinpath("tokenizer_2")) if getattr(model, "tokenizer_3", None) is not None: model.tokenizer_3.save_pretrained(output.joinpath("tokenizer_3")) if getattr(model, "feature_extractor", None) is not None: model.feature_extractor.save_pretrained(output.joinpath("feature_extractor")) model.save_config(output) models_and_neuron_configs = get_diffusion_models_for_export( pipeline=model, text_encoder_input_shapes=input_shapes["text_encoder"], unet_input_shapes=input_shapes.get("unet", None), transformer_input_shapes=input_shapes.get("transformer", None), vae_encoder_input_shapes=input_shapes["vae_encoder"], vae_decoder_input_shapes=input_shapes["vae_decoder"], lora_args=lora_args, dynamic_batch_size=dynamic_batch_size, output_hidden_states=output_hidden_states, controlnet_ids=controlnet_ids, controlnet_input_shapes=input_shapes.get("controlnet", None), image_encoder_input_shapes=input_shapes.get("image_encoder", None), ) output_model_names = { DIFFUSION_MODEL_VAE_ENCODER_NAME: os.path.join(DIFFUSION_MODEL_VAE_ENCODER_NAME, NEURON_FILE_NAME), DIFFUSION_MODEL_VAE_DECODER_NAME: os.path.join(DIFFUSION_MODEL_VAE_DECODER_NAME, NEURON_FILE_NAME), } if getattr(model, "text_encoder", None) is not None: output_model_names[DIFFUSION_MODEL_TEXT_ENCODER_NAME] = os.path.join( DIFFUSION_MODEL_TEXT_ENCODER_NAME, NEURON_FILE_NAME ) if getattr(model, "text_encoder_2", None) is not None: output_model_names[DIFFUSION_MODEL_TEXT_ENCODER_2_NAME] = os.path.join( DIFFUSION_MODEL_TEXT_ENCODER_2_NAME, NEURON_FILE_NAME ) if getattr(model, "unet", None) is not None: output_model_names[DIFFUSION_MODEL_UNET_NAME] = os.path.join(DIFFUSION_MODEL_UNET_NAME, NEURON_FILE_NAME) if getattr(model, "transformer", None) is not None: output_model_names[DIFFUSION_MODEL_TRANSFORMER_NAME] = os.path.join( DIFFUSION_MODEL_TRANSFORMER_NAME, NEURON_FILE_NAME ) if getattr(model, "image_encoder", None) is not None: output_model_names["image_encoder"] = os.path.join("image_encoder", NEURON_FILE_NAME) # ControlNet models if controlnet_ids: if isinstance(controlnet_ids, str): controlnet_ids = [controlnet_ids] for idx in range(len(controlnet_ids)): controlnet_name = DIFFUSION_MODEL_CONTROLNET_NAME + "_" + str(idx) output_model_names[controlnet_name] = os.path.join(controlnet_name, NEURON_FILE_NAME) del model return models_and_neuron_configs, output_model_names def _get_submodels_and_neuron_configs_for_encoder_decoder( model: "PreTrainedModel", input_shapes: Dict[str, int], tensor_parallel_size: int, task: str, output: Path, preprocessors: Optional[List] = None, dynamic_batch_size: bool = False, model_name_or_path: Optional[Union[str, Path]] = None, output_attentions: bool = False, output_hidden_states: bool = False, ): if is_neuron_available(): raise RuntimeError( "Encoder-decoder models export is not supported by neuron-cc on inf1, please use neuronx-cc on either inf2/trn1 instead." ) models_and_neuron_configs = get_encoder_decoder_models_for_export( model=model, task=task, tensor_parallel_size=tensor_parallel_size, dynamic_batch_size=dynamic_batch_size, input_shapes=input_shapes, output_attentions=output_attentions, output_hidden_states=output_hidden_states, model_name_or_path=model_name_or_path, preprocessors=preprocessors, ) output_model_names = { ENCODER_NAME: os.path.join(ENCODER_NAME, NEURON_FILE_NAME), DECODER_NAME: os.path.join(DECODER_NAME, NEURON_FILE_NAME), } model.config.save_pretrained(output) model.generation_config.save_pretrained(output) maybe_save_preprocessors(model_name_or_path, output) return models_and_neuron_configs, output_model_names def load_models_and_neuron_configs( model_name_or_path: str, output: Path, model: Optional[Union["PreTrainedModel", "ModelMixin"]], task: str, dynamic_batch_size: bool, cache_dir: Optional[str], trust_remote_code: bool, subfolder: str, revision: str, library_name: str, force_download: bool, local_files_only: bool, token: Optional[Union[bool, str]], submodels: Optional[Dict[str, Union[Path, str]]], torch_dtype: Optional[Union[str, torch.dtype]] = None, tensor_parallel_size: int = 1, controlnet_ids: Optional[Union[str, List[str]]] = None, lora_args: Optional[LoRAAdapterArguments] = None, ip_adapter_args: Optional[IPAdapterArguments] = None, output_attentions: bool = False, output_hidden_states: bool = False, **input_shapes, ): model_kwargs = { "task": task, "model_name_or_path": model_name_or_path, "subfolder": subfolder, "revision": revision, "cache_dir": cache_dir, "token": token, "local_files_only": local_files_only, "force_download": force_download, "trust_remote_code": trust_remote_code, "framework": "pt", "library_name": library_name, "torch_dtype": torch_dtype, } if model is None: model = TasksManager.get_model_from_task(**model_kwargs) # Load IP-Adapter if it exists if ip_adapter_args is not None and not all( getattr(ip_adapter_args, field.name) is None for field in fields(ip_adapter_args) ): model.load_ip_adapter( ip_adapter_args.model_id, subfolder=ip_adapter_args.subfolder, weight_name=ip_adapter_args.weight_name ) model.set_ip_adapter_scale(scale=ip_adapter_args.scale) models_and_neuron_configs, output_model_names = get_submodels_and_neuron_configs( model=model, input_shapes=input_shapes, tensor_parallel_size=tensor_parallel_size, task=task, library_name=library_name, output=output, subfolder=subfolder, trust_remote_code=trust_remote_code, dynamic_batch_size=dynamic_batch_size, model_name_or_path=model_name_or_path, submodels=submodels, output_attentions=output_attentions, output_hidden_states=output_hidden_states, controlnet_ids=controlnet_ids, lora_args=lora_args, ) return models_and_neuron_configs, output_model_names def main_export( model_name_or_path: str, output: Union[str, Path], compiler_kwargs: Dict[str, Any], torch_dtype: Optional[Union[str, torch.dtype]] = None, tensor_parallel_size: int = 1, model: Optional[Union["PreTrainedModel", "ModelMixin"]] = None, task: str = "auto", dynamic_batch_size: bool = False, atol: Optional[float] = None, cache_dir: Optional[str] = None, disable_neuron_cache: Optional[bool] = False, compiler_workdir: Optional[Union[str, Path]] = None, inline_weights_to_neff: bool = True, optlevel: str = "2", trust_remote_code: bool = False, subfolder: str = "", revision: str = "main", force_download: bool = False, local_files_only: bool = False, token: Optional[Union[bool, str]] = None, do_validation: bool = True, submodels: Optional[Dict[str, Union[Path, str]]] = None, output_attentions: bool = False, output_hidden_states: bool = False, library_name: Optional[str] = None, controlnet_ids: Optional[Union[str, List[str]]] = None, lora_args: Optional[LoRAAdapterArguments] = None, ip_adapter_args: Optional[IPAdapterArguments] = None, **input_shapes, ): output = Path(output) torch_dtype = map_torch_dtype(torch_dtype) if not output.parent.exists(): output.parent.mkdir(parents=True) task = TasksManager.map_from_synonym(task) if library_name is None: library_name = TasksManager.infer_library_from_model( model_name_or_path, revision=revision, cache_dir=cache_dir, token=token ) models_and_neuron_configs, output_model_names = load_models_and_neuron_configs( model_name_or_path=model_name_or_path, output=output, model=model, torch_dtype=torch_dtype, tensor_parallel_size=tensor_parallel_size, task=task, dynamic_batch_size=dynamic_batch_size, cache_dir=cache_dir, trust_remote_code=trust_remote_code, subfolder=subfolder, revision=revision, library_name=library_name, force_download=force_download, local_files_only=local_files_only, token=token, submodels=submodels, lora_args=lora_args, ip_adapter_args=ip_adapter_args, output_attentions=output_attentions, output_hidden_states=output_hidden_states, controlnet_ids=controlnet_ids, **input_shapes, ) _, neuron_outputs = export_models( models_and_neuron_configs=models_and_neuron_configs, task=task, output_dir=output, disable_neuron_cache=disable_neuron_cache, compiler_workdir=compiler_workdir, inline_weights_to_neff=inline_weights_to_neff, optlevel=optlevel, output_file_names=output_model_names, compiler_kwargs=compiler_kwargs, model_name_or_path=model_name_or_path, ) # Validate compiled model if do_validation and tensor_parallel_size > 1: # TODO: support the validation of tp models. logger.warning( "The validation is not yet supported for tensor parallel model, the validation will be turned off." ) do_validation = False if do_validation is True: try: validate_models_outputs( models_and_neuron_configs=models_and_neuron_configs, neuron_named_outputs=neuron_outputs, output_dir=output, atol=atol, neuron_files_subpaths=output_model_names, ) logger.info( f"The {NEURON_COMPILER} export succeeded and the exported model was saved at: {output.as_posix()}" ) except ShapeError as e: raise e except AtolError as e: logger.warning( f"The {NEURON_COMPILER} export succeeded with the warning: {e}.\n The exported model was saved at: " f"{output.as_posix()}" ) except OutputMatchError as e: logger.warning( f"The {NEURON_COMPILER} export succeeded with the warning: {e}.\n The exported model was saved at: " f"{output.as_posix()}" ) except Exception as e: logger.error( f"An error occurred with the error message: {e}.\n The exported model was saved at: {output.as_posix()}" ) def maybe_export_from_neuron_model_class( model: str, output: Union[str, Path], task: str = "auto", cache_dir: Optional[str] = None, subfolder: str = "", trust_remote_code: bool = False, **kwargs, ): """Export the model from the neuron model class if it exists.""" if task == "auto": task = infer_task(model) output = Path(output) # Remove None values from the kwargs kwargs = {key: value for key, value in kwargs.items() if value is not None} # Also remove some arguments that are not supported in this context kwargs.pop("disable_neuron_cache", None) kwargs.pop("inline_weights_neff", None) kwargs.pop("O1", None) kwargs.pop("O2", None) kwargs.pop("O3", None) kwargs.pop("disable_validation", None) kwargs.pop("dynamic_batch_size", None) kwargs.pop("output_hidden_states", None) kwargs.pop("output_attentions", None) kwargs.pop("tensor_parallel_size", None) # Fetch the model config config = AutoConfig.from_pretrained(model) # Check if we have an auto-model class for the model_type and task if not has_neuron_model_class(model_type=config.model_type, task=task, mode="inference"): return False neuron_model_class = get_neuron_model_class(model_type=config.model_type, task=task, mode="inference") neuron_model = neuron_model_class.from_pretrained( model_id=model, export=True, cache_dir=cache_dir, subfolder=subfolder, config=config, trust_remote_code=trust_remote_code, load_weights=False, # Reduce model size for nxd models **kwargs, ) if not output.parent.exists(): output.parent.mkdir(parents=True) neuron_model.save_pretrained(output) try: tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=trust_remote_code) tokenizer.save_pretrained(output) except Exception: logger.info(f"No tokenizer found while exporting {model}.") return True def main(): parser = ArgumentParser(f"Hugging Face Optimum {NEURON_COMPILER} exporter") parse_args_neuron(parser) # Retrieve CLI arguments args = parser.parse_args() task = infer_task(args.model) if args.task == "auto" else args.task library_name = TasksManager.infer_library_from_model(args.model, cache_dir=args.cache_dir) if library_name == "diffusers": input_shapes = normalize_stable_diffusion_input_shapes(args) submodels = {"unet": args.unet} elif library_name == "sentence_transformers": input_shapes = normalize_sentence_transformers_input_shapes(args) submodels = None else: # New export mode using dedicated neuron model classes kwargs = vars(args).copy() if maybe_export_from_neuron_model_class(**kwargs): return # Fallback to legacy export input_shapes = get_input_shapes(task, args) submodels = None disable_neuron_cache = args.disable_neuron_cache compiler_kwargs = infer_compiler_kwargs(args) optional_outputs = customize_optional_outputs(args) optlevel = parse_optlevel(args) lora_args = LoRAAdapterArguments( model_ids=getattr(args, "lora_model_ids", None), weight_names=getattr(args, "lora_weight_names", None), adapter_names=getattr(args, "lora_adapter_names", None), scales=getattr(args, "lora_scales", None), ) ip_adapter_args = IPAdapterArguments( model_id=getattr(args, "ip_adapter_id", None), subfolder=getattr(args, "ip_adapter_subfolder", None), weight_name=getattr(args, "ip_adapter_weight_name", None), scale=getattr(args, "ip_adapter_scale", None), ) main_export( model_name_or_path=args.model, output=args.output, compiler_kwargs=compiler_kwargs, torch_dtype=args.torch_dtype, tensor_parallel_size=args.tensor_parallel_size, task=task, dynamic_batch_size=args.dynamic_batch_size, atol=args.atol, cache_dir=args.cache_dir, disable_neuron_cache=disable_neuron_cache, compiler_workdir=args.compiler_workdir, inline_weights_to_neff=args.inline_weights_neff, optlevel=optlevel, trust_remote_code=args.trust_remote_code, subfolder=args.subfolder, do_validation=not args.disable_validation, submodels=submodels, library_name=library_name, controlnet_ids=getattr(args, "controlnet_ids", None), lora_args=lora_args, ip_adapter_args=ip_adapter_args, **optional_outputs, **input_shapes, ) if __name__ == "__main__": main()