import csv
import os
from datetime import datetime
from typing import Optional, Union

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

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.environ.get("HF_WRITE_TOKEN")

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 all possible tasks and their categories for coloring
TASK_CATEGORIES = {
    "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"},
    "text-to-image": {"color": "#ec4899", "category": "Multimodal"},
    "image-to-image": {"color": "#ec4899", "category": "Multimodal"},
    "inpaint": {"color": "#ec4899", "category": "Multimodal"},
    "zero-shot-image-classification": {"color": "#ec4899", "category": "Multimodal"},
    "sentence-similarity": {"color": "#06b6d4", "category": "Similarity"},
}

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

# Get all tasks for dropdown
ALL_TASKS = list(TASK_CATEGORIES.keys())

def create_task_tag(task: str) -> str:
    """Create a colored HTML tag for a task"""
    if task in TASK_CATEGORIES:
        color = TASK_CATEGORIES[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 neuron_export(model_id: str, task: str, oauth_token: gr.OAuthToken) -> str:
    if oauth_token.token is None:
        return "You must be logged in to use this space"

    if not model_id:
        return f"### Invalid input 🐞 Please specify a model name, got {model_id}"

    try:
        api = HfApi(token=oauth_token.token)

        error, result_info = convert(api=api, model_id=model_id, task=task, token=oauth_token.token)
        if error != "0":
            return error

        print("[result_info]", result_info)

        # Save in a private dataset if repo initialized
        if repo is not None:
            repo.git_pull(rebase=True)
            with open(os.path.join(DATADIR, DATA_FILE), "a") as csvfile:
                writer = csv.DictWriter(
                    csvfile, fieldnames=["model_id", "cache_info", "time"]
                )
                writer.writerow(
                    {
                        "model_id": model_id,
                        "cache_info": str(result_info),
                        "time": str(datetime.now()),
                    }
                )
            commit_url = repo.push_to_hub()
            print("[dataset]", commit_url)

        return f"#### Success 🔥 Yay! This model was successfully compiled and cached! The compilation artifacts have been synchronized to the optimum-neuron cache repository. You can now use this model with optimum-neuron and benefit from faster loading times. Model: {model_id}, Task: {task}"

    except Exception as e:
        return f"#### Error: {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 🏎️ (WIP)
</h1>
</div>
"""

DESCRIPTION = """
This Space allows you to automatically export 🤗 transformers models hosted on the Hugging Face Hub to AWS Neuron-optimized format for Inferentia/Trainium acceleration. It opens a PR on the target model, and it is up to the owner of the original model to merge the PR to allow people to leverage Neuron optimization!

**Features:**
- Automatically opens PR with Neuron-optimized model
- Preserves original model weights
- Adds proper tags to model card

**Requirements:**
- Model must be compatible with [Optimum Neuron](https://huggingface.co/docs/optimum-neuron)
- User must be logged in with write token
"""

# Custom CSS to fix dark mode compatibility and transparency issues
CUSTOM_CSS = """
/* Fix for HuggingfaceHubSearch component visibility in both light and dark modes */
.gradio-container .gr-form {
    background: var(--background-fill-primary) !important;
    border: 1px solid var(--border-color-primary) !important;
}

/* Ensure text is visible in both modes */
.gradio-container input[type="text"], 
.gradio-container textarea,
.gradio-container .gr-textbox input {
    color: var(--body-text-color) !important;
    background: var(--input-background-fill) !important;
    border: 1px solid var(--border-color-primary) !important;
}

/* Fix dropdown/search results visibility */
.gradio-container .gr-dropdown,
.gradio-container .gr-dropdown .gr-box,
.gradio-container [data-testid="textbox"] {
    background: var(--background-fill-primary) !important;
    color: var(--body-text-color) !important;
    border: 1px solid var(--border-color-primary) !important;
}

/* Fix for search component specifically */
.gradio-container .gr-form > div,
.gradio-container .gr-form input {
    background: var(--input-background-fill) !important;
    color: var(--body-text-color) !important;
}

/* Ensure proper contrast for placeholder text */
.gradio-container input::placeholder {
    color: var(--body-text-color-subdued) !important;
    opacity: 0.7;
}

/* Fix any remaining transparent backgrounds */
.gradio-container .gr-box,
.gradio-container .gr-panel {
    background: var(--background-fill-primary) !important;
}

/* Make sure search results are visible */
.gradio-container .gr-dropdown-item {
    color: var(--body-text-color) !important;
    background: var(--background-fill-primary) !important;
}

.gradio-container .gr-dropdown-item:hover {
    background: var(--background-fill-secondary) !important;
}

/* Task tag styling improvements */
.task-tags {
    line-height: 1.8;
}

.task-tags span {
    display: inline-block;
    margin: 2px;
}

/* Primary button styling with warm colors */
button[variant="primary"] {
    background: linear-gradient(135deg, #3B82F6, #10B981) !important;
    color: white !important;
    padding: 16px 32px !important;
    font-size: 1.1rem !important;
    font-weight: 700 !important;
    border: none !important;
    border-radius: 12px !important;
    box-shadow: 0 0 15px rgba(59, 130, 246, 0.5) !important;
    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;
}
button[variant="primary"]::before {
    content: "✨ ";
}
button[variant="primary"]:hover {
    transform: translateY(-5px) scale(1.05) !important;
    box-shadow: 0 10px 25px rgba(59, 130, 246, 0.7) !important;
}
@keyframes glow {
    from {
        box-shadow: 0 0 10px rgba(59, 130, 246, 0.5);
    }
    to {
        box-shadow: 0 0 20px rgba(59, 130, 246, 0.8), 0 0 30px rgba(16, 185, 129, 0.5);
    }
}


/* Login button styling with glow effect using warm colors */
#login-button {
    background: linear-gradient(135deg, #FF7A00, #FFD700) !important; /* Orange to Gold */
    color: white !important;
    font-weight: 700 !important;
    border: none !important;
    border-radius: 12px !important;
    box-shadow: 0 0 15px rgba(255, 165, 0, 0.5) !important; /* Warm 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: 300px !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(255, 140, 0, 0.7) !important; /* Deeper warm 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.2), transparent);
    transition: 0.5s;
}

#login-button:hover::after {
    left: 100%;
}
"""

with gr.Blocks(css=CUSTOM_CSS) as demo:
    # Login requirement notice and button
    gr.Markdown("**You must be logged in to use this space**")
    gr.LoginButton(elem_id="login-button", elem_classes="center-button", min_width=250)
    
    # Centered title and image
    gr.HTML(TITLE_IMAGE)
    gr.HTML(TITLE)
    
    # Full-width description
    gr.Markdown(DESCRIPTION)
    
    with gr.Tabs():
        with gr.Tab("Export Model"):
            # Input controls in a row
            with gr.Row():
                input_model = HuggingfaceHubSearch(
                    label="Hub model ID",
                    placeholder="Search for model ID on the hub",
                    search_type="model",
                )
                input_task = gr.Dropdown(
                    choices=ALL_TASKS,
                    value="auto",
                    label='Task (auto could infer task from model)',
                )
            
            # Export button below the inputs
            btn = gr.Button("Export to Neuron", size="lg", variant="primary")
            
            # Output section
            output = gr.Markdown(label="Output")

            btn.click(
                fn=neuron_export,
                inputs=[input_model, input_task],
                outputs=output,
            )
        
        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>
                </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__":
    def restart_space():
        if HF_TOKEN:
            HfApi().restart_space(repo_id="optimum/neuron-export", token=HF_TOKEN, factory_reboot=True)

    scheduler = BackgroundScheduler()
    scheduler.add_job(restart_space, "interval", seconds=21600)  # Restart every 6 hours
    scheduler.start()

    demo.launch()