Initial A360 migration

app.py

@@ -0,0 +1,4 @@
+from warp.gradio_app.app import demo
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,21 @@
+supabase
+gradio
+pandas
+pillow
+exifread
+requests
+pydantic
+boto3
+python-dotenv
+
+# Phase 2: Local background removal models
+rembg==2.0.67
+onnxruntime==1.23.2
+scikit-image==0.25.2
+
+# ML/AI dependencies for image upscaling
+torch>=2.0.0
+torchvision>=0.15.0
+transformers>=4.30.0
+diffusers>=0.21.0
+accelerate>=0.20.0

warp/data/__init__.py

@@ -0,0 +1,5 @@
+"""Data management for A360 WARP."""
+
+from .image_loader import ImageLoader, list_practices
+
+__all__ = ["ImageLoader", "list_practices"]

warp/data/image_loader.py

@@ -0,0 +1,227 @@
+"""Image loader for scraped medical practice images."""
+
+from pathlib import Path
+
+from PIL import Image
+
+# Default scraped images directory (now under top-level data/scrapedimages)
+DEFAULT_SCRAPED_IMAGES_DIR = Path(__file__).parent.parent.parent / "data" / "scrapedimages"
+
+
+class ImageLoader:
+    """Load and manage scraped medical practice images.
+
+    Attributes:
+        base_path: Root directory containing scraped images organized by practice
+        practices: List of available practice directories
+    """
+
+    def __init__(self, base_path: Path | str | None = None):
+        """Initialize the ImageLoader.
+
+        Args:
+            base_path: Root directory containing scraped images.
+                      Defaults to project's scrapedimages folder.
+        """
+        if base_path is None:
+            self.base_path = DEFAULT_SCRAPED_IMAGES_DIR
+        else:
+            self.base_path = Path(base_path)
+
+        if not self.base_path.exists():
+            raise ValueError(f"Image directory does not exist: {self.base_path}")
+
+        self._practices: list[str] | None = None
+
+    @property
+    def practices(self) -> list[str]:
+        """Get list of available practice directories.
+
+        Returns:
+            List of practice directory names (e.g., ['drleedy.com', 'drbirely.com'])
+        """
+        if self._practices is None:
+            self._practices = sorted(
+                [
+                    d.name
+                    for d in self.base_path.iterdir()
+                    if d.is_dir() and not d.name.startswith(".")
+                ]
+            )
+        return self._practices
+
+    def get_practice_path(self, practice_name: str) -> Path:
+        """Get the full path to a practice directory.
+
+        Args:
+            practice_name: Name of the practice (e.g., 'drleedy.com')
+
+        Returns:
+            Path object pointing to the practice directory
+
+        Raises:
+            ValueError: If practice does not exist
+        """
+        practice_path = self.base_path / practice_name
+        if not practice_path.exists():
+            raise ValueError(
+                f"Practice '{practice_name}' not found. "
+                f"Available practices: {', '.join(self.practices)}"
+            )
+        return practice_path
+
+    def list_images(self, practice_name: str, extensions: list[str] | None = None) -> list[Path]:
+        """List all images for a given practice.
+
+        Args:
+            practice_name: Name of the practice
+            extensions: List of file extensions to filter (e.g., ['.jpg', '.png'])
+                       If None, includes common image formats
+
+        Returns:
+            List of Path objects for all matching images
+        """
+        if extensions is None:
+            extensions = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"]
+
+        practice_path = self.get_practice_path(practice_name)
+        images: list[Path] = []
+
+        for ext in extensions:
+            images.extend(practice_path.glob(f"**/*{ext}"))
+            images.extend(practice_path.glob(f"**/*{ext.upper()}"))
+
+        return sorted(images)
+
+    def count_images(self, practice_name: str) -> int:
+        """Count total images for a practice.
+
+        Args:
+            practice_name: Name of the practice
+
+        Returns:
+            Number of images
+        """
+        return len(self.list_images(practice_name))
+
+    def load_image(self, image_path: Path | str) -> Image.Image:
+        """Load a single image.
+
+        Args:
+            image_path: Path to the image file
+
+        Returns:
+            PIL Image object
+
+        Raises:
+            FileNotFoundError: If image does not exist
+        """
+        image_path = Path(image_path)
+        if not image_path.exists():
+            raise FileNotFoundError(f"Image not found: {image_path}")
+
+        return Image.open(image_path)
+
+    def get_image_info(self, image_path: Path | str) -> dict:
+        """Get metadata about an image.
+
+        Args:
+            image_path: Path to the image file
+
+        Returns:
+            Dictionary with image metadata (size, format, mode, etc.)
+        """
+        image_path = Path(image_path)
+        img = self.load_image(image_path)
+
+        return {
+            "path": str(image_path),
+            "filename": image_path.name,
+            "practice": (
+                image_path.parent.name if image_path.is_relative_to(self.base_path) else None
+            ),
+            "size": img.size,
+            "width": img.width,
+            "height": img.height,
+            "format": img.format,
+            "mode": img.mode,
+            "file_size_bytes": image_path.stat().st_size,
+        }
+
+    def get_random_images(
+        self, practice_name: str, n: int = 5, seed: int | None = None
+    ) -> list[Path]:
+        """Get random sample of images from a practice.
+
+        Args:
+            practice_name: Name of the practice
+            n: Number of images to return
+            seed: Random seed for reproducibility
+
+        Returns:
+            List of n random image paths
+        """
+        import random
+
+        images = self.list_images(practice_name)
+
+        if seed is not None:
+            random.seed(seed)
+
+        return random.sample(images, min(n, len(images)))
+
+    def get_practice_stats(self, practice_name: str) -> dict:
+        """Get statistics for a practice's images.
+
+        Args:
+            practice_name: Name of the practice
+
+        Returns:
+            Dictionary with practice statistics
+        """
+        images = self.list_images(practice_name)
+        total_size = sum(img.stat().st_size for img in images)
+
+        # Get format distribution
+        formats: dict[str, int] = {}
+        for img_path in images:
+            ext = img_path.suffix.lower()
+            formats[ext] = formats.get(ext, 0) + 1
+
+        return {
+            "practice": practice_name,
+            "total_images": len(images),
+            "total_size_mb": total_size / (1024 * 1024),
+            "formats": formats,
+            "practice_path": str(self.get_practice_path(practice_name)),
+        }
+
+    def get_all_stats(self) -> dict:
+        """Get statistics for all practices.
+
+        Returns:
+            Dictionary with overall statistics
+        """
+        all_stats: dict = {"practices": {}, "total_images": 0, "total_size_mb": 0.0}
+
+        for practice in self.practices:
+            practice_stats = self.get_practice_stats(practice)
+            all_stats["practices"][practice] = practice_stats
+            all_stats["total_images"] += practice_stats["total_images"]
+            all_stats["total_size_mb"] += practice_stats["total_size_mb"]
+
+        return all_stats
+
+
+def list_practices(base_path: Path | str | None = None) -> list[str]:
+    """Convenience function to list all available practices.
+
+    Args:
+        base_path: Root directory containing scraped images.
+                  Defaults to project's scrapedimages folder.
+
+    Returns:
+        List of practice directory names
+    """
+    loader = ImageLoader(base_path)
+    return loader.practices

warp/gradio_app/__init__.py

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

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/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],
+    )