Create app.py

app.py

@@ -0,0 +1,492 @@
+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 optimum_neuron_export import convert
+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"},
+}
+
+# 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
+]
+
+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>'
+    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:  # diffusers
+        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="")
+
+def neuron_export(model_id: str, model_type: str, task_or_pipeline: str, 
+                  pr_destinations: List[str], custom_repo_id: str, custom_cache_repo: str, oauth_token: gr.OAuthToken):
+    
+    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
+        
+    if not model_id:
+        yield log("🚫 Invalid input. Please specify a model name from the hub.")
+        return
+        
+    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()
+        
+        yield log(f"🔑 Logging in ...")
+        try:
+            api.model_info(model_id, token=oauth_token.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=oauth_token.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 Exporter 🏎️
+</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 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.  
+
+### ⚙️ 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!  
+
+"""
+
+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%;
+}
+
+"""
+
+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)
+    
+    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
+                    )
+            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",
+                )
+                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,
+                    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>")
+
+if __name__ == "__main__":
+    demo.launch(debug=True)