Automatically check the device and supports running directly on Apple M-series chip devices

app.py

@@ -12,6 +12,12 @@ from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from PIL import Image, ImageDraw
 import numpy as np
 
+device = "cpu"
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():
+    device = "mps"
+
 config_file = hf_hub_download(
     "xinsir/controlnet-union-sdxl-1.0",
     filename="config_promax.json",
@@ -27,11 +33,11 @@ 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="cuda", dtype=torch.float16)
+model.to(device=device, dtype=torch.float16)
 
 vae = AutoencoderKL.from_pretrained(
     "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
-).to("cuda")
+).to(device)
 
 pipe = StableDiffusionXLFillPipeline.from_pretrained(
     "SG161222/RealVisXL_V5.0_Lightning",
@@ -39,7 +45,7 @@ pipe = StableDiffusionXLFillPipeline.from_pretrained(
     vae=vae,
     controlnet=model,
     variant="fp16",
-).to("cuda")
+).to(device)
 
 pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)
 
@@ -185,7 +191,7 @@ def infer(image, width, height, overlap_percentage, num_inference_steps, resize_
         negative_prompt_embeds,
         pooled_prompt_embeds,
         negative_pooled_prompt_embeds,
-    ) = pipe.encode_prompt(final_prompt, "cuda", True)
+    ) = pipe.encode_prompt(final_prompt, device, True)
 
     for image in pipe(
         prompt_embeds=prompt_embeds,