Update app.py

app.py

@@ -27,14 +27,14 @@ model_file = hf_hub_download(
     filename="diffusion_pytorch_model_promax.safetensors",
 )
 state_dict = load_state_dict(model_file)
-model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
+model = ControlNetModel_Union._load_pretrained_model(
     controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
 )
-model.to(device="cuda", dtype=torch.float16)
+model.to(device="cpu", dtype=torch.float16)
 
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
-).to("cuda")
+).to("cpu")
 
 pipe = StableDiffusionXLFillPipeline.from_pretrained(
     "SG161222/RealVisXL_V5.0_Lightning",
@@ -42,7 +42,7 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
     vae=vae,
     controlnet=model,
     variant="fp16",
-).to("cuda")
+).to("cpu")
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
@@ -91,7 +91,7 @@ def infer(image, model_selection, width, height, overlap_width, num_inference_st
         negative_prompt_embeds,
         pooled_prompt_embeds,
         negative_pooled_prompt_embeds,
-    ) = pipe.encode_prompt(final_prompt, "cuda", True)
+    ) = pipe.encode_prompt(final_prompt, "cpu", True)
 
     for image in pipe(
         prompt_embeds=prompt_embeds,