Restore lightweight warp package for HF Space runtime

.gitignore

@@ -0,0 +1,14 @@
+
+# keep warp lightweight
+warp/inference_outputs/
+warp/data/
+warp/**/checkpoints/
+warp/**/*.pt
+warp/**/*.pth
+warp/**/*.ckpt
+warp/**/*.safetensors
+warp/**/*.bin
+warp/**/*.npz
+*.zip
+*.7z
+*.tar

warp/gradio_app/__init__.py

@@ -0,0 +1 @@
+"""Gradio application for A360 WARP experimentation UI."""

warp/gradio_app/app-Nick.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+import gradio as gr
+from dotenv import load_dotenv
+
+if TYPE_CHECKING:
+    from supabase import Client as ClientType
+
+    from warp.data import ImageLoader as ImageLoaderType
+    from warp.gradio_app.models.upscaler import ImageUpscaler
+else:
+    ClientType = object
+    ImageLoaderType = object
+    ImageUpscaler = object
+
+try:
+    from supabase import Client, create_client
+except Exception:
+    create_client = None  # type: ignore
+    Client = None  # type: ignore
+
+try:
+    from warp.data import ImageLoader
+except ImportError:
+    ImageLoader = None  # type: ignore
+
+try:
+    from warp.gradio_app.models.upscaler import create_upscaler
+
+    UPSCALER_AVAILABLE = True
+except ImportError:
+    create_upscaler = None  # type: ignore
+    UPSCALER_AVAILABLE = False
+
+load_dotenv()
+SUPABASE_URL: str = os.getenv("SUPABASE_URL", "")
+SUPABASE_ANON_KEY: str = os.getenv("SUPABASE_ANON_KEY", "")
+
+supabase: ClientType | None = None
+if callable(create_client) and SUPABASE_URL and SUPABASE_ANON_KEY:
+    supabase = create_client(SUPABASE_URL, SUPABASE_ANON_KEY)
+
+# Initialize image loader
+image_loader: ImageLoaderType | None = None
+try:
+    if callable(ImageLoader):
+        image_loader = ImageLoader()
+        print(f"✓ Loaded {len(image_loader.practices)} practices with scraped images")
+except Exception as e:
+    print(f"Warning: Could not initialize ImageLoader: {e}")
+
+# Initialize upscaler (lazy load)
+upscaler: ImageUpscaler | None = None
+
+
+def load_practice_images(practice_name: str) -> tuple[list, str]:
+    """Load sample images from a practice.
+
+    Args:
+        practice_name: Name of the practice to load images from
+
+    Returns:
+        Tuple of (list of image paths, status message)
+    """
+    if not image_loader:
+        return [], "Image loader not available"
+
+    if not practice_name:
+        return [], "Please select a practice"
+
+    try:
+        # Get random sample of images
+        image_paths = image_loader.get_random_images(practice_name, n=10)
+        stats = image_loader.get_practice_stats(practice_name)
+        msg = (
+            f"Loaded {len(image_paths)} sample images from {practice_name} "
+            f"(Total: {stats['total_images']} images)"
+        )
+        return [str(p) for p in image_paths], msg
+    except Exception as e:
+        return [], f"Error loading images: {e}"
+
+
+def run_model(procedure: str | None, notes: str | None) -> str:
+    """Run a placeholder model execution.
+
+    Args:
+        procedure: The selected procedure type
+        notes: Additional context or parameters
+
+    Returns:
+        A formatted string with procedure and notes information
+    """
+    return f"Procedure={procedure or 'n/a'} | Notes={notes or 'n/a'}"
+
+
+def upscale_images(
+    before_img, after_img, prompt: str, num_steps: int, guidance: float
+) -> tuple[object, object, str]:
+    """Upscale before/after image pair.
+
+    Args:
+        before_img: Before image from Gradio
+        after_img: After image from Gradio
+        prompt: Quality prompt for upscaling
+        num_steps: Number of inference steps
+        guidance: Guidance scale
+
+    Returns:
+        Tuple of (upscaled_before, upscaled_after, status_message)
+    """
+    global upscaler
+
+    if not UPSCALER_AVAILABLE:
+        return (
+            None,
+            None,
+            "Upscaler not available. Install: pip install torch diffusers transformers",
+        )
+
+    if before_img is None or after_img is None:
+        return None, None, "Please upload both before and after images"
+
+    try:
+        # Lazy load upscaler
+        if upscaler is None and callable(create_upscaler):
+            upscaler = create_upscaler(model_type="sd-x4")
+
+        # Import PIL here to handle the images
+        from PIL import Image
+
+        # Convert Gradio images to PIL if needed
+        if not isinstance(before_img, Image.Image):
+            before_img = Image.fromarray(before_img)
+        if not isinstance(after_img, Image.Image):
+            after_img = Image.fromarray(after_img)
+
+        # Upscale the pair
+        before_upscaled, after_upscaled = upscaler.upscale_pair(
+            before_img,
+            after_img,
+            prompt=prompt,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance,
+        )
+
+        return (
+            before_upscaled,
+            after_upscaled,
+            f"✓ Successfully upscaled images 4x (Original: {before_img.size} → Upscaled: {before_upscaled.size})",
+        )
+
+    except Exception as e:
+        return None, None, f"Error during upscaling: {str(e)}"
+
+
+# Build the UI
+with gr.Blocks(title="A360 WARP — Gradio") as demo:
+    gr.Markdown("# A360 WARP — Experimentation UI (MVP)")
+    gr.Markdown("Load and experiment with before/after images from scraped medical practices.")
+
+    # Practice selection and image loading
+    with gr.Tab("Image Browser"):
+        with gr.Row():
+            practice_dropdown = gr.Dropdown(
+                label="Select Practice",
+                choices=image_loader.practices if image_loader else [],
+                value=None,
+            )
+            load_btn = gr.Button("Load Sample Images", variant="primary")
+
+        status_text = gr.Textbox(label="Status", interactive=False)
+        image_gallery = gr.Gallery(label="Sample Images", show_label=True, columns=5, height="auto")
+
+        load_btn.click(
+            fn=load_practice_images,
+            inputs=[practice_dropdown],
+            outputs=[image_gallery, status_text],
+        )
+
+    # Image Enhancement (Upscaling)
+    with gr.Tab("Image Enhancement"):
+        gr.Markdown(
+            "### Upscale Before/After Images\n"
+            "Upload medical before/after photos to upscale them 4x using AI. "
+            "This improves image quality and detail for better comparison."
+        )
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("#### Original Images")
+                before_input = gr.Image(label="Before Image", type="numpy")
+                after_input = gr.Image(label="After Image", type="numpy")
+
+            with gr.Column():
+                gr.Markdown("#### Upscaled Images (4x)")
+                before_output = gr.Image(label="Upscaled Before")
+                after_output = gr.Image(label="Upscaled After")
+
+        with gr.Row():
+            with gr.Column():
+                prompt_input = gr.Textbox(
+                    label="Quality Prompt",
+                    value="high quality medical photography, sharp details, professional lighting",
+                    placeholder="Describe desired image quality...",
+                )
+            with gr.Column():
+                num_steps = gr.Slider(
+                    minimum=20,
+                    maximum=100,
+                    value=50,
+                    step=5,
+                    label="Inference Steps (higher = better quality, slower)",
+                )
+                guidance_scale = gr.Slider(
+                    minimum=1.0, maximum=15.0, value=7.5, step=0.5, label="Guidance Scale"
+                )
+
+        upscale_btn = gr.Button("Upscale Images", variant="primary", size="lg")
+        upscale_status = gr.Textbox(label="Status", interactive=False)
+
+        upscale_btn.click(
+            fn=upscale_images,
+            inputs=[before_input, after_input, prompt_input, num_steps, guidance_scale],
+            outputs=[before_output, after_output, upscale_status],
+        )
+
+    # Model experimentation
+    with gr.Tab("Model Experiments"):
+        with gr.Row():
+            procedure = gr.Dropdown(
+                label="Procedure",
+                choices=[
+                    "breast-augmentation",
+                    "liposuction",
+                    "rhinoplasty",
+                    "ftm-top-surgery",
+                    "coolsculpting",
+                ],
+                value=None,
+            )
+            notes = gr.Textbox(label="Notes", placeholder="Run context / params…")
+        run = gr.Button("Run")
+        out = gr.Textbox(label="Output")
+
+        run.click(run_model, inputs=[procedure, notes], outputs=out)
+
+if __name__ == "__main__":
+    demo.launch()

warp/gradio_app/app.py

