Update optimum_neuron_export.py

optimum_neuron_export.py

@@ -162,17 +162,17 @@ def export(model_id: str, task_or_pipeline: str, model_type: str, folder: str):
             **inputs
         )
         neuron_model = NeuronModelForCausalLM.export(
-            model_id=model_id,  # Fixed variable name
+            model_id=model_id, 
             neuron_config=neuron_config,
             token=HF_TOKEN,
         )
-        neuron_model.save_pretrained(folder)  # Fixed variable name
+        neuron_model.save_pretrained(folder) 
 
     # DIFFUSION tasks
     elif task_or_pipeline in DIFFUSION_PIPELINE_MAPPING:
         model_class = DIFFUSION_PIPELINE_MAPPING.get(task_or_pipeline)
         model = model_class.from_pretrained(model_id)
-        # input_shapes = build_stable_diffusion_components_mandatory_shapes(**inputs)
+        input_shapes = build_stable_diffusion_components_mandatory_shapes(**inputs)
         compiler_kwargs = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
 
         result = main_export(
@@ -182,10 +182,12 @@ def export(model_id: str, task_or_pipeline: str, model_type: str, folder: str):
             torch_dtype= torch.bfloat16,
             token=HF_TOKEN,
             library_name=model_type,
+            tensor_parallel_size=4,
             cpu_backend=True,
             model=model,
             **inputs,
         )
+        
     else:
         raise ValueError(f"Unsupported task or pipeline: {task_or_pipeline}")