fix more

app.py

@@ -4,7 +4,7 @@ import gradio as gr
 from datasets import load_dataset
 from PIL import Image
 
-from model import get_sd_small, get_sd_tiny
+from model import get_sd_small, get_sd_tiny, get_sd_every
 from trans_google import google_translator
 
 from i18n import i18nTranslator
@@ -43,15 +43,18 @@ samplers = [
 rand = random.Random()
 translator = google_translator()
 
-pipe = get_sd_tiny()
-other_pipe = get_sd_small()
+tiny_pipe = get_sd_tiny()
+small_pipe = get_sd_small()
+every_pipe = get_sd_every()
 
 
 def get_pipe(width: int, height: int):
     if width == 512 and height == 512:
-        return pipe
+        return tiny_pipe
+    elif width == 256 and height == 256:
+        return small_pipe
     else:
-        return other_pipe
+        return every_pipe
 
 
 def infer(prompt: str, negative: str, width: int, height: int, sampler: str, steps: int, seed: int, scale):

model.py

@@ -28,10 +28,17 @@ def get_sd_21():
     return pipe
 
 
-def get_sd_small():
+def get_sd_every():
     model_id = 'OFA-Sys/small-stable-diffusion-v0'
     scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")
-    pipe = OVStableDiffusionPipeline.from_pretrained("OFA-Sys/small-stable-diffusion-v0", compile=False)
+    pipe = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-2-1-quantized", compile=False)
+    pipe.compile()
+    return pipe
+
+
+def get_sd_small():
+    pipe = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-2-1-quantized", compile=False)
+    pipe.reshape(batch_size=1, height=256, width=256, num_images_per_prompt=1)
     pipe.compile()
     return pipe