:tada: init

app.py

@@ -9,26 +9,34 @@ from diffusers.schedulers import (
     DDIMScheduler,
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
-    FlowMatchEulerDiscreteScheduler,
+    PNDMScheduler,
+    EulerDiscreteScheduler,
 )
 from diffusers.utils.torch_utils import randn_tensor
 
+MODEL_IDS = {
+    'Stable Diffusion v1.4': 'CompVis/stable-diffusion-v1-4',
+    'Stable Diffusion v3 medium': 'stabilityai/stable-diffusion-3-medium-diffusers',
+}
+
 SCHEDULERS = {
     "DDPMScheduler": DDPMScheduler,
     "DDIMScheduler": DDIMScheduler,
+    "PNDMScheduler": PNDMScheduler,
+    "EulerDiscreteScheduler": EulerDiscreteScheduler,
     "EulerAncestralDiscreteScheduler": EulerAncestralDiscreteScheduler,
     "DPMSolverMultistepScheduler": DPMSolverMultistepScheduler,
-    "FlowMatchEulerDiscreteScheduler": FlowMatchEulerDiscreteScheduler,
 }
 
 
 def inference(
         image_pil: Image.Image,
+        model_name: str,
         scheduler_name: str,
         per_step_time: int = 1,
         n_total_steps: int = 1000,
 ):
-    scheduler = SCHEDULERS[scheduler_name]()
+    scheduler = SCHEDULERS[scheduler_name].from_pretrained(MODEL_IDS[model_name], subfolder='scheduler')
     scheduler.set_timesteps(num_inference_steps=n_total_steps)
     timesteps = torch.flip(scheduler.timesteps, dims=[0])
 
@@ -40,6 +48,8 @@ def inference(
     noise = randn_tensor(image_tensor.shape, generator)
 
     for i, t in enumerate(timesteps):
+        t = torch.tensor([t])
+
         noised_image_tensor = scheduler.add_noise(image_tensor, noise, timesteps=t)
         noised_image_pil = image_processor.postprocess(noised_image_tensor)[0]
         time.sleep(per_step_time)
@@ -61,6 +71,7 @@ if __name__ == '__main__':
         fn=inference,
         inputs=[
             gr.Image(type='pil', label='Input Image'),
+            gr.Dropdown(list(MODEL_IDS.keys()), value='Stable Diffusion v1.4', label='Model ID'),
             gr.Dropdown(list(SCHEDULERS.keys()), value='DDPMScheduler', label='Scheduler'),
             gr.Radio(choices=[0, 1, 2], value=0, label='Per-Step time'),
             gr.Radio(choices=[10, 25, 50, 100, 1000], value=50, label='Total Steps'),
@@ -79,7 +90,7 @@ if __name__ == '__main__':
         """,
         cache_examples=True,
         examples=[
-            [Image.open("assets/corgi.png"), 'DDIMScheduler', 0, 50],
+            [Image.open("assets/corgi.png"), 'Stable Diffusion v1.4', 'DDIMScheduler', 0, 50],
         ],
     )
     demo.launch()