Update app.py

app.py

@@ -16,6 +16,7 @@ from pipeline_flux_control_removal import FluxControlRemovalPipeline
 
 torch.set_grad_enabled(False)
 os.environ['GRADIO_TEMP_DIR'] = './tmp'
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 image_path = mask_path = None
 image_examples = [...] 
 image_path = mask_path =None
@@ -78,7 +79,7 @@ def load_model(base_model_path, lora_path):
         base_model_path,
         transformer=transformer,
         torch_dtype=torch.bfloat16
-    ).to("cuda")
+    ).to(device)
     pipe.transformer.to(torch.bfloat16)
     gr.Info(str(f"Model loading: {int((80 / 100) * 100)}%"))
     gr.Info(str(f"Inject LoRA: {lora_path}"))
@@ -146,7 +147,7 @@ def predict(
         width=H,
         height=W,
         num_inference_steps=ddim_steps,
-        generator=torch.Generator("cuda").manual_seed(seed),
+        generator=torch.Generator(device).manual_seed(seed),
         guidance_scale=scale,
         max_sequence_length=512,
     ).images[0]