Update app.py

app.py

@@ -1,11 +1,10 @@
 import csv
 import os
 from datetime import datetime
-from typing import Optional, Union
-
+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
@@ -13,18 +12,16 @@ from apscheduler.schedulers.background import BackgroundScheduler
 DATASET_REPO_URL = "https://huggingface.co/datasets/optimum/neuron-exports"
 DATA_FILENAME = "exports.csv"
 DATA_FILE = os.path.join("data", DATA_FILENAME)
-
-HF_TOKEN = os.environ.get("HF_WRITE_TOKEN")
-
+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)
+# repo = Repository(local_dir=DATADIR, clone_from=DATASET_REPO_URL, token=HF_TOKEN)
 
-# Define all possible tasks and their categories for coloring
-TASK_CATEGORIES = {
+# 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"},
@@ -41,13 +38,28 @@ TASK_CATEGORIES = {
     "image-classification": {"color": "#ef4444", "category": "Vision"},
     "object-detection": {"color": "#ef4444", "category": "Vision"},
     "semantic-segmentation": {"color": "#ef4444", "category": "Vision"},
-    "text-to-image": {"color": "#ec4899", "category": "Multimodal"},
-    "image-to-image": {"color": "#ec4899", "category": "Multimodal"},
-    "inpaint": {"color": "#ec4899", "category": "Multimodal"},
     "zero-shot-image-classification": {"color": "#ec4899", "category": "Multimodal"},
     "sentence-similarity": {"color": "#06b6d4", "category": "Similarity"},
 }
 
+# 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"},
@@ -56,15 +68,37 @@ TAGS = {
     "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"},
 }
 
-# Get all tasks for dropdown
-ALL_TASKS = list(TASK_CATEGORIES.keys())
+# 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 TASK_CATEGORIES:
-        color = TASK_CATEGORIES[task]["color"]
+    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"]
@@ -77,46 +111,91 @@ 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"""
+    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="")
 
-def neuron_export(model_id: str, task: str, oauth_token: gr.OAuthToken) -> str:
-    if oauth_token.token is None:
-        return "You must be logged in to use this space"
-
+# 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:
-        return f"### Invalid input 🐞 Please specify a model name from the hub."
+        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(token=oauth_token.token)
-
-        error, commit_info = convert(api=api, model_id=model_id, task=task, token=oauth_token.token)
-        if error != "0":
-            return commit_info
-
-        print("[commit_info]", commit_info)
-
-         # Save in a private dataset if repo initialized
-        if repo is not None:
-            repo.git_pull(rebase=True)
-            with open(os.path.join(DATADIR, DATA_FILE), "a") as csvfile:
-                writer = csv.DictWriter(
-                    csvfile, fieldnames=["model_id", "pr_url", "time"]
-                )
-                writer.writerow(
-                    {
-                        "model_id": model_id,
-                        "pr_url": commit_info.pr_url,
-                        "time": str(datetime.now()),
-                    }
-                )
-            commit_url = repo.push_to_hub()
-            print("[dataset]", commit_url)
-
-        pr_revision = commit_info.pr_revision.replace("/", "%2F")
-        return f"#### Success 🔥 Yay! This model was successfully exported and a PR was opened using your token: [{commit_info.pr_url}]({commit_info.pr_url}). If you would like to use the exported model without waiting for the PR to be approved, head to https://huggingface.co/{model_id}/tree/{pr_revision}"
+        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:
-        return f"#### Error: {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%;">
@@ -127,205 +206,119 @@ TITLE_IMAGE = """
 TITLE = """
 <div style="text-align: center; max-width: 1400px; margin: 0 auto;">
 <h1 style="font-weight: 900; margin-bottom: 10px; margin-top: 10px; font-size: 2.2rem;">
-    🤗 Optimum Neuron Model Exporter 🏎️ (WIP)
+    🤗 Optimum Neuron Model Compiler 🏎️
 </h1>
 </div>
 """
 
