Merge pull request #63 from chenxwh/more-schedulers

add back schedulers for v2.1

Author: Jesse Andrews

README.md

@@ -1,8 +1,8 @@
 # Stable Diffusion v2 Cog model
 
-[![Replicate](https://replicate.com/cjwbw/stable-diffusion-v2/badge)](https://replicate.com/cjwbw/stable-diffusion-v2) 
+[![Replicate](https://replicate.com/stability-ai/stable-diffusion/badge)](https://replicate.com/stability-ai/stable-diffusion) 
 
-This is an implementation of the [Diffusers Stable Diffusion v2](https://huggingface.co/stabilityai/stable-diffusion-2) as a Cog model. [Cog packages machine learning models as standard containers.](https://github.com/replicate/cog)
+This is an implementation of the [Diffusers Stable Diffusion v2.1](https://huggingface.co/stabilityai/stable-diffusion-2-1) as a Cog model. [Cog packages machine learning models as standard containers.](https://github.com/replicate/cog)
 
 First, download the pre-trained weights:
 

predict.py

@@ -4,10 +4,13 @@
 import torch
 from cog import BasePredictor, Input, Path
 from diffusers import (
+    StableDiffusionPipeline,
+    PNDMScheduler,
+    LMSDiscreteScheduler,
     DDIMScheduler,
-    DPMSolverMultistepScheduler,
     EulerDiscreteScheduler,
-    StableDiffusionPipeline,
+    EulerAncestralDiscreteScheduler,
+    DPMSolverMultistepScheduler,
 )
 
 MODEL_ID = "stabilityai/stable-diffusion-2-1"
@@ -62,13 +65,16 @@ def predict(
             description="Scale for classifier-free guidance", ge=1, le=20, default=7.5
         ),
         scheduler: str = Input(
-            default="K_EULER",
+            default="DPMSolverMultistep",
             choices=[
                 "DDIM",
                 "K_EULER",
                 "DPMSolverMultistep",
+                "K_EULER_ANCESTRAL",
+                "PNDM",
+                "KLMS",
             ],
-            description="Choose a scheduler",
+            description="Choose a scheduler.",
         ),
         seed: int = Input(
             description="Random seed. Leave blank to randomize the seed", default=None
@@ -110,7 +116,10 @@ def predict(
 
 def make_scheduler(name, config):
     return {
+        "PNDM": PNDMScheduler.from_config(config),
+        "KLMS": LMSDiscreteScheduler.from_config(config),
         "DDIM": DDIMScheduler.from_config(config),
         "K_EULER": EulerDiscreteScheduler.from_config(config),
+        "K_EULER_ANCESTRAL": EulerAncestralDiscreteScheduler.from_config(config),
         "DPMSolverMultistep": DPMSolverMultistepScheduler.from_config(config),
     }[name]