Update optimum_neuron_export.py

optimum_neuron_export.py

@@ -55,12 +55,11 @@ from optimum.neuron import (
 
 from optimum.neuron.cache import synchronize_hub_cache
 from synchronizer import synchronize_hub_cache_with_pr
+from optimum.exporters.neuron import main_export, build_stable_diffusion_components_mandatory_shapes,
 
 SPACES_URL = "https://huggingface.co/spaces/optimum/neuron-export"
 CUSTOM_CACHE_REPO = os.getenv("CUSTOM_CACHE_REPO")
 HF_TOKEN = os.environ.get("HF_TOKEN")
-os.environ["NEURON_PLATFORM_TARGET_OVERRIDE"] = "inf2"
-
 
 # Task to NeuronModel mapping for transformers
 TASK_TO_MODEL_CLASS = {
@@ -84,31 +83,37 @@ TASK_TO_MODEL_CLASS = {
 
 # Diffusion pipeline mapping
 DIFFUSION_PIPELINE_MAPPING = {
-    "text-to-image": NeuronStableDiffusionPipeline,
-    "image-to-image": NeuronStableDiffusionImg2ImgPipeline,
-    "inpaint": NeuronStableDiffusionInpaintPipeline,
-    "instruct-pix2pix": NeuronStableDiffusionInstructPix2PixPipeline,
-    "latent-consistency": NeuronLatentConsistencyModelPipeline,
-    "stable-diffusion": NeuronStableDiffusionPipeline,
-    "stable-diffusion-xl": NeuronStableDiffusionXLPipeline,
-    "stable-diffusion-xl-img2img": NeuronStableDiffusionXLImg2ImgPipeline,
-    "stable-diffusion-xl-inpaint": NeuronStableDiffusionXLInpaintPipeline,
-    "controlnet": NeuronStableDiffusionControlNetPipeline,
-    "controlnet-xl": NeuronStableDiffusionXLControlNetPipeline,
-    "pixart-alpha": NeuronPixArtAlphaPipeline,
-    "pixart-sigma": NeuronPixArtSigmaPipeline,
-    "flux": NeuronFluxPipeline,
+    "text-to-image": StableDiffusionPipeline,
+    "image-to-image": StableDiffusionImg2ImgPipeline,
+    "inpaint": StableDiffusionInpaintPipeline,
+    "instruct-pix2pix": StableDiffusionInstructPix2PixPipeline,
+    "latent-consistency": LatentConsistencyModelPipeline,
+    "stable-diffusion": StableDiffusionPipeline,
+    "stable-diffusion-xl": StableDiffusionXLPipeline,
+    "stable-diffusion-xl-img2img": StableDiffusionXLImg2ImgPipeline,
+    "stable-diffusion-xl-inpaint": StableDiffusionXLInpaintPipeline,
+    "controlnet": StableDiffusionControlNetPipeline,
+    "controlnet-xl": StableDiffusionXLControlNetPipeline,
+    "pixart-alpha": PixArtAlphaPipeline,
+    "pixart-sigma": PixArtSigmaPipeline,
+    "flux": FluxPipeline,
 }
 
+ENCODER_TASKS = {"feature-extraction","sentence-transformers","fill-mask","question-answering","text-classification","token-classification","multiple-choice","image-classification","semantic-segmentation","object-detection","audio-classification","audio-frame-classification","automatic-speech-recognition","audio-xvector"}
+DECODER_TASKS = {"text-generation"}
+SEQ2SEQ_TAKS = {"text2text-generation"}
+
 def get_default_inputs(task_or_pipeline: str) -> Dict[str, int]:
     """Get default input shapes based on task type or diffusion pipeline type."""
-    if task_or_pipeline in ["feature-extraction", "sentence-transformers", "fill-mask", "question-answering", "text-classification", "token-classification","text-generation","text2text-generation"]:
+    if task_or_pipeline in ["feature-extraction", "sentence-transformers", "fill-mask", "question-answering", "text-classification", "token-classification","text-generation"]:
         return {"batch_size": 1, "sequence_length": 128}
     elif task_or_pipeline == "multiple-choice":
         return {"batch_size": 1, "num_choices": 4, "sequence_length": 128}
+    elif task_or_pipeline == "text2text-generation":
+        return {"batch_size": 1, "sequence_length": 128, "num_beams":4}
     elif task_or_pipeline in ["image-classification", "semantic-segmentation", "object-detection"]:
         return {"batch_size": 1, "num_channels": 3, "height": 224, "width": 224}
-    elif task_or_pipeline in ["audio-classification", "audio-frame-classification", "automatic-speech-recognition", "audio-xvector"]:
+    elif task_or_pipeline in ["audio-classification", "audio-frame-classification", "audio-xvector"]:
         return {"batch_size": 1, "audio_sequence_length": 16000}
     elif task_or_pipeline in DIFFUSION_PIPELINE_MAPPING:
         return {"batch_size": 1, "height": 1024, "width": 1024, "num_images_per_prompt": 1, "torch_dtype":torch.bfloat16}
@@ -116,33 +121,6 @@ def get_default_inputs(task_or_pipeline: str) -> Dict[str, int]:
         # Default to text-based shapes
         return {"batch_size": 1, "sequence_length": 128}
 
-def prepare_compiler_flags(
-    auto_cast: str = "matmul",
-    auto_cast_type: str = "bf16",
-    optlevel: str = "2",
-    instance_type: str = "inf2",
-):
-    if auto_cast is not None:
-        logger.info(f"Using Neuron: --auto-cast {auto_cast}")
-        auto_cast = "matmult" if auto_cast == "matmul" else auto_cast
-        compiler_args = ["--auto-cast", auto_cast]
-
-        logger.info(f"Using Neuron: --auto-cast-type {auto_cast_type}")
-        compiler_args.extend(["--auto-cast-type", auto_cast_type])
-    else:
-        compiler_args = ["--auto-cast", "none"]
-
-    compiler_args.extend(["--optlevel", optlevel])
-    logger.info(f"Using Neuron: --optlevel {optlevel}")
-
-    if instance_type == "trn2":
-        compiler_args.extend(["--target", "trn2"])
-    elif instance_type == "trn1":
-        compiler_args.extend(["--target", "trn1"])
-
-    compiler_args_str = " ".join(compiler_args)
-    return compiler_args_str
-
 def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:
     try:
         discussions = api.get_repo_discussions(repo_id=model_id)
@@ -157,30 +135,57 @@ def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discuss
             return discussion
     return None
 
-def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, folder: str, token: str) -> Any:
+def export(model_id: str, task_or_pipeline:str, model_type: str, folder: str):
 
     yield f"📦 Exporting model `{model_id}` for task `{task_or_pipeline}`..."
 
-    model_class = TASK_TO_MODEL_CLASS.get(task_or_pipeline) if model_type == "transformers" else DIFFUSION_PIPELINE_MAPPING.get(task_or_pipeline)
-    if model_class is None:
-        supported = list(TASK_TO_MODEL_CLASS.keys()) if model_type == "transformers" else list(DIFFUSION_PIPELINE_MAPPING.keys())
-        raise Exception(f"❌ Unsupported task/pipeline: {task_or_pipeline}. Supported: {supported}")
-
+    if model_type == "diffusers":
+        model_class = DIFFUSION_PIPELINE_MAPPING.get(task_or_pipeline)
+        
     inputs = get_default_inputs(task_or_pipeline)
-    compiler_args = prepare_compiler_flags()
     yield f"🔧 Using default inputs: {inputs}"
-    
-    try:
-        model = model_class.from_pretrained(
-            model_id,
-            export=True,
-            tensor_parallel_size=1,
-            token=HF_TOKEN,
+
+    if task_or_pipeline in ENCODER_TASKS or SEQ2SEQ_TAKS:   
+        result = main_export(
+            model_name_or_path=model_id,
+            output=folder,
+            token=HF_TOKEN
+            task=task_or_pipeline,
             cpu_backend=True,
-            compiler_args=compiler_args,
-            **inputs,
+            do_validation=False,  
+            input_shapes=**inputs            
+        )
+
+    if task_or_pipeline in DECODER_TASKS:
+        neuron_config = NeuronModelForCausalLM.get_neuron_config(model_name_or_path=model_id, **inputs)
+        neuron_model = NeuronModelForCausalLM.export(
+            model_id=export_decoder_id, 
+            neuron_config=neuron_config,
+            token = HF_TOKEN
         )
         model.save_pretrained(folder)
+
+    if task_or_pipeline in DIFFUSION_PIPELINE_MAPPING:
+        model = model_class.from_pretrained(model_id)
+        input_shapes = build_stable_diffusion_components_mandatory_shapes(**inputs)
+        compiler_kwargs = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
+
+        result = main_export(
+            model_name_or_path=model_id,
+            output=folder,
+            compiler_kwargs=compiler_kwargs,
+            token=HF_TOKEN,
+            library_name=model_type,
+            cpu_backend=True,
+            model=model,
+            **input_shapes,
+        )
+
+
+def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, folder: str, token: str) -> Any:
+    
+    try:
+        export(model_id, task_or_pipeline, model_type, folder)
         yield "✅ Export completed successfully."
     except Exception as e:
         yield f"❌ Export failed with error: {e}"