Spaces:
Paused
Paused
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) |