Update app.py

app.py

@@ -61,7 +61,7 @@ checkpoint = "microsoft/Phi-3.5-mini-instruct"
 vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16")
 #vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16")
 
-pipe = StableDiffusion3Pipeline.from_pretrained("ford442/stable-diffusion-3.5-medium-bf16").to(device).to(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"))
 #pipe = StableDiffusion3Pipeline.from_pretrained("ford442/RealVis_Medium_1.0b_bf16", torch_dtype=torch.bfloat16)
 #pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium", token=hftoken, torch_dtype=torch.float32, device_map='balanced')
@@ -75,7 +75,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(device).to(torch.bfloat16)
+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("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()
@@ -131,6 +131,8 @@ def infer(
     latent_file,  # Add latents file input
     progress=gr.Progress(track_tqdm=True),
 ):
+    upscaler_2.to(torch.device('cpu'))
+    pipe.to(torch.device('cuda'))
     torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
@@ -254,6 +256,8 @@ def infer(
     #torch.save(generated_latents, latent_path)
     #upload_to_ftp(latent_path)
     #refiner.scheduler.set_timesteps(num_inference_steps,device)
+    pipe.to(torch.device('cpu'))
+    refiner.to(torch.device('cuda'))
     refine = refiner(
             prompt=f"{enhanced_prompt_2}, high quality masterpiece, complex details",
             negative_prompt = negative_prompt,
@@ -265,9 +269,12 @@ def infer(
     refine_path = f"sd35_refine_{seed}.png"
     refine.save(refine_path,optimize=False,compress_level=0)
     upload_to_ftp(refine_path)
+    refiner.to(torch.device('cpu'))
+    upscaler_2.to(torch.device('cuda'))
     with torch.no_grad():
         upscale2 = upscaler_2(refine, tiling=True, tile_width=256, tile_height=256)
     print('-- got upscaled image --')
+    upscaler_2.to(torch.device('cpu'))
     downscale2 = upscale2.resize((upscale2.width // 4, upscale2.height // 4),Image.LANCZOS)
     upscale_path = f"sd35_upscale_{seed}.png"
     downscale2.save(upscale_path,optimize=False,compress_level=0)