Update optimum_neuron_export.py

optimum_neuron_export.py

@@ -53,6 +53,7 @@ from diffusers import (
     FluxPipeline,
     FluxInpaintPipeline,
     FluxImg2ImgPipeline,
+    FluxKontextPipeline,
 )
 
 from optimum.neuron.cache import synchronize_hub_cache
@@ -139,10 +140,10 @@ DIFFUSION_PIPELINE_MAPPING = {
         "tasks": ["inpaint"],
         "default_task": "inpaint"
     },
-    "flux-img2img": {
-        "class": FluxImg2ImgPipeline,
-        "tasks": ["image-to-image"],
-        "default_task": "image-to-image"
+    "flux-kont": {
+        "class": FluxKontextPipeline,
+        "tasks": ["text-to-image", "image-to-image"],
+        "default_task": "text-to-image"
     },
 }