Update app.py

app.py

@@ -58,8 +58,9 @@ torch_dtype = torch.bfloat16
 
 checkpoint = "microsoft/Phi-3.5-mini-instruct"
 #vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16")
 #vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16")
+#vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16")
+vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False) #, device_map='cpu') #.to(torch.bfloat16) #.to(device=device, dtype=torch.bfloat16)
 
 pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16") #.to(torch.bfloat16)
 #pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16").to(torch.device("cuda:0"))
@@ -75,7 +76,7 @@ pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-me
 #pipe = torch.compile(pipe)
 # pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear")
 
-refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("ford442/stable-diffusion-xl-refiner-1.0-bf16", vae=AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16"), requires_aesthetics_score=True).to(torch.bfloat16)
+refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("ford442/stable-diffusion-xl-refiner-1.0-bf16", vaeXL, requires_aesthetics_score=True) #.to(torch.bfloat16)
 #refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=vae, torch_dtype=torch.float32, requires_aesthetics_score=True, device_map='balanced')
 refiner.scheduler=EulerAncestralDiscreteScheduler.from_config(refiner.scheduler.config)
 #refiner.enable_model_cpu_offload()
@@ -257,7 +258,7 @@ def infer(
     #upload_to_ftp(latent_path)
     #refiner.scheduler.set_timesteps(num_inference_steps,device)
     pipe.to(torch.device('cpu'))
-    refiner.to(torch.device('cuda'))
+    refiner.to(device=device, dtype=torch.bfloat16)
     refine = refiner(
             prompt=f"{enhanced_prompt_2}, high quality masterpiece, complex details",
             negative_prompt = negative_prompt,