Update src/streamlit_app.py

src/streamlit_app.py

@@ -32,7 +32,6 @@ def load_birefnet_model():
 @st.cache_resource
 def load_vitmatte_model():
     """Option 3: The Refiner (Matting)"""
-    # VitMatte requires a rough mask first (we use RMBG for that)
     processor = AutoImageProcessor.from_pretrained("hustvl/vitmatte-small-composition-1k")
     model = VitMatteForImageMatting.from_pretrained("hustvl/vitmatte-small-composition-1k")
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -51,6 +50,12 @@ def load_upscaler(scale=2):
 
 # --- 2. HELPER FUNCTIONS ---
 
+def cleanup_memory():
+    """Forcibly clears memory."""
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
 def find_mask_tensor(output):
     """Recursively finds the mask tensor in complex model outputs."""
     if isinstance(output, torch.Tensor):
@@ -65,19 +70,13 @@ def find_mask_tensor(output):
     return None
 
 def generate_trimap(mask_tensor, erode_kernel_size=10, dilate_kernel_size=10):
-    """
-    Generates a trimap (Foreground, Background, Unknown) from a binary mask.
-    Values: 1=FG, 0=BG, 0.5=Unknown (Edge)
-    """
+    """Generates a trimap (Foreground, Background, Unknown) from a binary mask."""
     if mask_tensor.dim() == 3: mask_tensor = mask_tensor.unsqueeze(0)
     
     erode_k = erode_kernel_size
     dilate_k = dilate_kernel_size
     
-    # Dilation (Max Pooling)
     dilated = F.max_pool2d(mask_tensor, kernel_size=dilate_k, stride=1, padding=dilate_k//2)
-    
-    # Erosion (Negative Max Pooling)
     eroded = -F.max_pool2d(-mask_tensor, kernel_size=erode_k, stride=1, padding=erode_k//2)
     
     trimap = torch.full_like(mask_tensor, 0.5)
@@ -116,28 +115,42 @@ def inference_segmentation(model, image, device, resolution=1024):
 
 def inference_vitmatte(image, device):
     """
-    Runs pipeline: RMBG (Rough Mask) -> Trimap -> VitMatte (Refined Mask)
+    Runs pipeline: RMBG (Rough Mask) -> Trimap -> VitMatte (Refined Mask).
+    Includes memory safety downscaling.
     """
-    # 1. Get Rough Mask using RMBG (Fast)
+    cleanup_memory() # Clear RAM before starting
+    
+    original_size = image.size
+    
+    # --- MEMORY SAFETY CHECK ---
+    # If image is too large, downscale it for VitMatte processing
+    # 1536px is a sweet spot: good detail, safe RAM usage (~4-6GB peak)
+    max_dim = 1536
+    if max(image.size) > max_dim:
+        scale_ratio = max_dim / max(image.size)
+        new_w = int(image.size[0] * scale_ratio)
+        new_h = int(image.size[1] * scale_ratio)
+        # Create a smaller copy for processing
+        processing_image = image.resize((new_w, new_h), Image.LANCZOS)
+    else:
+        processing_image = image
+
+    # 1. Get Rough Mask using RMBG
     rmbg_model, _ = load_rmbg_model() 
-    rough_mask_pil = inference_segmentation(rmbg_model, image, device, resolution=1024)
+    rough_mask_pil = inference_segmentation(rmbg_model, processing_image, device, resolution=1024)
     
-    # 2. Create Trimap (Tensor)
+    # 2. Create Trimap
     mask_tensor = transforms.ToTensor()(rough_mask_pil).to(device)
     trimap_tensor = generate_trimap(mask_tensor, erode_kernel_size=25, dilate_kernel_size=25)
     
-    # --- 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)
+    # 3. Convert Trimap to PIL (Required for Processor)
     trimap_pil = transforms.ToPILImage()(trimap_tensor.squeeze().cpu())
     
     # 4. VitMatte Inference
     processor, model, _ = load_vitmatte_model()
     
-    # Pass PIL images for both
-    inputs = processor(images=image, trimaps=trimap_pil, return_tensors="pt").to(device)
-    # --- FIX END ---
+    # Pass PIL images
+    inputs = processor(images=processing_image, trimaps=trimap_pil, return_tensors="pt").to(device)
     
     with torch.no_grad():
         outputs = model(**inputs)
@@ -145,12 +158,18 @@ def inference_vitmatte(image, device):
     alphas = outputs.alphas
     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)
     
+    # 5. Restore Resolution
+    # If we downscaled, we must upscale the result mask back to match original
+    if original_size != processing_image.size:
+        alpha_pil = alpha_pil.resize(original_size, resample=Image.LANCZOS)
+    
+    cleanup_memory() # Cleanup after finish
     return alpha_pil
 
 @st.cache_data(show_spinner=False)
 def process_background_removal(image_bytes, method="RMBG-1.4"):
+    cleanup_memory() # Ensure clean state
     image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
     
     if method == "RMBG-1.4":
@@ -159,6 +178,7 @@ def process_background_removal(image_bytes, method="RMBG-1.4"):
         
     elif method == "BiRefNet (Heavy)":
         model, device = load_birefnet_model()
+        # BiRefNet handles 1024 internally well, generally safe on memory
         mask = inference_segmentation(model, image, device, resolution=1024)
         
     elif method == "VitMatte (Refiner)":
@@ -192,6 +212,7 @@ def upscale_chunk_logic(image, processor, model):
         return run_swin_inference(image, processor, model)
 
 def process_tiled_upscale(image, scale_factor, grid_n, progress_bar):
+    cleanup_memory()
     processor, model = load_upscaler(scale_factor)
     w, h = image.size
     rows = cols = grid_n
@@ -226,7 +247,7 @@ def process_tiled_upscale(image, scale_factor, grid_n, progress_bar):
             paste_y = target_upper * scale_factor
             full_image.paste(clean_tile, (paste_x, paste_y))
             del tile, upscaled_tile, clean_tile
-            gc.collect()
+            cleanup_memory()
             count += 1
             progress_bar.progress(count / total_tiles, text=f"Upscaling Tile {count}/{total_tiles}...")
     return full_image
@@ -283,7 +304,7 @@ def main():
         # 2. Upscaling
         if upscale_mode != "None":
             scale = 4 if "4x" in upscale_mode else 2
-            cache_key = f"{uploaded_file.name}_{bg_model}_{scale}_{grid_n}_v6"
+            cache_key = f"{uploaded_file.name}_{bg_model}_{scale}_{grid_n}_v7"
             
             if "upscale_cache" not in st.session_state:
                 st.session_state.upscale_cache = {}
@@ -306,10 +327,12 @@ def main():
         col1, col2 = st.columns(2)
         with col1:
             st.subheader("Original")
+            # Fixed deprecation warning for use_container_width
             st.image(Image.open(io.BytesIO(file_bytes)), use_container_width=True)
         
         with col2:
             st.subheader("Result")
+            # Fixed deprecation warning for use_container_width
             st.image(final_image, use_container_width=True)
 
         st.markdown("---")