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 optimum_neuron_export import convert | |
| from gradio_huggingfacehub_search import HuggingfaceHubSearch | |
| from apscheduler.schedulers.background import BackgroundScheduler | |
| # Define transformer tasks and their categories for coloring | |
| TRANSFORMER_TASKS = { | |
| "auto": {"color": "#6b7280", "category": "Auto"}, | |
| "feature-extraction": {"color": "#3b82f6", "category": "Feature Extraction"}, | |
| "fill-mask": {"color": "#8b5cf6", "category": "NLP"}, | |
| "multiple-choice": {"color": "#8b5cf6", "category": "NLP"}, | |
| "question-answering": {"color": "#8b5cf6", "category": "NLP"}, | |
| "text-classification": {"color": "#8b5cf6", "category": "NLP"}, | |
| "token-classification": {"color": "#8b5cf6", "category": "NLP"}, | |
| "text-generation": {"color": "#10b981", "category": "Text Generation"}, | |
| "text2text-generation": {"color": "#10b981", "category": "Text Generation"}, | |
| "audio-classification": {"color": "#f59e0b", "category": "Audio"}, | |
| "automatic-speech-recognition": {"color": "#f59e0b", "category": "Audio"}, | |
| "audio-frame-classification": {"color": "#f59e0b", "category": "Audio"}, | |
| "audio-xvector": {"color": "#f59e0b", "category": "Audio"}, | |
| "image-classification": {"color": "#ef4444", "category": "Vision"}, | |
| "object-detection": {"color": "#ef4444", "category": "Vision"}, | |
| "semantic-segmentation": {"color": "#ef4444", "category": "Vision"}, | |
| "zero-shot-image-classification": {"color": "#ec4899", "category": "Multimodal"}, | |
| "sentence-similarity": {"color": "#06b6d4", "category": "Similarity"}, | |
| } | |
| # Define diffusion pipeline types | |
| DIFFUSION_PIPELINES = { | |
| "text-to-image": {"color": "#ec4899", "category": "Stable Diffusion"}, | |
| "image-to-image": {"color": "#ec4899", "category": "Stable Diffusion"}, | |
| "inpaint": {"color": "#ec4899", "category": "Stable Diffusion"}, | |
| "instruct-pix2pix": {"color": "#ec4899", "category": "Stable Diffusion"}, | |
| "latent-consistency": {"color": "#8b5cf6", "category": "Latent Consistency"}, | |
| "stable-diffusion": {"color": "#10b981", "category": "Stable Diffusion"}, | |
| "stable-diffusion-xl": {"color": "#10b981", "category": "Stable Diffusion XL"}, | |
| "stable-diffusion-xl-img2img": {"color": "#10b981", "category": "Stable Diffusion XL"}, | |
| "stable-diffusion-xl-inpaint": {"color": "#10b981", "category": "Stable Diffusion XL"}, | |
| "controlnet": {"color": "#f59e0b", "category": "ControlNet"}, | |
| "controlnet-xl": {"color": "#f59e0b", "category": "ControlNet XL"}, | |
| "pixart-alpha": {"color": "#ef4444", "category": "PixArt"}, | |
| "pixart-sigma": {"color": "#ef4444", "category": "PixArt"}, | |
| "flux": {"color": "#06b6d4", "category": "Flux"}, | |
| } | |
| TAGS = { | |
| "Feature Extraction": {"color": "#3b82f6", "category": "Feature Extraction"}, | |
| "NLP": {"color": "#8b5cf6", "category": "NLP"}, | |
| "Text Generation": {"color": "#10b981", "category": "Text Generation"}, | |
| "Audio": {"color": "#f59e0b", "category": "Audio"}, | |
| "Vision": {"color": "#ef4444", "category": "Vision"}, | |
| "Multimodal": {"color": "#ec4899", "category": "Multimodal"}, | |
| "Similarity": {"color": "#06b6d4", "category": "Similarity"}, | |
| "Stable Diffusion": {"color": "#ec4899", "category": "Stable Diffusion"}, | |
| "Stable Diffusion XL": {"color": "#10b981", "category": "Stable Diffusion XL"}, | |
| "ControlNet": {"color": "#f59e0b", "category": "ControlNet"}, | |
| "ControlNet XL": {"color": "#f59e0b", "category": "ControlNet XL"}, | |
| "PixArt": {"color": "#ef4444", "category": "PixArt"}, | |
| "Latent Consistency": {"color": "#8b5cf6", "category": "Latent Consistency"}, | |
| "Flux": {"color": "#06b6d4", "category": "Flux"}, | |
| } | |
| # UPDATED: New choices for the Pull Request destination UI component | |
| DEST_NEW_NEURON_REPO = "Create new Neuron-optimized repository" | |
| DEST_CACHE_REPO = "Create a PR in the cache repository" | |
| DEST_CUSTOM_REPO = "Create a PR in a custom repository" | |
| PR_DESTINATION_CHOICES = [ | |
| DEST_NEW_NEURON_REPO, | |
| DEST_CACHE_REPO, | |
| DEST_CUSTOM_REPO | |
| ] | |
| DEFAULT_CACHE_REPO = "aws-neuron/optimum-neuron-cache" | |
| # Get all tasks and pipelines for dropdowns | |
| ALL_TRANSFORMER_TASKS = list(TRANSFORMER_TASKS.keys()) | |
| ALL_DIFFUSION_PIPELINES = list(DIFFUSION_PIPELINES.keys()) | |
| def create_task_tag(task: str) -> str: | |
| """Create a colored HTML tag for a task""" | |
| if task in TRANSFORMER_TASKS: | |
| color = TRANSFORMER_TASKS[task]["color"] | |
| return f'<span style="background-color: {color}; color: white; padding: 2px 6px; border-radius: 12px; font-size: 0.75rem; font-weight: 500; margin: 1px;">{task}</span>' | |
| elif task in DIFFUSION_PIPELINES: | |
| color = DIFFUSION_PIPELINES[task]["color"] | |
| return f'<span style="background-color: {color}; color: white; padding: 2px 6px; border-radius: 12px; font-size: 0.75rem; font-weight: 500; margin: 1px;">{task}</span>' | |
| elif task in TAGS: | |
| color = TAGS[task]["color"] | |
| return f'<span style="background-color: {color}; color: white; padding: 2px 6px; border-radius: 12px; font-size: 0.75rem; font-weight: 500; margin: 1px;">{task}</span>' | |
| else: | |
| return f'<span style="background-color: #6b7280; color: white; padding: 2px 6px; border-radius: 12px; font-size: 0.75rem; font-weight: 500; margin: 1px;">{task}</span>' | |
| def format_tasks_for_table(tasks_str: str) -> str: | |
| """Convert comma-separated tasks into colored tags""" | |
| tasks = [task.strip() for task in tasks_str.split(',')] | |
| return ' '.join([create_task_tag(task) for task in tasks]) | |
| def update_task_dropdown(model_type: str): | |
| """Update the task dropdown based on selected model type""" | |
| if model_type == "transformers": | |
| return gr.Dropdown( | |
| choices=ALL_TRANSFORMER_TASKS, | |
| value="auto", | |
| label="Task (auto can infer task from model)", | |
| visible=True | |
| ) | |
| else: # diffusers | |
| return gr.Dropdown( | |
| choices=ALL_DIFFUSION_PIPELINES, | |
| value="text-to-image", | |
| label="Pipeline Type", | |
| visible=True | |
| ) | |
| def toggle_custom_repo_box(pr_destinations: List[str]): | |
| """Show or hide the custom repo ID textbox based on checkbox selection.""" | |
| if DEST_CUSTOM_REPO in pr_destinations: | |
| return gr.Textbox(visible=True) | |
| else: | |
| return gr.Textbox(visible=False, value="") | |
| def neuron_export(model_id: str, model_type: str, task_or_pipeline: str, | |
| pr_destinations: List[str], custom_repo_id: str, custom_cache_repo: str, oauth_token: gr.OAuthToken): | |
| log_buffer = "" | |
| def log(msg): | |
| nonlocal log_buffer | |
| # Handle cases where the message from the backend is not a string | |
| if not isinstance(msg, str): | |
| msg = str(msg) | |
| log_buffer += msg + "\n" | |
| return log_buffer | |
| if oauth_token.token is None: | |
| yield log("You must be logged in to use this space") | |
| return | |
| if not model_id: | |
| yield log("🚫 Invalid input. Please specify a model name from the hub.") | |
| return | |
| try: | |
| api = HfApi(token=oauth_token.token) | |
| # Set custom cache repo as environment variable | |
| if custom_cache_repo: | |
| os.environ['CUSTOM_CACHE_REPO'] = custom_cache_repo.strip() | |
| yield log(f"🔑 Logging in ...") | |
| try: | |
| api.model_info(model_id, token=oauth_token.token) | |
| except Exception as e: | |
| yield log(f"❌ Could not access model `{model_id}`: {e}") | |
| return | |
| yield log(f"✅ Model `{model_id}` is accessible. Starting Neuron export...") | |
| # UPDATED: Build pr_options with new structure | |
| pr_options = { | |
| "create_neuron_repo": DEST_NEW_NEURON_REPO in pr_destinations, | |
| "create_cache_pr": DEST_CACHE_REPO in pr_destinations, | |
| "create_custom_pr": DEST_CUSTOM_REPO in pr_destinations, | |
| "custom_repo_id": custom_repo_id.strip() if custom_repo_id else "" | |
| } | |
| # The convert function is a generator, so we iterate through its messages | |
| for status_code, message in convert(api, model_id, task_or_pipeline, model_type, | |
| token=oauth_token.token, pr_options=pr_options): | |
| if isinstance(message, str): | |
| yield log(message) | |
| else: # It's the final result dictionary | |
| final_message = "🎉 Process finished.\n" | |
| if message.get("neuron_repo"): | |
| final_message += f"🏗️ New Neuron Repository: {message['neuron_repo']}\n" | |
| if message.get("readme_pr"): | |
| final_message += f"📝 README PR (Original Model): {message['readme_pr']}\n" | |
| if message.get("cache_pr"): | |
| final_message += f"🔗 Cache PR: {message['cache_pr']}\n" | |
| if message.get("custom_pr"): | |
| final_message += f"🔗 Custom PR: {message['custom_pr']}\n" | |
| yield log(final_message) | |
| except Exception as e: | |
| yield log(f"❗ An unexpected error occurred in the Gradio interface: {e}") | |
| TITLE_IMAGE = """ | |
| <div style="display: block; margin-left: auto; margin-right: auto; width: 50%;"> | |
| <img src="https://huggingface.co/spaces/optimum/neuron-export/resolve/main/huggingfaceXneuron.png"/> | |
| </div> | |
| """ | |
| TITLE = """ | |
| <div style="text-align: center; max-width: 1400px; margin: 0 auto;"> | |
| <h1 style="font-weight: 900; margin-bottom: 10px; margin-top: 10px; font-size: 2.2rem;"> | |
| 🤗 Optimum Neuron Model Exporter 🏎️ | |
| </h1> | |
| </div> | |
| """ | |
| # UPDATED: Description to reflect new workflow | |
| DESCRIPTION = """ | |
| This Space allows you to automatically export 🤗 transformers and 🧨 diffusion models to AWS Neuron-optimized format for Inferentia/Trainium acceleration. | |
| Simply provide a model ID from the Hugging Face Hub, and choose your desired output. | |
| ### ✨ Key Features | |
| * **🚀 Create a New Optimized Repo**: Automatically converts your model and uploads it to a new repository under your username (e.g., `your-username/model-name-neuron`). | |
| * **🔗 Link Back to Original**: Creates a Pull Request on the original model’s repository to add a link to your optimized version, making it easier for the community to discover. | |
| * **🛠️ PR to a Custom Repo**: For custom workflows, you can create a Pull Request to add the optimized files directly into an existing repository you own. | |
| * **📦 Contribute to Cache**: Contribute the generated compilation artifacts to a centralized cache repository (or your own private cache), helping avoid recompilation of already exported models. | |
| ### ⚙️ How to Use | |
| 1. **Model ID**: Enter the ID of the model you want to export (e.g., `bert-base-uncased` or `stabilityai/stable-diffusion-xl-base-1.0`) and choose the corresponding task. | |
| 2. **Export Options**: Select at least one option for where to save the exported model. You can provide your own cache repo ID or use the default (`aws-neuron/optimum-neuron-cache`). | |
| 3. **Convert & Upload**: Click the button and follow the logs to track progress! | |
| """ | |
| CUSTOM_CSS = """ | |
| /* Primary button styling with warm colors */ | |
| button.gradio-button.lg.primary { | |
| /* Changed the blue/green gradient to an orange/yellow one */ | |
| background: linear-gradient(135deg, #F97316, #FBBF24) !important; | |
| color: white !important; | |
| padding: 16px 32px !important; | |
| font-size: 1.1rem !important; | |
| font-weight: 700 !important; | |
| border: none !important; | |
| border-radius: 12px !important; | |
| /* Updated the shadow to match the new orange color */ | |
| box-shadow: 0 0 15px rgba(249, 115, 22, 0.5) !important; | |
| transition: all 0.3s cubic-bezier(0.25, 0.8, 0.25, 1) !important; | |
| position: relative; | |
| overflow: hidden; | |
| } | |
| /* Login button styling with glow effect using dark blue and violet colors */ | |
| #login-button { | |
| background: linear-gradient(135deg, #1a237e, #6a1b9a) !important; /* Dark Blue to Violet */ | |
| color: white !important; | |
| font-weight: 700 !important; | |
| border: none !important; | |
| border-radius: 12px !important; | |
| box-shadow: 0 0 15px rgba(106, 27, 154, 0.6) !important; /* Cool violet glow */ | |
| transition: all 0.3s cubic-bezier(0.25, 0.8, 0.25, 1) !important; | |
| position: relative; | |
| overflow: hidden; | |
| animation: glow 1.5s ease-in-out infinite alternate; | |
| max-width: 350px !important; | |
| margin: 0 auto !important; | |
| } | |
| #login-button::before { | |
| content: "🔑 "; | |
| display: inline-block !important; | |
| vertical-align: middle !important; | |
| margin-right: 5px !important; | |
| line-height: normal !important; | |
| } | |
| #login-button:hover { | |
| transform: translateY(-3px) scale(1.03) !important; | |
| box-shadow: 0 10px 25px rgba(26, 35, 126, 0.7) !important; /* Deeper blue glow */ | |
| } | |
| #login-button::after { | |
| content: ""; | |
| position: absolute; | |
| top: 0; | |
| left: -100%; | |
| width: 100%; | |
| height: 100%; | |
| background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.25), transparent); | |
| transition: 0.5s; | |
| } | |
| #login-button:hover::after { | |
| left: 100%; | |
| } | |
| """ | |
| with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("**You must be logged in to use this space**") | |
| gr.LoginButton(elem_id="login-button", elem_classes="center-button", min_width=250) | |
| gr.HTML(TITLE_IMAGE) | |
| gr.HTML(TITLE) | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Tabs(): | |
| with gr.Tab("Export Model"): | |
| with gr.Group(): | |
| with gr.Row(): | |
| pr_destinations_checkbox = gr.CheckboxGroup( | |
| choices=PR_DESTINATION_CHOICES, | |
| label="Export Destination", | |
| value=[DEST_NEW_NEURON_REPO], | |
| info="Select one or more destinations for the compiled model." | |
| ) | |
| custom_repo_id_textbox = gr.Textbox( | |
| label="Custom Repository ID", | |
| placeholder="e.g., your-username/your-repo-name", | |
| visible=False, | |
| interactive=True | |
| ) | |
| custom_cache_repo_textbox = gr.Textbox( | |
| label="Custom Cache Repository", | |
| placeholder="e.g., your-org/your-cache-repo", | |
| value=DEFAULT_CACHE_REPO, | |
| info=f"Repository to store and fetch from compilation cache artifacts (default: {DEFAULT_CACHE_REPO}) ", | |
| interactive=True | |
| ) | |
| with gr.Row(): | |
| model_type = gr.Radio( | |
| choices=["transformers", "diffusers"], | |
| value="transformers", | |
| label="Model Type", | |
| info="Choose the type of model you want to export" | |
| ) | |
| with gr.Row(): | |
| input_model = HuggingfaceHubSearch( | |
| label="Hub model ID", | |
| placeholder="Search for a model on the Hub...", | |
| search_type="model", | |
| ) | |
| task_dropdown = gr.Dropdown( | |
| choices=ALL_TRANSFORMER_TASKS, | |
| value="auto", | |
| label="Task (auto can infer from model)", | |
| ) | |
| btn = gr.Button("Export to Neuron", size="lg", variant="primary") | |
| log_box = gr.Textbox(label="Logs", lines=20, interactive=False, show_copy_button=True) | |
| # Event Handlers | |
| model_type.change( | |
| fn=update_task_dropdown, | |
| inputs=[model_type], | |
| outputs=[task_dropdown] | |
| ) | |
| pr_destinations_checkbox.change( | |
| fn=toggle_custom_repo_box, | |
| inputs=pr_destinations_checkbox, | |
| outputs=custom_repo_id_textbox | |
| ) | |
| btn.click( | |
| fn=neuron_export, | |
| inputs=[ | |
| input_model, | |
| model_type, | |
| task_dropdown, | |
| pr_destinations_checkbox, | |
| custom_repo_id_textbox, | |
| custom_cache_repo_textbox | |
| ], | |
| outputs=log_box, | |
| ) | |
| with gr.Tab("Supported Architectures"): | |
| gr.HTML(f""" | |
| <div style="margin-bottom: 20px;"> | |
| <h3>🎨 Task Categories Legend</h3> | |
| <div class="task-tags"> | |
| {create_task_tag("Feature Extraction")} | |
| {create_task_tag("NLP")} | |
| {create_task_tag("Text Generation")} | |
| {create_task_tag("Audio")} | |
| {create_task_tag("Vision")} | |
| {create_task_tag("Multimodal")} | |
| {create_task_tag("Similarity")} | |
| </div> | |
| </div> | |
| """) | |
| gr.HTML(f""" | |
| <h2>🤗 Transformers</h2> | |
| <table style="width: 100%; border-collapse: collapse; margin: 20px 0;"> | |
| <colgroup> | |
| <col style="width: 30%;"> | |
| <col style="width: 70%;"> | |
| </colgroup> | |
| <thead> | |
| <tr style="background-color: var(--background-fill-secondary);"> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Architecture</th> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Supported Tasks</th> | |
| </tr> | |
| </thead> | |
| <tbody> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ALBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">AST</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, audio-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">BERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">BLOOM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Beit</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CamemBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CLIP</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ConvBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ConvNext</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ConvNextV2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CvT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DeBERTa (INF2 only)</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DeBERTa-v2 (INF2 only)</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Deit</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DistilBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">DonutSwin</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Dpt</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ELECTRA</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ESM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">FlauBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">GPT2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Hubert</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Levit</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Llama, Llama 2, Llama 3</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Mistral</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Mixtral</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MobileBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MobileNetV2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification, semantic-segmentation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MobileViT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification, semantic-segmentation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ModernBERT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">MPNet</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">OPT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Phi</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">RoBERTa</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">RoFormer</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Swin</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">T5</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text2text-generation")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">UniSpeech</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">UniSpeech-SAT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification, audio-frame-classification, audio-xvector")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">ViT</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, image-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Wav2Vec2</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification, audio-frame-classification, audio-xvector")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">WavLM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, automatic-speech-recognition, audio-classification, audio-frame-classification, audio-xvector")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Whisper</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("automatic-speech-recognition")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">XLM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">XLM-RoBERTa</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, fill-mask, multiple-choice, question-answering, text-classification, token-classification")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Yolos</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, object-detection")}</td></tr> | |
| </tbody> | |
| </table> | |
| <h2>🧨 Diffusers</h2> | |
| <table style="width: 100%; border-collapse: collapse; margin: 20px 0;"> | |
| <colgroup> | |
| <col style="width: 30%;"> | |
| <col style="width: 70%;"> | |
| </colgroup> | |
| <thead> | |
| <tr style="background-color: var(--background-fill-secondary);"> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Architecture</th> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Supported Tasks</th> | |
| </tr> | |
| </thead> | |
| <tbody> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Stable Diffusion</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image, image-to-image, inpaint")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Stable Diffusion XL Base</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image, image-to-image, inpaint")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Stable Diffusion XL Refiner</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("image-to-image, inpaint")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">SDXL Turbo</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image, image-to-image, inpaint")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">LCM</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">PixArt-α</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">PixArt-Σ</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Flux</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("text-to-image")}</td></tr> | |
| </tbody> | |
| </table> | |
| <h2>🤖 Sentence Transformers</h2> | |
| <table style="width: 100%; border-collapse: collapse; margin: 20px 0;"> | |
| <colgroup> | |
| <col style="width: 30%;"> | |
| <col style="width: 70%;"> | |
| </colgroup> | |
| <thead> | |
| <tr style="background-color: var(--background-fill-secondary);"> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Architecture</th> | |
| <th style="border: 1px solid var(--border-color-primary); padding: 12px; text-align: left;">Supported Tasks</th> | |
| </tr> | |
| </thead> | |
| <tbody> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">Transformer</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, sentence-similarity")}</td></tr> | |
| <tr><td style="border: 1px solid var(--border-color-primary); padding: 8px; font-weight: bold;">CLIP</td><td style="border: 1px solid var(--border-color-primary); padding: 8px;" class="task-tags">{format_tasks_for_table("feature-extraction, zero-shot-image-classification")}</td></tr> | |
| </tbody> | |
| </table> | |
| <div style="margin-top: 20px;"> | |
| <p>💡 <strong>Note</strong>: Some architectures may have specific requirements or limitations. DeBERTa models are only supported on INF2 instances.</p> | |
| <p>For more details, check the <a href="https://huggingface.co/docs/optimum-neuron" target="_blank">Optimum Neuron documentation</a>.</p> | |
| </div> | |
| """) | |
| # Add spacing between tabs and content | |
| gr.Markdown("<br><br><br><br>") | |
| if __name__ == "__main__": | |
| demo.launch(debug=True) |