+# UPDATED: Description to reflect new workflow
 DESCRIPTION = """
-This Space allows you to automatically export 🤗 transformers models hosted on the Hugging Face Hub to AWS Neuron-optimized format for Inferentia/Trainium acceleration. It opens a PR on the target model, and it is up to the owner of the original model to merge the PR to allow people to leverage Neuron optimization!
-
-**Features:**
-- Automatically opens PR with Neuron-optimized model
-- Preserves original model weights
-- Adds proper tags to model card
-
-**Requirements:**
-- Model must be compatible with [Optimum Neuron](https://huggingface.co/docs/optimum-neuron)
-- User must be logged in with write token
-"""
-
-# Custom CSS to fix dark mode compatibility and transparency issues
-CUSTOM_CSS = """
-/* Fix for HuggingfaceHubSearch component visibility in both light and dark modes */
-.gradio-container .gr-form {
-    background: var(--background-fill-primary) !important;
-    border: 1px solid var(--border-color-primary) !important;
-}
-
-/* Ensure text is visible in both modes */
-.gradio-container input[type="text"], 
-.gradio-container textarea,
-.gradio-container .gr-textbox input {
-    color: var(--body-text-color) !important;
-    background: var(--input-background-fill) !important;
-    border: 1px solid var(--border-color-primary) !important;
-}
-
-/* Fix dropdown/search results visibility */
-.gradio-container .gr-dropdown,
-.gradio-container .gr-dropdown .gr-box,
-.gradio-container [data-testid="textbox"] {
-    background: var(--background-fill-primary) !important;
-    color: var(--body-text-color) !important;
-    border: 1px solid var(--border-color-primary) !important;
-}
-
-/* Fix for search component specifically */
-.gradio-container .gr-form > div,
-.gradio-container .gr-form input {
-    background: var(--input-background-fill) !important;
-    color: var(--body-text-color) !important;
-}
-
-/* Ensure proper contrast for placeholder text */
-.gradio-container input::placeholder {
-    color: var(--body-text-color-subdued) !important;
-    opacity: 0.7;
-}
+This Space allows you to automatically export 🤗 transformers and diffusion models to AWS Neuron-optimized format for Inferentia/Trainium acceleration. 
 
-/* Fix any remaining transparent backgrounds */
-.gradio-container .gr-box,
-.gradio-container .gr-panel {
-    background: var(--background-fill-primary) !important;
-}
+Simply provide a model ID from the Hugging Face Hub, and choose your desired output.
 
-/* Make sure search results are visible */
-.gradio-container .gr-dropdown-item {
-    color: var(--body-text-color) !important;
-    background: var(--background-fill-primary) !important;
-}
+### ✨ Key Features
 
-.gradio-container .gr-dropdown-item:hover {
-    background: var(--background-fill-secondary) !important;
-}
+* **🚀 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.
 
-/* Task tag styling improvements */
-.task-tags {
-    line-height: 1.8;
-}
+### ⚙️ How to Use
 
-.task-tags span {
-    display: inline-block;
-    margin: 2px;
-}
+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[variant="primary"] {
-    background: linear-gradient(135deg, #3B82F6, #10B981) !important;
+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;
-    box-shadow: 0 0 15px rgba(59, 130, 246, 0.5) !important;
-    transition: all 0.3s cubic-bezier(0.25, 0.8, 0.25, 1) !important;
-    position: relative;
-    overflow: hidden;
-    animation: glow 1.5s ease-in-out infinite alternate;
-}
-button[variant="primary"]::before {
-    content: "✨ ";
-}
-button[variant="primary"]:hover {
-    transform: translateY(-5px) scale(1.05) !important;
-    box-shadow: 0 10px 25px rgba(59, 130, 246, 0.7) !important;
-}
-@keyframes glow {
-    from {
-        box-shadow: 0 0 10px rgba(59, 130, 246, 0.5);
-    }
-    to {
-        box-shadow: 0 0 20px rgba(59, 130, 246, 0.8), 0 0 30px rgba(16, 185, 129, 0.5);
-    }
-}
-
-
-/* Login button styling with glow effect using warm colors */
-#login-button {
-    background: linear-gradient(135deg, #FF7A00, #FFD700) !important; /* Orange to Gold */
-    color: white !important;
-    font-weight: 700 !important;
-    border: none !important;
-    border-radius: 12px !important;
-    box-shadow: 0 0 15px rgba(255, 165, 0, 0.5) !important; /* Warm glow */
+    /* 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;
-    animation: glow 1.5s ease-in-out infinite alternate;
-    max-width: 300px !important;
-    margin: 0 auto !important;
-}
-
-#login-button::before {
-    content: "🔑 ";
-    display: inline-block !important;
-    vertical-align: middle !important;
-    margin-right: 5px !important;
-    line-height: normal !important;
-}
-
-#login-button:hover {
-    transform: translateY(-3px) scale(1.03) !important;
-    box-shadow: 0 10px 25px rgba(255, 140, 0, 0.7) !important; /* Deeper warm glow */
-}
-
-#login-button::after {
-    content: "";
-    position: absolute;
-    top: 0;
-    left: -100%;
-    width: 100%;
-    height: 100%;
-    background: linear-gradient(90deg, transparent, rgba(255, 255, 255, 0.2), transparent);
-    transition: 0.5s;
-}
-
-#login-button:hover::after {
-    left: 100%;
 }
 """
 
-with gr.Blocks(css=CUSTOM_CSS) as demo:
-    # Login requirement notice and button
-    gr.Markdown("**You must be logged in to use this space**")
-    gr.LoginButton(elem_id="login-button", elem_classes="center-button", min_width=250)
-    
-    # Centered title and image
+with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
     gr.HTML(TITLE_IMAGE)
     gr.HTML(TITLE)
-    
-    # Full-width description
     gr.Markdown(DESCRIPTION)
     
     with gr.Tabs():
         with gr.Tab("Export Model"):
-            # Input controls in a row
+            with gr.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 model ID on the hub",
+                    placeholder="Search for a model on the Hub...",
                     search_type="model",
                 )
-                input_task = gr.Dropdown(
-                    choices=ALL_TASKS,
+                task_dropdown = gr.Dropdown(
+                    choices=ALL_TRANSFORMER_TASKS,
                     value="auto",
-                    label='Task (auto could infer task from model)',
+                    label="Task (auto can infer from model)",
                 )
             
-            # Export button below the inputs
             btn = gr.Button("Export to Neuron", size="lg", variant="primary")
             
-            # Output section
-            output = gr.Markdown(label="Output")
-
+            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, input_task],
-                outputs=output,
+                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;">
@@ -404,7 +397,6 @@ with gr.Blocks(css=CUSTOM_CSS) as demo:
                     <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>
@@ -425,9 +417,10 @@ with gr.Blocks(css=CUSTOM_CSS) as demo:
                     <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>
@@ -445,7 +438,6 @@ with gr.Blocks(css=CUSTOM_CSS) as demo:
                     <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>
@@ -456,12 +448,4 @@ with gr.Blocks(css=CUSTOM_CSS) as demo:
     gr.Markdown("<br><br><br><br>")
 
 if __name__ == "__main__":
-    def restart_space():
-        if HF_TOKEN:
-            HfApi().restart_space(repo_id="optimum/neuron-export", token=HF_TOKEN, factory_reboot=True)
-
-    scheduler = BackgroundScheduler()
-    scheduler.add_job(restart_space, "interval", seconds=21600)  # Restart every 6 hours
-    scheduler.start()
-
-    demo.launch()
+    demo.launch(debug=True)