runtime fix

app.py

@@ -33,6 +33,10 @@ def set_seed(seed):
     torch.cuda.manual_seed_all(seed)
     random.seed(seed)
 
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
 
 def remove_prefix(text, prefix):
     if text.startswith(prefix):
@@ -176,9 +180,6 @@ class HandDiffOpts:
     num_workers: int = 10
     n_val_samples: int = 4
 
-if not torch.cuda.is_available():
-    raise ValueError("No GPU")
-
 # load models
 if NEW_MODEL:
     opts = HandDiffOpts()
@@ -202,15 +203,15 @@ if NEW_MODEL:
         latent_dim=opts.latent_dim,
         in_channels=opts.latent_dim+opts.n_keypoints+opts.n_mask,
         learn_sigma=True,
-    ).cuda()
+    ).to(device)
     # ckpt_state_dict = torch.load(model_path)['model_state_dict']
-    ckpt_state_dict = torch.load(model_path, map_location=torch.device('cuda'))['ema_state_dict']
+    ckpt_state_dict = torch.load(model_path, map_location=torch.device(device))['ema_state_dict']
     missing_keys, extra_keys = model.load_state_dict(ckpt_state_dict, strict=False)
     model.eval()
     print(missing_keys, extra_keys)
     assert len(missing_keys) == 0
     vae_state_dict = torch.load(vae_path)['state_dict']
-    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False).cuda()
+    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False).to(device)
     missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
     autoencoder.eval()
     assert len(missing_keys) == 0
@@ -225,18 +226,18 @@ else:
         latent_dim=opts.latent_dim,
         in_channels=opts.latent_dim+opts.n_keypoints+opts.n_mask,
         learn_sigma=True,
-    ).cuda()
+    ).to(device)
     ckpt_state_dict = torch.load(model_path)['state_dict']
     dit_state_dict = {remove_prefix(k, 'diffusion_backbone.'): v for k, v in ckpt_state_dict.items() if k.startswith('diffusion_backbone')}
     vae_state_dict = {remove_prefix(k, 'autoencoder.'): v for k, v in ckpt_state_dict.items() if k.startswith('autoencoder')}
     missing_keys, extra_keys = model.load_state_dict(dit_state_dict, strict=False)
     model.eval()
     assert len(missing_keys) == 0 and len(extra_keys) == 0
-    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False).cuda()
+    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False).to(device)
     missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
     autoencoder.eval()
     assert len(missing_keys) == 0 and len(extra_keys) == 0
-sam_predictor = init_sam(ckpt_path="./sam_vit_h_4b8939.pth")
+sam_predictor = init_sam(ckpt_path="./sam_vit_h_4b8939.pth", device=device)
 
 
 print("Mediapipe hand detector and SAM ready...")
@@ -312,7 +313,7 @@ def get_ref_anno(ref):
         img,
         keypts,
         hand_mask,
-        device="cuda",
+        device=device,
         target_size=(256, 256),
         latent_size=(32, 32),
     ):
@@ -348,7 +349,7 @@ def get_ref_anno(ref):
         img,
         keypts,
         hand_mask,
-        device="cuda",
+        device=device,
         target_size=opts.image_size,
         latent_size=opts.latent_size,
     )
@@ -405,7 +406,7 @@ def get_target_anno(target):
             )
             * kpts_valid[:, None, None],
             dtype=torch.float,
-            device="cuda",
+            device=device,
         )[None, ...]
         target_cond = torch.cat(
             [target_heatmaps, torch.zeros_like(target_heatmaps)[:, :1]], 1
@@ -525,12 +526,12 @@ def sample_diff(ref_cond, target_cond, target_keypts, num_gen, seed, cfg):
     set_seed(seed)
     z = torch.randn(
         (num_gen, opts.latent_dim, opts.latent_size[0], opts.latent_size[1]),
-        device="cuda",
+        device=device,
     )
     target_cond = target_cond.repeat(num_gen, 1, 1, 1)
     ref_cond = ref_cond.repeat(num_gen, 1, 1, 1)
     # novel view synthesis mode = off
-    nvs = torch.zeros(num_gen, dtype=torch.int, device="cuda")
+    nvs = torch.zeros(num_gen, dtype=torch.int, device=device)
     z = torch.cat([z, z], 0)
     model_kwargs = dict(
         target_cond=torch.cat([target_cond, torch.zeros_like(target_cond)]),
@@ -546,7 +547,7 @@ def sample_diff(ref_cond, target_cond, target_keypts, num_gen, seed, cfg):
         clip_denoised=False,
         model_kwargs=model_kwargs,
         progress=True,
-        device="cuda",
+        device=device,
     ).chunk(2)
     sampled_images = autoencoder.decode(samples / opts.latent_scaling_factor)
     sampled_images = torch.clamp(sampled_images, min=-1.0, max=1.0)
@@ -635,14 +636,14 @@ def ready_sample(img_ori, inpaint_mask, keypts):
             inpaint_mask, dsize=opts.latent_size, interpolation=cv2.INTER_NEAREST
         ),
         dtype=torch.float,
-        device="cuda",
+        device=device,
     ).unsqueeze(0)[None, ...]
 
     def make_ref_cond(
         img,
         keypts,
         hand_mask,
-        device="cuda",
+        device=device,
         target_size=(256, 256),
         latent_size=(32, 32),
     ):
@@ -678,7 +679,7 @@ def ready_sample(img_ori, inpaint_mask, keypts):
         img,
         keypts,
         hand_mask * (1 - inpaint_mask),
-        device="cuda",
+        device=device,
         target_size=opts.image_size,
         latent_size=opts.latent_size,
     )
@@ -736,12 +737,12 @@ def sample_inpaint(
     jump_n_sample = quality
     cfg_scale = cfg
     z = torch.randn(
-        (N, opts.latent_dim, opts.latent_size[0], opts.latent_size[1]), device="cuda"
+        (N, opts.latent_dim, opts.latent_size[0], opts.latent_size[1]), device=device
     )
     target_cond_N = target_cond.repeat(N, 1, 1, 1)
     ref_cond_N = ref_cond.repeat(N, 1, 1, 1)
     # novel view synthesis mode = off
-    nvs = torch.zeros(N, dtype=torch.int, device="cuda")
+    nvs = torch.zeros(N, dtype=torch.int, device=device)
     z = torch.cat([z, z], 0)
     model_kwargs = dict(
         target_cond=torch.cat([target_cond_N, torch.zeros_like(target_cond_N)]),
@@ -759,7 +760,7 @@ def sample_inpaint(
         clip_denoised=False,
         model_kwargs=model_kwargs,
         progress=True,
-        device="cuda",
+        device=device,
         jump_length=jump_length,
         jump_n_sample=jump_n_sample,
     ).chunk(2)