Add euler scheduler. Also, make it default for sd2.1.

nod-ai · Dec 12, 2022 · 8ae76d1 · 8ae76d1
1 parent e5be179
commit 8ae76d1
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/shark/examples/shark_inference/stable_diffusion/main.py b/shark/examples/shark_inference/stable_diffusion/main.py
@@ -6,6 +6,7 @@
     PNDMScheduler,
     DDIMScheduler,
     DPMSolverMultistepScheduler,
+    EulerDiscreteScheduler,
 )
 from tqdm.auto import tqdm
 import numpy as np
@@ -77,6 +78,15 @@ def end_profiling(device):
             subfolder="scheduler",
         )
 
+    if args.version == "v2.1base":
+        tokenizer = CLIPTokenizer.from_pretrained(
+            "stabilityai/stable-diffusion-2-1-base", subfolder="tokenizer"
+        )
+
+        scheduler = EulerDiscreteScheduler.from_pretrained(
+            "stabilityai/stable-diffusion-2-1-base",
+            subfolder="scheduler",
+        )
     start = time.time()
 
     text_input = tokenizer(