off cache

app.py

@@ -1,733 +1,724 @@
-"""
-Background Removal - Image and Video
-- BRIA RMBG 1.4: Fast general-purpose background removal
-- BRIA RMBG 2.0: High-quality BiRefNet-based (slower, better quality)
-- IS-Net/DIS: Ultra-fine edge detection (hair, fur, transparency)
-- Robust Video Matting: Video background removal with ONNX Runtime (2x faster!)
-
-CPU-only HuggingFace Space
-"""
-
-import os
-import sys
-import tempfile
-import argparse
-from typing import Tuple, Optional
-
-# Set OMP threads to match cpu-basic (2 vCPU)
-os.environ["OMP_NUM_THREADS"] = "2"
-os.environ["OMP_WAIT_POLICY"] = "ACTIVE"
-
-import gradio as gr
-import numpy as np
-from PIL import Image
-import torch
-import torch.nn.functional as F
-from torchvision.transforms.functional import normalize
-from huggingface_hub import hf_hub_download
-import cv2
-import onnxruntime as ort
-
-# Force CPU for PyTorch (image model)
-device = torch.device("cpu")
-
-# =============================================================================
-# Image Background Removal (BRIA RMBG 1.4)
-# =============================================================================
-
-class BriaRMBGModel:
-    """Wrapper for BRIA-RMBG 1.4 model"""
-
-    def __init__(self):
-        self.model = None
-        self.input_size = [1024, 1024]
-
-    def load_model(self):
-        """Load the BRIA-RMBG 1.4 model from HuggingFace"""
-        if self.model is not None:
-            return
-
-        print("Loading BRIA-RMBG 1.4 model...")
-
-        # Download all required files from BRIA repo
-        repo_id = "briaai/RMBG-1.4"
-        model_path = hf_hub_download(repo_id=repo_id, filename="model.pth")
-        briarmbg_path = hf_hub_download(repo_id=repo_id, filename="briarmbg.py")
-        myconfig_path = hf_hub_download(repo_id=repo_id, filename="MyConfig.py")
-
-        # Setup package directory
-        briarmbg_dir = os.path.dirname(briarmbg_path)
-
-        # Create __init__.py to make it a proper package
-        init_path = os.path.join(briarmbg_dir, "__init__.py")
-        if not os.path.exists(init_path):
-            with open(init_path, "w") as f:
-                f.write("# Auto-generated for package import\n")
-
-        # Add parent directory to path so we can import as package
-        parent_dir = os.path.dirname(briarmbg_dir)
-        package_name = os.path.basename(briarmbg_dir)
-
-        if parent_dir not in sys.path:
-            sys.path.insert(0, parent_dir)
-
-        # Import using the package structure
-        import importlib
-        briarmbg_module = importlib.import_module(f"{package_name}.briarmbg")
-        BriaRMBG = briarmbg_module.BriaRMBG
-
-        self.model = BriaRMBG()
-        self.model.load_state_dict(torch.load(model_path, map_location=device))
-        self.model.to(device)
-        self.model.eval()
-        print("BRIA-RMBG model loaded!")
-
-    def preprocess(self, image: Image.Image):
-        """Preprocess image for model input"""
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-
-        orig_size = image.size
-        im = np.array(image)
-
-        if len(im.shape) < 3:
-            im = im[:, :, np.newaxis]
-        if im.shape[2] == 1:
-            im = np.repeat(im, 3, axis=2)
-
-        im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1)
-        im_tensor = F.interpolate(
-            im_tensor.unsqueeze(0),
-            self.input_size,
-            mode="bilinear",
-            align_corners=False
-        )
-
-        image_tensor = im_tensor / 255.0
-        image_tensor = normalize(image_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
-
-        return image_tensor, orig_size
-
-    def postprocess(self, result, orig_size):
-        """Postprocess model output to mask"""
-        result = F.interpolate(
-            result,
-            size=(orig_size[1], orig_size[0]),
-            mode="bilinear",
-            align_corners=False
-        )
-
-        ma = torch.max(result)
-        mi = torch.min(result)
-        result = (result - mi) / (ma - mi + 1e-8)
-
-        mask = (result.squeeze().cpu().numpy() * 255).astype(np.uint8)
-        return mask
-
-    @torch.no_grad()
-    def remove_background(self, image: Image.Image):
-        """Remove background from image"""
-        self.load_model()
-
-        input_tensor, orig_size = self.preprocess(image)
-        input_tensor = input_tensor.to(device)
-
-        result = self.model(input_tensor)
-
-        if isinstance(result, (list, tuple)):
-            result = result[0]
-            if isinstance(result, (list, tuple)):
-                result = result[0]
-
-        mask = self.postprocess(result, orig_size)
-
-        pil_mask = Image.fromarray(mask).convert("L")
-
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-
-        output = image.copy()
-        output.putalpha(pil_mask)
-
-        return output, pil_mask
-
-
-# =============================================================================
-# Image Background Removal (IS-Net/DIS - Ultra-fine edges) - ONNX Runtime
-# =============================================================================
-
-class ISNetModel:
-    """Wrapper for IS-Net/DIS ONNX model - better for hair, fur, transparency"""
-
-    def __init__(self):
-        self.session = None
-        self.input_size = [1024, 1024]
-
-    def load_model(self):
-        """Load IS-Net ONNX model"""
-        if self.session is not None:
-            return
-
-        print("Loading IS-Net/DIS ONNX model...")
-        model_path = os.path.join(os.path.dirname(__file__), "isnet_dis_fp32.onnx")
-
-        # Setup ONNX Runtime with optimizations
-        sess_options = ort.SessionOptions()
-        sess_options.intra_op_num_threads = 2
-        sess_options.inter_op_num_threads = 1
-        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-
-        self.session = ort.InferenceSession(
-            model_path,
-            sess_options,
-            providers=['CPUExecutionProvider']
-        )
-        print("IS-Net/DIS ONNX model loaded!")
-
-    def preprocess(self, image: Image.Image):
-        """Preprocess image for model input"""
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-
-        orig_size = image.size
-        # Resize to 1024x1024
-        im_resized = image.resize((1024, 1024), Image.BILINEAR)
-        im = np.array(im_resized).astype(np.float32)
-
-        # Normalize: (x/255 - 0.5) / 1.0 = x/255 - 0.5
-        im = im / 255.0
-        im = (im - 0.5) / 1.0
-
-        # HWC -> NCHW
-        im = np.transpose(im, (2, 0, 1))[np.newaxis, ...]
-
-        return im.astype(np.float32), orig_size
-
-    def postprocess(self, result, orig_size):
-        """Postprocess model output to mask"""
-        # Result is already sigmoid (0-1), resize to original
-        mask_1024 = result[0, 0]  # 1024x1024
-
-        # Resize to original size
-        mask_pil = Image.fromarray((mask_1024 * 255).astype(np.uint8))
-        mask_pil = mask_pil.resize(orig_size, Image.BILINEAR)
-
-        return np.array(mask_pil)
-
-    def remove_background(self, image: Image.Image):
-        """Remove background from image using IS-Net ONNX"""
-        self.load_model()
-
-        input_tensor, orig_size = self.preprocess(image)
-
-        # Run ONNX inference - d0 is the main output
-        outputs = self.session.run(None, {'input': input_tensor})
-        result = outputs[0]  # d0: main segmentation mask
-
-        mask = self.postprocess(result, orig_size)
-
-        pil_mask = Image.fromarray(mask).convert("L")
-
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-
-        output = image.copy()
-        output.putalpha(pil_mask)
-
-        return output, pil_mask
-
-
-# =============================================================================
-# Image Background Removal (BRIA RMBG 2.0 - BiRefNet based, high quality)
-# =============================================================================
-
-class RMBG2Model:
-    """Wrapper for BRIA-RMBG 2.0 model (BiRefNet architecture) - higher quality, slower"""
-
-    def __init__(self):
-        self.model = None
-        self.input_size = 1024  # Native resolution for best quality
-
-    def load_model(self):
-        """Load RMBG-2.0 model from HuggingFace"""
-        if self.model is not None:
-            return
-
-        print("Loading BRIA-RMBG 2.0 model (BiRefNet)...")
-
-        # Download model files from ungated repo
-        repo_id = "1038lab/RMBG-2.0"
-        model_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors")
-        birefnet_path = hf_hub_download(repo_id=repo_id, filename="birefnet.py")
-        config_path = hf_hub_download(repo_id=repo_id, filename="BiRefNet_config.py")
-
-        # Setup import path
-        model_dir = os.path.dirname(birefnet_path)
-        if model_dir not in sys.path:
-            sys.path.insert(0, model_dir)
-
-        # Import model classes
-        from safetensors.torch import load_file
-        import importlib.util
-
-        # Load BiRefNet_config
-        spec = importlib.util.spec_from_file_location("BiRefNet_config", config_path)
-        config_module = importlib.util.module_from_spec(spec)
-        sys.modules["BiRefNet_config"] = config_module
-        spec.loader.exec_module(config_module)
-
-        # Fix relative import in birefnet.py - read, fix, and exec
-        with open(birefnet_path, 'r', encoding='utf-8') as f:
-            birefnet_code = f.read()
-
-        # Replace relative import with absolute
-        birefnet_code = birefnet_code.replace(
-            'from .BiRefNet_config import BiRefNetConfig',
-            'from BiRefNet_config import BiRefNetConfig'
-        )
-
-        # Execute the fixed code
-        birefnet_module = type(sys)("birefnet")
-        exec(compile(birefnet_code, birefnet_path, 'exec'), birefnet_module.__dict__)
-
-        BiRefNet = birefnet_module.BiRefNet
-        BiRefNetConfig = config_module.BiRefNetConfig
-
-        # Create and load model
-        config = BiRefNetConfig(bb_pretrained=False)
-        self.model = BiRefNet(config)
-
-        state_dict = load_file(model_path)
-        self.model.load_state_dict(state_dict)
-        self.model.to(device)
-        self.model.eval()
-
-        print(f"BRIA-RMBG 2.0 loaded! ({sum(p.numel() for p in self.model.parameters())/1e6:.0f}M params)")
-
-    def preprocess(self, image: Image.Image):
-        """Preprocess image for model input"""
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-
-        orig_size = image.size
-
-        # Resize to model input size
-        im_resized = image.resize((self.input_size, self.input_size), Image.BILINEAR)
-        im = np.array(im_resized).astype(np.float32)
-
-        # Normalize
-        im = im / 255.0
-        im = (im - 0.5) / 1.0
-
-        # HWC -> NCHW
-        im_tensor = torch.from_numpy(im).permute(2, 0, 1).unsqueeze(0)
-
-        return im_tensor.float(), orig_size
-
-    def postprocess(self, result, orig_size):
-        """Postprocess model output to mask"""
-        # Get final output (BiRefNet returns list for deep supervision)
-        if isinstance(result, (list, tuple)):
-            result = result[-1]  # Last output is the refined mask
-
-        # Sigmoid and resize
-        mask = torch.sigmoid(result).squeeze().cpu().numpy()
-
-        # Resize to original size
-        mask_pil = Image.fromarray((mask * 255).astype(np.uint8))
-        mask_pil = mask_pil.resize(orig_size, Image.BILINEAR)
-
-        return np.array(mask_pil)
-
-    @torch.no_grad()
-    def remove_background(self, image: Image.Image):
-        """Remove background from image using RMBG-2.0"""
-        self.load_model()
-
-        input_tensor, orig_size = self.preprocess(image)
-        input_tensor = input_tensor.to(device)
-
-        result = self.model(input_tensor)
-        mask = self.postprocess(result, orig_size)
-
-        pil_mask = Image.fromarray(mask).convert("L")
-
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-
-        output = image.copy()
-        output.putalpha(pil_mask)
-
-        return output, pil_mask
-
-
-# =============================================================================
-# Video Background Removal (Robust Video Matting - ONNX Runtime)
-# =============================================================================
-
-rvm_session = None
-
-
-def load_rvm_onnx():
-    """Load RobustVideoMatting ONNX model for fast CPU inference."""
-    global rvm_session
-    if rvm_session is not None:
-        return rvm_session
-
-    print("Loading Robust Video Matting ONNX model...")
-
-    # Load ONNX model from local Space directory
-    model_path = os.path.join(os.path.dirname(__file__), "rvm_mobilenetv3_fp32.onnx")
-
-    # Setup ONNX Runtime with optimizations (2 vCPU on cpu-basic)
-    sess_options = ort.SessionOptions()
-    sess_options.intra_op_num_threads = 2
-    sess_options.inter_op_num_threads = 1
-    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-
-    rvm_session = ort.InferenceSession(
-        model_path,
-        sess_options,
-        providers=['CPUExecutionProvider']
-    )
-
-    print("RVM ONNX model loaded! (2x faster than PyTorch)")
-    return rvm_session
-
-
-# Output resolution options
-RESOLUTION_OPTIONS = {
-    "Original": None,
-    "1080p": 1080,
-    "720p": 720,
-    "480p": 480,
-}
-
-
-def process_video(input_video, output_type="green_screen", max_resolution="720p", downsample_ratio=0.25, progress=None):
-    """
-    Process video to remove background using ONNX Runtime.
-
-    Args:
-        input_video: Path to input video
-        output_type: "green_screen" or "foreground_only"
-        max_resolution: Output video resolution (affects speed)
-        downsample_ratio: Internal model quality (affects edge accuracy)
-        progress: Gradio Progress tracker (optional, for web UI)
-
-    Returns:
-        Path to processed video
-    """
-    if input_video is None:
-        return None
-
-    # Get max height from resolution option
-    max_h = RESOLUTION_OPTIONS.get(max_resolution)
-
-    session = load_rvm_onnx()
-    output_names = [o.name for o in session.get_outputs()]
-
-    cap = cv2.VideoCapture(input_video)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    orig_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    orig_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-
-    # Resize based on quality preset
-    if max_h and orig_height > max_h:
-        scale = max_h / orig_height
-        width = int(orig_width * scale)
-        height = max_h
-        print(f"Resizing: {orig_width}x{orig_height} -> {width}x{height}")
-    else:
-        width, height = orig_width, orig_height
-
-    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
-        output_path = tmp.name
-
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-
-    green_bg = np.zeros((height, width, 3), dtype=np.uint8)
-    green_bg[:] = (0, 255, 0)
-
-    # Initialize recurrent states (zeros for first frame)
-    r1 = np.zeros((1, 1, 1, 1), dtype=np.float32)
-    r2 = np.zeros((1, 1, 1, 1), dtype=np.float32)
-    r3 = np.zeros((1, 1, 1, 1), dtype=np.float32)
-    r4 = np.zeros((1, 1, 1, 1), dtype=np.float32)
-
-    frame_idx = 0
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-
-        # Resize if needed
-        if frame.shape[0] != height or frame.shape[1] != width:
-            frame = cv2.resize(frame, (width, height))
-
-        # Preprocess frame
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        src = frame_rgb.astype(np.float32) / 255.0
-        src = np.transpose(src, (2, 0, 1))[np.newaxis, ...]  # NCHW
-
-        # ONNX inference
-        outputs = session.run(output_names, {
-            'src': src,
-            'r1i': r1,
-            'r2i': r2,
-            'r3i': r3,
-            'r4i': r4,
-            'downsample_ratio': np.array([downsample_ratio], dtype=np.float32)
-        })
-
-        fgr, pha, r1, r2, r3, r4 = outputs
-
-        # Postprocess
-        fgr_np = np.transpose(fgr[0], (1, 2, 0))  # CHW -> HWC
-        pha_np = pha[0, 0]
-
-        if fgr_np.shape[:2] != (height, width):
-            fgr_np = cv2.resize(fgr_np, (width, height))
-            pha_np = cv2.resize(pha_np, (width, height))
-
-        pha_3ch = np.stack([pha_np] * 3, axis=-1)
-
-        if output_type == "green_screen":
-            result = (fgr_np * pha_3ch + (green_bg / 255.0) * (1 - pha_3ch))
-        else:
-            result = fgr_np * pha_3ch
-
-        result_bgr = cv2.cvtColor((result * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
-        out.write(result_bgr)
-
-        frame_idx += 1
-        if frame_count > 0:
-            if progress is not None:
-                progress(frame_idx / frame_count, desc=f"Processing frame {frame_idx}/{frame_count}")
-            elif frame_idx % 30 == 0:
-                print(f"Processing frame {frame_idx}/{frame_count}")
-
-    cap.release()
-    out.release()
-
-    return output_path
-
-
-# =============================================================================
-# Gradio Interface
-# =============================================================================
-
-# Initialize image models (lazy loading)
-bria_model = BriaRMBGModel()
-rmbg2_model = RMBG2Model()
-isnet_model = ISNetModel()
-
-IMAGE_MODELS = {
-    "BRIA-RMBG 1.4 (fast, general)": bria_model,
-    "BRIA-RMBG 2.0 (high quality, slow)": rmbg2_model,
-    "IS-Net/DIS (fine edges, hair)": isnet_model,
-}
-
-
-IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif', '.tiff'}
-VIDEO_EXTENSIONS = {'.mp4', '.avi', '.mov', '.mkv', '.webm', '.flv', '.wmv'}
-
-
-def detect_file_type(file_path):
-    """Detect if file is image or video based on extension."""
-    if file_path is None:
-        return None
-    ext = os.path.splitext(file_path)[1].lower()
-    if ext in IMAGE_EXTENSIONS:
-        return "image"
-    elif ext in VIDEO_EXTENSIONS:
-        return "video"
-    return None
-
-
-def on_file_upload(file_path):
-    """Handle file upload - detect type and update preview."""
-    if file_path is None:
-        return gr.Image(visible=False), gr.Video(visible=False), None
-
-    file_type = detect_file_type(file_path)
-    if file_type == "image":
-        return gr.Image(value=file_path, visible=True), gr.Video(visible=False), "image"
-    elif file_type == "video":
-        return gr.Image(visible=False), gr.Video(value=file_path, visible=True), "video"
-    return gr.Image(visible=False), gr.Video(visible=False), None
-
-
-def process_media(file_path, file_type, image_model_choice, output_type, max_resolution, edge_quality, progress=gr.Progress()):
-    """Process based on stored file type."""
-    if file_path is None:
-        return None, None
-
-    # Auto-detect file type if not provided (for API calls)
-    if file_type is None:
-        file_type = detect_file_type(file_path)
-
-    if file_type is None:
-        return None, None
-
-    if file_type == "image":
-        pil_image = Image.open(file_path)
-        model = IMAGE_MODELS.get(image_model_choice, bria_model)
-        output_img, mask = model.remove_background(pil_image)
-        return output_img, None
-
-    elif file_type == "video":
-        output_path = process_video(file_path, output_type, max_resolution, edge_quality, progress)
-        return None, output_path
-
-    return None, None
-
-
-def build_ui():
-    """Build compact unified Gradio interface."""
-    with gr.Blocks(title="Background Removal") as demo:
-        gr.Markdown("## Background Removal\nUpload image or video. Image models: RMBG 1.4 (fast) | RMBG 2.0 (quality) | IS-Net (fine edges) | Video: RVM")
-
-        # State to track file type
-        file_type_state = gr.State(value=None)
-
-        with gr.Row():
-            with gr.Column(scale=1):
-                input_file = gr.File(
-                    label="Upload Image or Video",
-                    file_types=["image", "video"],
-                    type="filepath"
-                )
-                # Preview - only one visible at a time
-                preview_image = gr.Image(label="Preview", interactive=False, visible=False)
-                preview_video = gr.Video(label="Preview", interactive=False, visible=False)
-
-                # Image model selector
-                image_model_dropdown = gr.Dropdown(
-                    choices=list(IMAGE_MODELS.keys()),
-                    value="BRIA-RMBG 1.4 (fast, general)",
-                    label="Image Model",
-                    info="1.4: fast | 2.0: quality (slow) | IS-Net: fine edges"
-                )
-
-                with gr.Accordion("Video Settings", open=False):
-                    output_type = gr.Radio(
-                        choices=["green_screen", "foreground_only"],
-                        value="green_screen",
-                        label="Output Type"
-                    )
-                    max_res = gr.Dropdown(
-                        choices=list(RESOLUTION_OPTIONS.keys()),
-                        value="720p",
-                        label="Output Resolution"
-                    )
-                    edge_quality = gr.Slider(
-                        minimum=0.1,
-                        maximum=1.0,
-                        value=0.25,
-                        step=0.05,
-                        label="Edge Quality"
-                    )
-
-                process_btn = gr.Button("Remove Background", variant="primary")
-
-            with gr.Column(scale=1):
-                # Output components
-                output_image = gr.Image(label="Result (Image)", type="pil")
-                output_video = gr.Video(label="Result (Video)")
-
-        # Wrapper for examples (auto-detects file type)
-        def process_example(file_path, image_model_choice, output_type, max_resolution, edge_quality):
-            file_type = detect_file_type(file_path)
-            return process_media(file_path, file_type, image_model_choice, output_type, max_resolution, edge_quality)
-
-        # Example files with cached results
-        gr.Examples(
-            examples=[
-                ["examples/cat.jpg", "BRIA-RMBG 1.4 (fast, general)", "green_screen", "720p", 0.25],
-            ],
-            inputs=[input_file, image_model_dropdown, output_type, max_res, edge_quality],
-            outputs=[output_image, output_video],
-            fn=process_example,
-            cache_examples=True,
-            cache_mode="lazy",
-            label="Examples"
-        )
-
-        # On upload: detect type and show preview
-        input_file.change(
-            fn=on_file_upload,
-            inputs=[input_file],
-            outputs=[preview_image, preview_video, file_type_state]
-        )
-
-        # On process: fill the already-visible output
-        process_btn.click(
-            fn=process_media,
-            inputs=[input_file, file_type_state, image_model_dropdown, output_type, max_res, edge_quality],
-            outputs=[output_image, output_video]
-        )
-
-    return demo
-
-
-# =============================================================================
-# CLI Interface
-# =============================================================================
-
-def cli_main():
-    """CLI entry point."""
-    parser = argparse.ArgumentParser(description="Background Removal - Image and Video")
-    subparsers = parser.add_subparsers(dest="command", help="Commands")
-
-    # Image command
-    img_parser = subparsers.add_parser("image", help="Remove background from image")
-    img_parser.add_argument("input", type=str, help="Input image path")
-    img_parser.add_argument("--output", "-o", type=str, default="output.png", help="Output path")
-    img_parser.add_argument("--mask", "-m", type=str, help="Save mask to path")
-    img_parser.add_argument("--model", type=str, default="BRIA-RMBG 1.4 (fast, general)",
-                           choices=list(IMAGE_MODELS.keys()),
-                           help="Image model: BRIA (fast) or IS-Net (fine edges)")
-
-    # Video command
-    vid_parser = subparsers.add_parser("video", help="Remove background from video")
-    vid_parser.add_argument("input", type=str, help="Input video path")
-    vid_parser.add_argument("--output", "-o", type=str, default="output.mp4", help="Output path")
-    vid_parser.add_argument("--type", type=str, default="green_screen",
-                           choices=["green_screen", "foreground_only"])
-    vid_parser.add_argument("--downsample", type=float, default=0.25, help="Downsample ratio")
-
-    args = parser.parse_args()
-
-    if args.command == "image":
-        print(f"Processing image: {args.input}")
-        input_img = Image.open(args.input)
-        # Use selected model (default: BRIA)
-        model = IMAGE_MODELS.get(args.model, bria_model)
-        output_img, mask = model.remove_background(input_img)
-        output_img.save(args.output)
-        print(f"Saved: {args.output}")
-        if args.mask:
-            mask.save(args.mask)
-            print(f"Saved mask: {args.mask}")
-
-    elif args.command == "video":
-        print(f"Processing video: {args.input}")
-        output_path = process_video(args.input, args.type, args.downsample)
-        if output_path:
-            import shutil
-            shutil.move(output_path, args.output)
-            print(f"Saved: {args.output}")
-
-    else:
-        parser.print_help()
-
-
-if __name__ == "__main__":
-    if len(sys.argv) > 1 and sys.argv[1] in ["image", "video", "-h", "--help"]:
-        cli_main()
-    else:
-        demo = build_ui()
-        demo.launch(theme=gr.themes.Soft())
+"""
+Background Removal - Image and Video
+- BRIA RMBG 1.4: Fast general-purpose background removal
+- BRIA RMBG 2.0: High-quality BiRefNet-based (slower, better quality)
+- IS-Net/DIS: Ultra-fine edge detection (hair, fur, transparency)
+- Robust Video Matting: Video background removal with ONNX Runtime (2x faster!)
+
+CPU-only HuggingFace Space
+"""
+
+import os
+import sys
+import tempfile
+import argparse
+from typing import Tuple, Optional
+
+# Set OMP threads to match cpu-basic (2 vCPU)
+os.environ["OMP_NUM_THREADS"] = "2"
+os.environ["OMP_WAIT_POLICY"] = "ACTIVE"
+
+import gradio as gr
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+from huggingface_hub import hf_hub_download
+import cv2
+import onnxruntime as ort
+
+# Force CPU for PyTorch (image model)
+device = torch.device("cpu")
+
+# =============================================================================
+# Image Background Removal (BRIA RMBG 1.4)
+# =============================================================================
+
+class BriaRMBGModel:
+    """Wrapper for BRIA-RMBG 1.4 model"""
+
+    def __init__(self):
+        self.model = None
+        self.input_size = [1024, 1024]
+
+    def load_model(self):
+        """Load the BRIA-RMBG 1.4 model from HuggingFace"""
+        if self.model is not None:
+            return
+
+        print("Loading BRIA-RMBG 1.4 model...")
+
+        # Download all required files from BRIA repo
+        repo_id = "briaai/RMBG-1.4"
+        model_path = hf_hub_download(repo_id=repo_id, filename="model.pth")
+        briarmbg_path = hf_hub_download(repo_id=repo_id, filename="briarmbg.py")
+        myconfig_path = hf_hub_download(repo_id=repo_id, filename="MyConfig.py")
+
+        # Setup package directory
+        briarmbg_dir = os.path.dirname(briarmbg_path)
+
+        # Create __init__.py to make it a proper package
+        init_path = os.path.join(briarmbg_dir, "__init__.py")
+        if not os.path.exists(init_path):
+            with open(init_path, "w") as f:
+                f.write("# Auto-generated for package import\n")
+
+        # Add parent directory to path so we can import as package
+        parent_dir = os.path.dirname(briarmbg_dir)
+        package_name = os.path.basename(briarmbg_dir)
+
+        if parent_dir not in sys.path:
+            sys.path.insert(0, parent_dir)
+
+        # Import using the package structure
+        import importlib
+        briarmbg_module = importlib.import_module(f"{package_name}.briarmbg")
+        BriaRMBG = briarmbg_module.BriaRMBG
+
+        self.model = BriaRMBG()
+        self.model.load_state_dict(torch.load(model_path, map_location=device))
+        self.model.to(device)
+        self.model.eval()
+        print("BRIA-RMBG model loaded!")
+
+    def preprocess(self, image: Image.Image):
+        """Preprocess image for model input"""
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        orig_size = image.size
+        im = np.array(image)
+
+        if len(im.shape) < 3:
+            im = im[:, :, np.newaxis]
+        if im.shape[2] == 1:
+            im = np.repeat(im, 3, axis=2)
+
+        im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1)
+        im_tensor = F.interpolate(
+            im_tensor.unsqueeze(0),
+            self.input_size,
+            mode="bilinear",
+            align_corners=False
+        )
+
+        image_tensor = im_tensor / 255.0
+        image_tensor = normalize(image_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+
+        return image_tensor, orig_size
+
+    def postprocess(self, result, orig_size):
+        """Postprocess model output to mask"""
+        result = F.interpolate(
+            result,
+            size=(orig_size[1], orig_size[0]),
+            mode="bilinear",
+            align_corners=False
+        )
+
+        ma = torch.max(result)
+        mi = torch.min(result)
+        result = (result - mi) / (ma - mi + 1e-8)
+
+        mask = (result.squeeze().cpu().numpy() * 255).astype(np.uint8)
+        return mask
+
+    @torch.no_grad()
+    def remove_background(self, image: Image.Image):
+        """Remove background from image"""
+        self.load_model()
+
+        input_tensor, orig_size = self.preprocess(image)
+        input_tensor = input_tensor.to(device)
+
+        result = self.model(input_tensor)
+
+        if isinstance(result, (list, tuple)):
+            result = result[0]
+            if isinstance(result, (list, tuple)):
+                result = result[0]
+
+        mask = self.postprocess(result, orig_size)
+
+        pil_mask = Image.fromarray(mask).convert("L")
+
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        output = image.copy()
+        output.putalpha(pil_mask)
+
+        return output, pil_mask
+
+
+# =============================================================================
+# Image Background Removal (IS-Net/DIS - Ultra-fine edges) - ONNX Runtime
+# =============================================================================
+
+class ISNetModel:
+    """Wrapper for IS-Net/DIS ONNX model - better for hair, fur, transparency"""
+
+    def __init__(self):
+        self.session = None
+        self.input_size = [1024, 1024]
+
+    def load_model(self):
+        """Load IS-Net ONNX model"""
+        if self.session is not None:
+            return
+
+        print("Loading IS-Net/DIS ONNX model...")
+        model_path = os.path.join(os.path.dirname(__file__), "isnet_dis_fp32.onnx")
+
+        # Setup ONNX Runtime with optimizations
+        sess_options = ort.SessionOptions()
+        sess_options.intra_op_num_threads = 2
+        sess_options.inter_op_num_threads = 1
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        self.session = ort.InferenceSession(
+            model_path,
+            sess_options,
+            providers=['CPUExecutionProvider']
+        )
+        print("IS-Net/DIS ONNX model loaded!")
+
+    def preprocess(self, image: Image.Image):
+        """Preprocess image for model input"""
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        orig_size = image.size
+        # Resize to 1024x1024
+        im_resized = image.resize((1024, 1024), Image.BILINEAR)
+        im = np.array(im_resized).astype(np.float32)
+
+        # Normalize: (x/255 - 0.5) / 1.0 = x/255 - 0.5
+        im = im / 255.0
+        im = (im - 0.5) / 1.0
+
+        # HWC -> NCHW
+        im = np.transpose(im, (2, 0, 1))[np.newaxis, ...]
+
+        return im.astype(np.float32), orig_size
+
+    def postprocess(self, result, orig_size):
+        """Postprocess model output to mask"""
+        # Result is already sigmoid (0-1), resize to original
+        mask_1024 = result[0, 0]  # 1024x1024
+
+        # Resize to original size
+        mask_pil = Image.fromarray((mask_1024 * 255).astype(np.uint8))
+        mask_pil = mask_pil.resize(orig_size, Image.BILINEAR)
+
+        return np.array(mask_pil)
+
+    def remove_background(self, image: Image.Image):
+        """Remove background from image using IS-Net ONNX"""
+        self.load_model()
+
+        input_tensor, orig_size = self.preprocess(image)
+
+        # Run ONNX inference - d0 is the main output
+        outputs = self.session.run(None, {'input': input_tensor})
+        result = outputs[0]  # d0: main segmentation mask
+
+        mask = self.postprocess(result, orig_size)
+
+        pil_mask = Image.fromarray(mask).convert("L")
+
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        output = image.copy()
+        output.putalpha(pil_mask)
+
+        return output, pil_mask
+
+
+# =============================================================================
+# Image Background Removal (BRIA RMBG 2.0 - BiRefNet based, high quality)
+# =============================================================================
+
+class RMBG2Model:
+    """Wrapper for BRIA-RMBG 2.0 model (BiRefNet architecture) - higher quality, slower"""
+
+    def __init__(self):
+        self.model = None
+        self.input_size = 1024  # Native resolution for best quality
+
+    def load_model(self):
+        """Load RMBG-2.0 model from HuggingFace"""
+        if self.model is not None:
+            return
+
+        print("Loading BRIA-RMBG 2.0 model (BiRefNet)...")
+
+        # Download model files from ungated repo
+        repo_id = "1038lab/RMBG-2.0"
+        model_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors")
+        birefnet_path = hf_hub_download(repo_id=repo_id, filename="birefnet.py")
+        config_path = hf_hub_download(repo_id=repo_id, filename="BiRefNet_config.py")
+
+        # Setup import path
+        model_dir = os.path.dirname(birefnet_path)
+        if model_dir not in sys.path:
+            sys.path.insert(0, model_dir)
+
+        # Import model classes
+        from safetensors.torch import load_file
+        import importlib.util
+
+        # Load BiRefNet_config
+        spec = importlib.util.spec_from_file_location("BiRefNet_config", config_path)
+        config_module = importlib.util.module_from_spec(spec)
+        sys.modules["BiRefNet_config"] = config_module
+        spec.loader.exec_module(config_module)
+
+        # Fix relative import in birefnet.py - read, fix, and exec
+        with open(birefnet_path, 'r', encoding='utf-8') as f:
+            birefnet_code = f.read()
+
+        # Replace relative import with absolute
+        birefnet_code = birefnet_code.replace(
+            'from .BiRefNet_config import BiRefNetConfig',
+            'from BiRefNet_config import BiRefNetConfig'
+        )
+
+        # Execute the fixed code
+        birefnet_module = type(sys)("birefnet")
+        exec(compile(birefnet_code, birefnet_path, 'exec'), birefnet_module.__dict__)
+
+        BiRefNet = birefnet_module.BiRefNet
+        BiRefNetConfig = config_module.BiRefNetConfig
+
+        # Create and load model
+        config = BiRefNetConfig(bb_pretrained=False)
+        self.model = BiRefNet(config)
+
+        state_dict = load_file(model_path)
+        self.model.load_state_dict(state_dict)
+        self.model.to(device)
+        self.model.eval()
+
+        print(f"BRIA-RMBG 2.0 loaded! ({sum(p.numel() for p in self.model.parameters())/1e6:.0f}M params)")
+
+    def preprocess(self, image: Image.Image):
+        """Preprocess image for model input"""
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        orig_size = image.size
+
+        # Resize to model input size
+        im_resized = image.resize((self.input_size, self.input_size), Image.BILINEAR)
+        im = np.array(im_resized).astype(np.float32)
+
+        # Normalize
+        im = im / 255.0
+        im = (im - 0.5) / 1.0
+
+        # HWC -> NCHW
+        im_tensor = torch.from_numpy(im).permute(2, 0, 1).unsqueeze(0)
+
+        return im_tensor.float(), orig_size
+
+    def postprocess(self, result, orig_size):
+        """Postprocess model output to mask"""
+        # Get final output (BiRefNet returns list for deep supervision)
+        if isinstance(result, (list, tuple)):
+            result = result[-1]  # Last output is the refined mask
+
+        # Sigmoid and resize
+        mask = torch.sigmoid(result).squeeze().cpu().numpy()
+
+        # Resize to original size
+        mask_pil = Image.fromarray((mask * 255).astype(np.uint8))
+        mask_pil = mask_pil.resize(orig_size, Image.BILINEAR)
+
+        return np.array(mask_pil)
+
+    @torch.no_grad()
+    def remove_background(self, image: Image.Image):
+        """Remove background from image using RMBG-2.0"""
+        self.load_model()
+
+        input_tensor, orig_size = self.preprocess(image)
+        input_tensor = input_tensor.to(device)
+
+        result = self.model(input_tensor)
+        mask = self.postprocess(result, orig_size)
+
+        pil_mask = Image.fromarray(mask).convert("L")
+
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        output = image.copy()
+        output.putalpha(pil_mask)
+
+        return output, pil_mask
+
+
+# =============================================================================
+# Video Background Removal (Robust Video Matting - ONNX Runtime)
+# =============================================================================
+
+rvm_session = None
+
+
+def load_rvm_onnx():
+    """Load RobustVideoMatting ONNX model for fast CPU inference."""
+    global rvm_session
+    if rvm_session is not None:
+        return rvm_session
+
+    print("Loading Robust Video Matting ONNX model...")
+
+    # Load ONNX model from local Space directory
+    model_path = os.path.join(os.path.dirname(__file__), "rvm_mobilenetv3_fp32.onnx")
+
+    # Setup ONNX Runtime with optimizations (2 vCPU on cpu-basic)
+    sess_options = ort.SessionOptions()
+    sess_options.intra_op_num_threads = 2
+    sess_options.inter_op_num_threads = 1
+    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+    rvm_session = ort.InferenceSession(
+        model_path,
+        sess_options,
+        providers=['CPUExecutionProvider']
+    )
+
+    print("RVM ONNX model loaded! (2x faster than PyTorch)")
+    return rvm_session
+
+
+# Output resolution options
+RESOLUTION_OPTIONS = {
+    "Original": None,
+    "1080p": 1080,
+    "720p": 720,
+    "480p": 480,
+}
+
+
+def process_video(input_video, output_type="green_screen", max_resolution="720p", downsample_ratio=0.25, progress=None):
+    """
+    Process video to remove background using ONNX Runtime.
+
+    Args:
+        input_video: Path to input video
+        output_type: "green_screen" or "foreground_only"
+        max_resolution: Output video resolution (affects speed)
+        downsample_ratio: Internal model quality (affects edge accuracy)
+        progress: Gradio Progress tracker (optional, for web UI)
+
+    Returns:
+        Path to processed video
+    """
+    if input_video is None:
+        return None
+
+    # Get max height from resolution option
+    max_h = RESOLUTION_OPTIONS.get(max_resolution)
+
+    session = load_rvm_onnx()
+    output_names = [o.name for o in session.get_outputs()]
+
+    cap = cv2.VideoCapture(input_video)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    orig_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    orig_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    # Resize based on quality preset
+    if max_h and orig_height > max_h:
+        scale = max_h / orig_height
+        width = int(orig_width * scale)
+        height = max_h
+        print(f"Resizing: {orig_width}x{orig_height} -> {width}x{height}")
+    else:
+        width, height = orig_width, orig_height
+
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
+        output_path = tmp.name
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    green_bg = np.zeros((height, width, 3), dtype=np.uint8)
+    green_bg[:] = (0, 255, 0)
+
+    # Initialize recurrent states (zeros for first frame)
+    r1 = np.zeros((1, 1, 1, 1), dtype=np.float32)
+    r2 = np.zeros((1, 1, 1, 1), dtype=np.float32)
+    r3 = np.zeros((1, 1, 1, 1), dtype=np.float32)
+    r4 = np.zeros((1, 1, 1, 1), dtype=np.float32)
+
+    frame_idx = 0
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Resize if needed
+        if frame.shape[0] != height or frame.shape[1] != width:
+            frame = cv2.resize(frame, (width, height))
+
+        # Preprocess frame
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        src = frame_rgb.astype(np.float32) / 255.0
+        src = np.transpose(src, (2, 0, 1))[np.newaxis, ...]  # NCHW
+
+        # ONNX inference
+        outputs = session.run(output_names, {
+            'src': src,
+            'r1i': r1,
+            'r2i': r2,
+            'r3i': r3,
+            'r4i': r4,
+            'downsample_ratio': np.array([downsample_ratio], dtype=np.float32)
+        })
+
+        fgr, pha, r1, r2, r3, r4 = outputs
+
+        # Postprocess
+        fgr_np = np.transpose(fgr[0], (1, 2, 0))  # CHW -> HWC
+        pha_np = pha[0, 0]
+
+        if fgr_np.shape[:2] != (height, width):
+            fgr_np = cv2.resize(fgr_np, (width, height))
+            pha_np = cv2.resize(pha_np, (width, height))
+
+        pha_3ch = np.stack([pha_np] * 3, axis=-1)
+
+        if output_type == "green_screen":
+            result = (fgr_np * pha_3ch + (green_bg / 255.0) * (1 - pha_3ch))
+        else:
+            result = fgr_np * pha_3ch
+
+        result_bgr = cv2.cvtColor((result * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
+        out.write(result_bgr)
+
+        frame_idx += 1
+        if frame_count > 0:
+            if progress is not None:
+                progress(frame_idx / frame_count, desc=f"Processing frame {frame_idx}/{frame_count}")
+            elif frame_idx % 30 == 0:
+                print(f"Processing frame {frame_idx}/{frame_count}")
+
+    cap.release()
+    out.release()
+
+    return output_path
+
+
+# =============================================================================
+# Gradio Interface
+# =============================================================================
+
+# Initialize image models (lazy loading)
+bria_model = BriaRMBGModel()
+rmbg2_model = RMBG2Model()
+isnet_model = ISNetModel()
+
+IMAGE_MODELS = {
+    "BRIA-RMBG 1.4 (fast, general)": bria_model,
+    "BRIA-RMBG 2.0 (high quality, slow)": rmbg2_model,
+    "IS-Net/DIS (fine edges, hair)": isnet_model,
+}
+
+
+IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif', '.tiff'}
+VIDEO_EXTENSIONS = {'.mp4', '.avi', '.mov', '.mkv', '.webm', '.flv', '.wmv'}
+
+
+def detect_file_type(file_path):
+    """Detect if file is image or video based on extension."""
+    if file_path is None:
+        return None
+    ext = os.path.splitext(file_path)[1].lower()
+    if ext in IMAGE_EXTENSIONS:
+        return "image"
+    elif ext in VIDEO_EXTENSIONS:
+        return "video"
+    return None
+
+
+def on_file_upload(file_path):
+    """Handle file upload - detect type and update preview."""
+    if file_path is None:
+        return gr.Image(visible=False), gr.Video(visible=False), None
+
+    file_type = detect_file_type(file_path)
+    if file_type == "image":
+        return gr.Image(value=file_path, visible=True), gr.Video(visible=False), "image"
+    elif file_type == "video":
+        return gr.Image(visible=False), gr.Video(value=file_path, visible=True), "video"
+    return gr.Image(visible=False), gr.Video(visible=False), None
+
+
+def process_media(file_path, file_type, image_model_choice, output_type, max_resolution, edge_quality, progress=gr.Progress()):
+    """Process based on stored file type."""
+    if file_path is None:
+        return None, None
+
+    # Auto-detect file type if not provided (for API calls)
+    if file_type is None:
+        file_type = detect_file_type(file_path)
+
+    if file_type is None:
+        return None, None
+
+    if file_type == "image":
+        pil_image = Image.open(file_path)
+        model = IMAGE_MODELS.get(image_model_choice, bria_model)
+        output_img, mask = model.remove_background(pil_image)
+        return output_img, None
+
+    elif file_type == "video":
+        output_path = process_video(file_path, output_type, max_resolution, edge_quality, progress)
+        return None, output_path
+
+    return None, None
+
+
+def build_ui():
+    """Build compact unified Gradio interface."""
+    with gr.Blocks(title="Background Removal") as demo:
+        gr.Markdown("## Background Removal\nUpload image or video. Image models: RMBG 1.4 (fast) | RMBG 2.0 (quality) | IS-Net (fine edges) | Video: RVM")
+
+        # State to track file type
+        file_type_state = gr.State(value=None)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_file = gr.File(
+                    label="Upload Image or Video",
+                    file_types=["image", "video"],
+                    type="filepath"
+                )
+                # Preview - only one visible at a time
+                preview_image = gr.Image(label="Preview", interactive=False, visible=False)
+                preview_video = gr.Video(label="Preview", interactive=False, visible=False)
+
+                # Image model selector
+                image_model_dropdown = gr.Dropdown(
+                    choices=list(IMAGE_MODELS.keys()),
+                    value="BRIA-RMBG 1.4 (fast, general)",
+                    label="Image Model",
+                    info="1.4: fast | 2.0: quality (slow) | IS-Net: fine edges"
+                )
+
+                with gr.Accordion("Video Settings", open=False):
+                    output_type = gr.Radio(
+                        choices=["green_screen", "foreground_only"],
+                        value="green_screen",
+                        label="Output Type"
+                    )
+                    max_res = gr.Dropdown(
+                        choices=list(RESOLUTION_OPTIONS.keys()),
+                        value="720p",
+                        label="Output Resolution"
+                    )
+                    edge_quality = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        value=0.25,
+                        step=0.05,
+                        label="Edge Quality"
+                    )
+
+                process_btn = gr.Button("Remove Background", variant="primary")
+
+            with gr.Column(scale=1):
+                # Output components
+                output_image = gr.Image(label="Result (Image)", type="pil")
+                output_video = gr.Video(label="Result (Video)")
+
+        # Example files (click to populate inputs, then press button)
+        gr.Examples(
+            examples=[
+                ["examples/cat.jpg", "BRIA-RMBG 1.4 (fast, general)", "green_screen", "720p", 0.25],
+            ],
+            inputs=[input_file, image_model_dropdown, output_type, max_res, edge_quality],
+            label="Examples"
+        )
+
+        # On upload: detect type and show preview
+        input_file.change(
+            fn=on_file_upload,
+            inputs=[input_file],
+            outputs=[preview_image, preview_video, file_type_state]
+        )
+
+        # On process: fill the already-visible output
+        process_btn.click(
+            fn=process_media,
+            inputs=[input_file, file_type_state, image_model_dropdown, output_type, max_res, edge_quality],
+            outputs=[output_image, output_video]
+        )
+
+    return demo
+
+
+# =============================================================================
+# CLI Interface
+# =============================================================================
+
+def cli_main():
+    """CLI entry point."""
+    parser = argparse.ArgumentParser(description="Background Removal - Image and Video")
+    subparsers = parser.add_subparsers(dest="command", help="Commands")
+
+    # Image command
+    img_parser = subparsers.add_parser("image", help="Remove background from image")
+    img_parser.add_argument("input", type=str, help="Input image path")
+    img_parser.add_argument("--output", "-o", type=str, default="output.png", help="Output path")
+    img_parser.add_argument("--mask", "-m", type=str, help="Save mask to path")
+    img_parser.add_argument("--model", type=str, default="BRIA-RMBG 1.4 (fast, general)",
+                           choices=list(IMAGE_MODELS.keys()),
+                           help="Image model: BRIA (fast) or IS-Net (fine edges)")
+
+    # Video command
+    vid_parser = subparsers.add_parser("video", help="Remove background from video")
+    vid_parser.add_argument("input", type=str, help="Input video path")
+    vid_parser.add_argument("--output", "-o", type=str, default="output.mp4", help="Output path")
+    vid_parser.add_argument("--type", type=str, default="green_screen",
+                           choices=["green_screen", "foreground_only"])
+    vid_parser.add_argument("--downsample", type=float, default=0.25, help="Downsample ratio")
+
+    args = parser.parse_args()
+
+    if args.command == "image":
+        print(f"Processing image: {args.input}")
+        input_img = Image.open(args.input)
+        # Use selected model (default: BRIA)
+        model = IMAGE_MODELS.get(args.model, bria_model)
+        output_img, mask = model.remove_background(input_img)
+        output_img.save(args.output)
+        print(f"Saved: {args.output}")
+        if args.mask:
+            mask.save(args.mask)
+            print(f"Saved mask: {args.mask}")
+
+    elif args.command == "video":
+        print(f"Processing video: {args.input}")
+        output_path = process_video(args.input, args.type, args.downsample)
+        if output_path:
+            import shutil
+            shutil.move(output_path, args.output)
+            print(f"Saved: {args.output}")
+
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    if len(sys.argv) > 1 and sys.argv[1] in ["image", "video", "-h", "--help"]:
+        cli_main()
+    else:
+        demo = build_ui()
+        demo.launch(theme=gr.themes.Soft())