Update app.py

app.py

@@ -4,7 +4,7 @@ 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 optimum_neuron_export import convert, DIFFUSION_PIPELINE_MAPPING
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from apscheduler.schedulers.background import BackgroundScheduler
 
@@ -30,22 +30,16 @@ TRANSFORMER_TASKS = {
     "sentence-similarity": {"color": "#06b6d4", "category": "Similarity"},
 }
 
-# Define diffusion pipeline types
+# Define diffusion pipeline types - updated structure
 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"},
+    "stable-diffusion": {"color": "#ec4899", "category": "Stable Diffusion", "tasks": ["text-to-image", "image-to-image", "inpaint"]},
+    "stable-diffusion-xl": {"color": "#10b981", "category": "Stable Diffusion XL", "tasks": ["text-to-image", "image-to-image", "inpaint"]},
+    "sdxl-turbo": {"color": "#f59e0b", "category": "SDXL Turbo", "tasks": ["text-to-image", "image-to-image", "inpaint"]},
+    "lcm": {"color": "#8b5cf6", "category": "LCM", "tasks": ["text-to-image"]},
+    "pixart-alpha": {"color": "#ef4444", "category": "PixArt", "tasks": ["text-to-image"]},
+    "pixart-sigma": {"color": "#ef4444", "category": "PixArt", "tasks": ["text-to-image"]},
+    "flux": {"color": "#06b6d4", "category": "Flux", "tasks": ["text-to-image", "inpaint"]},
+    "flux-kont": {"color": "#06b6d4", "category": "Flux Kont", "tasks": ["text-to-image", "image-to-image"]},
 }
 
 TAGS = {
@@ -101,22 +95,46 @@ def format_tasks_for_table(tasks_str: str) -> str:
     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"""
+def update_pipeline_and_task_dropdowns(model_type: str):
+    """Update the pipeline and task dropdowns 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
+        return (
+            gr.Dropdown(visible=False),  # pipeline dropdown hidden
+            gr.Dropdown(
+                choices=ALL_TRANSFORMER_TASKS,
+                value="auto",
+                label="Task (auto can infer task from model)",
+                visible=True
+            )
         )
     else:  # diffusers
+        # Show pipeline dropdown, hide task dropdown initially
+        return (
+            gr.Dropdown(
+                choices=ALL_DIFFUSION_PIPELINES,
+                value="stable-diffusion",
+                label="Pipeline Type",
+                visible=True
+            ),
+            gr.Dropdown(
+                choices=DIFFUSION_PIPELINES["stable-diffusion"]["tasks"],
+                value=DIFFUSION_PIPELINES["stable-diffusion"]["tasks"][0],
+                label="Task",
+                visible=True
+            )
+        )
+
+def update_task_dropdown_for_pipeline(pipeline_name: str):
+    """Update task dropdown based on selected pipeline"""
+    if pipeline_name in DIFFUSION_PIPELINES:
+        tasks = DIFFUSION_PIPELINES[pipeline_name]["tasks"]
         return gr.Dropdown(
-            choices=ALL_DIFFUSION_PIPELINES,
-            value="text-to-image",
-            label="Pipeline Type",
+            choices=tasks,
+            value=tasks[0] if tasks else None,
+            label="Task",
             visible=True
         )
+    return gr.Dropdown(visible=False)
 
 def toggle_custom_repo_box(pr_destinations: List[str]):
     """Show or hide the custom repo ID textbox based on checkbox selection."""
@@ -125,7 +143,7 @@ def toggle_custom_repo_box(pr_destinations: List[str]):
     else:
         return gr.Textbox(visible=False, value="")
 
-def neuron_export(model_id: str, model_type: str, task_or_pipeline: str, 
+def neuron_export(model_id: str, model_type: str, pipeline_name: str, task_or_pipeline: str, 
                   pr_destinations: List[str], custom_repo_id: str, custom_cache_repo: str, oauth_token: gr.OAuthToken):
     
     log_buffer = ""
@@ -169,8 +187,11 @@ def neuron_export(model_id: str, model_type: str, task_or_pipeline: str,
         }
 
         # 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):
+        for status_code, message in convert(
+            api, model_id, task_or_pipeline, model_type, 
+            token=oauth_token.token, pr_options=pr_options, 
+            pipeline_name=pipeline_name if model_type == "diffusers" else None
+        ):
             if isinstance(message, str):
                 yield log(message)
             else:  # It's the final result dictionary
@@ -211,7 +232,7 @@ Simply provide a model ID from the Hugging Face Hub, and choose your desired out
 ### ✨ 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.  
+* **🔗 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.  
 
@@ -324,6 +345,12 @@ with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
                     placeholder="Search for a model on the Hub...",
                     search_type="model",
                 )
+                pipeline_dropdown = gr.Dropdown(
+                    choices=ALL_DIFFUSION_PIPELINES,
+                    value="stable-diffusion",
+                    label="Pipeline Type",
+                    visible=False
+                )
                 task_dropdown = gr.Dropdown(
                     choices=ALL_TRANSFORMER_TASKS,
                     value="auto",
@@ -336,8 +363,14 @@ with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
             
             # Event Handlers
             model_type.change(
-                fn=update_task_dropdown,
+                fn=update_pipeline_and_task_dropdowns,
                 inputs=[model_type],
+                outputs=[pipeline_dropdown, task_dropdown]
+            )
+            
+            pipeline_dropdown.change(
+                fn=update_task_dropdown_for_pipeline,
+                inputs=[pipeline_dropdown],
                 outputs=[task_dropdown]
             )
             
@@ -351,7 +384,8 @@ with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
                 fn=neuron_export,
                 inputs=[
                     input_model, 
-                    model_type, 
+                    model_type,
+                    pipeline_dropdown,
                     task_dropdown,
                     pr_destinations_checkbox,
                     custom_repo_id_textbox,