app.py

import csv
import os
from datetime import datetime
from typing import Optional, Union, List
import gradio as gr
from huggingface_hub import HfApi, Repository
from huggingface_hub import login
from optimum_neuron_export import convert
from gradio_huggingfacehub_search import HuggingfaceHubSearch
from apscheduler.schedulers.background import BackgroundScheduler

DATASET_REPO_URL = "https://huggingface.co/datasets/optimum/neuron-exports"
DATA_FILENAME = "exports.csv"
DATA_FILE = os.path.join("data", DATA_FILENAME)
HF_TOKEN = os.getenv("HF_TOKEN") # It's better to use environment variables
DATADIR = "neuron_exports_data"
repo: Optional[Repository] = None

# Uncomment if you want to push to dataset repo with token
# if HF_TOKEN:
# repo = Repository(local_dir=DATADIR, clone_from=DATASET_REPO_URL, token=HF_TOKEN)

# 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"},
}

# Define diffusion pipeline types
DIFFUSION_PIPELINES = {
    "text-to-image": {"color": "#ec4899", "category": "Stable Diffusion"},
    "image-to-image": {"color": "#ec4899", "category": "Stable Diffusion"},
    "inpaint": {"color": "#ec4899", "category": "Stable Diffusion"},
    "instruct-pix2pix": {"color": "#ec4899", "category": "Stable Diffusion"},
    "latent-consistency": {"color": "#8b5cf6", "category": "Latent Consistency"},
    "stable-diffusion": {"color": "#10b981", "category": "Stable Diffusion"},
    "stable-diffusion-xl": {"color": "#10b981", "category": "Stable Diffusion XL"},
    "stable-diffusion-xl-img2img": {"color": "#10b981", "category": "Stable Diffusion XL"},
    "stable-diffusion-xl-inpaint": {"color": "#10b981", "category": "Stable Diffusion XL"},
    "controlnet": {"color": "#f59e0b", "category": "ControlNet"},
    "controlnet-xl": {"color": "#f59e0b", "category": "ControlNet XL"},
    "pixart-alpha": {"color": "#ef4444", "category": "PixArt"},
    "pixart-sigma": {"color": "#ef4444", "category": "PixArt"},
    "flux": {"color": "#06b6d4", "category": "Flux"},
}

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"},
}

# 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
]

# 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>'
    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_task_dropdown(model_type: str):
    """Update the task dropdown based on selected model type"""
    if model_type == "transformers":
        return gr.Dropdown(
            choices=ALL_TRANSFORMER_TASKS,
            value="auto",
            label="Task (auto can infer task from model)",
            visible=True
        )
    else:  # diffusion
        return gr.Dropdown(
            choices=ALL_DIFFUSION_PIPELINES,
            value="text-to-image",
            label="Pipeline Type",
            visible=True
        )

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="")

# UPDATED: Modified function to handle new repository creation workflow
def neuron_export(model_id: str, model_type: str, task_or_pipeline: str, 
                  pr_destinations: List[str], custom_repo_id: str):
    if not model_id:
        yield "🚫 Invalid input. Please specify a model name from the hub."
        return

    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

    try:
        api = HfApi()
        yield log(f"🔑 Logging in with provided token...")
        if not HF_TOKEN:
            yield log("❌ HF_TOKEN not found. Please set it as an environment variable in the Space secrets.")
            return
        
        login(token=HF_TOKEN)
        yield log("✅ Login successful.")
        yield log(f"🔍 Checking access to `{model_id}`...")
        try:
            api.model_info(model_id, token=HF_TOKEN)
        except Exception as e:
            yield log(f"❌ Could not access model `{model_id}`: {e}")
            return

        yield log(f"✅ Model `{model_id}` is accessible. Starting Neuron export...")
        
        # 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 ""
        }

        # 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=HF_TOKEN, pr_options=pr_options):
            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)

    except Exception as e:
        yield log(f"❗ An unexpected error occurred in the Gradio interface: {e}")

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>
"""

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 Compiler 🏎️
</h1>
</div>
"""

# 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. 

Simply provide a model ID from the Hugging Face Hub, and choose your desired output.

### ✨ Key Features

* **🚀 Create a New Optimized Repo**: Automatically converts the 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 new optimized version, making it easily discoverable by the community.
* **🛠️ PR to a Custom Repo**: For custom workflows, you can create a Pull Request with the optimized files directly into an existing repository you own.
* **📦 Contribute to Cache**: You can also contribute the generated compilation artifacts to a centralized cache repository, which helps speed up future compilations for everyone.

### ⚙️ 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`).
2.  **Export Options**: Select at least one option for where to save the exported model.
3.  **Convert & Upload**: Click the button and follow the logs for progress!
"""

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;
}
"""

with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
    gr.HTML(TITLE_IMAGE)
    gr.HTML(TITLE)
    gr.Markdown(DESCRIPTION)
    
    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
                    )
            with gr.Row():
                model_type = gr.Radio(
                    choices=["transformers", "diffusion"],
                    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",
                )
                task_dropdown = gr.Dropdown(
                    choices=ALL_TRANSFORMER_TASKS,
                    value="auto",
                    label="Task (auto can infer from model)",
                )
            
            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_task_dropdown,
                inputs=[model_type],
                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, 
                    task_dropdown,
                    pr_destinations_checkbox,
                    custom_repo_id_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>")

if __name__ == "__main__":
    demo.launch(debug=True)