From 8ae76d18b5f5a1f3b2f72ba73a349ce87fe0d5f5 Mon Sep 17 00:00:00 2001
From: Prashant Kumar <prashant@nod-labs.com>
Date: Mon, 12 Dec 2022 15:54:59 +0000
Subject: [PATCH] Add euler scheduler. Also, make it default for sd2.1.

---
 .../examples/shark_inference/stable_diffusion/main.py  | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/shark/examples/shark_inference/stable_diffusion/main.py b/shark/examples/shark_inference/stable_diffusion/main.py
index a18f1105a1..d380d70820 100644
--- 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(