@@ -0,0 +1,359 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+import gradio as gr
+from dotenv import load_dotenv
+
+if TYPE_CHECKING:
+    from supabase import Client as ClientType
+
+    from warp.data import ImageLoader as ImageLoaderType
+    from warp.gradio_app.models.upscaler import ImageUpscaler
+else:
+    ClientType = object
+    ImageLoaderType = object
+    ImageUpscaler = object
+
+try:
+    from supabase import Client, create_client
+except Exception:
+    create_client = None  # type: ignore
+    Client = None  # type: ignore
+
+try:
+    from warp.data import ImageLoader
+except ImportError:
+    ImageLoader = None  # type: ignore
+
+try:
+    from warp.gradio_app.models.upscaler import create_upscaler
+
+    UPSCALER_AVAILABLE = True
+    print("✓ Upscaler module loaded successfully")
+except ImportError as e:
+    create_upscaler = None  # type: ignore
+    UPSCALER_AVAILABLE = False
+    print(f"✗ Upscaler import failed: {e}")
+
+# Temporarily disable Advanced Upscaling tab due to import issues
+# Will re-enable after fixing module resolution
+COMPARE_TAB_AVAILABLE = False
+build_upscale_compare = None
+
+load_dotenv()
+SUPABASE_URL: str = os.getenv("SUPABASE_URL", "")
+SUPABASE_ANON_KEY: str = os.getenv("SUPABASE_ANON_KEY", "")
+
+supabase: ClientType | None = None
+if callable(create_client) and SUPABASE_URL and SUPABASE_ANON_KEY:
+    supabase = create_client(SUPABASE_URL, SUPABASE_ANON_KEY)
+
+# Initialize image loader
+image_loader: ImageLoaderType | None = None
+try:
+    if callable(ImageLoader):
+        image_loader = ImageLoader()
+        print(f"✓ Loaded {len(image_loader.practices)} practices with scraped images")
+except Exception as e:
+    # If initialization fails for any reason (including mocked import errors),
+    # fall back to no image loader so the rest of the app can still import.
+    image_loader = None
+    print(f"Warning: Could not initialize ImageLoader: {e}")
+
+# Initialize upscaler (lazy load)
+upscaler: ImageUpscaler | None = None
+
+
+def load_practice_images(practice_name: str) -> tuple[list, str]:
+    """Load sample images from a practice.
+
+    Args:
+        practice_name: Name of the practice to load images from
+
+    Returns:
+        Tuple of (list of image paths, status message)
+    """
+    if not image_loader:
+        return [], "Image loader not available"
+
+    if not practice_name:
+        return [], "Please select a practice"
+
+    try:
+        # Get random sample of images
+        image_paths = image_loader.get_random_images(practice_name, n=10)
+        stats = image_loader.get_practice_stats(practice_name)
+        msg = (
+            f"Loaded {len(image_paths)} sample images from {practice_name} "
+            f"(Total: {stats['total_images']} images)"
+        )
+        return [str(p) for p in image_paths], msg
+    except Exception as e:
+        return [], f"Error loading images: {e}"
+
+
+def run_model(procedure: str | None, notes: str | None) -> str:
+    """Run a placeholder model execution.
+
+    Args:
+        procedure: The selected procedure type
+        notes: Additional context or parameters
+
+    Returns:
+        A formatted string with procedure and notes information
+    """
+    return f"Procedure={procedure or 'n/a'} | Notes={notes or 'n/a'}"
+
+
+def upscale_images(
+    before_img, after_img, prompt: str, num_steps: int, guidance: float, progress=gr.Progress()
+):
+    """Upscale before/after image pair (synchronous helper).
+
+    This version is a simple function (not a generator) so tests can call it
+    and make assertions about the returned tuple. The Gradio UI wraps this in
+    a streaming function that updates the progress bar and status text.
+
+    Returns:
+        Tuple of (upscaled_before, upscaled_after, status_message)
+    """
+    global upscaler
+
+    # Handle missing upscaler dependency
+    if not UPSCALER_AVAILABLE:
+        return (
+            None,
+            None,
+            "Upscaler not available. Install: pip install torch diffusers transformers",
+        )
+
+    # Validate inputs
+    if before_img is None or after_img is None:
+        return None, None, "Please upload both before and after images"
+
+    try:
+        # Lazy load upscaler on first use
+        if upscaler is None and callable(create_upscaler):
+            upscaler = create_upscaler(model_type="sd-x4")
+
+        # Import PIL here to handle the images
+        from PIL import Image
+
+        # Convert numpy arrays to PIL Images if needed
+        if not isinstance(before_img, Image.Image):
+            before_img = Image.fromarray(before_img)
+        if not isinstance(after_img, Image.Image):
+            after_img = Image.fromarray(after_img)
+
+        orig_size = before_img.size
+
+        # Use the pair upscaling helper with a callback that updates the
+        # Gradio progress bar more granularly during diffusion steps.
+        callback_state = {"phase": "before", "last_step": -1}
+
+        def progress_callback(step, timestep, latents):  # type: ignore[unused-argument]
+            """Update progress bar for each diffusion step.
+
+            We see steps 0..num_steps-1 for the "before" image first, then
+            again for the "after" image. When the step counter resets, we
+            switch to the "after" phase and map progress into [0.5, 0.9].
+            """
+            try:
+                # Detect phase change when step counter resets
+                if step < callback_state["last_step"]:
+                    callback_state["phase"] = "after"
+                callback_state["last_step"] = step
+
+                frac = step / max(num_steps, 1)
+                if callback_state["phase"] == "before":
+                    # Map to [0.1, 0.5]
+                    pct = 0.1 + 0.4 * frac
+                    desc = f"Upscaling BEFORE image ({step}/{num_steps})"
+                else:
+                    # Map to [0.5, 0.9]
+                    pct = 0.5 + 0.4 * frac
+                    desc = f"Upscaling AFTER image ({step}/{num_steps})"
+
+                try:
+                    progress(pct, desc=desc)
+                except Exception:
+                    # In tests or non-Gradio contexts, progress may be a no-op
+                    pass
+            except Exception:
+                # Never allow progress UI issues to break the core upscaling
+                pass
+
+        before_upscaled, after_upscaled = upscaler.upscale_pair(
+            before_img,
+            after_img,
+            prompt=prompt,
+            num_inference_steps=num_steps,
+            guidance_scale=guidance,
+            callback=progress_callback,
+            callback_steps=1,
+        )
+
+        status = (
+            "Successfully upscaled both images 4x\n"
+            f"Original: {orig_size[0]}×{orig_size[1]} → "
+            f"Upscaled: {before_upscaled.size[0]}×{before_upscaled.size[1]}"
+        )
+        return before_upscaled, after_upscaled, status
+
+    except Exception as e:
+        # Graceful error handling for tests and UI
+        return None, None, f"Error during upscaling: {str(e)}"
+
+
+def upscale_images_stream(
+    before_img, after_img, prompt: str, num_steps: int, guidance: float, progress=gr.Progress()
+):
+    """Streaming wrapper for ``upscale_images`` used by the Gradio UI.
+
+    Yields intermediate status updates so the user sees a live progress bar
+    and status text while the heavy model runs.
+    """
+    # Handle missing upscaler dependency
+    if not UPSCALER_AVAILABLE:
+        yield (
+            None,
+            None,
+            "Upscaler not available. Install: pip install torch diffusers transformers",
+        )
+        return
+
+    # Validate inputs
+    if before_img is None or after_img is None:
+        yield None, None, "Please upload both before and after images"
+        return
+
+    try:
+        # Initial progress
+        try:
+            progress(0.0, desc="Initializing upscaler...")
+        except Exception:
+            pass
+        yield None, None, "Initializing upscaler..."
+
+        # Coarse progress while running the model
+        try:
+            progress(0.3, desc="Upscaling images...")
+        except Exception:
+            pass
+
+        before_upscaled, after_upscaled, status = upscale_images(
+            before_img, after_img, prompt, num_steps, guidance, progress
+        )
+
+        try:
+            progress(1.0, desc="Complete")
+        except Exception:
+            pass
+
+        yield before_upscaled, after_upscaled, status
+
+    except Exception as e:
+        yield None, None, f"Error during upscaling: {str(e)}"
+
+
+# Build the UI
+with gr.Blocks(title="A360 WARP — Gradio") as demo:
+    gr.Markdown("# A360 WARP — Experimentation UI (MVP)")
+    gr.Markdown("Load and experiment with before/after images from scraped medical practices.")
+
+    # Practice selection and image loading
+    with gr.Tab("Image Browser"):
+        with gr.Row():
+            practice_dropdown = gr.Dropdown(
+                label="Select Practice",
+                choices=image_loader.practices if image_loader else [],
+                value=None,
+            )
+            load_btn = gr.Button("Load Sample Images", variant="primary")
+
+        status_text = gr.Textbox(label="Status", interactive=False)
+        image_gallery = gr.Gallery(label="Sample Images", show_label=True, columns=5, height="auto")
+
+        load_btn.click(
+            fn=load_practice_images,
+            inputs=[practice_dropdown],
+            outputs=[image_gallery, status_text],
+        )
+
+    # Image Enhancement (Upscaling)
+    with gr.Tab("Image Enhancement"):
+        gr.Markdown(
+            "### Upscale Before/After Images\n"
+            "Upload medical before/after photos to upscale them 4x using AI. "
+            "This improves image quality and detail for better comparison."
+        )
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("#### Original Images")
+                before_input = gr.Image(label="Before Image", type="numpy")
+                after_input = gr.Image(label="After Image", type="numpy")
+
+            with gr.Column():
+                gr.Markdown("#### Upscaled Images (4x)")
+                before_output = gr.Image(label="Upscaled Before")
+                after_output = gr.Image(label="Upscaled After")
+
+        with gr.Row():
+            with gr.Column():
+                prompt_input = gr.Textbox(
+                    label="Quality Prompt",
+                    value="high quality medical photography, sharp details, professional lighting",
+                    placeholder="Describe desired image quality...",
+                )
+            with gr.Column():
+                num_steps = gr.Slider(
+                    minimum=20,
+                    maximum=100,
+                    value=50,
+                    step=5,
+                    label="Inference Steps (higher = better quality, slower)",
+                )
+                guidance_scale = gr.Slider(
+                    minimum=1.0, maximum=15.0, value=7.5, step=0.5, label="Guidance Scale"
+                )
+
+        upscale_btn = gr.Button("Upscale Images", variant="primary", size="lg")
+        upscale_status = gr.Textbox(label="Status", interactive=False)
+
+        # Use the streaming wrapper so users see live progress/status updates
+        upscale_btn.click(
+            fn=upscale_images_stream,
+            inputs=[before_input, after_input, prompt_input, num_steps, guidance_scale],
+            outputs=[before_output, after_output, upscale_status],
+        )
+
+    # Advanced Upscaling with Comparison
+    if COMPARE_TAB_AVAILABLE and build_upscale_compare:
+        with gr.Tab("Advanced Upscaling"):
+            build_upscale_compare()
+
+    # Model experimentation
+    with gr.Tab("Model Experiments"):
+        with gr.Row():
+            procedure = gr.Dropdown(
+                label="Procedure",
+                choices=[
+                    "breast-augmentation",
+                    "liposuction",
+                    "rhinoplasty",
+                    "ftm-top-surgery",
+                    "coolsculpting",
+                ],
+                value=None,
+            )
+            notes = gr.Textbox(label="Notes", placeholder="Run context / params…")
+        run = gr.Button("Run")
+        out = gr.Textbox(label="Output")
+
+        run.click(run_model, inputs=[procedure, notes], outputs=out)
+
+if __name__ == "__main__":
+    demo.launch()

warp/gradio_app/model_comparison.py

