Update app.py

app.py

@@ -1,526 +1,807 @@
-import csv
 import os
-from datetime import datetime
-from typing import Optional, Union, List
+import shutil
+from tempfile import TemporaryDirectory, NamedTemporaryFile
+from typing import List, Union, Optional, Tuple, Dict, Any, Generator
+from pathlib import Path
+import torch
 import gradio as gr
-from huggingface_hub import HfApi, Repository
-from optimum_neuron_export import convert, DIFFUSION_PIPELINE_MAPPING
-from gradio_huggingfacehub_search import HuggingfaceHubSearch
-from apscheduler.schedulers.background import BackgroundScheduler
-
-# Define transformer tasks and their categories for coloring
-TRANSFORMER_TASKS = {
-    "auto": {"color": "#6b7280", "category": "Auto"},
-    "feature-extraction": {"color": "#3b82f6", "category": "Feature Extraction"},
-    "fill-mask": {"color": "#8b5cf6", "category": "NLP"},
-    "multiple-choice": {"color": "#8b5cf6", "category": "NLP"},
-    "question-answering": {"color": "#8b5cf6", "category": "NLP"},
-    "text-classification": {"color": "#8b5cf6", "category": "NLP"},
-    "token-classification": {"color": "#8b5cf6", "category": "NLP"},
-    "text-generation": {"color": "#10b981", "category": "Text Generation"},
-    "text2text-generation": {"color": "#10b981", "category": "Text Generation"},
-    "audio-classification": {"color": "#f59e0b", "category": "Audio"},
-    "automatic-speech-recognition": {"color": "#f59e0b", "category": "Audio"},
-    "audio-frame-classification": {"color": "#f59e0b", "category": "Audio"},
-    "audio-xvector": {"color": "#f59e0b", "category": "Audio"},
-    "image-classification": {"color": "#ef4444", "category": "Vision"},
-    "object-detection": {"color": "#ef4444", "category": "Vision"},
-    "semantic-segmentation": {"color": "#ef4444", "category": "Vision"},
-    "zero-shot-image-classification": {"color": "#ec4899", "category": "Multimodal"},
-    "sentence-similarity": {"color": "#06b6d4", "category": "Similarity"},
-}
+from huggingface_hub import (
+    CommitOperationAdd,
+    HfApi,
+    ModelCard,
+    Discussion,
+    CommitInfo,
+    create_repo,
+    RepoUrl,
+)
+from huggingface_hub.file_download import repo_folder_name
+from optimum.exporters.tasks import TasksManager
+from optimum.exporters.neuron.model_configs import *
+from optimum.exporters.neuron import build_stable_diffusion_components_mandatory_shapes
+from optimum.exporters.neuron.model_configs import *
+from optimum.exporters.neuron import get_submodels_and_neuron_configs, export_models
+from optimum.neuron import (
+    NeuronModelForFeatureExtraction,
+    NeuronModelForSentenceTransformers,
+    NeuronModelForMaskedLM,
+    NeuronModelForQuestionAnswering,
+    NeuronModelForSequenceClassification,
+    NeuronModelForTokenClassification,
+    NeuronModelForMultipleChoice,
+    NeuronModelForImageClassification,
+    NeuronModelForSemanticSegmentation,
+    NeuronModelForObjectDetection,
+    NeuronModelForAudioClassification,
+    NeuronModelForAudioFrameClassification,
+    NeuronModelForCTC,
+    NeuronModelForXVector,
+    NeuronModelForCausalLM,
+    NeuronModelForSeq2SeqLM,
+)
+
+# Import diffusers pipelines
+from diffusers import (
+    StableDiffusionPipeline,
+    StableDiffusionImg2ImgPipeline,
+    StableDiffusionInpaintPipeline,
+    StableDiffusionXLPipeline,
+    StableDiffusionXLImg2ImgPipeline,
+    StableDiffusionXLInpaintPipeline,
+    LatentConsistencyModelPipeline,
+    PixArtAlphaPipeline,
+    PixArtSigmaPipeline,
+    FluxPipeline,
+    FluxInpaintPipeline,
+    FluxImg2ImgPipeline,
+)
+
+from optimum.neuron.cache import synchronize_hub_cache
+from synchronizer import synchronize_hub_cache_with_pr
 
