Update app.py

app.py

@@ -10,7 +10,7 @@ import numpy as np
 import PIL.Image
 import spaces
 import torch
-from diffusers import AutoencoderKL, DiffusionPipeline
+from diffusers import StableDiffusionKDiffusionPipeline, AutoencoderKL, DiffusionPipeline
 
 DESCRIPTION = "# SDXL"
 if not torch.cuda.is_available():
@@ -55,8 +55,8 @@ def generate(
     lora = 'amazonaws-la/juliette',
 ) -> PIL.Image.Image:
     if torch.cuda.is_available():
-        pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
-
+        pipe = StableDiffusionKDiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
+        pipe.set_scheduler('sample_dpmpp_2m_sde')
         if use_lora:
             pipe.load_lora_weights(lora)
             pipe.fuse_lora(lora_scale=0.7)
@@ -90,32 +90,9 @@ def generate(
             guidance_scale=guidance_scale_base,
             num_inference_steps=num_inference_steps_base,
             generator=generator,
+            use_karras_sigmas=True,
             output_type="pil",
         ).images[0]
-    else:
-        latents = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            prompt_2=prompt_2,
-            negative_prompt_2=negative_prompt_2,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale_base,
-            num_inference_steps=num_inference_steps_base,
-            generator=generator,
-            output_type="latent",
-        ).images
-        image = refiner(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            prompt_2=prompt_2,
-            negative_prompt_2=negative_prompt_2,
-            guidance_scale=guidance_scale_refiner,
-            num_inference_steps=num_inference_steps_refiner,
-            image=latents,
-            generator=generator,
-        ).images[0]
-        return image
 
 
 examples = [