neuron-export / app.py
badaoui's picture
badaoui HF Staff
Update app.py
959ee6e verified
raw
history blame
36.1 kB
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)