Update app.py

app.py

@@ -438,12 +438,12 @@ def register(state, drawpad, model):
 
     if drawpad['composite'] is None:
         user_input = np.zeros((opt.height, opt.width, 4))
-        foreground_mask = torch.zeros((1, 1, opt.height, opt.width))
-        user_input = torch.tensor(user_input[..., :-1]) # (H, W, 3)
+        foreground_mask = torch.zeros((1, 1, opt.height, opt.width), device=model.device)
+        user_input = torch.tensor(user_input[..., :-1], device=model.device) # (H, W, 3)
     else:
         user_input = np.asarray(drawpad['composite']) # (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))
@@ -599,7 +599,7 @@ def draw(state, drawpad):
 
 css = f"""
 body {{
-    background: unset;
+    background: rgb(11, 15, 25) !important
 }}
 
 #run-button {{