-# Define diffusion pipeline types - updated structure
-DIFFUSION_PIPELINES = {
-    "stable-diffusion": {"color": "#ec4899", "category": "Stable Diffusion", "tasks": ["text-to-image", "image-to-image", "inpaint"]},
-    "stable-diffusion-xl": {"color": "#10b981", "category": "Stable Diffusion XL", "tasks": ["text-to-image", "image-to-image", "inpaint"]},
-    "sdxl-turbo": {"color": "#f59e0b", "category": "SDXL Turbo", "tasks": ["text-to-image", "image-to-image", "inpaint"]},
-    "lcm": {"color": "#8b5cf6", "category": "LCM", "tasks": ["text-to-image"]},
-    "pixart-alpha": {"color": "#ef4444", "category": "PixArt", "tasks": ["text-to-image"]},
-    "pixart-sigma": {"color": "#ef4444", "category": "PixArt", "tasks": ["text-to-image"]},
-    "flux": {"color": "#06b6d4", "category": "Flux", "tasks": ["text-to-image", "inpaint"]},
-    "flux-kont": {"color": "#06b6d4", "category": "Flux Kont", "tasks": ["text-to-image", "image-to-image"]},
+SPACES_URL = "https://huggingface.co/spaces/optimum/neuron-export"
+CUSTOM_CACHE_REPO = os.getenv("CUSTOM_CACHE_REPO")
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# Task to NeuronModel mapping for transformers
+TASK_TO_MODEL_CLASS = {
+    "feature-extraction": NeuronModelForFeatureExtraction,
+    "sentence-transformers": NeuronModelForSentenceTransformers,
+    "fill-mask": NeuronModelForMaskedLM,
+    "question-answering": NeuronModelForQuestionAnswering,
+    "text-classification": NeuronModelForSequenceClassification,
+    "token-classification": NeuronModelForTokenClassification,
+    "multiple-choice": NeuronModelForMultipleChoice,
+    "image-classification": NeuronModelForImageClassification,
+    "semantic-segmentation": NeuronModelForSemanticSegmentation,
+    "object-detection": NeuronModelForObjectDetection,
+    "audio-classification": NeuronModelForAudioClassification,
+    "audio-frame-classification": NeuronModelForAudioFrameClassification,
+    "automatic-speech-recognition": NeuronModelForCTC,
+    "audio-xvector": NeuronModelForXVector,
+    "text-generation": NeuronModelForCausalLM,
+    "text2text-generation": NeuronModelForSeq2SeqLM,
 }
 
-TAGS = {
-    "Feature Extraction": {"color": "#3b82f6", "category": "Feature Extraction"},
-    "NLP": {"color": "#8b5cf6", "category": "NLP"},
-    "Text Generation": {"color": "#10b981", "category": "Text Generation"},
-    "Audio": {"color": "#f59e0b", "category": "Audio"},
-    "Vision": {"color": "#ef4444", "category": "Vision"},
-    "Multimodal": {"color": "#ec4899", "category": "Multimodal"},
-    "Similarity": {"color": "#06b6d4", "category": "Similarity"},
-    "Stable Diffusion": {"color": "#ec4899", "category": "Stable Diffusion"},
-    "Stable Diffusion XL": {"color": "#10b981", "category": "Stable Diffusion XL"},
-    "ControlNet": {"color": "#f59e0b", "category": "ControlNet"},
-    "ControlNet XL": {"color": "#f59e0b", "category": "ControlNet XL"},
-    "PixArt": {"color": "#ef4444", "category": "PixArt"},
-    "Latent Consistency": {"color": "#8b5cf6", "category": "Latent Consistency"},
-    "Flux": {"color": "#06b6d4", "category": "Flux"},
+# Diffusion pipeline mapping with their corresponding diffusers classes and supported tasks
+DIFFUSION_PIPELINE_MAPPING = {
+    "stable-diffusion": {
+        "class": StableDiffusionPipeline,
+        "tasks": ["text-to-image"],
+        "default_task": "text-to-image"
+    },
+    "stable-diffusion-img2img": {
+        "class": StableDiffusionImg2ImgPipeline,
+        "tasks": ["image-to-image"],
+        "default_task": "image-to-image"
+    },
+    "stable-diffusion-inpaint": {
+        "class": StableDiffusionInpaintPipeline,
+        "tasks": ["inpaint"],
+        "default_task": "inpaint"
+    },
+    "stable-diffusion-xl": {
+        "class": StableDiffusionXLPipeline,
+        "tasks": ["text-to-image"],
+        "default_task": "text-to-image"
+    },
+    "stable-diffusion-xl-img2img": {
+        "class": StableDiffusionXLImg2ImgPipeline,
+        "tasks": ["image-to-image"],
+        "default_task": "image-to-image"
+    },
+    "stable-diffusion-xl-inpaint": {
+        "class": StableDiffusionXLInpaintPipeline,
+        "tasks": ["inpaint"],
+        "default_task": "inpaint"
+    },
+    "lcm": {
+        "class": LatentConsistencyModelPipeline,
+        "tasks": ["text-to-image"],
+        "default_task": "text-to-image"
+    },
+    "pixart-alpha": {
+        "class": PixArtAlphaPipeline,
+        "tasks": ["text-to-image"],
+        "default_task": "text-to-image"
+    },
+    "pixart-sigma": {
+        "class": PixArtSigmaPipeline,
+        "tasks": ["text-to-image"],
+        "default_task": "text-to-image"
+    },
+    "flux": {
+        "class": FluxPipeline,
+        "tasks": ["text-to-image"],
+        "default_task": "text-to-image"
+    },
+    "flux-inpaint": {
+        "class": FluxInpaintPipeline,
+        "tasks": ["inpaint"],
+        "default_task": "inpaint"
+    },
+    "flux-img2img": {
+        "class": FluxImg2ImgPipeline,
+        "tasks": ["image-to-image"],
+        "default_task": "image-to-image"
+    },
 }
 
-# UPDATED: New choices for the Pull Request destination UI component
-DEST_NEW_NEURON_REPO = "Create new Neuron-optimized repository"
-DEST_CACHE_REPO = "Create a PR in the cache repository"
-DEST_CUSTOM_REPO = "Create a PR in a custom repository"
-
-PR_DESTINATION_CHOICES = [
-    DEST_NEW_NEURON_REPO,
-    DEST_CACHE_REPO,
-    DEST_CUSTOM_REPO
-]
-
-DEFAULT_CACHE_REPO = "aws-neuron/optimum-neuron-cache"
-
-# Get all tasks and pipelines for dropdowns
-ALL_TRANSFORMER_TASKS = list(TRANSFORMER_TASKS.keys())
-ALL_DIFFUSION_PIPELINES = list(DIFFUSION_PIPELINES.keys())
-
-def create_task_tag(task: str) -> str:
-    """Create a colored HTML tag for a task"""
-    if task in TRANSFORMER_TASKS:
-        color = TRANSFORMER_TASKS[task]["color"]
-        return f'<span style="background-color: {color}; color: white; padding: 2px 6px; border-radius: 12px; font-size: 0.75rem; font-weight: 500; margin: 1px;">{task}</span>'
-    elif task in DIFFUSION_PIPELINES:
-        color = DIFFUSION_PIPELINES[task]["color"]
-        return f'<span style="background-color: {color}; color: white; padding: 2px 6px; border-radius: 12px; font-size: 0.75rem; font-weight: 500; margin: 1px;">{task}</span>'
-    elif task in TAGS:
-        color = TAGS[task]["color"]
-        return f'<span style="background-color: {color}; color: white; padding: 2px 6px; border-radius: 12px; font-size: 0.75rem; font-weight: 500; margin: 1px;">{task}</span>'
+def get_default_inputs(task_or_pipeline: str, pipeline_name: str = None) -> 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"]:
+        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", "audio-xvector"]:
+        return {"batch_size": 1, "audio_sequence_length": 16000}
+    elif pipeline_name and pipeline_name in DIFFUSION_PIPELINE_MAPPING:
+        # For diffusion models, use appropriate sizes based on pipeline
+        if "xl" in pipeline_name.lower():
+            return {"batch_size": 1, "height": 1024, "width": 1024, "num_images_per_prompt": 1}
+        else:
+            return {"batch_size": 1, "height": 512, "width": 512, "num_images_per_prompt": 1}
     else:
-        return f'<span style="background-color: #6b7280; color: white; padding: 2px 6px; border-radius: 12px; font-size: 0.75rem; font-weight: 500; margin: 1px;">{task}</span>'
-
-def format_tasks_for_table(tasks_str: str) -> str:
-    """Convert comma-separated tasks into colored tags"""
-    tasks = [task.strip() for task in tasks_str.split(',')]
-    return ' '.join([create_task_tag(task) for task in tasks])
-
-def update_pipeline_and_task_dropdowns(model_type: str):
-    """Update the pipeline and task dropdowns based on selected model type"""
-    if model_type == "transformers":
-        return (
-            gr.Dropdown(visible=False),  # pipeline dropdown hidden
-            gr.Dropdown(
-                choices=ALL_TRANSFORMER_TASKS,
-                value="auto",
-                label="Task (auto can infer task from model)",
-                visible=True
-            )
-        )
-    else:  # diffusers
-        # Show pipeline dropdown, hide task dropdown initially
-        return (
-            gr.Dropdown(
-                choices=ALL_DIFFUSION_PIPELINES,
-                value="stable-diffusion",
-                label="Pipeline Type",
-                visible=True
-            ),
-            gr.Dropdown(
-                choices=DIFFUSION_PIPELINES["stable-diffusion"]["tasks"],
-                value=DIFFUSION_PIPELINES["stable-diffusion"]["tasks"][0],
-                label="Task",
-                visible=True
-            )
-        )
+        # Default to text-based shapes
+        return {"batch_size": 1, "sequence_length": 128}
 
-def update_task_dropdown_for_pipeline(pipeline_name: str):
-    """Update task dropdown based on selected pipeline"""
-    if pipeline_name in DIFFUSION_PIPELINES:
-        tasks = DIFFUSION_PIPELINES[pipeline_name]["tasks"]
-        return gr.Dropdown(
-            choices=tasks,
-            value=tasks[0] if tasks else None,
-            label="Task",
-            visible=True
+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 (for transformers)
+                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)
+    except Exception:
+        return None
+    for discussion in discussions:
+        if (
+            discussion.status == "open"
+            and discussion.is_pull_request
+            and discussion.title == pr_title
+        ):
+            return discussion
+    return None
+
+def export_diffusion_model(model_id: str, pipeline_name: str, task: str, folder: str, token: str) -> Generator:
+    """Export diffusion model using optimum.exporters.neuron"""
+    
+    yield f"📦 Exporting diffusion model `{model_id}` with pipeline `{pipeline_name}` for task `{task}`..."
+    
+    if pipeline_name not in DIFFUSION_PIPELINE_MAPPING:
+        supported = list(DIFFUSION_PIPELINE_MAPPING.keys())
+        raise Exception(f"❌ Unsupported pipeline: {pipeline_name}. Supported: {supported}")
+    
+    pipeline_config = DIFFUSION_PIPELINE_MAPPING[pipeline_name]
+    pipeline_class = pipeline_config["class"]
+    
+    # Get default inputs
+    inputs = get_default_inputs(task, pipeline_name)
+    yield f"🔧 Using default inputs: {inputs}"
+    
+    try:
+        # Load the pipeline
+        yield "📥 Loading diffusion pipeline..."
+        model = pipeline_class.from_pretrained(model_id, token=token)
+        
+        # Build input shapes for compilation
+        input_shapes = build_stable_diffusion_components_mandatory_shapes(**inputs)
+        
+        # Compiler arguments
+        compiler_kwargs = {
+            "auto_cast": "matmul",
+            "auto_cast_type": "bf16",
+        }
+        
+        yield "🔨 Starting compilation process..."
+        
+        # Get submodels and neuron configs
+        models_and_neuron_configs, output_model_names = get_submodels_and_neuron_configs(
+            model=model,
+            input_shapes=input_shapes,
+            task=task,
+            library_name="diffusers",
+            output=Path(folder),
+            model_name_or_path=model_id,
         )
-    return gr.Dropdown(visible=False)
+        
+        # Export models
+        _, neuron_outputs = export_models(
+            models_and_neuron_configs=models_and_neuron_configs,
+            task=task,
+            output_dir=Path(folder),
+            output_file_names=output_model_names,
+            compiler_kwargs=compiler_kwargs,
+        )
+        
+        yield f"✅ Diffusion model export completed. Files saved to {folder}"
+        
+    except Exception as e:
+        yield f"❌ Export failed with error: {e}"
+        raise
 
-def toggle_custom_repo_box(pr_destinations: List[str]):
-    """Show or hide the custom repo ID textbox based on checkbox selection."""
-    if DEST_CUSTOM_REPO in pr_destinations:
-        return gr.Textbox(visible=True)
-    else:
-        return gr.Textbox(visible=False, value="")
+def export_transformer_model(model_id: str, task: str, folder: str, token: str) -> Generator:
+    """Export transformer model using optimum.neuron"""
+    
+    yield f"📦 Exporting transformer model `{model_id}` for task `{task}`..."
+    
+    model_class = TASK_TO_MODEL_CLASS.get(task)
+    if model_class is None:
+        supported = list(TASK_TO_MODEL_CLASS.keys())
+        raise Exception(f"❌ Unsupported task: {task}. Supported: {supported}")
 
-def neuron_export(model_id: str, model_type: str, pipeline_name: str, task_or_pipeline: str, 
-                  pr_destinations: List[str], custom_repo_id: str, custom_cache_repo: str, oauth_token: gr.OAuthToken):
+    inputs = get_default_inputs(task)
+    compiler_configs = {"auto_cast": "matmul", "auto_cast_type": "bf16", "instance_type": "inf2"}
+    yield f"🔧 Using default inputs: {inputs}"
     
-    log_buffer = ""
-    def log(msg):
-        nonlocal log_buffer
-        # Handle cases where the message from the backend is not a string
-        if not isinstance(msg, str):
-            msg = str(msg)
-        log_buffer += msg + "\n"
-        return log_buffer
-
-    if oauth_token.token is None:
-        yield log("You must be logged in to use this space")
-        return
+    # Clear any old cache artifacts before export
+    cache_base_dir = "/var/tmp/neuron-compile-cache"
+    
+    try:
+        # Trigger the export/compilation
+        model = model_class.from_pretrained(
+            model_id,
+            export=True,
+            tensor_parallel_size=4,
+            token=token,
+            **compiler_configs,
+            **inputs,
+        )
         
-    if not model_id:
-        yield log("🚫 Invalid input. Please specify a model name from the hub.")
-        return
+        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
+
+def export_and_git_add(model_id: str, task_or_pipeline: str, model_type: str, folder: str, token: str, pipeline_name: str = None) -> Any:
+
+    operations = []
+    
     try:
-        api = HfApi(token=oauth_token.token)
-        # Set custom cache repo as environment variable
-        if custom_cache_repo:
-            os.environ['CUSTOM_CACHE_REPO'] = custom_cache_repo.strip()
+        if model_type == "diffusers":
+            # For diffusion models, use the new export function
+            export_gen = export_diffusion_model(model_id, pipeline_name, task_or_pipeline, folder, token)
+            for message in export_gen:
+                yield message
+        else:
+            # For transformer models, use the existing function
+            export_gen = export_transformer_model(model_id, task_or_pipeline, folder, token)
+            for message in export_gen:
+                yield message
+        
+        # Create operations from exported files
+        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))
         
