Update app.py

app.py

@@ -159,11 +159,24 @@ def infer_inp(prompt, audio_path, spec_with_mask, progress=gr.Progress(track_tqd
     raw_image = image_add_color(torch_to_pil(norm_spec))
     
     # Add Mask
-    mask = torch.zeros_like(norm_spec)[:1,...]
-    mask[:, :, width_start:width_start+width] = 1
-    mask_image = torch_to_pil(mask)
+    #mask = torch.zeros_like(norm_spec)[:1,...]
+    #mask[:, :, width_start:width_start+width] = 1
+    #mask_image = torch_to_pil(mask)
+
+    # Load the mask image (input from user)
+    mask_pil = spec_with_mask['layers'][0]
+    
+    # Convert to tensor and normalize
+    mask_tensor = transforms.ToTensor()(mask_pil)  # Shape: (1, H, W), values in [0, 1]
+    
+    # Ensure the shape matches expected input (add batch dimension if needed)
+    mask_tensor = mask_tensor[:1, :, :]  # Keep only one channel (grayscale)
+    mask_tensor = mask_tensor.to(device, dtype)  # Send to correct device and dtype
+    
+    # Convert to PIL image if needed for visualization
+    mask_image = torch_to_pil(mask_tensor)
     
-    mask, masked_spec = prepare_mask_and_masked_image(norm_spec, mask)
+    mask, masked_spec = prepare_mask_and_masked_image(norm_spec, mask_tensor)
     masked_spec_image = torch_to_pil(masked_spec)
 
     # color masked spec and paint masked area to black