fix more

app.py

@@ -44,6 +44,14 @@ rand = random.Random()
 translator = google_translator()
 
 pipe = get_sd_tiny()
+other_pipe = get_sd_small()
+
+
+def get_pipe(width: int, height: int):
+    if width == 512 and height == 512:
+        return pipe
+    else:
+        return other_pipe
 
 
 def infer(prompt: str, negative: str, width: int, height: int, sampler: str, steps: int, seed: int, scale):
@@ -54,37 +62,39 @@ def infer(prompt: str, negative: str, width: int, height: int, sampler: str, ste
     else:
         seed = int(seed)
 
+    pipeline = get_pipe(width, height)
+
     images = []
     if torch.cuda.is_available():
         generator = torch.Generator(device="cuda").manual_seed(seed)
     else:
         generator = None
     if sampler == "EulerDiscrete":
-        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = EulerDiscreteScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "EulerAncestralDiscrete":
-        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "KDPM2Discrete":
-        pipe.scheduler = KDPM2DiscreteScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = KDPM2DiscreteScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "KDPM2AncestralDiscrete":
-        pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "UniPCMultistep":
-        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = UniPCMultistepScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "DPMSolverSinglestep":
-        pipe.scheduler = DPMSolverSinglestepScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = DPMSolverSinglestepScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "DPMSolverMultistep":
-        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "HeunDiscrete":
-        pipe.scheduler = HeunDiscreteScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = HeunDiscreteScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "DEISMultistep":
-        pipe.scheduler = DEISMultistepScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = DEISMultistepScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "PNDM":
-        pipe.scheduler = PNDMScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = PNDMScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "DDPM":
-        pipe.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = DDPMScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "DDIM":
-        pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
     elif sampler == "LMSDiscrete":
-        pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
+        pipeline.scheduler = LMSDiscreteScheduler.from_config(pipeline.scheduler.config)
 
     try:
         translate_prompt = translator.translate(prompt, lang_tgt='en')
@@ -94,13 +104,13 @@ def infer(prompt: str, negative: str, width: int, height: int, sampler: str, ste
         translate_prompt = prompt
         translate_negative = negative
 
-    image = pipe(prompt=translate_prompt,
-                 negative_prompt=translate_negative,
-                 guidance_scale=scale,
-                 num_inference_steps=steps,
-                 generator=generator,
-                 height=height,
-                 width=width).images[0]
+    image = pipeline(prompt=translate_prompt,
+                     negative_prompt=translate_negative,
+                     guidance_scale=scale,
+                     num_inference_steps=steps,
+                     generator=generator,
+                     height=height,
+                     width=width).images[0]
 
     buffered = BytesIO()
     image.save(buffered, format="JPEG")
@@ -133,7 +143,7 @@ css = """
             padding-top: 1.5rem;
         }
         #prompt-column {
-            min-height: 500px
+            min-height: 520px
         }
         #gallery {
             min-height: 22rem;
@@ -360,7 +370,10 @@ with block:
               <p style="margin-bottom: 10px; font-size: 94%; line-height: 23px;">
                 Stable Diffusion 2.1 Demo App. <br />
                 Click "Generate image" Button to generate image. <br />
-                Also Change params to have a try
+                Also Change params to have a try <br />
+                512*512 is optimized, every image will cost 30s. <br />
+                other size may cost more time. <br />
+                It's just a simplified demo, you can use more advanced features optimize image quality <br />
               </p>
             </div>
         """

model.py

@@ -46,4 +46,4 @@ def get_sd_tiny():
     pipe = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-2-1-quantized", compile=False)
     pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
     pipe.compile()
-    return pipe
+    return pipe

requirements.txt

@@ -1,6 +1,6 @@
 python-dotenv
 diffusers
-transformers
+transformers<5
 accelerate
 scipy
 safetensors