Update optimum_neuron_export.py

optimum_neuron_export.py

@@ -140,7 +140,7 @@ DIFFUSION_PIPELINE_MAPPING = {
         "tasks": ["inpaint"],
         "default_task": "inpaint"
     },
-    "flux-kont": {
+    "flux-kontext": {
         "class": FluxKontextPipeline,
         "tasks": ["text-to-image", "image-to-image"],
         "default_task": "text-to-image"
@@ -245,6 +245,7 @@ def export_diffusion_model(model_id: str, pipeline_name: str, task: str, folder:
             input_shapes=input_shapes,
             task=task,
             library_name="diffusers",
+            tensor_parallel_size=4,
             output=Path(folder),
             model_name_or_path=model_id,
         )