@@ -0,0 +1,288 @@
+"""
+Model Comparison App - Gradio interface for testing and comparing background removal models.
+"""
+
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import gradio as gr
+import pandas as pd
+from dotenv import load_dotenv
+from PIL import Image
+
+# Load environment variables from .env file
+env_path = Path(__file__).parent.parent.parent / ".env"
+load_dotenv(env_path)
+
+from warp.data import ImageLoader
+from warp.inference.background_removal import BackgroundRemovalEngine, BackgroundRemovalResult
+from warp.models import list_model_names
+
+# Initialize components
+engine = BackgroundRemovalEngine()
+loader = ImageLoader()
+
+
+def load_test_image(practice: str) -> Optional[Image.Image]:
+    """Load a random test image from a practice."""
+    try:
+        images = loader.get_random_images(practice, n=1)
+        if images:
+            return loader.load_image(images[0])
+    except Exception as e:
+        print(f"Error loading image: {e}")
+    return None
+
+
+def run_comparison(
+    image: Image.Image, selected_models: List[str]
+) -> Tuple[Dict, pd.DataFrame, str]:
+    """
+    Run background removal comparison across selected models.
+
+    Args:
+        image: Input PIL Image
+        selected_models: List of model names to test
+
+    Returns:
+        Tuple of (results_dict, metrics_df, summary_text)
+    """
+    if not image:
+        return {}, pd.DataFrame(), "❌ No image provided"
+
+    if not selected_models:
+        return {}, pd.DataFrame(), "❌ No models selected"
+
+    print(f"\n{'='*60}")
+    print(f"Running comparison with {len(selected_models)} models...")
+    print(f"{'='*60}")
+
+    # Run comparisons
+    results = {}
+    for model_name in selected_models:
+        result = engine.remove_background(image, model_name)
+        results[model_name] = result
+
+    # Prepare outputs for gallery (list of tuples with image and caption)
+    output_images = [(res.output_image, f"{name}\n{res.processing_time_ms}ms") 
+                     for name, res in results.items() if res.success]
+
+    # Create metrics DataFrame
+    metrics_data = []
+    for name, result in results.items():
+        metrics_data.append(
+            {
+                "Model": name,
+                "Status": "✓ Success" if result.success else "✗ Failed",
+                "Time (ms)": result.processing_time_ms,
+                "Edge Quality": f"{result.edge_quality:.3f}" if result.edge_quality else "-",
+                "SSIM": f"{result.ssim:.3f}" if result.ssim else "-",
+                "PSNR (dB)": f"{result.psnr:.1f}" if result.psnr else "-",
+                "Transparency": f"{result.transparency_coverage:.1%}" if result.transparency_coverage else "-",
+                "Quality Score": f"{result.weighted_quality_score:.3f}" if result.weighted_quality_score else "-",
+                "Error": result.error_message or "-",
+            }
+        )
+
+    metrics_df = pd.DataFrame(metrics_data)
+
+    # Create summary
+    successful = sum(1 for r in results.values() if r.success)
+    avg_time = (
+        sum(r.processing_time_ms for r in results.values() if r.success) / successful
+        if successful > 0
+        else 0
+    )
+
+    summary = f"""
+## Comparison Summary
+
+- **Models Tested**: {len(selected_models)}
+- **Successful**: {successful}
+- **Failed**: {len(selected_models) - successful}
+- **Average Time**: {avg_time:.0f}ms
+    """.strip()
+
+    return output_images, metrics_df, summary
+
+
+def run_single_model(
+    image: Image.Image, model_name: str
+) -> Tuple[Optional[Image.Image], str]:
+    """Run a single model and return output + info."""
+    if not image:
+        return None, "❌ No image provided"
+
+    result = engine.remove_background(image, model_name)
+
+    if result.success:
+        info = f"""
+### ✓ Success
+
+- **Model**: {result.model_name}
+- **Processing Time**: {result.processing_time_ms}ms
+- **Input Size**: {result.input_size[0]}x{result.input_size[1]}
+- **Output Size**: {result.output_size[0]}x{result.output_size[1]}
+        """.strip()
+        return result.output_image, info
+    else:
+        info = f"""
+### ✗ Failed
+
+- **Model**: {result.model_name}
+- **Error**: {result.error_message}
+- **Time Elapsed**: {result.processing_time_ms}ms
+        """.strip()
+        return image, info
+
+
+# ============================================================================
+# Gradio Interface
+# ============================================================================
+
+
+def create_comparison_tab():
+    """Create the model comparison tab."""
+    with gr.Column():
+        gr.Markdown("# 🔬 Model Comparison")
+        gr.Markdown("Compare multiple background removal models side-by-side")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(
+                    type="pil", label="Input Image", height=400
+                )
+
+                # Model selection
+                available_models = list_model_names("bg_removal")
+                model_selector = gr.CheckboxGroup(
+                    choices=available_models,
+                    value=[available_models[0]] if available_models else [],
+                    label="Select Models to Compare",
+                )
+
+                # Quick load options
+                with gr.Row():
+                    practice_dropdown = gr.Dropdown(
+                        choices=loader.practices if hasattr(loader, "practices") else [],
+                        label="Load Random Image From",
+                        value=None,
+                    )
+                    load_btn = gr.Button("📁 Load Sample", size="sm")
+
+                run_btn = gr.Button("▶️ Run Comparison", variant="primary", size="lg")
+
+            with gr.Column(scale=2):
+                summary_md = gr.Markdown("### Ready to compare models")
+                metrics_table = gr.DataFrame(label="Performance Metrics")
+
+        # Output gallery
+        gr.Markdown("### Model Outputs")
+        output_gallery = gr.Gallery(
+            label="Results", columns=3, height="auto", object_fit="contain"
+        )
+
+        # Wire up events
+        def load_sample(practice):
+            if practice:
+                img = load_test_image(practice)
+                return img
+            return None
+
+        load_btn.click(fn=load_sample, inputs=[practice_dropdown], outputs=[input_image])
+
+        run_btn.click(
+            fn=run_comparison,
+            inputs=[input_image, model_selector],
+            outputs=[output_gallery, metrics_table, summary_md],
+        )
+
+
+def create_single_model_tab():
+    """Create the single model testing tab."""
+    with gr.Column():
+        gr.Markdown("# 🎯 Single Model Test")
+        gr.Markdown("Test a single model with detailed results")
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(type="pil", label="Input Image")
+
+                model_dropdown = gr.Dropdown(
+                    choices=list_model_names("bg_removal"),
+                    value=list_model_names("bg_removal")[0],
+                    label="Select Model",
+                )
+
+                run_btn = gr.Button("▶️ Remove Background", variant="primary")
+
+            with gr.Column(scale=1):
+                output_image = gr.Image(type="pil", label="Output Image")
+                result_info = gr.Markdown("### Waiting for input...")
+
+        run_btn.click(
+            fn=run_single_model,
+            inputs=[input_image, model_dropdown],
+            outputs=[output_image, result_info],
+        )
+
+
+def create_app():
+    """Create the full Gradio app."""
+    with gr.Blocks(title="WARP Model Comparison", theme=gr.themes.Soft()) as app:
+        gr.Markdown(
+            """
+        # 🚀 WARP Model Test Harness
+        
+        **AI-Powered Image Processing Pipeline**
+        
+        Test and compare background removal models with real-time performance metrics.
+        """
+        )
+
+        with gr.Tabs():
+            with gr.Tab("Model Comparison"):
+                create_comparison_tab()
+
+            with gr.Tab("Single Model Test"):
+                create_single_model_tab()
+
+            with gr.Tab("Model Registry"):
+                gr.Markdown("## 📚 Available Models")
+
+                # Display model registry
+                from warp.models import get_models_by_type
+
+                bg_models = get_models_by_type("bg_removal")
+                model_info_data = []
+
+                for name, config in bg_models.items():
+                    model_info_data.append(
+                        {
+                            "Name": config.name,
+                            "Display Name": config.display_name,
+                            "Model ID": config.model_id,
+                            "Est. Time": f"{config.estimated_time_ms}ms",
+                            "Default": "✓" if config.is_default else "",
+                            "Description": config.description,
+                        }
+                    )
+
+                gr.DataFrame(value=pd.DataFrame(model_info_data), label="Background Removal Models")
+
+    return app
+
+
+# ============================================================================
+# Main Entry Point
+# ============================================================================
+
+if __name__ == "__main__":
+    app = create_app()
+    app.launch(
+        server_name="127.0.0.1",  # Bind to localhost so the browser URL is valid on Windows
+        server_port=7860,
+        share=False,
+        show_error=True,
+    )

warp/gradio_app/models/__init__.py

@@ -0,0 +1,5 @@
+"""Model registry for A360 WARP."""
+
+from .registry import MODELS
+
+__all__ = ["MODELS"]

warp/gradio_app/models/registry-Nick.py

@@ -0,0 +1,9 @@
+MODELS = {
+    "CLIP": "openai/clip-vit-base-patch32",
+    "BLIP-2": "Salesforce/blip2-flan-t5-xl",
+    "DINOv2": "facebook/dinov2-base",
+}
+
+UPSCALER_MODELS = {
+    "SD-X4-Upscaler": "stabilityai/stable-diffusion-x4-upscaler",
+}

warp/gradio_app/models/registry.py

@@ -0,0 +1,5 @@
+MODELS = {
+    "CLIP": "openai/clip-vit-base-patch32",
+    "BLIP-2": "Salesforce/blip2-flan-t5-xl",
+    "DINOv2": "facebook/dinov2-base",
+}

warp/gradio_app/models/upscaler.py

