Update app.py

app.py

@@ -7,6 +7,8 @@ is_shared_ui = True if "fffiloni/sdxl-control-loras" in os.environ['SPACE_ID'] e
 hf_token = os.environ.get("HF_TOKEN")
 login(token=hf_token)
 
+device="cuda" if torch.cuda.is_available() else "cpu"
+
 from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
 from diffusers.utils import load_image
 from PIL import Image
@@ -30,7 +32,7 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
     use_safetensors=True
 )
 
-pipe.to("cuda")
+pipe.to(device)
 
 
 
@@ -60,7 +62,7 @@ def resize_image(input_path, output_path, target_height):
 def infer(use_custom_model, model_name, custom_lora_weight, image_in, prompt, negative_prompt, preprocessor, controlnet_conditioning_scale, guidance_scale, inf_steps, seed, progress=gr.Progress(track_tqdm=True)):
     prompt = prompt
     negative_prompt = negative_prompt
-    generator = torch.Generator(device="cuda").manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
     if image_in == None:
         raise gr.Error("You forgot to upload a source image.")