diff --git a/optimum/neuron/modeling_diffusion.py b/optimum/neuron/modeling_diffusion.py index 4257e25ba..3712ffe45 100644 --- a/optimum/neuron/modeling_diffusion.py +++ b/optimum/neuron/modeling_diffusion.py @@ -14,6 +14,7 @@ # limitations under the License. """NeuroStableDiffusionPipeline class for inference of diffusion models on neuron devices.""" +import copy import importlib import logging import os @@ -698,6 +699,7 @@ def _export( save_dir = TemporaryDirectory() save_dir_path = Path(save_dir.name) # 1. Fetch all model configs + input_shapes_copy = copy.deepcopy(input_shapes) models_and_neuron_configs, _ = load_models_and_neuron_configs( model_name_or_path=model_id, output=save_dir_path, @@ -717,7 +719,7 @@ def _export( lora_weight_names=lora_weight_names, lora_adapter_names=lora_adapter_names, lora_scales=lora_scales, - **input_shapes, + **input_shapes_copy, ) # 2. Build compilation config