Update app.py

app.py

@@ -423,7 +423,11 @@ def register(state, drawpad, model):
     seed_everything(state.seed if state.seed >=0 else np.random.randint(2147483647))
     print('Generate!')
 
-    background = drawpad['background'].convert('RGBA')
+    background = drawpad['background']
+    if background is None:
+        background = Image.new(size=(opt.width, opt.height), mode='RGB', color=(255, 255, 255))
+    else:
+        background = background.convert('RGBA')
     inpainting_mode = np.asarray(background).sum() != 0
     if not inpainting_mode:
         background = Image.new(size=(opt.width, opt.height), mode='RGB', color=(255, 255, 255))
@@ -432,9 +436,14 @@ def register(state, drawpad, model):
         background_prompt = None
     print('Inpainting mode: ', inpainting_mode)
 
-    user_input = np.asarray(drawpad['layers'][0]) # (H, W, 4)
-    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)
+    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)
+    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)
 
     palette = torch.tensor([
         tuple(int(s[i+1:i+3], 16) for i in (0, 2, 4))