Update optimum_neuron_export.py

optimum_neuron_export.py

@@ -1,4 +1,3 @@
-
 import os
 import shutil
 from tempfile import TemporaryDirectory, NamedTemporaryFile
@@ -36,25 +35,25 @@ from optimum.neuron import (
     NeuronModelForCausalLM,
     NeuronModelForSeq2SeqLM,
 )
-from diffusers import (
-    StableDiffusionPipeline,
-    StableDiffusionImg2ImgPipeline,
-    StableDiffusionInpaintPipeline,
-    StableDiffusionInstructPix2PixPipeline,
-    LatentConsistencyModelPipeline,
-    StableDiffusionXLPipeline,
-    StableDiffusionXLImg2ImgPipeline,
-    StableDiffusionXLInpaintPipeline,
-    StableDiffusionControlNetPipeline,
-    StableDiffusionXLControlNetPipeline,
-    PixArtAlphaPipeline,
-    PixArtSigmaPipeline,
-    FluxPipeline
+from optimum.neuron import (
+    NeuronDiffusionPipelineBase,
+    NeuronStableDiffusionPipeline,
+    NeuronStableDiffusionImg2ImgPipeline,
+    NeuronStableDiffusionInpaintPipeline,
+    NeuronStableDiffusionInstructPix2PixPipeline,
+    NeuronLatentConsistencyModelPipeline,
+    NeuronStableDiffusionXLPipeline,
+    NeuronStableDiffusionXLImg2ImgPipeline,
+    NeuronStableDiffusionXLInpaintPipeline,
+    NeuronStableDiffusionControlNetPipeline,
+    NeuronStableDiffusionXLControlNetPipeline,
+    NeuronPixArtAlphaPipeline,
+    NeuronPixArtSigmaPipeline,
+    NeuronFluxPipeline
 )
 
 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")
@@ -82,26 +81,22 @@ TASK_TO_MODEL_CLASS = {
 
 # Diffusion pipeline mapping
 DIFFUSION_PIPELINE_MAPPING = {
-    "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,
+    "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,
 }
 
-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"]:
@@ -120,6 +115,30 @@ def get_default_inputs(task_or_pipeline: str) -> Dict[str, int]:
         # Default to text-based shapes
         return {"batch_size": 1, "sequence_length": 128}
 
+def find_neuron_cache_artifacts(cache_base_dir: str = "/var/tmp/neuron-compile-cache") -> Optional[str]:
+    """
+    Find the most recently created Neuron cache artifacts directory.
+    Returns the path to the MODULE directory containing the compiled artifacts.
+    """
+    if not os.path.exists(cache_base_dir):
+        return None
+    
+    # Find all MODULE directories
+    module_dirs = []
+    for root, dirs, files in os.walk(cache_base_dir):
+        for d in dirs:
+            if d.startswith("MODULE_"):
+                full_path = os.path.join(root, d)
+                # Check if it contains the expected files
+                if os.path.exists(os.path.join(full_path, "model.neuron")):
+                    module_dirs.append(full_path)
+    
+    if not module_dirs:
+        return None
+    
+    # Return the most recently modified directory
+    return max(module_dirs, key=os.path.getmtime)
+
 def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]:
     try:
         discussions = api.get_repo_discussions(repo_id=model_id)
@@ -134,102 +153,69 @@ def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discuss
             return discussion
     return None
 
-def export(model_id: str, task_or_pipeline: str, model_type: str, folder: str):
-    """Export model to Neuron format. This is NOT a generator."""
-    
-    print(f"📦 Exporting model `{model_id}` for task `{task_or_pipeline}`...")
-    
-    inputs = get_default_inputs(task_or_pipeline)
-    print(f"🔧 Using default inputs: {inputs}")
+def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, folder: str, token: str) -> Any:
 
-    # ENCODER and SEQ2SEQ tasks
-    if task_or_pipeline in ENCODER_TASKS or task_or_pipeline in SEQ2SEQ_TAKS:   
-        result = main_export(
-            model_name_or_path=model_id,
-            output=folder,
-            token=HF_TOKEN,
-            task=task_or_pipeline,
-            cpu_backend=True,
-            do_validation=False,  
-            compiler_kwargs={},
-            **inputs,            
-        )
+    yield f"📦 Exporting model `{model_id}` for task `{task_or_pipeline}`..."
 
-    # DECODER tasks
-    elif 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=model_id, 
-            neuron_config=neuron_config,
-            token=HF_TOKEN,
-        )
-        neuron_model.save_pretrained(folder) 
+    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}")
 
