Update app.py

app.py

@@ -1,16 +1,14 @@
-from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
+from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler, AutoencoderKL
 import torch
 import gradio as gr
 import spaces
 
 
 lora_path = "OedoSoldier/detail-tweaker-lora"
+vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to("cuda")
 
 @spaces.GPU
 def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_scale=7.0,model="Real6.0"):
-    """
-    Generate an image using Stable Diffusion based on the input prompt
-    """
         
     if model == "Real5.0":
         model_id = "SG161222/Realistic_Vision_V5.0_noVAE"
@@ -22,7 +20,7 @@ def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_sca
         model_id = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
         
         
-    pipe = DiffusionPipeline.from_pretrained(model_id).to("cuda")
+    pipe = DiffusionPipeline.from_pretrained(model_id, vae=vae).to("cuda")
     
     if model == "Real6.0":
         pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))
@@ -43,8 +41,8 @@ def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_sca
         cross_attention_kwargs = {"scale":1},
         num_inference_steps = num_inference_steps,
         guidance_scale = guidance_scale,
-        width = 960,
-        height = 960
+        width = 800,
+        height = 800
     ).images[0]
     
     return image