-        yield log(f"🔑 Logging in ...")
+        # Update model card
+        try:
+            card = ModelCard.load(model_id, token=token)
+            if not hasattr(card.data, "tags") or card.data.tags is None:
+                card.data.tags = []
+            if "neuron" not in card.data.tags:
+                card.data.tags.append("neuron")
+            
+            readme_path = os.path.join(folder, "README.md")
+            card.save(readme_path)
+            
+            # Check if README.md is already in operations, if so update, else add
+            readme_op = next((op for op in operations if op.path_in_repo == "README.md"), None)
+            if readme_op:
+                readme_op.path_or_fileobj = readme_path
+            else:
+                operations.append(CommitOperationAdd(path_in_repo="README.md", path_or_fileobj=readme_path))
+
+        except Exception as e:
+            yield f"⚠️ Warning: Could not update model card: {e}"
+
+    except Exception as e:
+        yield f"❌ Export failed with error: {e}"
+        raise
+
+    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."""
+    requesting_user = api.whoami(token=token)["name"]
+    base_name = original_model_id.replace('/', '-')
+    return f"{requesting_user}/{base_name}-neuron"
+
+def create_neuron_repo_and_upload(
+    operations: List[CommitOperationAdd],
+    original_model_id: str,
+    model_type: str,
+    task_or_pipeline: str,
+    requesting_user: str,
+    token: str,
+    pipeline_name: str = None,
+) -> Generator[Union[str, RepoUrl], None, None]:
+    """
+    Creates a new repository with Neuron files and uploads them.
+    """
+    api = HfApi(token=token)
+
+    if task_or_pipeline == "auto" and model_type == "transformers":
         try:
-            api.model_info(model_id, token=oauth_token.token)
+            task_or_pipeline = TasksManager.infer_task_from_model(original_model_id, token=token)
         except Exception as e:
-            yield log(f"❌ Could not access model `{model_id}`: {e}")
-            return
+            raise Exception(f"❌ Could not infer task for model {original_model_id}: {e}")
+    
+    # Generate repository name
+    neuron_repo_name = generate_neuron_repo_name(api, original_model_id, task_or_pipeline, token)
+    
+    try:
+        # Create the repository
+        repo_url = create_repo(
+            repo_id=neuron_repo_name,
+            token=token,
+            repo_type="model",
+            private=False,
+            exist_ok=True,  
+        )
+                
+        # Get the appropriate class name for the Python example
+        if model_type == "transformers":
+            model_class = TASK_TO_MODEL_CLASS.get(task_or_pipeline)
+            model_class_name = model_class.__name__ if model_class else "NeuronModel"
+            usage_example = f"""```python
+from optimum.neuron import {model_class_name}
+
+model = {model_class_name}.from_pretrained("{neuron_repo_name}")
+```"""
+        else:
+            # For diffusion models
+            pipeline_config = DIFFUSION_PIPELINE_MAPPING.get(pipeline_name, {})
+            pipeline_class = pipeline_config.get("class")
+            if pipeline_class:
+                class_name = pipeline_class.__name__.replace("Pipeline", "")
+                model_class_name = f"Neuron{class_name}Pipeline"
+            else:
+                model_class_name = "NeuronStableDiffusionPipeline"
+            
+            usage_example = f"""```python
+from optimum.neuron import {model_class_name}
 
-        yield log(f"✅ Model `{model_id}` is accessible. Starting Neuron export...")
+pipeline = {model_class_name}.from_pretrained("{neuron_repo_name}")
+```"""
         
-        # UPDATED: Build pr_options with new structure
-        pr_options = {
-            "create_neuron_repo": DEST_NEW_NEURON_REPO in pr_destinations,
-            "create_cache_pr": DEST_CACHE_REPO in pr_destinations,
-            "create_custom_pr": DEST_CUSTOM_REPO in pr_destinations,
-            "custom_repo_id": custom_repo_id.strip() if custom_repo_id else ""
-        }
+        # Create enhanced model card for the Neuron repo
+        neuron_readme_content = f"""---
+tags:
+- neuron
+- optimized
+- aws-neuron
+- {task_or_pipeline}
+base_model: {original_model_id}
+---
 
