Update app.py

app.py

@@ -12,7 +12,7 @@ import spaces
 import torch
 from PIL import Image
 from io import BytesIO
-from diffusers import AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, DPMSolverMultistepScheduler
+from diffusers import AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, DPMSolverMultistepScheduler, StableDiffusionKDiffusionPipeline
 
 DESCRIPTION = "# Run any LoRA or SD Model"
 if not torch.cuda.is_available():
@@ -61,24 +61,20 @@ def generate(
         if not use_img2img:
             pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
             pipe.enable_model_cpu_offload()
-            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
             
             if use_vae:  
                 vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
                 pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
                 pipe.enable_model_cpu_offload()
-                pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 
         if use_img2img:
             pipe = AutoPipelineForImage2Image.from_pretrained(model, torch_dtype=torch.float16)
             pipe.enable_model_cpu_offload()
-            pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
             
             if use_vae:  
                 vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
                 pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
                 pipe.enable_model_cpu_offload()
-                pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 
             response = requests.get(url)
             init_image = Image.open(BytesIO(response.content)).convert("RGB")
@@ -101,6 +97,7 @@ def generate(
         negative_prompt_2 = None  # type: ignore
 
     if not use_img2img:
+        pipe.set_scheduler('sample_dpmpp_2m')
         return pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -111,9 +108,11 @@ 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:
+        pipe.set_scheduler('sample_dpmpp_2m')
         images = pipe(
             prompt=prompt,
             image=init_image,
@@ -126,6 +125,7 @@ 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]
         return images