-    # 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)
-        compiler_kwargs = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
-
-        result = main_export(
-            model_name_or_path=model_id,
-            output=folder,
-            compiler_kwargs=compiler_kwargs,
-            torch_dtype= torch.bfloat16,
-            token=HF_TOKEN,
-            library_name=model_type,
-            tensor_parallel_size=4,
-            model=model,
-            **inputs_shapes,
-        )
-        
-    else:
-        raise ValueError(f"Unsupported task or pipeline: {task_or_pipeline}")
+    inputs = get_default_inputs(task_or_pipeline)
+    compiler_configs = {"auto_cast": "matmul", "auto_cast_type": "bf16", "instance_type": "inf2"}
+    yield f"🔧 Using default inputs: {inputs}"
     
-    print(f"✅ Export completed successfully to {folder}")
-
-
-def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, folder: str, token: str):
-    """Export model and prepare git operations. This IS a generator."""
+    # Clear any old cache artifacts before export
+    cache_base_dir = "/var/tmp/neuron-compile-cache"
+    initial_cache_dirs = set()
+    if os.path.exists(cache_base_dir):
+        for root, dirs, _ in os.walk(cache_base_dir):
+            for d in dirs:
+                if d.startswith("MODULE_"):
+                    initial_cache_dirs.add(os.path.join(root, d))
     
     try:
-        # Actually execute the export (not a generator anymore)
-        export(model_id, task_or_pipeline, model_type, folder)
-        yield "✅ Export completed successfully."
+        # Trigger the export/compilation
+        model = model_class.from_pretrained(
+            model_id,
+            export=True,
+            tensor_parallel_size=4,
+            token=HF_TOKEN,
+            **compiler_configs,
+            **inputs,
+        )
+        
+        yield "✅ Export/compilation completed successfully."
+        
+        # Find the newly created cache artifacts
+        yield "🔍 Locating compiled artifacts in Neuron cache..."
+        cache_artifact_dir = find_neuron_cache_artifacts(cache_base_dir)
+        
+        if not cache_artifact_dir:
+            raise Exception("❌ Could not find compiled artifacts in Neuron cache")
+        
+        yield f"📂 Found artifacts at: {cache_artifact_dir}"
+        
+        # Copy artifacts from cache to our target folder
+        yield f"📋 Copying artifacts to export folder..."
+        if os.path.exists(folder):
+            shutil.rmtree(folder)
+        shutil.copytree(cache_artifact_dir, folder)
+        
+        yield f"✅ Artifacts successfully copied to {folder}"
+        
     except Exception as e:
         yield f"❌ Export failed with error: {e}"
         raise
 
-    # Verify that files were actually created
-    if not os.path.exists(folder) or not os.listdir(folder):
-        error_msg = f"❌ Export folder is empty or doesn't exist: {folder}"
-        yield error_msg
-        raise Exception(error_msg)
-    
-    yield f"📁 Found exported files in {folder}"
-
-    # Collect all files for git operations
     operations = []
-    file_count = 0
     for root, _, files in os.walk(folder):
         for filename in files:
             file_path = os.path.join(root, filename)
             repo_path = os.path.relpath(file_path, folder)
             operations.append(CommitOperationAdd(path_in_repo=repo_path, path_or_fileobj=file_path))
-            file_count += 1
-    
-    yield f"📦 Prepared {file_count} files for upload"
     
-    if file_count == 0:
-        error_msg = "❌ No files found to upload after export"
-        yield error_msg
-        raise Exception(error_msg)
-    
-    # Update model card
     try:
         card = ModelCard.load(model_id, token=token)
         if not hasattr(card.data, "tags") or card.data.tags is None:
@@ -246,18 +232,14 @@ def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, fo
             readme_op.path_or_fileobj = readme_path
         else:
             operations.append(CommitOperationAdd(path_in_repo="README.md", path_or_fileobj=readme_path))
-        
-        yield "📝 Updated model card with neuron tag"
 
     except Exception as e:
         yield f"⚠️ Warning: Could not update model card: {e}"
 
-    # Return the operations
     yield ("__RETURN__", operations)
 
 def generate_neuron_repo_name(api, original_model_id: str, task_or_pipeline: str, token:str) -> str:
     """Generate a name for the Neuron-optimized repository."""
-    # Replace '©' with '-' and add neuron suffix
     requesting_user = api.whoami(token=token)["name"]
     base_name = original_model_id.replace('/', '-')
     return f"{requesting_user}/{base_name}-neuron"
@@ -475,7 +457,6 @@ Generated using: [Optimum Neuron Compiler Space]({SPACES_URL})
         yield f"❌ Failed to create README PR: {e}"
         raise
 
-# --- Updated upload_to_custom_repo function ---
 def upload_to_custom_repo(
     operations: List[CommitOperationAdd],
     custom_repo_id: str,
@@ -684,4 +665,4 @@ These files were generated to accelerate model loading on AWS Neuron devices.
         except Exception as e:
             yield "1", f"❌ Conversion failed with a critical error: {e}"
             # Re-raise the exception to be caught by the outer try-except in the Gradio app if needed
-            raise
+            raise