Update app.py

app.py

@@ -62,6 +62,10 @@ def generate(
 
         if not use_img2img:
             pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
+
+            if use_vae:  
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
             
         if use_img2img:
             pipe = AutoPipelineForImage2Image.from_pretrained(model, torch_dtype=torch.float16)
@@ -70,6 +74,10 @@ def generate(
             init_image = Image.open(BytesIO(response.content)).convert("RGB")
             init_image = init_image.resize((width, height))
 
+            if use_vae:  
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+                pipe = AutoPipelineForImage2Image.from_pretrained(model, torch_dtype=torch.float16)
+
         if use_lora:
             pipe.load_lora_weights(lora)
             pipe.fuse_lora(lora_scale)