import streamlit as st from PIL import Image import torch import torch.nn.functional as F from torchvision import transforms from transformers import AutoModelForImageSegmentation, AutoImageProcessor, Swin2SRForImageSuperResolution, VitMatteForImageMatting import io import numpy as np import gc # Page Configuration st.set_page_config(layout="wide", page_title="AI Image Lab") # --- 1. MODEL LOADING (Cached) --- @st.cache_resource def load_rmbg_model(): """Option 1: The Lightweight Specialist""" model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4", trust_remote_code=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) return model, device @st.cache_resource def load_birefnet_model(): """Option 2: The Heavyweight Generalist""" model = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet", trust_remote_code=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) return model, device @st.cache_resource def load_vitmatte_model(): """Option 3: The Refiner (Matting)""" 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") model.to(device) return processor, model, device @st.cache_resource def load_upscaler(scale=2): if scale == 4: model_id = "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr" else: model_id = "caidas/swin2SR-classical-sr-x2-64" processor = AutoImageProcessor.from_pretrained(model_id) model = Swin2SRForImageSuperResolution.from_pretrained(model_id) return processor, model # --- 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): if output.dim() == 4 and output.shape[1] == 1: return output elif output.dim() == 3 and output.shape[0] == 1: return output return None if hasattr(output, "logits"): return find_mask_tensor(output.logits) elif isinstance(output, (list, tuple)): for item in output: found = find_mask_tensor(item) if found is not None: return found 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.""" if mask_tensor.dim() == 3: mask_tensor = mask_tensor.unsqueeze(0) erode_k = erode_kernel_size dilate_k = dilate_kernel_size dilated = F.max_pool2d(mask_tensor, kernel_size=dilate_k, stride=1, padding=dilate_k//2) eroded = -F.max_pool2d(-mask_tensor, kernel_size=erode_k, stride=1, padding=erode_k//2) trimap = torch.full_like(mask_tensor, 0.5) trimap[eroded > 0.5] = 1.0 trimap[dilated < 0.5] = 0.0 return trimap # --- 3. INFERENCE LOGIC --- def inference_segmentation(model, image, device, resolution=1024): """Generic inference for RMBG and BiRefNet.""" w, h = image.size transform = transforms.Compose([ transforms.Resize((resolution, resolution)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) input_tensor = transform(image).unsqueeze(0).to(device) with torch.no_grad(): outputs = model(input_tensor) result_tensor = find_mask_tensor(outputs) if result_tensor is None: result_tensor = outputs[0] if isinstance(outputs, (list, tuple)) else outputs if not isinstance(result_tensor, torch.Tensor): if isinstance(result_tensor, (list, tuple)): result_tensor = result_tensor[0] pred = result_tensor.squeeze().cpu() if pred.max() > 1 or pred.min() < 0: pred = pred.sigmoid() pred_pil = transforms.ToPILImage()(pred) mask = pred_pil.resize((w, h), resample=Image.LANCZOS) return mask def inference_vitmatte(image, device): """ Runs pipeline: RMBG (Rough Mask) -> Trimap -> VitMatte (Refined Mask). Includes memory safety downscaling. """ 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, processing_image, device, resolution=1024) # 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) # 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 inputs = processor(images=processing_image, trimaps=trimap_pil, return_tensors="pt").to(device) with torch.no_grad(): outputs = model(**inputs) alphas = outputs.alphas alpha_np = alphas.squeeze().cpu().numpy() alpha_pil = Image.fromarray((alpha_np * 255).astype("uint8"), mode="L") # 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": model, device = load_rmbg_model() mask = inference_segmentation(model, image, device) 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)": device = torch.device("cuda" if torch.cuda.is_available() else "cpu") mask = inference_vitmatte(image, device) else: return image image.putalpha(mask) return image # --- Upscaling Logic --- def run_swin_inference(image, processor, model): inputs = processor(image, return_tensors="pt") with torch.no_grad(): outputs = model(**inputs) output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy() output = np.moveaxis(output, 0, -1) output = (output * 255.0).round().astype(np.uint8) return Image.fromarray(output) def upscale_chunk_logic(image, processor, model): if image.mode == 'RGBA': r, g, b, a = image.split() rgb_image = Image.merge('RGB', (r, g, b)) upscaled_rgb = run_swin_inference(rgb_image, processor, model) upscaled_a = a.resize(upscaled_rgb.size, Image.Resampling.LANCZOS) return Image.merge('RGBA', (*upscaled_rgb.split(), upscaled_a)) else: 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 tile_w = w // cols tile_h = h // rows overlap = 32 full_image = Image.new(image.mode, (w * scale_factor, h * scale_factor)) total_tiles = rows * cols count = 0 for y in range(rows): for x in range(cols): target_left = x * tile_w target_upper = y * tile_h target_right = w if x == cols - 1 else (x + 1) * tile_w target_lower = h if y == rows - 1 else (y + 1) * tile_h source_left = max(0, target_left - overlap) source_upper = max(0, target_upper - overlap) source_right = min(w, target_right + overlap) source_lower = min(h, target_lower + overlap) tile = image.crop((source_left, source_upper, source_right, source_lower)) upscaled_tile = upscale_chunk_logic(tile, processor, model) target_w = target_right - target_left target_h = target_lower - target_upper extra_left = target_left - source_left extra_upper = target_upper - source_upper crop_x = extra_left * scale_factor crop_y = extra_upper * scale_factor crop_w = target_w * scale_factor crop_h = target_h * scale_factor clean_tile = upscaled_tile.crop((crop_x, crop_y, crop_x + crop_w, crop_y + crop_h)) paste_x = target_left * scale_factor paste_y = target_upper * scale_factor full_image.paste(clean_tile, (paste_x, paste_y)) del tile, upscaled_tile, clean_tile cleanup_memory() count += 1 progress_bar.progress(count / total_tiles, text=f"Upscaling Tile {count}/{total_tiles}...") return full_image def convert_image_to_bytes(img): buf = io.BytesIO() img.save(buf, format="PNG") return buf.getvalue() # --- 4. MAIN APP --- def main(): st.title("✨ AI Image Lab: Ultimate Edition") st.markdown("Features: **Multi-Model Background** | **Swin2SR** | **Progress Bar**") # --- Sidebar --- st.sidebar.header("1. Background Removal") remove_bg = st.sidebar.checkbox("Remove Background", value=False) if remove_bg: bg_model = st.sidebar.selectbox( "Select AI Model", ["RMBG-1.4", "BiRefNet (Heavy)", "VitMatte (Refiner)"], index=0, help="RMBG: Fast.\nBiRefNet: Better.\nVitMatte: Best for hair/transparency." ) else: bg_model = "None" st.sidebar.header("2. AI Upscaling") upscale_mode = st.sidebar.radio("Magnification", ["None", "2x", "4x"]) if upscale_mode != "None": grid_n = st.sidebar.slider("Grid Split", 2, 8, 4, help="Higher = Safer RAM usage") else: grid_n = 2 st.sidebar.header("3. Geometry") rotate_angle = st.sidebar.slider("Rotate", -180, 180, 0, 1) # --- Main Logic --- uploaded_file = st.file_uploader("Upload Image", type=["png", "jpg", "jpeg", "webp"]) if uploaded_file is not None: file_bytes = uploaded_file.getvalue() # 1. Background if remove_bg: with st.spinner(f"Removing background using {bg_model}..."): processed_image = process_background_removal(file_bytes, bg_model) else: processed_image = Image.open(io.BytesIO(file_bytes)).convert("RGB") # 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}_v7" if "upscale_cache" not in st.session_state: st.session_state.upscale_cache = {} if cache_key in st.session_state.upscale_cache: processed_image = st.session_state.upscale_cache[cache_key] st.info("✅ Loaded upscaled image from cache") else: progress_bar = st.progress(0, text="Initializing AI models...") processed_image = process_tiled_upscale(processed_image, scale, grid_n, progress_bar) progress_bar.empty() st.session_state.upscale_cache[cache_key] = processed_image # 3. Geometry final_image = processed_image.copy() if rotate_angle != 0: final_image = final_image.rotate(rotate_angle, expand=True) # --- Display --- 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("---") st.download_button( label="💾 Download Result (PNG)", data=convert_image_to_bytes(final_image), file_name="processed_image.png", mime="image/png" ) if __name__ == "__main__": main()