@@ -0,0 +1,211 @@
+"""Image upscaling using HuggingFace models."""
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+from PIL import Image
+
+if TYPE_CHECKING:
+    from diffusers import StableDiffusionUpscalePipeline as PipelineType
+else:
+    PipelineType = object
+
+try:
+    import torch
+    from diffusers import StableDiffusionUpscalePipeline
+
+    TORCH_AVAILABLE = True
+except ImportError:
+    # If either torch or diffusers is missing, mark them unavailable. The
+    # tests patch these symbols as needed, and the runtime gracefully degrades
+    # by raising a clear ImportError from ImageUpscaler.__init__.
+    torch = None  # type: ignore[assignment]
+    StableDiffusionUpscalePipeline = None  # type: ignore[assignment]
+    TORCH_AVAILABLE = False
+
+
+class ImageUpscaler:
+    """Handle image upscaling using HuggingFace models."""
+
+    def __init__(
+        self, model_id: str = "stabilityai/stable-diffusion-x4-upscaler", device: str | None = None
+    ):
+        """Initialize the upscaler.
+
+        Args:
+            model_id: HuggingFace model identifier
+            device: Device to run model on ('cuda', 'cpu', or None for auto)
+        """
+        if not TORCH_AVAILABLE:
+            raise ImportError(
+                "torch and diffusers are required for upscaling. "
+                "Install with: pip install torch diffusers transformers"
+            )
+
+        self.model_id = model_id
+
+        # Auto-detect device
+        if device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = device
+
+        self.pipeline: PipelineType | None = None
+        self._load_model()
+
+    def _load_model(self) -> None:
+        """Load the upscaling model."""
+        print(f"Loading upscaler model: {self.model_id} on {self.device}...")
+
+        # Determine torch dtype based on device
+        torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
+
+        self.pipeline = StableDiffusionUpscalePipeline.from_pretrained(
+            self.model_id, torch_dtype=torch_dtype
+        )
+        self.pipeline = self.pipeline.to(self.device)
+
+        # Enable memory optimizations if on CUDA
+        if self.device == "cuda":
+            self.pipeline.enable_attention_slicing()
+
+        print(f"✓ Model loaded successfully on {self.device}")
+
+    def upscale(
+        self,
+        image: Image.Image | str | Path,
+        prompt: str = "high quality, detailed, sharp",
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        callback=None,
+        callback_steps: int = 1,
+    ) -> Image.Image:
+        """Upscale an image 4x.
+
+        Args:
+            image: PIL Image or path to image file
+            prompt: Text prompt to guide upscaling (helps with quality)
+            num_inference_steps: Number of denoising steps (higher = better quality, slower)
+            guidance_scale: How closely to follow the prompt (7.5 is good default)
+            callback: Optional callback function(step, timestep, latents) called each step
+            callback_steps: How often to call the callback (default: every step)
+
+        Returns:
+            Upscaled PIL Image
+        """
+        # Load image if path is provided
+        if isinstance(image, (str, Path)):
+            image = Image.open(image).convert("RGB")
+
+        # Ensure RGB mode
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        # Run upscaling
+        if self.pipeline is None:
+            raise RuntimeError("Pipeline not initialized")
+        result = self.pipeline(
+            prompt=prompt,
+            image=image,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            callback=callback,
+            callback_steps=callback_steps,
+        )
+        upscaled: Image.Image = result.images[0]
+
+        return upscaled
+
+    def upscale_pair(
+        self,
+        before_image: Image.Image | str | Path,
+        after_image: Image.Image | str | Path,
+        prompt: str = "high quality medical photography, sharp details, professional lighting",
+        **kwargs,
+    ) -> tuple[Image.Image, Image.Image]:
+        """Upscale a before/after image pair.
+
+        Args:
+            before_image: Before image (PIL Image or path)
+            after_image: After image (PIL Image or path)
+            prompt: Text prompt for upscaling quality
+            **kwargs: Additional arguments for upscale()
+
+        Returns:
+            Tuple of (upscaled_before, upscaled_after)
+        """
+        print("Upscaling before image...")
+        before_upscaled = self.upscale(before_image, prompt=prompt, **kwargs)
+
+        print("Upscaling after image...")
+        after_upscaled = self.upscale(after_image, prompt=prompt, **kwargs)
+
+        return before_upscaled, after_upscaled
+
+    def batch_upscale(
+        self,
+        images: list[Image.Image | str | Path],
+        prompt: str = "high quality, detailed, sharp",
+        **kwargs,
+    ) -> list[Image.Image]:
+        """Upscale multiple images.
+
+        Args:
+            images: List of PIL Images or paths
+            prompt: Text prompt for upscaling
+            **kwargs: Additional arguments for upscale()
+
+        Returns:
+            List of upscaled PIL Images
+        """
+        results = []
+        for i, img in enumerate(images, 1):
+            print(f"Upscaling image {i}/{len(images)}...")
+            upscaled = self.upscale(img, prompt=prompt, **kwargs)
+            results.append(upscaled)
+        return results
+
+
+def create_upscaler(
+    model_type: Literal["sd-x4", "fast"] = "sd-x4", device: str | None = None
+) -> ImageUpscaler:
+    """Factory function to create an upscaler.
+
+    Args:
+        model_type: Type of upscaler model
+            - "sd-x4": Stable Diffusion 4x upscaler (high quality, slower)
+            - "fast": Faster alternative (to be implemented)
+        device: Device to run on ('cuda', 'cpu', or None for auto)
+
+    Returns:
+        Initialized ImageUpscaler
+    """
+    model_map = {
+        "sd-x4": "stabilityai/stable-diffusion-x4-upscaler",
+        # Can add more models here later
+    }
+
+    model_id = model_map.get(model_type, model_map["sd-x4"])
+    return ImageUpscaler(model_id=model_id, device=device)
+
+
+# NOTE:
+# -----
+# When this module is imported as a submodule of ``warp.gradio_app.models``
+# (e.g. via ``from warp.gradio_app.models import upscaler``), Python normally
+# caches it as an attribute on the parent package. That caching can interfere
+# with tests that manipulate ``sys.modules`` to simulate import failures
+# (like removing ``torch``/``diffusers`` and re-importing this module).
+#
+# To ensure those tests can reliably exercise the fallback path, we avoid
+# permanently caching this submodule on the parent package by removing the
+# attribute if it exists. The module itself remains available via
+# ``sys.modules['warp.gradio_app.models.upscaler']``.
+try:  # Best-effort; never fail import because of this cleanup.
+    import sys as _sys
+
+    _parent_pkg = _sys.modules.get("warp.gradio_app.models")
+    if _parent_pkg is not None and hasattr(_parent_pkg, "upscaler"):
+        delattr(_parent_pkg, "upscaler")
+except Exception:
+    pass

warp/gradio_app/upscale_compare_tab.py

