Update app.py

app.py

@@ -1,4 +1,4 @@
-from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler, AutoencoderKL
+from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler, AutoencoderKL, UNet2DConditionModel
 from transformers import CLIPTextModel, CLIPTokenizer
 import torch
 import gradio as gr
@@ -7,8 +7,6 @@ import spaces
 
 lora1 = "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",num_images=1, width=512, height=512):
         
@@ -20,7 +18,12 @@ def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_sca
         
     else:
         model_id = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
-        
+
+
+    vae = AutoencoderKL.from_pretrained(
+        model_id,
+        subfolder="vae"
+    ).to("cuda")
 
     text_encoder = CLIPTextModel.from_pretrained(
         model_id,
@@ -31,6 +34,11 @@ def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_sca
         model_id,
         subfolder="tokenizer"
     )
+
+    unet = UNet2DConditionModel.from_pretrained(
+        model_id,
+        subfolder="unet"
+    ).to("cuda")
         
     pipe = DiffusionPipeline.from_pretrained(
         model_id,