-        # The convert function is a generator, so we iterate through its messages
-        for status_code, message in convert(
-            api, model_id, task_or_pipeline, model_type, 
-            token=oauth_token.token, pr_options=pr_options, 
-            pipeline_name=pipeline_name if model_type == "diffusers" else None
-        ):
-            if isinstance(message, str):
-                yield log(message)
-            else:  # It's the final result dictionary
-                final_message = "🎉 Process finished.\n"
-                if message.get("neuron_repo"):
-                    final_message += f"🏗️ New Neuron Repository: {message['neuron_repo']}\n"
-                if message.get("readme_pr"):
-                    final_message += f"📝 README PR (Original Model): {message['readme_pr']}\n"
-                if message.get("cache_pr"):
-                    final_message += f"🔗 Cache PR: {message['cache_pr']}\n"
-                if message.get("custom_pr"):
-                    final_message += f"🔗 Custom PR: {message['custom_pr']}\n"
-                yield log(final_message)
+# Neuron-Optimized {original_model_id}
 
-    except Exception as e:
-        yield log(f"❗ An unexpected error occurred in the Gradio interface: {e}")
+This repository contains AWS Neuron-optimized files for [{original_model_id}](https://huggingface.co/{original_model_id}).
+
+## Model Details
 
-TITLE_IMAGE = """
-<div style="display: block; margin-left: auto; margin-right: auto; width: 50%;">
-<img src="https://huggingface.co/spaces/optimum/neuron-export/resolve/main/huggingfaceXneuron.png"/>
-</div>
+- **Base Model**: [{original_model_id}](https://huggingface.co/{original_model_id})
+- **Task**: {task_or_pipeline}
+- **Optimization**: AWS Neuron compilation
+- **Generated by**: [{requesting_user}](https://huggingface.co/{requesting_user})
+- **Generated using**: [Optimum Neuron Compiler Space]({SPACES_URL})
+
+## Usage
+
+This model has been optimized for AWS Neuron devices (Inferentia/Trainium). To use it:
+
+{usage_example}
+
+## Performance
+
+These files are pre-compiled for AWS Neuron devices and should provide improved inference performance compared to the original model when deployed on Inferentia or Trainium instances.
+
+## Original Model
+
+For the original model, training details, and more information, please visit: [{original_model_id}](https://huggingface.co/{original_model_id})
 """
+        
+        # Update the README in operations
+        readme_op = next((op for op in operations if op.path_in_repo == "README.md"), None)
+        if readme_op:
+            # Create a temporary file with the new content
+            with NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
+                f.write(neuron_readme_content)
+                readme_op.path_or_fileobj = f.name
+        else:
+            # Add new README operation
+            with NamedTemporaryFile(mode='w', suffix='.md', delete=False) as f:
+                f.write(neuron_readme_content)
+                operations.append(CommitOperationAdd(path_in_repo="README.md", path_or_fileobj=f.name))
+        
+        # Upload files to the new repository
+        commit_message = f"Add Neuron-optimized files for {original_model_id}"
+        commit_description = f"""
+🤖 Neuron Export Bot: Adding AWS Neuron-optimized model files.
+
+Original model: [{original_model_id}](https://huggingface.co/{original_model_id})
+Task: {task_or_pipeline}
+Generated by: [{requesting_user}](https://huggingface.co/{requesting_user})
+Generated using: [Optimum Neuron Compiler Space]({SPACES_URL})
 
-TITLE = """
-<div style="text-align: center; max-width: 1400px; margin: 0 auto;">
-<h1 style="font-weight: 900; margin-bottom: 10px; margin-top: 10px; font-size: 2.2rem;">
-    🤗 Optimum Neuron Model Exporter 🏎️
-</h1>
-</div>
+These files have been pre-compiled for AWS Neuron devices (Inferentia/Trainium) and should provide improved inference performance.
 """
+        
+        commit_info = api.create_commit(
+            repo_id=neuron_repo_name,
+            operations=operations,
+            commit_message=commit_message,
+            commit_description=commit_description,
+            token=token,
+        )
+        yield f"✅ Repository created: {repo_url}"
+
+    except Exception as e:
+        yield f"❌ Failed to create/upload to Neuron repository: {e}"
+        raise
+
+def create_readme_pr_for_original_model(
+    original_model_id: str,
+    neuron_repo_name: str,
+    task_or_pipeline: str,
+    requesting_user: str,
+    token: str,
+) -> Generator[Union[str, CommitInfo], None, None]:
+    """
+    Creates a PR on the original model repository to add a link to the Neuron-optimized version.
+    """
+    api = HfApi(token=token)
 
-# UPDATED: Description to reflect new workflow
-DESCRIPTION = """
-This Space allows you to automatically export 🤗 transformers and 🧨 diffusion models to AWS Neuron-optimized format for Inferentia/Trainium acceleration. 
+    yield f"📝 Creating PR to add Neuron repo link in {original_model_id}..."
 
-Simply provide a model ID from the Hugging Face Hub, and choose your desired output.
+    try:
+        # Check if there's already an open PR
+        pr_title = "Add link to Neuron-optimized version"
+        existing_pr = previous_pr(api, original_model_id, pr_title)
+
+        if existing_pr:
+            yield f"⚠️ PR already exists: https://huggingface.co/{original_model_id}/discussions/{existing_pr.num}"
+            return
+
+        # Get the current README
+        try:
+            current_readme_path = api.hf_hub_download(
+                repo_id=original_model_id,
+                filename="README.md",
+                token=token,
+            )
+            with open(current_readme_path, 'r', encoding='utf-8') as f:
+                readme_content = f.read()
+        except Exception:
+            # If README doesn't exist, create a basic one
+            readme_content = f"# {original_model_id}\n\n"
 
-### ✨ Key Features
+        # Add Neuron optimization section, separated by a horizontal rule
+        neuron_section = f"""
+---
+## 🚀 AWS Neuron Optimized Version Available
 
-* **🚀 Create a New Optimized Repo**: Automatically converts your model and uploads it to a new repository under your username (e.g., `your-username/model-name-neuron`).  
-* **🔗 Link Back to Original**: Creates a Pull Request on the original model's repository to add a link to your optimized version, making it easier for the community to discover.  
-* **🛠️ PR to a Custom Repo**: For custom workflows, you can create a Pull Request to add the optimized files directly into an existing repository you own.  
-* **📦 Contribute to Cache**: Contribute the generated compilation artifacts to a centralized cache repository (or your own private cache), helping avoid recompilation of already exported models.  
+A Neuron-optimized version of this model is available for improved performance on AWS Inferentia/Trainium instances:
 
-### ⚙️ How to Use  
-1. **Model ID**: Enter the ID of the model you want to export (e.g., `bert-base-uncased` or `stabilityai/stable-diffusion-xl-base-1.0`) and choose the corresponding task.  
-2. **Export Options**: Select at least one option for where to save the exported model. You can provide your own cache repo ID or use the default (`aws-neuron/optimum-neuron-cache`).  
-3. **Convert & Upload**: Click the button and follow the logs to track progress!  
+**[{neuron_repo_name}](https://huggingface.co/{neuron_repo_name})**
 
+The Neuron-optimized version provides:
+- Pre-compiled artifacts for faster loading
+- Optimized performance on AWS Neuron devices
+- Same model capabilities with improved inference speed
 """
 
-CUSTOM_CSS = """
-/* Primary button styling with warm colors */
-button.gradio-button.lg.primary {
-    /* Changed the blue/green gradient to an orange/yellow one */
-    background: linear-gradient(135deg, #F97316, #FBBF24) !important;
-    color: white !important;
-    padding: 16px 32px !important;
-    font-size: 1.1rem !important;
-    font-weight: 700 !important;
-    border: none !important;
-    border-radius: 12px !important;
-    /* Updated the shadow to match the new orange color */
-    box-shadow: 0 0 15px rgba(249, 115, 22, 0.5) !important;
-    transition: all 0.3s cubic-bezier(0.25, 0.8, 0.25, 1) !important;
-    position: relative;
-    overflow: hidden;
-}
-/* Login button styling with glow effect using dark blue and violet colors */
-#login-button {
-    background: linear-gradient(135deg, #1a237e, #6a1b9a) !important; /* Dark Blue to Violet */
-    color: white !important;
-    font-weight: 700 !important;
-    border: none !important;
-    border-radius: 12px !important;
-    box-shadow: 0 0 15px rgba(106, 27, 154, 0.6) !important; /* Cool violet glow */
-    transition: all 0.3s cubic-bezier(0.25, 0.8, 0.25, 1) !important;
-    position: relative;
-    overflow: hidden;
-    animation: glow 1.5s ease-in-out infinite alternate;
-    max-width: 350px !important;
-    margin: 0 auto !important;
-}
-#login-button::before {
-    content: "🔑 ";
-    display: inline-block !important;
-    vertical-align: middle !important;
-    margin-right: 5px !important;
-    line-height: normal !important;
-}
-#login-button:hover {
-    transform: translateY(-3px) scale(1.03) !important;
-    box-shadow: 0 10px 25px rgba(26, 35, 126, 0.7) !important; /* Deeper blue glow */
-}
-#login-button::after {
-    content: "";
-    position: absolute;
-    top: 0;
-    left: -100%;
-    width: 100%;
-    height: 100%;
-    background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.25), transparent);
-    transition: 0.5s;
-}
-#login-button:hover::after {
-    left: 100%;
-}
+        # Append the Neuron section to the end of the README
+        updated_readme = readme_content.rstrip() + "\n" + neuron_section
+
+        # Create temporary file with updated README
+        with NamedTemporaryFile(mode='w', suffix='.md', delete=False, encoding="utf-8") as f:
+            f.write(updated_readme)
+            temp_readme_path = f.name
+
+        # Create the PR
+        operations = [CommitOperationAdd(path_in_repo="README.md", path_or_fileobj=temp_readme_path)]
+
+        commit_description = f"""
+🤖 Neuron Export Bot: Adding link to Neuron-optimized version.
 
+A Neuron-optimized version of this model has been created at [{neuron_repo_name}](https://huggingface.co/{neuron_repo_name}).
+
+The optimized version provides improved performance on AWS Inferentia/Trainium instances with pre-compiled artifacts.
+
+Generated by: [{requesting_user}](https://huggingface.co/{requesting_user})
+Generated using: [Optimum Neuron Compiler Space]({SPACES_URL})
 """
 
-with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
-    gr.Markdown("**You must be logged in to use this space**")
-    gr.LoginButton(elem_id="login-button", elem_classes="center-button", min_width=250)
-    gr.HTML(TITLE_IMAGE)
-    gr.HTML(TITLE)
-    gr.Markdown(DESCRIPTION)
+        pr = api.create_commit(
+            repo_id=original_model_id,
+            operations=operations,
+            commit_message=pr_title,
+            commit_description=commit_description,
+            create_pr=True,
+            token=token,
+        )
+
+        yield f"✅ README PR created: https://huggingface.co/{original_model_id}/discussions/{pr.pr_num}"
+
+        # Clean up temporary file
+        os.unlink(temp_readme_path)
+
+    except Exception as e:
+        yield f"❌ Failed to create README PR: {e}"
+        raise
+
+def upload_to_custom_repo(
+    operations: List[CommitOperationAdd],
+    custom_repo_id: str,
+    original_model_id: str,
+    requesting_user: str,
+    token: str,
+) -> Generator[Union[str, CommitInfo], None, None]:
+    """
+    Uploads neuron files to a custom repository and creates a PR.
+    """
+    api = HfApi(token=token)
     
-    with gr.Tabs():
-        with gr.Tab("Export Model"):
-            with gr.Group():          
-                with gr.Row():
-                    pr_destinations_checkbox = gr.CheckboxGroup(
-                        choices=PR_DESTINATION_CHOICES,
-                        label="Export Destination",
-                        value=[DEST_NEW_NEURON_REPO],
-                        info="Select one or more destinations for the compiled model."
-                    )
-                    custom_repo_id_textbox = gr.Textbox(
-                        label="Custom Repository ID",
-                        placeholder="e.g., your-username/your-repo-name",
-                        visible=False,  
-                        interactive=True
-                    )
-                    custom_cache_repo_textbox = gr.Textbox(
-                        label="Custom Cache Repository",
-                        placeholder="e.g., your-org/your-cache-repo",
-                        value=DEFAULT_CACHE_REPO,
-                        info=f"Repository to store and fetch from compilation cache artifacts (default: {DEFAULT_CACHE_REPO}) ",
-                        interactive=True
+    try:
+        # Ensure the custom repo exists
+        api.repo_info(repo_id=custom_repo_id, repo_type="model")
+    except Exception as e:
+        yield f"❌ Could not access custom repository `{custom_repo_id}`. Please ensure it exists and you have write access. Error: {e}"
+        raise
+
+    pr_title = f"Add Neuron-optimized files for {original_model_id}"
+    commit_description = f"""
+🤖 Neuron Export Bot: On behalf of [{requesting_user}](https://huggingface.co/{requesting_user}), adding AWS Neuron-optimized model files for `{original_model_id}`.
+
+These files were generated using the [Optimum Neuron Compiler Space](https://huggingface.co/spaces/optimum/neuron-export).
+"""
+
+    try:
+        custom_pr = api.create_commit(
+            repo_id=custom_repo_id,
+            operations=operations,
+            commit_message=pr_title,
+            commit_description=commit_description,
+            create_pr=True,
+            token=token,
+        )
+        yield f"✅ Custom PR created successfully: https://huggingface.co/{custom_repo_id}/discussions/{custom_pr.pr_num}"
+        yield custom_pr
+
+    except Exception as e:
+        yield f"❌ Failed to create PR in custom repository: {e}"
+        raise
+
+def convert(
+    api: "HfApi",
+    model_id: str,
+    task_or_pipeline: str,
+    model_type: str = "transformers",
+    token: str = None,
+    pr_options: Dict = None,
+    pipeline_name: str = None,
+) -> Generator[Tuple[str, Any], None, None]:
+    if pr_options is None:
+        pr_options = {}
+    
+    info = api.model_info(model_id, token=token)
+    filenames = {s.rfilename for s in info.siblings}
+    requesting_user = api.whoami(token=token)["name"]
+
+    if not any(pr_options.values()):
+        yield "1", "⚠️ No option selected. Please choose at least one option."
+        return
+
+    if pr_options.get("create_custom_pr") and not pr_options.get("custom_repo_id"):
+        yield "1", "⚠️ Custom PR selected but no repository ID was provided."
+        return
+
+    yield "0", f"🚀 Starting export process with options: {pr_options}..."
+
+    if task_or_pipeline == "auto" and model_type == "transformers":
+        try:
+            task_or_pipeline = TasksManager.infer_task_from_model(model_id, token=token)
+        except Exception as e:
+            raise Exception(f"❌ Could not infer task for model {model_id}: {e}")
+
+    with TemporaryDirectory() as temp_dir:
+        export_folder = os.path.join(temp_dir, "export")
+        cache_mirror_dir = os.path.join(temp_dir, "cache_mirror")
+        os.makedirs(export_folder, exist_ok=True)
+        os.makedirs(cache_mirror_dir, exist_ok=True)
+        
+        result_info = {}
+
+        try:
+            # --- Export Logic ---
+            export_gen = export_and_git_add(model_id, task_or_pipeline, model_type, export_folder, token=token, pipeline_name=pipeline_name)
+            operations = None
+            for message in export_gen:
+                if isinstance(message, tuple) and message[0] == "__RETURN__":
+                    operations = message[1]
+                    break
+                else:
+                    yield "0", message
+            
+            if not operations:
+                raise Exception("Export process did not produce any files to commit.")
+
+            # --- Cache Handling ---
+            if pr_options.get("create_cache_pr"):
+                yield "0", f"📤 Creating a Pull Request for the cache repository ..."
+                
+                try:
+                    pr_title = f"Add Neuron cache artifacts for {model_id}"
+                    custom_pr_description = f"""
+🤖 **Neuron Cache Sync Bot**
+
+This PR adds newly compiled cache artifacts for the model:
+- **Original Model ID:** `{model_id}`
+- **Task:** `{task_or_pipeline}`
+
+These files were generated to accelerate model loading on AWS Neuron devices.
+""" 
+                    
+                    # 1. Create an instance of your generator
+                    commit_message = f"Synchronizing local compiler cache of {model_id}"
+                    inputs = get_default_inputs(task_or_pipeline, pipeline_name)
+                    commit_description = f"""
+🤖 **Neuron Cache Sync Bot**
+
+This commit adds newly compiled cache artifacts for the model:
+- **Original Model ID:** `{model_id}`
+- **Task:** `{task_or_pipeline}`
+- **Compilation inputs:** {inputs}
+- **Generated by:** [{requesting_user}](https://huggingface.co/{requesting_user})
+- **Generated using:** [Optimum Neuron Model Exporter]({SPACES_URL})
+
+These files were generated to accelerate model loading on AWS Neuron devices.
+"""
+             
+                    pr_generator = synchronize_hub_cache_with_pr(
+                        cache_repo_id=CUSTOM_CACHE_REPO,
+                        commit_message=commit_message,
+                        commit_description=commit_description,
+                        token=token,
                     )
-            with gr.Row():
-                model_type = gr.Radio(
-                    choices=["transformers", "diffusers"],
-                    value="transformers",
-                    label="Model Type",
-                    info="Choose the type of model you want to export"
-                ) 
-            with gr.Row():
-                input_model = HuggingfaceHubSearch(
-                    label="Hub model ID",
-                    placeholder="Search for a model on the Hub...",
-                    search_type="model",
-                )
-                pipeline_dropdown = gr.Dropdown(
-                    choices=ALL_DIFFUSION_PIPELINES,
-                    value="stable-diffusion",
-                    label="Pipeline Type",
-                    visible=False
+
+                    pr_url = None
+                    # 2. Loop to process yielded status messages and capture the final return value
+                    while True:
+                        try:
+                            # Get the next status message from your generator
+                            status_message = next(pr_generator)
+                            yield "0", status_message
+                        except StopIteration as e:
+                            # The generator is finished. Its `return` value is in e.value.
+                            pr_url = e.value
+                            break # Exit the loop
+
+                    # 3. Process the final result
+                    if pr_url:
+                        yield "0", f"✅ Successfully captured PR URL."
+                        result_info["cache_pr"] = pr_url
+                    else:
+                        yield "0", "⚠️ PR process finished, but no URL was returned. This may be expected in non-blocking mode."
+
+                except Exception as e:
+                    yield "0", f"❌ Failed to create cache PR: {e}"
+            
+            # --- New Repository Creation (Replaces Model PR) ---
+            if pr_options.get("create_neuron_repo"):
+                yield "0", "🏗️ Creating new Neuron-optimized repository..."
+                neuron_repo_url = None
+                # Generate the repo name first so we can use it consistently
+                neuron_repo_name = generate_neuron_repo_name(api, model_id, task_or_pipeline, token)
+                
+                repo_creation_gen = create_neuron_repo_and_upload(
+                    operations, model_id, model_type, task_or_pipeline, requesting_user, token, pipeline_name
                 )
-                task_dropdown = gr.Dropdown(
-                    choices=ALL_TRANSFORMER_TASKS,
-                    value="auto",
-                    label="Task (auto can infer from model)",
+                
+                for msg in repo_creation_gen:
+                    if isinstance(msg, str):
+                        yield "0", msg
+                    else:
+                        neuron_repo_url = msg
+                
+                result_info["neuron_repo"] = f"https://huggingface.co/{neuron_repo_name}"
+                
+                # Automatically create a PR on the original model to add a link
+                readme_pr = None
+                readme_pr_gen = create_readme_pr_for_original_model(
+                    model_id, neuron_repo_name, task_or_pipeline, requesting_user, token
                 )
+                for msg in readme_pr_gen:
+                    if isinstance(msg, str):
+                        yield "0", msg
+                    else:
+                        readme_pr = msg
+                
+                if readme_pr:
+                    result_info["readme_pr"] = f"https://huggingface.co/{model_id}/discussions/{readme_pr.pr_num}"
             
-            btn = gr.Button("Export to Neuron", size="lg", variant="primary")
-            
-            log_box = gr.Textbox(label="Logs", lines=20, interactive=False, show_copy_button=True)
-            
-            # Event Handlers
-            model_type.change(
-                fn=update_pipeline_and_task_dropdowns,
-                inputs=[model_type],
-                outputs=[pipeline_dropdown, task_dropdown]
-            )
-            
-            pipeline_dropdown.change(
-                fn=update_task_dropdown_for_pipeline,
-                inputs=[pipeline_dropdown],
-                outputs=[task_dropdown]
-            )
-            
-            pr_destinations_checkbox.change(
-                fn=toggle_custom_repo_box,
-                inputs=pr_destinations_checkbox,
-                outputs=custom_repo_id_textbox
-            )
-            
-            btn.click(
-                fn=neuron_export,
-                inputs=[
-                    input_model, 
-                    model_type,
-                    pipeline_dropdown,
-                    task_dropdown,
-                    pr_destinations_checkbox,
-                    custom_repo_id_textbox,
-                    custom_cache_repo_textbox
-                ],
-                outputs=log_box,
-            )
-            
-        with gr.Tab("Supported Architectures"):
-            gr.HTML(f"""
-            <div style="margin-bottom: 20px;">
-                <h3>🎨 Task Categories Legend</h3>
-                <div class="task-tags">
-                    {create_task_tag("Feature Extraction")} 
-                    {create_task_tag("NLP")}
-                    {create_task_tag("Text Generation")}
-                    {create_task_tag("Audio")}
-                    {create_task_tag("Vision")}
-                    {create_task_tag("Multimodal")}
-                    {create_task_tag("Similarity")}
-                </div>
-            </div>
-            """)
-            
-            gr.HTML(f"""
-            <h2>🤗 Transformers</h2>
-            <table style="width: 100%; border-collapse: collapse; margin: 20px 0;">
-                <colgroup>
-                    <col style="width: 30%;">
-                    <col style="width: 70%;">
-                </colgroup>
-                <thead>
-                    <tr style="background-color: var(--background-fill-secondary);">
-                        <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Architecture</th>
-                        <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Supported Tasks</th>
-                    </tr>
-                </thead>
-                <tbody>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ALBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">AST</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, audio-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">BERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">BLOOM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Beit</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CamemBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CLIP</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ConvBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ConvNext</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ConvNextV2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CvT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DeBERTa (INF2 only)</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DeBERTa-v2  (INF2 only)</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Deit</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DistilBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DonutSwin</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Dpt</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ELECTRA</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ESM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">FlauBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">GPT2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Hubert</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Levit</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Llama, Llama 2, Llama 3</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Mistral</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Mixtral</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MobileBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MobileNetV2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification, semantic-segmentation")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MobileViT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification, semantic-segmentation")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ModernBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MPNet</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">OPT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Phi</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">RoBERTa</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">RoFormer</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Swin</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">T5</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text2text-generation")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">UniSpeech</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">UniSpeech-SAT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification, audio-frame-classification, audio-xvector")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ViT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Wav2Vec2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification, audio-frame-classification, audio-xvector")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">WavLM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification, audio-frame-classification, audio-xvector")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Whisper</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("automatic-speech-recognition")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">XLM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">XLM-RoBERTa</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Yolos</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, object-detection")}</td></tr>
-                </tbody>
-            </table>
-            <h2>🧨 Diffusers</h2>
-            <table style="width: 100%; border-collapse: collapse; margin: 20px 0;">
-                <colgroup>
-                    <col style="width: 30%;">
-                    <col style="width: 70%;">
-                </colgroup>
-                <thead>
-                    <tr style="background-color: var(--background-fill-secondary);">
-                        <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Architecture</th>
-                        <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Supported Tasks</th>
-                    </tr>
-                </thead>
-                <tbody>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Stable Diffusion</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image, image-to-image, inpaint")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Stable Diffusion XL Base</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image, image-to-image, inpaint")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Stable Diffusion XL Refiner</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("image-to-image, inpaint")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">SDXL Turbo</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image, image-to-image, inpaint")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">LCM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">PixArt-α</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">PixArt-Σ</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Flux</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr>
-
-                </tbody>
-            </table>
-            <h2>🤖 Sentence Transformers</h2>
-            <table style="width: 100%; border-collapse: collapse; margin: 20px 0;">
-                <colgroup>
-                    <col style="width: 30%;">
-                    <col style="width: 70%;">
-                </colgroup>
-                <thead>
-                    <tr style="background-color: var(--background-fill-secondary);">
-                        <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Architecture</th>
-                        <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Supported Tasks</th>
-                    </tr>
-                </thead>
-                <tbody>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Transformer</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, sentence-similarity")}</td></tr>
-                    <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CLIP</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, zero-shot-image-classification")}</td></tr>
-                </tbody>
-            </table>
-            <div style="margin-top: 20px;">
-                <p>💡 <strong>Note</strong>: Some architectures may have specific requirements or limitations. DeBERTa models are only supported on INF2 instances.</p>
-                <p>For more details, check the <a href="https://huggingface.co/docs/optimum-neuron" target="_blank">Optimum Neuron documentation</a>.</p>
-            </div>
-            """)
-    
-    # Add spacing between tabs and content
-    gr.Markdown("<br><br><br><br>")
+            # --- Custom Repository PR ---
+            if pr_options.get("create_custom_pr"):
+                custom_repo_id = pr_options["custom_repo_id"]
+                yield "0", f"📤 Creating PR in custom repository: {custom_repo_id}..."
+                custom_pr = None
+                custom_upload_gen = upload_to_custom_repo(operations, custom_repo_id, model_id, requesting_user, token)
+                for msg in custom_upload_gen:
+                    if isinstance(msg, str):
+                        yield "0", msg
+                    else:
+                        custom_pr = msg
+                if custom_pr:
+                    result_info["custom_pr"] = f"https://huggingface.co/{custom_repo_id}/discussions/{custom_pr.pr_num}"
 
-if __name__ == "__main__":
-    demo.launch(debug=True)
+            yield "0", result_info
+
+        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