Update app.py

app.py

@@ -1,9 +1,22 @@
 import gradio as gr
 import numpy as np
 import random
+from typing import Optional
 
 # import spaces #[uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline
+
+from diffusers import (
+    DPMSolverMultistepScheduler,
+    DPMSolverSinglestepScheduler,
+    KDPM2DiscreteScheduler,
+    KDPM2AncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    EulerAncestralDiscreteScheduler,
+    HeunDiscreteScheduler,
+    LMSDiscreteScheduler,
+)
+
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -29,10 +42,47 @@ def infer(
     height,
     guidance_scale,
     num_inference_steps,
+    scheduler: Optional[str] = None,
     progress=gr.Progress(track_tqdm=True),
 ):
 
     pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
+    match scheduler:
+        case None:
+            pass
+        case "DPMSolverMultistepScheduler":
+            if DPMSolverMultistepScheduler in pipe.scheduler.compatibles:
+                scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+                pipe.scheduler = scheduler
+        case "DPMSolverSinglestepScheduler":
+            if DPMSolverSinglestepScheduler in pipe.scheduler.compatibles:
+                scheduler = DPMSolverSinglestepScheduler.from_config(pipe.scheduler.config)
+                pipe.scheduler = scheduler
+        case "KDPM2DiscreteScheduler":
+            if KDPM2DiscreteScheduler in pipe.scheduler.compatibles:
+                scheduler = KDPM2DiscreteScheduler.from_config(pipe.scheduler.config)
+                pipe.scheduler = scheduler
+        case "KDPM2AncestralDiscreteScheduler":
+            if KDPM2AncestralDiscreteScheduler in pipe.scheduler.compatibles:
+                scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+                pipe.scheduler = scheduler
+        case "EulerDiscreteScheduler":
+            if EulerDiscreteScheduler in pipe.scheduler.compatibles:
+                scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
+                pipe.scheduler = scheduler
+        case "EulerAncestralDiscreteScheduler":
+            if EulerAncestralDiscreteScheduler in pipe.scheduler.compatibles:
+                scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+                pipe.scheduler = scheduler
+        case "HeunDiscreteScheduler":
+            if HeunDiscreteScheduler in pipe.scheduler.compatibles:
+                scheduler = HeunDiscreteScheduler.from_config(pipe.scheduler.config)
+                pipe.scheduler = scheduler
+        case "LMSDiscreteScheduler":
+            if LMSDiscreteScheduler in pipe.scheduler.compatibles:
+                scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
+                pipe.scheduler = scheduler
+    
     pipe = pipe.to(device)
     
     if randomize_seed:
@@ -74,7 +124,6 @@ with gr.Blocks(css=css) as demo:
                 ["stabilityai/sdxl-turbo", "lightx2v/Qwen-Image-Lightning", "tencent/HunyuanImage-2.1", "black-forest-labs/FLUX.1-dev"],
                 label="Image-to-text model",
                 visible=True,
-                value=model_repo_id
             )
 
         with gr.Row():
@@ -91,6 +140,7 @@ with gr.Blocks(css=css) as demo:
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
+
             
             negative_prompt = gr.Text(
                 label="Negative prompt",
@@ -109,6 +159,15 @@ with gr.Blocks(css=css) as demo:
 
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
+            scheduler = gr.Dropdown(
+                [None, "DPMSolverMultistepScheduler", "DPMSolverSinglestepScheduler", 
+    "KDPM2DiscreteScheduler", "KDPM2AncestralDiscreteScheduler",
+    "EulerDiscreteScheduler", "EulerAncestralDiscreteScheduler",
+    "HeunDiscreteScheduler", "LMSDiscreteScheduler",],
+                label="Scheduler",
+                visible=True
+            )
+
             with gr.Row():
                 width = gr.Slider(
                     label="Width",
@@ -148,6 +207,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[
+            model_id,
             prompt,
             negative_prompt,
             seed,
@@ -156,6 +216,7 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            scheduler,
         ],
         outputs=[result, seed],
     )