Update src/streamlit_app.py

src/streamlit_app.py

@@ -24,7 +24,6 @@ def load_rmbg_model():
 @st.cache_resource
 def load_birefnet_model():
     """Option 2: The Heavyweight Generalist"""
-    # This requires 'timm' installed
     model = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet", trust_remote_code=True)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
@@ -67,33 +66,21 @@ def find_mask_tensor(output):
 
 def generate_trimap(mask_tensor, erode_kernel_size=10, dilate_kernel_size=10):
     """
-    Generates a trimap (Foreground, Background, Unknown) from a binary mask 
-    using Pure PyTorch (No OpenCV required).
+    Generates a trimap (Foreground, Background, Unknown) from a binary mask.
     Values: 1=FG, 0=BG, 0.5=Unknown (Edge)
     """
-    # Ensure mask is Bx1xHxW
     if mask_tensor.dim() == 3: mask_tensor = mask_tensor.unsqueeze(0)
     
-    # Create kernels
     erode_k = erode_kernel_size
     dilate_k = dilate_kernel_size
     
-    # Dilation (Max Pooling) - Expands the white area
-    # We pad to keep size same
+    # Dilation (Max Pooling)
     dilated = F.max_pool2d(mask_tensor, kernel_size=dilate_k, stride=1, padding=dilate_k//2)
     
-    # Erosion (Negative Max Pooling) - Shrinks the white area
+    # Erosion (Negative Max Pooling)
     eroded = -F.max_pool2d(-mask_tensor, kernel_size=erode_k, stride=1, padding=erode_k//2)
     
-    # Trimap construction
-    # Pixels that are 1 in eroded are definitely FG (1.0)
-    # Pixels that are 0 in dilated are definitely BG (0.0)
-    # Everything else is the "Unknown" zone (0.5)
-    
-    # Start with Unknown (0.5)
     trimap = torch.full_like(mask_tensor, 0.5)
-    
-    # Set definites
     trimap[eroded > 0.5] = 1.0
     trimap[dilated < 0.5] = 0.0
     
@@ -120,11 +107,9 @@ def inference_segmentation(model, image, device, resolution=1024):
     if not isinstance(result_tensor, torch.Tensor):
          if isinstance(result_tensor, (list, tuple)): result_tensor = result_tensor[0]
 
-    # Get binary-ish mask (logits or sigmoid)
     pred = result_tensor.squeeze().cpu()
     if pred.max() > 1 or pred.min() < 0: pred = pred.sigmoid()
     
-    # Resize back to original
     pred_pil = transforms.ToPILImage()(pred)
     mask = pred_pil.resize((w, h), resample=Image.LANCZOS)
     return mask
@@ -134,36 +119,36 @@ def inference_vitmatte(image, device):
     Runs pipeline: RMBG (Rough Mask) -> Trimap -> VitMatte (Refined Mask)
     """
     # 1. Get Rough Mask using RMBG (Fast)
-    rmbg_model, _ = load_rmbg_model() # Re-use loaded model
+    rmbg_model, _ = load_rmbg_model() 
     rough_mask_pil = inference_segmentation(rmbg_model, image, device, resolution=1024)
     
-    # 2. Create Trimap
-    # Convert PIL mask to Tensor
+    # 2. Create Trimap (Tensor)
     mask_tensor = transforms.ToTensor()(rough_mask_pil).to(device)
-    # Generate trimap (1=FG, 0=BG, 0.5=Unknown)
     trimap_tensor = generate_trimap(mask_tensor, erode_kernel_size=25, dilate_kernel_size=25)
     
-    # 3. VitMatte Inference
+    # --- FIX START ---
+    # 3. Convert Trimap Tensor to PIL Image
+    # VitMatte Processor crashes on raw tensors. It wants a PIL Image.
+    # We take the tensor (0.0 to 1.0), move to CPU, and convert to PIL (0 to 255)
+    trimap_pil = transforms.ToPILImage()(trimap_tensor.squeeze().cpu())
+    
+    # 4. VitMatte Inference
     processor, model, _ = load_vitmatte_model()
     
-    # VitMatte expects inputs: pixel_values (image) and mask_labels (trimap)
-    inputs = processor(images=image, trimaps=trimap_tensor, return_tensors="pt").to(device)
+    # Pass PIL images for both
+    inputs = processor(images=image, trimaps=trimap_pil, return_tensors="pt").to(device)
+    # --- FIX END ---
     
     with torch.no_grad():
         outputs = model(**inputs)
     
-    # Output is the refined alphas
     alphas = outputs.alphas
-    
-    # 4. Post-process
-    # Extract alpha, resize to original
     alpha_np = alphas.squeeze().cpu().numpy()
     alpha_pil = Image.fromarray((alpha_np * 255).astype("uint8"), mode="L")
     alpha_pil = alpha_pil.resize(image.size, resample=Image.LANCZOS)
     
     return alpha_pil
 
-
 @st.cache_data(show_spinner=False)
 def process_background_removal(image_bytes, method="RMBG-1.4"):
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
@@ -177,19 +162,16 @@ def process_background_removal(image_bytes, method="RMBG-1.4"):
         mask = inference_segmentation(model, image, device, resolution=1024)
         
     elif method == "VitMatte (Refiner)":
-        # VitMatte needs GPU ideally, works on CPU but slow
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         mask = inference_vitmatte(image, device)
     
     else:
-        # Fallback
         return image
 
-    # Apply mask
     image.putalpha(mask)
     return image
 
-# --- Upscaling Logic (Same as before) ---
+# --- Upscaling Logic ---
 def run_swin_inference(image, processor, model):
     inputs = processor(image, return_tensors="pt")
     with torch.no_grad():
@@ -264,13 +246,12 @@ def main():
     st.sidebar.header("1. Background Removal")
     remove_bg = st.sidebar.checkbox("Remove Background", value=False)
     
-    # NEW: Model Selector
     if remove_bg:
         bg_model = st.sidebar.selectbox(
             "Select AI Model",
             ["RMBG-1.4", "BiRefNet (Heavy)", "VitMatte (Refiner)"],
             index=0,
-            help="RMBG: Fast, Standard Quality.\nBiRefNet: Slower, Better Edges.\nVitMatte: Slowest, Best for Hair/Transparency."
+            help="RMBG: Fast.\nBiRefNet: Better.\nVitMatte: Best for hair/transparency."
         )
     else:
         bg_model = "None"
@@ -294,7 +275,6 @@ def main():
         
         # 1. Background
         if remove_bg:
-            # We add the model name to the spinner text so user knows what's happening
             with st.spinner(f"Removing background using {bg_model}..."):
                 processed_image = process_background_removal(file_bytes, bg_model)
         else:
@@ -303,9 +283,7 @@ def main():
         # 2. Upscaling
         if upscale_mode != "None":
             scale = 4 if "4x" in upscale_mode else 2
-            
-            # Cache Key includes model name now
-            cache_key = f"{uploaded_file.name}_{bg_model}_{scale}_{grid_n}_v5"
+            cache_key = f"{uploaded_file.name}_{bg_model}_{scale}_{grid_n}_v6"
             
             if "upscale_cache" not in st.session_state:
                 st.session_state.upscale_cache = {}