@@ -0,0 +1,231 @@
+"""Advanced upscaling tab with before/after comparison and detailed metrics."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import gradio as gr
+import numpy as np
+from PIL import Image, ImageDraw
+
+if TYPE_CHECKING:
+    from .models.upscaler import ImageUpscaler
+
+try:
+    from .models.upscaler import create_upscaler
+
+    UPSCALER_AVAILABLE = True
+except (ImportError, ModuleNotFoundError):
+    create_upscaler = None  # type: ignore[assignment]
+    UPSCALER_AVAILABLE = False
+
+
+# Global upscaler instance (lazy load)
+_upscaler: ImageUpscaler | None = None
+
+# Configuration
+MAX_INPUT_WIDTH = 1024
+MAX_INPUT_HEIGHT = 1024
+UPSCALE_FACTOR = 4
+QUALITY_PROMPT = "ultra realistic, natural contrast, high clarity, clean skin texture, professional medical photography"
+
+
+def _get_upscaler() -> ImageUpscaler | None:
+    """Lazy load upscaler on first use."""
+    global _upscaler
+    if _upscaler is None and UPSCALER_AVAILABLE:
+        try:
+            _upscaler = create_upscaler(model_type="sd-x4")
+        except Exception as e:
+            raise RuntimeError(f"Failed to load upscaler: {e}") from e
+    return _upscaler
+
+
+def _validate_and_resize_image(
+    img: np.ndarray | Image.Image, max_w: int = MAX_INPUT_WIDTH, max_h: int = MAX_INPUT_HEIGHT
+) -> Image.Image:
+    """Convert and validate image, resize if needed."""
+    # Convert numpy to PIL if needed
+    if isinstance(img, np.ndarray):
+        img = Image.fromarray(img.astype("uint8"))
+    elif not isinstance(img, Image.Image):
+        raise ValueError(f"Invalid image type: {type(img)}")
+
+    # Convert to RGB
+    if img.mode != "RGB":
+        img = img.convert("RGB")
+
+    # Check and resize if too large
+    w, h = img.size
+    if w > max_w or h > max_h:
+        img.thumbnail((max_w, max_h), Image.Resampling.LANCZOS)
+
+    return img
+
+
+def _create_comparison_grid(
+    before_orig: Image.Image, after_orig: Image.Image, before_up: Image.Image, after_up: Image.Image
+) -> Image.Image:
+    """Create a 2x2 grid showing before/after, original/upscaled."""
+
+    # All upscaled images should be the same size (4x original)
+    # For display, we'll scale them down to fit alongside originals
+
+    orig_w, orig_h = before_orig.size
+    up_display_w, up_display_h = before_up.size  # Should be 4x larger
+
+    # Create display versions of upscaled (scaled down slightly for display)
+    display_scale = 0.5  # Show upscaled at 2x (half of 4x)
+    display_w = int(up_display_w * display_scale)
+    display_h = int(up_display_h * display_scale)
+
+    before_up_display = before_up.resize((display_w, display_h), Image.Resampling.LANCZOS)
+    after_up_display = after_up.resize((display_w, display_h), Image.Resampling.LANCZOS)
+
+    # Create grid background
+    grid_w = display_w * 2 + 40  # padding
+    grid_h = display_h * 2 + 80  # padding + title space
+
+    grid = Image.new("RGB", (grid_w, grid_h), color=(30, 30, 30))
+    draw = ImageDraw.Draw(grid)
+
+    # Add labels (simple text, no font to avoid system dependencies)
+    label_y = 10
+    draw.text((10, label_y), "BEFORE (Orig → Upscaled 2x)", fill=(255, 150, 0))
+    draw.text((display_w + 20, label_y), "AFTER (Orig → Upscaled 2x)", fill=(255, 150, 0))
+
+    # Paste images
+    paste_y = 40
+    grid.paste(before_orig, (10, paste_y))
+    grid.paste(before_up_display, (10, paste_y + orig_h + 10))
+
+    grid.paste(after_orig, (display_w + 20, paste_y))
+    grid.paste(after_up_display, (display_w + 20, paste_y + orig_h + 10))
+
+    return grid
+
+
+def upscale_and_compare(
+    before_img, after_img, prompt: str = QUALITY_PROMPT, num_steps: int = 50, guidance: float = 7.5
+) -> tuple[Image.Image, Image.Image, Image.Image, str]:
+    """Upscale before/after pair with detailed comparison.
+
+    Args:
+        before_img: Before image (numpy or PIL)
+        after_img: After image (numpy or PIL)
+        prompt: Quality prompt for upscaling
+        num_steps: Inference steps (20-100)
+        guidance: Guidance scale (1.0-15.0)
+
+    Returns:
+        Tuple of (before_upscaled, after_upscaled, comparison_grid, status_message)
+    """
+
+    if not UPSCALER_AVAILABLE:
+        return (
+            None,
+            None,
+            None,
+            "❌ Upscaler not available. Install: pip install torch diffusers transformers",
+        )  # type: ignore[return-value]
+
+    if before_img is None or after_img is None:
+        return None, None, None, "❌ Please upload both before and after images"  # type: ignore[return-value]
+
+    try:
+        # Get upscaler instance
+        upscaler = _get_upscaler()
+        if upscaler is None:
+            return None, None, None, "❌ Upscaler not available"  # type: ignore[return-value]
+
+        # Validate and resize inputs
+        before_pil = _validate_and_resize_image(before_img)
+        after_pil = _validate_and_resize_image(after_img)
+
+        orig_before_size = before_pil.size
+        orig_after_size = after_pil.size
+
+        # Upscale both images
+        print(f"Upscaling before image ({orig_before_size})...")
+        before_upscaled = upscaler.upscale(
+            before_pil, prompt=prompt, num_inference_steps=num_steps, guidance_scale=guidance
+        )
+
+        print(f"Upscaling after image ({orig_after_size})...")
+        after_upscaled = upscaler.upscale(
+            after_pil, prompt=prompt, num_inference_steps=num_steps, guidance_scale=guidance
+        )
+
+        # Create comparison grid
+        comparison = _create_comparison_grid(before_pil, after_pil, before_upscaled, after_upscaled)
+
+        # Build status message
+        status = (
+            f"✅ Successfully upscaled both images!\n\n"
+            f"Before: {orig_before_size} → {before_upscaled.size}\n"
+            f"After: {orig_after_size} → {after_upscaled.size}\n"
+            f"Upscale Factor: {UPSCALE_FACTOR}x\n"
+            f"Steps: {num_steps} | Guidance: {guidance}\n"
+            f"\nNote: Comparison shows upscaled images at 2x (50% of 4x for display)"
+        )
+
+        return before_upscaled, after_upscaled, comparison, status
+
+    except Exception as e:
+        error_msg = f"❌ Error during upscaling: {str(e)}"
+        print(error_msg)
+        return None, None, None, error_msg  # type: ignore[return-value]
+
+
+def build_ui() -> None:
+    """Build the advanced upscaling UI tab."""
+
+    gr.Markdown("### Upscale Before/After Images (Max Detail & Clarity)")
+    gr.Markdown(
+        "Upload medical before/after photos to upscale them 4x using Stable Diffusion x4 Upscaler. "
+        "Both images are processed with identical parameters for fair comparison.\n\n"
+        f"⚠️ **Note:** Processing takes 30-60 seconds per image (CPU) or 5-10 seconds (GPU). "
+        f"Maximum input size: {MAX_INPUT_WIDTH}x{MAX_INPUT_HEIGHT}px (automatically resized if larger)."
+    )
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("#### Original Images")
+            before_input = gr.Image(label="Before Image", type="numpy")
+            after_input = gr.Image(label="After Image", type="numpy")
+
+            # Parameters
+            prompt_input = gr.Textbox(
+                label="Quality Prompt",
+                value=QUALITY_PROMPT,
+                placeholder="Describe desired image quality...",
+                lines=3,
+            )
+
+        with gr.Column():
+            gr.Markdown("#### Upscaled Results (4x)")
+            before_output = gr.Image(label="Upscaled Before", type="pil")
+            after_output = gr.Image(label="Upscaled After", type="pil")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            num_steps = gr.Slider(
+                minimum=20, maximum=100, value=50, step=5, label="Inference Steps"
+            )
+        with gr.Column(scale=1):
+            guidance_scale = gr.Slider(
+                minimum=1.0, maximum=15.0, value=7.5, step=0.5, label="Guidance Scale"
+            )
+
+    upscale_btn = gr.Button("🚀 Upscale Both", variant="primary", size="lg")
+    upscale_status = gr.Textbox(label="Status", interactive=False, lines=4)
+
+    gr.Markdown("#### Side-by-Side Comparison")
+    comparison_output = gr.Image(label="Comparison Grid", type="pil")
+
+    # Button click handler
+    upscale_btn.click(
+        fn=upscale_and_compare,
+        inputs=[before_input, after_input, prompt_input, num_steps, guidance_scale],
+        outputs=[before_output, after_output, comparison_output, upscale_status],
+    )

warp/inference/__init__.py

@@ -0,0 +1,10 @@
+"""WARP Inference - AI model inference interfaces."""
+
+from .hf_client import HuggingFaceAPIError, HuggingFaceClient, create_client, infer_image
+
+__all__ = [
+    "HuggingFaceClient",
+    "HuggingFaceAPIError",
+    "create_client",
+    "infer_image",
+]

warp/inference/background_removal.py

@@ -0,0 +1,226 @@
+"""
+Background Removal - Unified interface for background removal models with performance tracking.
+"""
+
+import time
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from PIL import Image
+
+from warp.inference.hf_client import HuggingFaceClient
+from warp.inference.local_client import LocalBackgroundRemovalClient
+from warp.inference.metrics import calculate_comprehensive_metrics, calculate_weighted_quality_score
+from warp.models import get_model, list_model_names
+
+
+@dataclass
+class BackgroundRemovalResult:
+    """Result of a background removal operation."""
+
+    output_image: Image.Image
+    model_name: str
+    processing_time_ms: int
+    success: bool = True
+    error_message: Optional[str] = None
+    input_size: tuple = (0, 0)
+    output_size: tuple = (0, 0)
+    
+    # Quality metrics
+    edge_quality: Optional[float] = None
+    ssim: Optional[float] = None
+    psnr: Optional[float] = None
+    transparency_coverage: Optional[float] = None
+    mask_accuracy: Optional[float] = None
+    weighted_quality_score: Optional[float] = None
+
+
+class BackgroundRemovalEngine:
+    """Engine for running background removal with multiple models."""
+
+    def __init__(self, api_key: Optional[str] = None):
+        """
+        Initialize background removal engine.
+
+        Args:
+            api_key: Hugging Face API key (optional, defaults to env var)
+        """
+        self.hf_client = None
+        self.local_clients = {}  # Cache local model sessions
+        self.api_key = api_key
+        self.available_models = list_model_names("bg_removal")
+
+    def remove_background(
+        self, image: Image.Image, model_name: Optional[str] = None
+    ) -> BackgroundRemovalResult:
+        """
+        Remove background from an image using the specified model.
+
+        Args:
+            image: Input PIL Image
+            model_name: Model name from registry (uses default if None)
+
+        Returns:
+            BackgroundRemovalResult with output image and metrics
+
+        Raises:
+            ValueError: If model not found
+        """
+        # Get model config
+        if model_name is None:
+            from warp.models import get_default_model
+
+            model_config = get_default_model("bg_removal")
+            if not model_config:
+                raise ValueError("No default background removal model found")
+        else:
+            model_config = get_model(model_name)
+            if not model_config:
+                raise ValueError(f"Model '{model_name}' not found in registry")
+
+        input_size = image.size
+        start_time = time.time()
+
+        try:
+            # Route to appropriate client based on provider
+            if model_config.provider == "local":
+                # Use local rembg
+                if model_config.model_id not in self.local_clients:
+                    self.local_clients[model_config.model_id] = LocalBackgroundRemovalClient(
+                        model_name=model_config.model_id
+                    )
+                client = self.local_clients[model_config.model_id]
+                output_image = client.remove_background(image)
+            else:
+                # Use Hugging Face API
+                if self.hf_client is None:
+                    self.hf_client = HuggingFaceClient(api_key=self.api_key)
+                output_image = self.hf_client.infer_image(
+                    model_id=model_config.model_id,
+                    image=image,
+                    parameters=model_config.default_parameters,
+                )
+
+            processing_time_ms = int((time.time() - start_time) * 1000)
+            
+            # Calculate quality metrics
+            metrics = calculate_comprehensive_metrics(
+                output_image=output_image,
+                input_image=image
+            )
+            weighted_score = calculate_weighted_quality_score(metrics)
+
+            return BackgroundRemovalResult(
+                output_image=output_image,
+                model_name=model_config.name,
+                processing_time_ms=processing_time_ms,
+                success=True,
+                input_size=input_size,
+                output_size=output_image.size,
+                edge_quality=metrics.get('edge_quality'),
+                ssim=metrics.get('ssim'),
+                psnr=metrics.get('psnr'),
+                transparency_coverage=metrics.get('transparency_coverage'),
+                mask_accuracy=metrics.get('mask_accuracy'),
+                weighted_quality_score=weighted_score,
+            )
+
+        except Exception as e:
+            processing_time_ms = int((time.time() - start_time) * 1000)
+            return BackgroundRemovalResult(
+                output_image=image,  # Return original on error
+                model_name=model_config.name,
+                processing_time_ms=processing_time_ms,
+                success=False,
+                error_message=str(e),
+                input_size=input_size,
+                output_size=image.size,
+            )
+
+    def compare_models(
+        self, image: Image.Image, model_names: Optional[list] = None
+    ) -> Dict[str, BackgroundRemovalResult]:
+        """
+        Compare multiple background removal models on the same image.
+
+        Args:
+            image: Input PIL Image
+            model_names: List of model names to compare (uses all if None)
+
+        Returns:
+            Dictionary mapping model names to results
+        """
+        if model_names is None:
+            model_names = self.available_models
+
+        results = {}
+        for model_name in model_names:
+            print(f"Processing with {model_name}...")
+            result = self.remove_background(image, model_name)
+            results[model_name] = result
+
+            if result.success:
+                print(f"  ✓ {result.processing_time_ms}ms")
+            else:
+                print(f"  ✗ Error: {result.error_message}")
+
+        return results
+
+
+# ============================================================================
+# Convenience Functions
+# ============================================================================
+
+
+def remove_background(
+    image: Image.Image, model_name: Optional[str] = None, api_key: Optional[str] = None
+) -> Image.Image:
+    """
+    Quick background removal function.
+
+    Args:
+        image: Input PIL Image
+        model_name: Model name (uses default if None)
+        api_key: Optional API key
+
+    Returns:
+        Output PIL Image with background removed
+    """
+    engine = BackgroundRemovalEngine(api_key=api_key)
+    result = engine.remove_background(image, model_name)
+
+    if not result.success:
+        raise RuntimeError(f"Background removal failed: {result.error_message}")
+
+    return result.output_image
+
+
+if __name__ == "__main__":
+    # Demo/test code
+    import os
+
+    print("=== Background Removal Engine ===\n")
+
+    engine = BackgroundRemovalEngine()
+    print(f"Available models: {', '.join(engine.available_models)}\n")
+
+    # Test with a sample image if available
+    test_image_path = "data/scrapedimages/drleedy.com"
+    if os.path.exists(test_image_path):
+        from warp.data import ImageLoader
+
+        loader = ImageLoader()
+        images = loader.get_random_images("drleedy.com", n=1)
+
+        if images:
+            print(f"Testing with image: {images[0]}")
+            test_image = loader.load_image(images[0])
+
+            # Test single model
+            result = engine.remove_background(test_image, model_name="rmbg-1.4")
+            print(f"\nTest result:")
+            print(f"  Model: {result.model_name}")
+            print(f"  Time: {result.processing_time_ms}ms")
+            print(f"  Success: {result.success}")
+    else:
+        print("No test images found. Run with actual images to test.")

warp/inference/hf_client.py

@@ -0,0 +1,253 @@
+"""
+Hugging Face API Client - Unified interface for calling Hugging Face Inference API
+with retry logic, rate limiting, and error handling.
+"""
+
+import os
+import time
+from io import BytesIO
+from typing import Any, Dict, Optional
+
+import requests
+from PIL import Image
+
+# Default configuration
+DEFAULT_API_URL = "https://api-inference.huggingface.co/models"
+DEFAULT_TIMEOUT = 120  # seconds
+DEFAULT_MAX_RETRIES = 3
+DEFAULT_RETRY_DELAY = 2  # seconds
+
+
+class HuggingFaceAPIError(Exception):
+    """Custom exception for Hugging Face API errors."""
+
+    pass
+
+
+class HuggingFaceClient:
+    """Client for Hugging Face Inference API."""
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        retry_delay: int = DEFAULT_RETRY_DELAY,
+        timeout: int = DEFAULT_TIMEOUT,
+    ):
+        """
+        Initialize Hugging Face API client.
+
+        Args:
+            api_key: Hugging Face API key (defaults to HF_API_KEY env var)
+            max_retries: Maximum number of retry attempts
+            retry_delay: Delay between retries in seconds
+            timeout: Request timeout in seconds
+        """
+        self.api_key = api_key or os.getenv("HF_API_KEY")
+        if not self.api_key:
+            raise ValueError(
+                "Hugging Face API key not provided. "
+                "Set HF_API_KEY environment variable or pass api_key parameter."
+            )
+
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        self.timeout = timeout
+
+        self.headers = {"Authorization": f"Bearer {self.api_key}"}
+
+    def _build_url(self, model_id: str) -> str:
+        """Build the full API URL for a model."""
+        return f"{DEFAULT_API_URL}/{model_id}"
+
+    def _handle_response(self, response: requests.Response) -> Any:
+        """
+        Handle API response and raise appropriate errors.
+
+        Args:
+            response: Response from API
+
+        Returns:
+            Response content (bytes or dict)
+
+        Raises:
+            HuggingFaceAPIError: If API returns an error
+        """
+        if response.status_code == 200:
+            # Check content type
+            content_type = response.headers.get("content-type", "")
+            if "application/json" in content_type:
+                return response.json()
+            else:
+                return response.content
+
+        # Handle errors
+        error_msg = f"API request failed with status {response.status_code}"
+        try:
+            error_detail = response.json()
+            if "error" in error_detail:
+                error_msg = f"{error_msg}: {error_detail['error']}"
+        except Exception:
+            error_msg = f"{error_msg}: {response.text[:200]}"
+
+        raise HuggingFaceAPIError(error_msg)
+
+    def infer(
+        self,
+        model_id: str,
+        inputs: Any,
+        parameters: Optional[Dict] = None,
+        return_json: bool = False,
+    ) -> Any:
+        """
+        Call Hugging Face Inference API with retry logic.
+
+        Args:
+            model_id: Hugging Face model ID (e.g., 'briaai/RMBG-2.0')
+            inputs: Input data (can be bytes, PIL Image, or dict)
+            parameters: Additional parameters for the model
+            return_json: Whether to expect JSON response
+
+        Returns:
+            API response (bytes for images, dict for JSON)
+
+        Raises:
+            HuggingFaceAPIError: If all retries fail
+        """
+        url = self._build_url(model_id)
+        payload = self._prepare_payload(inputs, parameters)
+
+        for attempt in range(1, self.max_retries + 1):
+            try:
+                response = requests.post(
+                    url, headers=self.headers, data=payload, timeout=self.timeout
+                )
+
+                # Handle model loading
+                if response.status_code == 503:
+                    error_data = response.json()
+                    if "estimated_time" in error_data:
+                        wait_time = error_data["estimated_time"]
+                        print(f"Model loading, waiting {wait_time}s...")
+                        time.sleep(wait_time + 1)
+                        continue
+
+                return self._handle_response(response)
+
+            except requests.exceptions.RequestException as e:
+                if attempt == self.max_retries:
+                    raise HuggingFaceAPIError(f"Request failed after {attempt} attempts: {e}")
+
+                print(f"Attempt {attempt} failed: {e}. Retrying in {self.retry_delay}s...")
+                time.sleep(self.retry_delay)
+
+        raise HuggingFaceAPIError(f"All {self.max_retries} retries exhausted")
+
+    def _prepare_payload(self, inputs: Any, parameters: Optional[Dict] = None) -> bytes:
+        """
+        Prepare payload for API request.
+
+        Args:
+            inputs: Input data (bytes, PIL Image, or dict)
+            parameters: Additional parameters
+
+        Returns:
+            Bytes payload
+        """
+        # Convert PIL Image to bytes
+        if isinstance(inputs, Image.Image):
+            buffer = BytesIO()
+            inputs.save(buffer, format="PNG")
+            inputs = buffer.getvalue()
+
+        # If already bytes, return as-is
+        if isinstance(inputs, bytes):
+            return inputs
+
+        # For dict/json inputs, serialize
+        if isinstance(inputs, dict):
+            import json
+
+            return json.dumps(inputs).encode("utf-8")
+
+        raise ValueError(f"Unsupported input type: {type(inputs)}")
+
+    def infer_image(self, model_id: str, image: Image.Image, parameters: Optional[Dict] = None) -> Image.Image:
+        """
+        Call inference API with a PIL Image and return a PIL Image.
+
+        Args:
+            model_id: Hugging Face model ID
+            image: Input PIL Image
+            parameters: Additional parameters
+
+        Returns:
+            Output PIL Image
+
+        Raises:
+            HuggingFaceAPIError: If inference fails
+        """
+        response_bytes = self.infer(model_id, image, parameters)
+        return Image.open(BytesIO(response_bytes))
+
+    def health_check(self, model_id: str) -> bool:
+        """
+        Check if a model is available and loaded.
+
+        Args:
+            model_id: Hugging Face model ID
+
+        Returns:
+            True if model is accessible
+        """
+        try:
+            url = self._build_url(model_id)
+            response = requests.get(url, headers=self.headers, timeout=10)
+            return response.status_code in [200, 503]  # 503 means loading
+        except Exception:
+            return False
+
+
+# ============================================================================
+# Convenience Functions
+# ============================================================================
+
+
+def create_client(api_key: Optional[str] = None) -> HuggingFaceClient:
+    """Create a HuggingFace client with default settings."""
+    return HuggingFaceClient(api_key=api_key)
+
+
+def infer_image(
+    model_id: str, image: Image.Image, api_key: Optional[str] = None, parameters: Optional[Dict] = None
+) -> Image.Image:
+    """
+    One-shot image inference with a Hugging Face model.
+
+    Args:
+        model_id: Hugging Face model ID
+        image: Input PIL Image
+        api_key: Optional API key (defaults to env var)
+        parameters: Additional parameters
+
+    Returns:
+        Output PIL Image
+    """
+    client = create_client(api_key)
+    return client.infer_image(model_id, image, parameters)
+
+
+if __name__ == "__main__":
+    # Demo/test code
+    client = create_client()
+
+    # Test model availability
+    test_model = "briaai/RMBG-1.4"
+    print(f"Testing connection to {test_model}...")
+    is_available = client.health_check(test_model)
+    print(f"Model available: {is_available}")
+
+    if is_available:
+        print("\n✓ Hugging Face API client ready")
+    else:
+        print("\n✗ Failed to connect to model")

warp/inference/local_client.py

@@ -0,0 +1,64 @@
+"""
+Local Background Removal Client - Uses rembg for local processing.
+"""
+
+from typing import Optional
+from PIL import Image
+from rembg import remove, new_session
+
+
+class LocalBackgroundRemovalClient:
+    """Client for local background removal using rembg."""
+    
+    # Available models in rembg
+    MODELS = {
+        "u2net": "General purpose model (fast)",
+        "u2netp": "Lightweight version of u2net",
+        "u2net_human_seg": "Optimized for human segmentation",
+        "u2net_cloth_seg": "Optimized for cloth segmentation",
+        "silueta": "General segmentation",
+        "isnet-general-use": "Improved segmentation (recommended)",
+        "isnet-anime": "Anime/illustration specific",
+    }
+    
+    def __init__(self, model_name: str = "isnet-general-use"):
+        """
+        Initialize local client with a specific model.
+        
+        Args:
+            model_name: Name of the rembg model to use
+        """
+        self.model_name = model_name
+        self.session = None
+    
+    def _get_session(self):
+        """Lazy-load the model session."""
+        if self.session is None:
+            self.session = new_session(self.model_name)
+        return self.session
+    
+    def remove_background(self, image: Image.Image) -> Image.Image:
+        """
+        Remove background from an image.
+        
+        Args:
+            image: Input PIL Image
+            
+        Returns:
+            PIL Image with background removed (RGBA)
+        """
+        session = self._get_session()
+        output = remove(image, session=session)
+        return output
+    
+    @classmethod
+    def list_models(cls):
+        """List available models."""
+        return cls.MODELS
+
+
+if __name__ == "__main__":
+    # Test
+    print("Available rembg models:")
+    for name, desc in LocalBackgroundRemovalClient.MODELS.items():
+        print(f"  - {name}: {desc}")

warp/inference/metrics.py

@@ -0,0 +1,256 @@
+"""
+Quality Metrics - Calculate image quality metrics for background removal evaluation.
+"""
+
+import numpy as np
+from PIL import Image
+from typing import Dict, Tuple, Optional
+import cv2
+
+
+def calculate_edge_quality(image: Image.Image, mask: Optional[Image.Image] = None) -> float:
+    """
+    Calculate edge quality score based on gradient strength.
+    
+    Higher score = sharper edges.
+    
+    Args:
+        image: PIL Image (RGBA for background removal)
+        mask: Optional mask (if available separately)
+        
+    Returns:
+        Edge quality score (0.0 to 1.0)
+    """
+    # Convert to numpy array
+    if image.mode == 'RGBA':
+        # Use alpha channel as mask
+        img_array = np.array(image)
+        mask_array = img_array[:, :, 3]
+    elif mask:
+        mask_array = np.array(mask.convert('L'))
+    else:
+        # Convert to grayscale if no alpha
+        mask_array = np.array(image.convert('L'))
+    
+    # Calculate gradients using Sobel
+    grad_x = cv2.Sobel(mask_array, cv2.CV_64F, 1, 0, ksize=3)
+    grad_y = cv2.Sobel(mask_array, cv2.CV_64F, 0, 1, ksize=3)
+    
+    # Gradient magnitude
+    gradient_magnitude = np.sqrt(grad_x**2 + grad_y**2)
+    
+    # Normalize to 0-1
+    if gradient_magnitude.max() > 0:
+        edge_score = np.mean(gradient_magnitude) / gradient_magnitude.max()
+    else:
+        edge_score = 0.0
+    
+    return float(edge_score)
+
+
+def calculate_mask_accuracy(pred_mask: Image.Image, gt_mask: Image.Image) -> float:
+    """
+    Calculate mask accuracy (IoU - Intersection over Union).
+    
+    Only applicable when ground truth mask is available.
+    
+    Args:
+        pred_mask: Predicted mask (alpha channel or grayscale)
+        gt_mask: Ground truth mask
+        
+    Returns:
+        IoU score (0.0 to 1.0)
+    """
+    # Convert to binary masks
+    pred_array = np.array(pred_mask.convert('L')) > 127
+    gt_array = np.array(gt_mask.convert('L')) > 127
+    
+    # Calculate IoU
+    intersection = np.logical_and(pred_array, gt_array).sum()
+    union = np.logical_or(pred_array, gt_array).sum()
+    
+    if union == 0:
+        return 0.0
+    
+    iou = intersection / union
+    return float(iou)
+
+
+def calculate_ssim(image1: Image.Image, image2: Image.Image) -> float:
+    """
+    Calculate Structural Similarity Index (SSIM) between two images.
+    
+    Used to compare processed image with reference.
+    
+    Args:
+        image1: First image
+        image2: Second image
+        
+    Returns:
+        SSIM score (-1.0 to 1.0, higher is better)
+    """
+    from skimage.metrics import structural_similarity as ssim
+    
+    # Convert to grayscale numpy arrays
+    img1_gray = np.array(image1.convert('L'))
+    img2_gray = np.array(image2.convert('L'))
+    
+    # Ensure same size
+    if img1_gray.shape != img2_gray.shape:
+        # Resize to match
+        img2_gray = cv2.resize(img2_gray, (img1_gray.shape[1], img1_gray.shape[0]))
+    
+    # Calculate SSIM
+    score = ssim(img1_gray, img2_gray)
+    return float(score)
+
+
+def calculate_psnr(image1: Image.Image, image2: Image.Image) -> float:
+    """
+    Calculate Peak Signal-to-Noise Ratio (PSNR) between two images.
+    
+    Higher PSNR = better quality (less noise/distortion).
+    
+    Args:
+        image1: First image
+        image2: Second image
+        
+    Returns:
+        PSNR in dB (higher is better, typically 20-50)
+    """
+    # Convert to numpy arrays
+    img1_array = np.array(image1.convert('RGB')).astype(float)
+    img2_array = np.array(image2.convert('RGB')).astype(float)
+    
+    # Ensure same size
+    if img1_array.shape != img2_array.shape:
+        img2_array = cv2.resize(img2_array, (img1_array.shape[1], img1_array.shape[0]))
+    
+    # Calculate MSE
+    mse = np.mean((img1_array - img2_array) ** 2)
+    
+    if mse == 0:
+        return 100.0  # Perfect match
+    
+    # Calculate PSNR
+    max_pixel = 255.0
+    psnr = 20 * np.log10(max_pixel / np.sqrt(mse))
+    
+    return float(psnr)
+
+
+def calculate_transparency_coverage(image: Image.Image) -> float:
+    """
+    Calculate percentage of image that is transparent (for RGBA images).
+    
+    Useful for background removal evaluation.
+    
+    Args:
+        image: RGBA PIL Image
+        
+    Returns:
+        Transparency coverage (0.0 to 1.0)
+    """
+    if image.mode != 'RGBA':
+        return 0.0
+    
+    img_array = np.array(image)
+    alpha_channel = img_array[:, :, 3]
+    
+    # Count transparent pixels (alpha < 10)
+    transparent_pixels = (alpha_channel < 10).sum()
+    total_pixels = alpha_channel.size
+    
+    coverage = transparent_pixels / total_pixels
+    return float(coverage)
+
+
+def calculate_comprehensive_metrics(
+    output_image: Image.Image,
+    input_image: Optional[Image.Image] = None,
+    ground_truth_mask: Optional[Image.Image] = None
+) -> Dict[str, float]:
+    """
+    Calculate all available quality metrics for an image.
+    
+    Args:
+        output_image: Processed image (typically RGBA after bg removal)
+        input_image: Original input image (for SSIM/PSNR comparison)
+        ground_truth_mask: Ground truth mask (if available, for accuracy)
+        
+    Returns:
+        Dictionary of metric names to scores
+    """
+    metrics = {}
+    
+    # Edge quality (always available)
+    metrics['edge_quality'] = calculate_edge_quality(output_image)
+    
+    # Transparency coverage (for RGBA images)
+    if output_image.mode == 'RGBA':
+        metrics['transparency_coverage'] = calculate_transparency_coverage(output_image)
+    
+    # Comparison metrics (if input provided)
+    if input_image:
+        try:
+            metrics['ssim'] = calculate_ssim(input_image, output_image)
+        except Exception as e:
+            print(f"SSIM calculation failed: {e}")
+            metrics['ssim'] = 0.0
+        
+        try:
+            metrics['psnr'] = calculate_psnr(input_image, output_image)
+        except Exception as e:
+            print(f"PSNR calculation failed: {e}")
+            metrics['psnr'] = 0.0
+    
+    # Mask accuracy (if ground truth provided)
+    if ground_truth_mask and output_image.mode == 'RGBA':
+        try:
+            # Extract alpha channel as predicted mask
+            alpha = output_image.split()[3]
+            metrics['mask_accuracy'] = calculate_mask_accuracy(alpha, ground_truth_mask)
+        except Exception as e:
+            print(f"Mask accuracy calculation failed: {e}")
+            metrics['mask_accuracy'] = 0.0
+    
+    return metrics
+
+
+def calculate_weighted_quality_score(metrics: Dict[str, float]) -> float:
+    """
+    Calculate overall weighted quality score from individual metrics.
+    
+    Formula: (mask_accuracy × 0.6 + edge_quality × 0.4) × success_rate
+    
+    If mask_accuracy unavailable, uses edge_quality only.
+    
+    Args:
+        metrics: Dictionary of metric scores
+        
+    Returns:
+        Weighted quality score (0.0 to 1.0)
+    """
+    edge_quality = metrics.get('edge_quality', 0.0)
+    mask_accuracy = metrics.get('mask_accuracy')
+    
+    if mask_accuracy is not None:
+        # Use weighted combination
+        score = (mask_accuracy * 0.6 + edge_quality * 0.4)
+    else:
+        # Use edge quality only
+        score = edge_quality
+    
+    return float(score)
+
+
+if __name__ == "__main__":
+    # Demo
+    print("=== Quality Metrics Module ===\n")
+    print("Available metrics:")
+    print("  - Edge Quality: Gradient-based edge sharpness")
+    print("  - Mask Accuracy: IoU with ground truth (if available)")
+    print("  - SSIM: Structural similarity (0-1)")
+    print("  - PSNR: Peak signal-to-noise ratio (dB)")
+    print("  - Transparency Coverage: % of transparent pixels")
+    print("  - Weighted Quality Score: Combined metric for ranking")

warp/inference/upscaler.py

@@ -0,0 +1,117 @@
+"""
+Upscaler Interface - Unified interface for image upscaling models.
+
+Placeholder for Phase 3 upscaling integration.
+"""
+
+import time
+from dataclass import dataclass
+from typing import Optional
+from PIL import Image
+
+from warp.models import get_model
+
+
+@dataclass
+class UpscaleResult:
+    """Result of an upscaling operation."""
+    
+    output_image: Image.Image
+    model_name: str
+    processing_time_ms: int
+    success: bool = True
+    error_message: Optional[str] = None
+    input_size: tuple = (0, 0)
+    output_size: tuple = (0, 0)
+    scale_factor: int = 4
+    
+
+class UpscalerEngine:
+    """Engine for running upscaling models."""
+    
+    def __init__(self, api_key: Optional[str] = None):
+        """
+        Initialize upscaler engine.
+        
+        Args:
+            api_key: API key for external services (if needed)
+        """
+        self.api_key = api_key
+        # TODO: Initialize clients when upscaling models are added
+        
+    def upscale(
+        self, 
+        image: Image.Image, 
+        model_name: Optional[str] = None,
+        scale_factor: int = 4
+    ) -> UpscaleResult:
+        """
+        Upscale an image using the specified model.
+        
+        Args:
+            image: Input PIL Image
+            model_name: Model name from registry (uses default if None)
+            scale_factor: Upscaling factor (2x, 4x, etc.)
+            
+        Returns:
+            UpscaleResult with output image and metrics
+            
+        Raises:
+            ValueError: If model not found
+            NotImplementedError: If upscaling not yet implemented
+        """
+        # Get model config
+        if model_name is None:
+            from warp.models import get_default_model
+            model_config = get_default_model("upscale")
+            if not model_config:
+                raise ValueError("No default upscaling model found")
+        else:
+            model_config = get_model(model_name)
+            if not model_config:
+                raise ValueError(f"Model '{model_name}' not found in registry")
+        
+        # TODO: Implement actual upscaling in Phase 3
+        raise NotImplementedError(
+            f"Upscaling with {model_config.name} not yet implemented. "
+            "This will be added in Phase 3."
+        )
+
+
+# Convenience function
+def upscale_image(
+    image: Image.Image,
+    model_name: Optional[str] = None,
+    scale_factor: int = 4,
+    api_key: Optional[str] = None
+) -> Image.Image:
+    """
+    Quick upscaling function.
+    
+    Args:
+        image: Input PIL Image
+        model_name: Model name (uses default if None)
+        scale_factor: Upscaling factor
+        api_key: Optional API key
+        
+    Returns:
+        Upscaled PIL Image
+    """
+    engine = UpscalerEngine(api_key=api_key)
+    result = engine.upscale(image, model_name, scale_factor)
+    
+    if not result.success:
+        raise RuntimeError(f"Upscaling failed: {result.error_message}")
+    
+    return result.output_image
+
+
+if __name__ == "__main__":
+    print("=== Upscaler Module ===\\n")
+    print("Status: Placeholder for Phase 3")
+    print("\\nPlanned upscaling models:")
+    print("  - Real-ESRGAN 4x")
+    print("  - GFPGAN (face restoration)")
+    print("  - CodeFormer (face enhancement)")
+    print("  - Swin2SR")
+    print("\\nUpscaling will be integrated in Phase 3 (Week 4-6)")

warp/models/__init__.py

@@ -0,0 +1,25 @@
+"""WARP Models - Model registry and configuration."""
+
+from .registry import (
+    BACKGROUND_REMOVAL_MODELS,
+    MODEL_REGISTRY,
+    UPSCALING_MODELS,
+    ModelConfig,
+    get_default_model,
+    get_model,
+    get_models_by_type,
+    list_all_models,
+    list_model_names,
+)
+
+__all__ = [
+    "ModelConfig",
+    "MODEL_REGISTRY",
+    "BACKGROUND_REMOVAL_MODELS",
+    "UPSCALING_MODELS",
+    "get_model",
+    "get_models_by_type",
+    "get_default_model",
+    "list_all_models",
+    "list_model_names",
+]

warp/models/registry.py

@@ -0,0 +1,248 @@
+"""
+Model Registry - Centralized configuration for all AI models in WARP.
+"""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+
+@dataclass
+class ModelConfig:
+    """Configuration for an AI model."""
+
+    name: str
+    display_name: str
+    model_id: str  # Hugging Face model ID or path
+    provider: str = "huggingface"  # 'huggingface', 'replicate', 'local'
+    operation_type: str = "bg_removal"  # 'bg_removal', 'upscale', 'color_correct'
+    description: str = ""
+    version: Optional[str] = None
+
+    # Model capabilities
+    input_formats: List[str] = field(default_factory=lambda: ["png", "jpg", "jpeg", "webp"])
+    output_format: str = "png"
+    max_input_size: int = 2048  # Max dimension in pixels
+    requires_gpu: bool = False
+
+    # Default parameters
+    default_parameters: Dict = field(default_factory=dict)
+
+    # Performance hints
+    estimated_time_ms: int = 3000
+    is_default: bool = False
+    is_active: bool = True
+
+
+# ============================================================================
+# Background Removal Models
+# ============================================================================
+
+BACKGROUND_REMOVAL_MODELS = {
+    # Local models (using rembg)
+    "isnet-general": ModelConfig(
+        name="isnet-general",
+        display_name="ISNet General (Local)",
+        model_id="isnet-general-use",
+        provider="local",
+        operation_type="bg_removal",
+        description="Improved segmentation model - highest quality, slower",
+        estimated_time_ms=678,  # Updated from benchmark results
+        is_default=False,
+        default_parameters={},
+    ),
+    "u2net": ModelConfig(
+        name="u2net",
+        display_name="U2Net (Local)",
+        model_id="u2net",
+        provider="local",
+        operation_type="bg_removal",
+        description="General purpose background removal - BEST SPEED/QUALITY BALANCE",
+        estimated_time_ms=416,  # Updated from benchmark results
+        is_default=True,  # Updated based on Phase 2 benchmark results
+        default_parameters={},
+    ),
+    "u2net-human": ModelConfig(
+        name="u2net-human",
+        display_name="U2Net Human (Local)",
+        model_id="u2net_human_seg",
+        provider="local",
+        operation_type="bg_removal",
+        description="Optimized for human segmentation and portraits",
+        estimated_time_ms=436,  # Updated from benchmark results
+        is_default=False,  # Changed: u2net outperformed in benchmarks
+        default_parameters={},
+    ),
+    "isnet-anime": ModelConfig(
+        name="isnet-anime",
+        display_name="ISNet Anime (Local)",
+        model_id="isnet-anime",
+        provider="local",
+        operation_type="bg_removal",
+        description="Specialized for anime/illustrations - not suitable for medical photos",
+        estimated_time_ms=727,  # Updated from benchmark results
+        is_default=False,
+        default_parameters={},
+    ),
+    # HuggingFace models (currently unavailable)
+    "rmbg-1.4": ModelConfig(
+        name="rmbg-1.4",
+        display_name="RMBG 1.4 (HF)",
+        model_id="briaai/RMBG-1.4",
+        operation_type="bg_removal",
+        description="Fast and accurate background removal by Bria AI (v1.4) - Currently unavailable",
+        version="1.4",
+        estimated_time_ms=2000,
+        is_default=False,
+        is_active=False,  # Disabled due to HF API deprecation
+        default_parameters={},
+    ),
+}
+
+# ============================================================================
+# Upscaling Models
+# ============================================================================
+
+UPSCALING_MODELS = {
+    "realesrgan": ModelConfig(
+        name="realesrgan",
+        display_name="Real-ESRGAN 4x",
+        model_id="ai-forever/Real-ESRGAN",
+        operation_type="upscale",
+        description="Real-ESRGAN 4x upscaling for general images",
+        estimated_time_ms=8000,
+        is_default=True,
+        requires_gpu=True,
+        default_parameters={"scale": 4},
+    ),
+    "gfpgan": ModelConfig(
+        name="gfpgan",
+        display_name="GFPGAN",
+        model_id="TencentARC/GFPGAN",
+        operation_type="upscale",
+        description="Face restoration and enhancement",
+        estimated_time_ms=7000,
+        is_default=False,
+        requires_gpu=True,
+        default_parameters={"version": "1.3"},
+    ),
+    "codeformer": ModelConfig(
+        name="codeformer",
+        display_name="CodeFormer",
+        model_id="sczhou/CodeFormer",
+        operation_type="upscale",
+        description="Face restoration with better quality preservation",
+        estimated_time_ms=9000,
+        is_default=False,
+        requires_gpu=True,
+        default_parameters={"w": 0.5},
+    ),
+    "swin2sr": ModelConfig(
+        name="swin2sr",
+        display_name="Swin2SR",
+        model_id="caidas/swin2SR-classical-sr-x4-64",
+        operation_type="upscale",
+        description="Swin Transformer for image super-resolution",
+        estimated_time_ms=6000,
+        is_default=False,
+        requires_gpu=True,
+        default_parameters={"scale": 4},
+    ),
+}
+
+# ============================================================================
+# Combined Registry
+# ============================================================================
+
+MODEL_REGISTRY = {
+    **BACKGROUND_REMOVAL_MODELS,
+    **UPSCALING_MODELS,
+}
+
+
+# ============================================================================
+# Helper Functions
+# ============================================================================
+
+
+def get_model(name: str) -> Optional[ModelConfig]:
+    """Get model configuration by name."""
+    return MODEL_REGISTRY.get(name)
+
+
+def get_models_by_type(operation_type: str) -> Dict[str, ModelConfig]:
+    """Get all models for a specific operation type."""
+    return {
+        name: config
+        for name, config in MODEL_REGISTRY.items()
+        if config.operation_type == operation_type and config.is_active
+    }
+
+
+def get_default_model(operation_type: str) -> Optional[ModelConfig]:
+    """Get the default model for an operation type."""
+    models = get_models_by_type(operation_type)
+    for model in models.values():
+        if model.is_default:
+            return model
+    # If no default set, return first active model
+    return next(iter(models.values())) if models else None
+
+
+def list_all_models() -> Dict[str, ModelConfig]:
+    """Get all active models."""
+    return {name: config for name, config in MODEL_REGISTRY.items() if config.is_active}
+
+
+def list_model_names(operation_type: Optional[str] = None) -> List[str]:
+    """Get list of model names, optionally filtered by operation type."""
+    if operation_type:
+        return list(get_models_by_type(operation_type).keys())
+    return list(MODEL_REGISTRY.keys())
+
+
+# ============================================================================
+# Model Information Display
+# ============================================================================
+
+
+def get_model_info(name: str) -> str:
+    """Get formatted information about a model."""
+    model = get_model(name)
+    if not model:
+        return f"Model '{name}' not found in registry."
+
+    info = f"""
+Model: {model.display_name}
+ID: {model.model_id}
+Type: {model.operation_type}
+Provider: {model.provider}
+Description: {model.description}
+Est. Time: {model.estimated_time_ms}ms
+Default: {'Yes' if model.is_default else 'No'}
+Active: {'Yes' if model.is_active else 'No'}
+    """.strip()
+
+    return info
+
+
+if __name__ == "__main__":
+    # Demo/test code
+    print("=== WARP Model Registry ===\n")
+
+    print("Background Removal Models:")
+    for name in list_model_names("bg_removal"):
+        print(f"  - {name}")
+
+    print("\nUpscaling Models:")
+    for name in list_model_names("upscale"):
+        print(f"  - {name}")
+
+    print("\nDefault Background Removal Model:")
+    default_bg = get_default_model("bg_removal")
+    if default_bg:
+        print(f"  {default_bg.display_name} ({default_bg.name})")
+
+    print("\nDefault Upscaling Model:")
+    default_up = get_default_model("upscale")
+    if default_up:
+        print(f"  {default_up.display_name} ({default_up.name})")