Get device automatically

app.py

@@ -52,7 +52,7 @@ is_attention_slicing_enabled = True
 
 # Load model
 dtype = torch.float16
-# device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
 
 model_path = "stabilityai/stable-diffusion-xl-base-1.0"
@@ -63,7 +63,7 @@ pipeline = DiffusionPipeline.from_pretrained(
     variant="fp16",
     use_safetensors=True,
     torch_dtype=dtype,
-).to('cuda')
+).to(device)
 
 if is_attention_slicing_enabled:
     pipeline.enable_attention_slicing()
@@ -75,13 +75,13 @@ if is_cpu_offload_enabled:
 @spaces.GPU
 def remove(gradio_image, rm_guidance_scale=9, num_inference_steps=50, seed=42, strength=0.8):
     try:
-        generator = torch.Generator('cuda').manual_seed(seed)
+        generator = torch.Generator(device).manual_seed(seed)
         prompt = "" # Set prompt to null
 
         source_image_pure = gradio_image["background"]
         mask_image_pure = gradio_image["layers"][0]
-        source_image = preprocess_image(source_image_pure.convert('RGB'), 'cuda')
-        mask = preprocess_mask(mask_image_pure, 'cuda')
+        source_image = preprocess_image(source_image_pure.convert('RGB'), device)
+        mask = preprocess_mask(mask_image_pure, device)
         
         START_STEP = 0 # AAS start step
         END_STEP = int(strength * num_inference_steps) # AAS end step