Update app.py

app.py

@@ -14,9 +14,10 @@ from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlowMatch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.float16
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
 repo = "stabilityai/stable-diffusion-3-medium-diffusers"
-pipe = StableDiffusion3Pipeline.from_pretrained(repo, torch_dtype=torch.float16).to(device)
+pipe = StableDiffusion3Pipeline.from_pretrained(repo, torch_dtype=torch.float16, vae=vae).to(device)
 
 help_text = """
 To optimize image results:
@@ -44,7 +45,6 @@ def set_timesteps_patched(self, num_inference_steps: int, device = None):
 # Image Editor
 edit_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl_edit.safetensors")
 EDMEulerScheduler.set_timesteps = set_timesteps_patched
-vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 pipe_edit = StableDiffusionXLInstructPix2PixPipeline.from_single_file(
     edit_file, num_in_channels=8, is_cosxl_edit=True, vae=vae, torch_dtype=torch.float16,
 )