change device

app.py

@@ -12,6 +12,7 @@ from src.utils import (
 )
 from diffusers import StableDiffusionXLPipeline
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def get_model_param_summary(model, verbose=False):
     params_dict = dict()
@@ -29,7 +30,6 @@ def get_model_param_summary(model, verbose=False):
 @dataclass
 class GradioArgs:
     ckpt: str = "./mask/ff.pt"
-    device: str = "cuda:0"
     seed: list = None
     prompt: str = None
     mix_precision: str = "bf16"
@@ -95,9 +95,8 @@ def binary_mask_eval(args):
     # load sdxl model
     pipe = StableDiffusionXLPipeline.from_pretrained(
         "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
-    ).to(args.device)
+    ).to(device)
 
-    device = args.device
     torch_dtype = torch.bfloat16 if args.mix_precision == "bf16" else torch.float32
     mask_pipe, hookers = create_pipeline(
         pipe,
@@ -132,7 +131,7 @@ def binary_mask_eval(args):
     # reload the original model
     pipe = StableDiffusionXLPipeline.from_pretrained(
         "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.bfloat16
-    ).to(args.device)
+    ).to(device)
 
     # get model param summary
     print(f"original model param: {get_model_param_summary(pipe.unet)['overall']}")
@@ -143,9 +142,9 @@ def binary_mask_eval(args):
 @spaces.GPU
 def generate_images(prompt, seed, steps, pipe, pruned_pipe):
     # Run the model and return images directly
-    g_cpu = torch.Generator("cuda:0").manual_seed(seed)
+    g_cpu = torch.Generator(device).manual_seed(seed)
     original_image = pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
-    g_cpu = torch.Generator("cuda:0").manual_seed(seed)
+    g_cpu = torch.Generator(device).manual_seed(seed)
     ecodiff_image = pruned_pipe(prompt=prompt, generator=g_cpu, num_inference_steps=steps).images[0]
     return original_image, ecodiff_image