added cuda/cpu useabiltiy

app.py

@@ -28,19 +28,20 @@ state_dict = load_state_dict(model_file)
 model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
     controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
 )
-model.to(device="cpu", dtype=torch.float16)
+device = torch.cuda.is_available()
 
+model.to(device=device, dtype=torch.float16 if device == 'cuda' else torch.float32)
 vae = AutoencoderKL.from_pretrained(
-    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
-).to("cpu")
+    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16 if device == 'cuda' else torch.float32
+).to(device)
 
 pipe = StableDiffusionXLFillPipeline.from_pretrained(
     "SG161222/RealVisXL_V5.0_Lightning",
-    torch_dtype=torch.float16,
+    torch_dtype=torch.float16 if device == 'cuda' else torch.float32,
     vae=vae,
     controlnet=model,
-    variant="fp16",
-).to("cpu")
+    variant="fp16" if device == 'cuda' else None,
+).to(device)
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
@@ -50,7 +51,7 @@ prompt = "high quality"
     negative_prompt_embeds,
     pooled_prompt_embeds,
     negative_pooled_prompt_embeds,
-) = pipe.encode_prompt(prompt, "cpu", True)
+) = pipe.encode_prompt(prompt, device, True)
 
 
 def fill_image(image, model_selection):