Update app.py

app.py

@@ -421,7 +421,6 @@ def import_state(state, json_text):
 
 ### Main worker
 
-
 def register(state, drawpad, model):
     seed_everything(state.seed if state.seed >=0 else np.random.randint(2147483647))
     print('Generate!')
@@ -436,13 +435,13 @@ def register(state, drawpad, model):
     print('Inpainting mode: ', inpainting_mode)
 
     user_input = np.asarray(drawpad['layers'][0]) # (H, W, 4)
-    foreground_mask = torch.tensor(user_input[..., -1])[None, None] # (1, 1, H, W)
-    user_input = torch.tensor(user_input[..., :-1]) # (H, W, 3)
+    foreground_mask = torch.tensor(user_input[..., -1], device=model.device)[None, None] # (1, 1, H, W)
+    user_input = torch.tensor(user_input[..., :-1], device=model.device) # (H, W, 3)
 
     palette = torch.tensor([
         tuple(int(s[i+1:i+3], 16) for i in (0, 2, 4))
         for s in opt.colors[1:]
-    ]) # (N, 3)
+    ], device=model.device) # (N, 3)
     masks = (palette[:, None, None, :] == user_input[None]).all(dim=-1)[:, None, ...] # (N, 1, H, W)
     # has_masks = [i for i, m in enumerate(masks.sum(dim=(1, 2, 3)) == 0) if not m]
     has_masks = list(range(opt.max_palettes))
@@ -542,13 +541,13 @@ def draw(state, drawpad):
     # conn = Client(opt.address, authkey=opt.authkey)
 
     user_input = np.asarray(drawpad['layers'][0]) # (H, W, 4)
-    foreground_mask = torch.tensor(user_input[..., -1])[None, None] # (1, 1, H, W)
-    user_input = torch.tensor(user_input[..., :-1]) # (H, W, 3)
+    foreground_mask = torch.tensor(user_input[..., -1], device=model.device)[None, None] # (1, 1, H, W)
+    user_input = torch.tensor(user_input[..., :-1], device=model.device) # (H, W, 3)
 
     palette = torch.tensor([
         tuple(int(s[i+1:i+3], 16) for i in (0, 2, 4))
         for s in opt.colors[1:]
-    ]) # (N, 3)
+    ], device=model.device) # (N, 3)
     masks = (palette[:, None, None, :] == user_input[None]).all(dim=-1)[:, None, ...] # (N, 1, H, W)
     # has_masks = [i for i, m in enumerate(masks.sum(dim=(1, 2, 3)) == 0) if not m]
     has_masks = list(range(opt.max_palettes))