Add GitHub Actions workflows and comprehensive test suite

- Add CI/CD workflows for tests, code quality, and Docker builds
- Add comprehensive Makefile with test, lint, format, and Docker targets
- Add new test files for CLI, fixtures, UI components, and UI events
- Refactor batch analysis into main analysis module
- Update model manager and inference modules
- Add .dockerignore for cleaner Docker builds
- Update dependencies in uv.lock

.dockerignore

@@ -0,0 +1,79 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+*.egg-info/
+dist/
+build/
+*.egg
+
+# Virtual environments
+.venv/
+venv/
+ENV/
+env/
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+.tox/
+*.cover
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# Git
+.git/
+.gitignore
+.gitattributes
+
+# CI/CD
+.github/
+.gitlab-ci.yml
+
+# Documentation
+docs/
+*.md
+!README.md
+
+# Data and outputs
+data/
+output/
+*.svs
+*.tiff
+*.tif
+*.png
+*.jpg
+*.jpeg
+
+# Logs
+*.log
+logs/
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Project specific
+tests/
+*.csv
+profile.stats
+benchmark_output/
+profile_output/
+
+# Lock files (we use uv.lock)
+poetry.lock
+Pipfile.lock
+requirements*.txt
+
+# Makefile and CI configs
+Makefile
+.dockerignore
+Dockerfile*

.github/workflows/code-quality.yml

@@ -0,0 +1,80 @@
+name: Code Quality
+
+on:
+  push:
+    branches: [ main, dev ]
+  pull_request:
+    branches: [ main, dev ]
+  workflow_dispatch:
+
+jobs:
+  format-check:
+    name: Check Code Formatting
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v4
+      with:
+        enable-cache: true
+
+    - name: Install dependencies
+      run: |
+        uv sync
+
+    - name: Check formatting with black
+      run: |
+        make format-check
+
+    - name: Format Summary
+      if: always()
+      run: |
+        echo "## Code Formatting :art:" >> $GITHUB_STEP_SUMMARY
+        echo "" >> $GITHUB_STEP_SUMMARY
+        echo "Status: ${{ job.status }}" >> $GITHUB_STEP_SUMMARY
+        if [ "${{ job.status }}" == "failure" ]; then
+          echo "" >> $GITHUB_STEP_SUMMARY
+          echo "Run \`make format\` to auto-fix formatting issues." >> $GITHUB_STEP_SUMMARY
+        fi
+
+  lint:
+    name: Lint Code
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v4
+      with:
+        enable-cache: true
+
+    - name: Install dependencies
+      run: |
+        uv sync
+
+    - name: Lint with pylint
+      run: |
+        make lint
+      continue-on-error: true  # Don't fail CI on pylint warnings
+
+    - name: Lint Summary
+      if: always()
+      run: |
+        echo "## Linting Results :mag:" >> $GITHUB_STEP_SUMMARY
+        echo "" >> $GITHUB_STEP_SUMMARY
+        echo "Status: ${{ job.status }}" >> $GITHUB_STEP_SUMMARY

.github/workflows/docker.yml

@@ -0,0 +1,73 @@
+name: Docker Build
+
+on:
+  push:
+    branches: [ main, dev ]
+    tags:
+      - 'v*'
+  pull_request:
+    branches: [ main ]
+  workflow_dispatch:
+
+env:
+  REGISTRY: ghcr.io
+  IMAGE_NAME: ${{ github.repository }}
+
+jobs:
+  build:
+    name: Build Docker Image
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+
+    - name: Set up Docker Buildx
+      uses: docker/setup-buildx-action@v3
+
+    - name: Log in to Container Registry
+      if: github.event_name != 'pull_request'
+      uses: docker/login-action@v3
+      with:
+        registry: ${{ env.REGISTRY }}
+        username: ${{ github.actor }}
+        password: ${{ secrets.GITHUB_TOKEN }}
+
+    - name: Extract metadata
+      id: meta
+      uses: docker/metadata-action@v5
+      with:
+        images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
+        tags: |
+          type=ref,event=branch
+          type=ref,event=pr
+          type=semver,pattern={{version}}
+          type=semver,pattern={{major}}.{{minor}}
+          type=sha,prefix={{branch}}-
+
+    - name: Build and push Docker image
+      uses: docker/build-push-action@v5
+      with:
+        context: .
+        push: ${{ github.event_name != 'pull_request' }}
+        tags: ${{ steps.meta.outputs.tags }}
+        labels: ${{ steps.meta.outputs.labels }}
+        cache-from: type=gha
+        cache-to: type=gha,mode=max
+        secret-files: |
+          "github_token=${{ secrets.GITHUB_TOKEN }}"
+
+    - name: Docker Summary
+      run: |
+        echo "## Docker Build :whale:" >> $GITHUB_STEP_SUMMARY
+        echo "" >> $GITHUB_STEP_SUMMARY
+        echo "Registry: ${{ env.REGISTRY }}" >> $GITHUB_STEP_SUMMARY
+        echo "Image: ${{ env.IMAGE_NAME }}" >> $GITHUB_STEP_SUMMARY
+        echo "" >> $GITHUB_STEP_SUMMARY
+        echo "### Tags" >> $GITHUB_STEP_SUMMARY
+        echo '```' >> $GITHUB_STEP_SUMMARY
+        echo "${{ steps.meta.outputs.tags }}" >> $GITHUB_STEP_SUMMARY
+        echo '```' >> $GITHUB_STEP_SUMMARY

.github/workflows/tests.yml

@@ -0,0 +1,74 @@
+name: Tests
+
+on:
+  push:
+    branches: [ main, dev ]
+  pull_request:
+    branches: [ main, dev ]
+  workflow_dispatch:  # Allow manual trigger
+
+jobs:
+  test:
+    name: Run Tests (Python ${{ matrix.python-version }})
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.11"]
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0  # Full history for better coverage reports
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v4
+      with:
+        enable-cache: true
+        cache-dependency-glob: "uv.lock"
+
+    - name: Install dependencies
+      run: |
+        uv sync
+
+    - name: Run tests with coverage
+      run: |
+        make test-coverage
+
+    - name: Generate coverage badge
+      if: matrix.python-version == '3.10'
+      run: |
+        COVERAGE=$(uv run coverage report | grep TOTAL | awk '{print $NF}' | sed 's/%//')
+        echo "COVERAGE=$COVERAGE" >> $GITHUB_ENV
+        echo "Coverage: $COVERAGE%"
+
+    - name: Upload coverage reports to Codecov
+      if: matrix.python-version == '3.10'
+      uses: codecov/codecov-action@v4
+      with:
+        file: ./coverage.xml
+        fail_ci_if_error: false
+        token: ${{ secrets.CODECOV_TOKEN }}
+      continue-on-error: true
+
+    - name: Upload coverage HTML report
+      if: matrix.python-version == '3.10'
+      uses: actions/upload-artifact@v4
+      with:
+        name: coverage-report
+        path: htmlcov/
+        retention-days: 30
+
+    - name: Test Summary
+      if: always()
+      run: |
+        echo "## Test Results :test_tube:" >> $GITHUB_STEP_SUMMARY
+        echo "" >> $GITHUB_STEP_SUMMARY
+        echo "Python Version: ${{ matrix.python-version }}" >> $GITHUB_STEP_SUMMARY
+        echo "Status: ${{ job.status }}" >> $GITHUB_STEP_SUMMARY

MAKEFILE_QUICK_REF.md

@@ -0,0 +1,93 @@
+# Makefile Quick Reference
+
+## Most Common Commands
+
+```bash
+# Setup
+make install-dev              # Install dev dependencies
+make help                     # Show all available commands
+
+# Testing
+make test                     # Run tests with coverage
+make test-fast                # Run tests quickly (no coverage)
+make test-ui                  # Test UI components only
+make test-cli                 # Test CLI only
+
+# Code Quality
+make format                   # Format code with black
+make format-check             # Check formatting
+make quality                  # Run all quality checks
+
+# Running
+make run-ui                   # Launch web interface
+make run-single SLIDE=x.svs OUTPUT=out/  # Process single slide
+make run-batch CSV=s.csv OUTPUT=out/     # Process batch
+
+# Docker
+make docker-build             # Build image
+make docker-run               # Run web UI in container
+make docker-shell             # Shell into container
+
+# Cleanup
+make clean                    # Remove cache files
+make clean-all                # Remove everything
+```
+
+## Development Workflow
+
+```bash
+# 1. Initial setup
+make install-dev
+
+# 2. Make changes to code
+# ... edit files ...
+
+# 3. Format and test
+make format
+make test
+
+# 4. Before committing
+make quality
+make test-coverage
+
+# 5. Optional: Install pre-commit hooks
+make pre-commit-install
+```
+
+## Docker Workflow
+
+```bash
+# Build and test locally
+make docker-build
+make docker-run
+
+# Process slides with Docker
+make docker-run-single SLIDE=my_slide.svs
+make docker-run-batch CSV=settings.csv
+
+# Push to registry
+make docker-tag DOCKER_REGISTRY=myregistry.com/user
+make docker-push DOCKER_REGISTRY=myregistry.com/user
+```
+
+## CI/CD
+
+```bash
+# Run all CI checks
+make ci-test                  # Tests + format check (fast)
+make ci-test-strict           # Tests + format check + pylint (slow)
+make ci-docker                # Build Docker for CI
+```
+
+## Tips
+
+- Use `make help` to see all available commands
+- Use `make test-specific TEST=path/to/test` for debugging
+- Use `make test-verbose` to see print statements
+- Use `make info` to see project information
+- Set environment variables to customize Docker:
+  ```bash
+  export DOCKER_REGISTRY=myregistry.com/user
+  export DOCKER_TAG=v1.0.0
+  make docker-build
+  ```

MAKEFILE_USAGE.md

@@ -0,0 +1,459 @@
+# Makefile Usage Guide
+
+This document provides detailed information about the Makefile targets available in the Mosaic project.
+
+## Quick Start
+
+```bash
+# See all available commands
+make help
+
+# Setup development environment
+make install-dev
+
+# Run tests
+make test
+
+# Launch web interface
+make run-ui
+```
+
+## Development Setup
+
+### `make install`
+Install production dependencies only (no dev tools).
+```bash
+make install
+```
+
+### `make install-dev`
+Install all dependencies including development tools (pytest, ruff, etc.).
+```bash
+make install-dev
+```
+
+## Testing
+
+### `make test`
+Run full test suite with coverage reporting.
+```bash
+make test
+```
+
+### `make test-fast`
+Run tests without coverage (faster execution).
+```bash
+make test-fast
+```
+
+### `make test-coverage`
+Run tests with detailed coverage report (terminal + HTML).
+```bash
+make test-coverage
+# View HTML report at: htmlcov/index.html
+```
+
+### `make test-ui`
+Run only UI-related tests.
+```bash
+make test-ui
+```
+
+### `make test-cli`
+Run only CLI-related tests.
+```bash
+make test-cli
+```
+
+### `make test-verbose`
+Run tests with verbose output and show print statements.
+```bash
+make test-verbose
+```
+
+### `make test-specific`
+Run a specific test file, class, or method.
+```bash
+# Run specific test file
+make test-specific TEST=tests/test_cli.py
+
+# Run specific test class
+make test-specific TEST=tests/test_cli.py::TestArgumentParsing
+
+# Run specific test method
+make test-specific TEST=tests/test_cli.py::TestArgumentParsing::test_no_arguments_launches_web_interface
+```
+
+## Code Quality
+
+### `make lint`
+Check code for linting issues using pylint (src only for speed).
+```bash
+make lint
+```
+
+### `make lint-strict`
+Run pylint on both src and tests (slower but comprehensive).
+```bash
+make lint-strict
+```
+
+### `make format`
+Format code using black formatter.
+```bash
+make format
+```
+
+### `make format-check`
+Check if code is properly formatted without making changes.
+```bash
+make format-check
+```
+
+### `make quality`
+Run all code quality checks (format-check + lint).
+```bash
+make quality
+```
+
+## Running the Application
+
+### `make run-ui`
+Launch the Gradio web interface locally.
+```bash
+make run-ui
+# Open browser to http://localhost:7860
+```
+
+### `make run-ui-public`
+Launch Gradio web interface with public sharing enabled.
+```bash
+make run-ui-public
+# Returns a public gradio.app URL for sharing
+```
+
+### `make run-single`
+Process a single slide from the command line.
+```bash
+make run-single SLIDE=data/my_slide.svs OUTPUT=output/
+```
+
+### `make run-batch`
+Process multiple slides from a CSV file.
+```bash
+make run-batch CSV=data/settings.csv OUTPUT=output/
+```
+
+## Docker
+
+### `make docker-build`
+Build Docker image for Mosaic.
+```bash
+make docker-build
+
+# Build with custom tag
+make docker-build DOCKER_TAG=v1.0.0
+
+# Build with custom image name
+make docker-build DOCKER_IMAGE_NAME=my-mosaic DOCKER_TAG=latest
+```
+
+### `make docker-build-no-cache`
+Build Docker image without using cache (useful for clean builds).
+```bash
+make docker-build-no-cache
+```
+
+### `make docker-run`
+Run Docker container in web UI mode.
+```bash
+make docker-run
+# Access at http://localhost:7860
+```
+
+### `make docker-run-single`
+Run Docker container to process a single slide.
+```bash
+# Place your slide in ./data directory first
+make docker-run-single SLIDE=my_slide.svs
+# Results will be in ./output directory
+```
+
+### `make docker-run-batch`
+Run Docker container for batch processing.
+```bash
+# Place CSV and slides in ./data directory
+make docker-run-batch CSV=settings.csv
+# Results will be in ./output directory
+```
+
+### `make docker-shell`
+Open an interactive shell inside the Docker container.
+```bash
+make docker-shell
+```
+
+### `make docker-tag`
+Tag Docker image for pushing to a registry.
+```bash
+make docker-tag DOCKER_REGISTRY=docker.io/myusername
+```
+
+### `make docker-push`
+Push Docker image to registry.
+```bash
+# Set your registry first
+make docker-push DOCKER_REGISTRY=docker.io/myusername DOCKER_TAG=latest
+```
+
+### `make docker-clean`
+Remove local Docker image.
+```bash
+make docker-clean
+```
+
+### `make docker-prune`
+Clean up Docker build cache to free space.
+```bash
+make docker-prune
+```
+
+## Cleanup
+
+### `make clean`
+Remove Python cache files and build artifacts.
+```bash
+make clean
+```
+
+### `make clean-outputs`
+Remove generated output files (masks, CSVs).
+```bash
+make clean-outputs
+```
+
+### `make clean-all`
+Remove all artifacts, cache, and Docker images.
+```bash
+make clean-all
+```
+
+## Model Management
+
+### `make download-models`
+Explicitly download required models from HuggingFace.
+```bash
+make download-models
+# Note: Models are automatically downloaded on first run
+```
+
+## CI/CD
+
+### `make ci-test`
+Run complete CI test suite (install deps, test with coverage, lint).
+```bash
+make ci-test
+```
+
+### `make ci-docker`
+Build Docker image for CI pipeline.
+```bash
+make ci-docker
+```
+
+## Development Utilities
+
+### `make shell`
+Open Python shell with project in path.
+```bash
+make shell
+```
+
+### `make ipython`
+Open IPython shell with project in path.
+```bash
+make ipython
+```
+
+### `make notebook`
+Start Jupyter notebook server.
+```bash
+make notebook
+```
+
+### `make check-deps`
+Check for outdated dependencies.
+```bash
+make check-deps
+```
+
+### `make update-deps`
+Update all dependencies (use with caution).
+```bash
+make update-deps
+```
+
+### `make lock`
+Update uv.lock file.
+```bash
+make lock
+```
+
+## Git Hooks
+
+### `make pre-commit-install`
+Install pre-commit hooks that run lint, format-check, and test-fast before each commit.
+```bash
+make pre-commit-install
+```
+
+### `make pre-commit-uninstall`
+Remove pre-commit hooks.
+```bash
+make pre-commit-uninstall
+```
+
+## Information
+
+### `make info`
+Display project information and key commands.
+```bash
+make info
+```
+
+### `make version`
+Show version information.
+```bash
+make version
+```
+
+### `make tree`
+Show project directory structure (requires `tree` command).
+```bash
+make tree
+```
+
+## Performance
+
+### `make profile`
+Profile single slide analysis to identify performance bottlenecks.
+```bash
+make profile SLIDE=tests/testdata/948176.svs
+# Creates profile.stats file with profiling data
+```
+
+### `make benchmark`
+Run performance benchmarks on test slide.
+```bash
+make benchmark
+# Times full analysis pipeline
+```
+
+## Common Workflows
+
+### Setting up for development
+```bash
+# 1. Install dependencies
+make install-dev
+
+# 2. Run tests to ensure everything works
+make test
+
+# 3. Install pre-commit hooks
+make pre-commit-install
+```
+
+### Before committing changes
+```bash
+# Run quality checks
+make quality
+
+# Run tests
+make test
+
+# Clean up
+make clean
+```
+
+### Preparing a release
+```bash
+# Run full CI suite
+make ci-test
+
+# Build Docker image
+make docker-build DOCKER_TAG=v1.0.0
+
+# Test Docker image
+make docker-run DOCKER_TAG=v1.0.0
+
+# Push to registry
+make docker-push DOCKER_REGISTRY=your-registry DOCKER_TAG=v1.0.0
+```
+
+### Processing slides
+```bash
+# Web UI (recommended for exploration)
+make run-ui
+
+# Single slide (CLI)
+make run-single SLIDE=data/sample.svs OUTPUT=results/
+
+# Batch processing (CLI)
+make run-batch CSV=data/batch_settings.csv OUTPUT=results/
+
+# Using Docker
+make docker-build
+make docker-run-batch CSV=batch_settings.csv
+```
+
+## Customization
+
+You can customize Makefile behavior by setting environment variables or editing the Makefile:
+
+```bash
+# Custom Docker registry
+export DOCKER_REGISTRY=my-registry.com/username
+
+# Custom image name
+export DOCKER_IMAGE_NAME=my-custom-mosaic
+
+# Then use make commands as normal
+make docker-build
+make docker-push
+```
+
+## Troubleshooting
+
+### Tests fail
+```bash
+# Run with verbose output
+make test-verbose
+
+# Run specific failing test
+make test-specific TEST=tests/test_file.py::test_name
+```
+
+### Docker build fails
+```bash
+# Build without cache
+make docker-build-no-cache
+
+# Check Docker logs
+docker logs <container-id>
+```
+
+### Permission errors
+```bash
+# Clean and rebuild
+make clean-all
+make install-dev
+```
+
+### Out of disk space
+```bash
+# Clean Docker cache
+make docker-prune
+
+# Clean project artifacts
+make clean
+```

Makefile

@@ -0,0 +1,257 @@
+.PHONY: help install install-dev test test-coverage test-verbose lint format clean docker-build docker-run docker-push docker-clean run-ui run-cli
+
+# Default target
+.DEFAULT_GOAL := help
+
+# Variables
+DOCKER_IMAGE_NAME := mosaic
+DOCKER_TAG := latest
+DOCKER_REGISTRY := # Set your registry here (e.g., docker.io/username)
+PYTHON := uv run python
+PYTEST := uv run pytest
+BLACK := uv run black
+PYLINT := uv run pylint
+
+##@ General
+
+help: ## Display this help message
+	@awk 'BEGIN {FS = ":.*##"; printf "\nUsage:\n  make \033[36m<target>\033[0m\n"} /^[a-zA-Z_-]+:.*?##/ { printf "  \033[36m%-20s\033[0m %s\n", $$1, $$2 } /^##@/ { printf "\n\033[1m%s\033[0m\n", substr($$0, 5) } ' $(MAKEFILE_LIST)
+
+##@ Development Setup
+
+install: ## Install production dependencies using uv
+	uv sync --no-dev
+
+install-dev: ## Install development dependencies using uv
+	uv sync
+
+##@ Testing
+
+test: ## Run all tests
+	$(PYTEST) tests/ -v
+
+test-fast: ## Run tests without coverage (faster)
+	$(PYTEST) tests/ -v --no-cov
+
+test-coverage: ## Run tests with detailed coverage report
+	$(PYTEST) tests/ -v --cov=src/mosaic --cov-report=term-missing --cov-report=html
+
+test-ui: ## Run only UI tests
+	$(PYTEST) tests/test_ui_components.py tests/test_ui_events.py -v
+
+test-cli: ## Run only CLI tests
+	$(PYTEST) tests/test_cli.py -v
+
+test-verbose: ## Run tests with verbose output and show print statements
+	$(PYTEST) tests/ -vv -s
+
+test-specific: ## Run specific test (usage: make test-specific TEST=tests/test_cli.py::TestClass::test_method)
+	$(PYTEST) $(TEST) -v
+
+test-watch: ## Run tests in watch mode (requires pytest-watch)
+	$(PYTEST) tests/ --watch
+
+##@ Code Quality
+
+lint: ## Run linting checks with pylint
+	$(PYLINT) src/mosaic/
+
+lint-strict: ## Run pylint on both src and tests
+	$(PYLINT) src/mosaic/ tests/
+
+format: ## Format code with black
+	$(BLACK) src/ tests/
+
+format-check: ## Check code formatting without making changes
+	$(BLACK) --check src/ tests/
+
+quality: format-check lint ## Run all code quality checks
+
+##@ Application
+
+run-ui: ## Launch Gradio web interface
+	$(PYTHON) -m mosaic.gradio_app
+
+run-ui-public: ## Launch Gradio web interface with public sharing
+	$(PYTHON) -m mosaic.gradio_app --share
+
+run-single: ## Run single slide analysis (usage: make run-single SLIDE=path/to/slide.svs OUTPUT=output_dir)
+	$(PYTHON) -m mosaic.gradio_app --slide-path $(SLIDE) --output-dir $(OUTPUT)
+
+run-batch: ## Run batch analysis from CSV (usage: make run-batch CSV=settings.csv OUTPUT=output_dir)
+	$(PYTHON) -m mosaic.gradio_app --slide-csv $(CSV) --output-dir $(OUTPUT)
+
+##@ Docker
+
+docker-build: ## Build Docker image
+	docker build -t $(DOCKER_IMAGE_NAME):$(DOCKER_TAG) .
+
+docker-build-no-cache: ## Build Docker image without cache
+	docker build --no-cache -t $(DOCKER_IMAGE_NAME):$(DOCKER_TAG) .
+
+docker-run: ## Run Docker container (web UI mode)
+	docker run -it --rm \
+		--gpus all \
+		-p 7860:7860 \
+		-v $(PWD)/data:/app/data \
+		-v $(PWD)/output:/app/output \
+		$(DOCKER_IMAGE_NAME):$(DOCKER_TAG)
+
+docker-run-single: ## Run Docker container (single slide mode)
+	docker run -it --rm \
+		--gpus all \
+		-v $(PWD)/data:/app/data \
+		-v $(PWD)/output:/app/output \
+		$(DOCKER_IMAGE_NAME):$(DOCKER_TAG) \
+		--slide-path /app/data/$(SLIDE) \
+		--output-dir /app/output
+
+docker-run-batch: ## Run Docker container (batch mode)
+	docker run -it --rm \
+		--gpus all \
+		-v $(PWD)/data:/app/data \
+		-v $(PWD)/output:/app/output \
+		$(DOCKER_IMAGE_NAME):$(DOCKER_TAG) \
+		--slide-csv /app/data/$(CSV) \
+		--output-dir /app/output
+
+docker-shell: ## Open shell in Docker container
+	docker run -it --rm \
+		--gpus all \
+		-v $(PWD)/data:/app/data \
+		-v $(PWD)/output:/app/output \
+		$(DOCKER_IMAGE_NAME):$(DOCKER_TAG) \
+		/bin/bash
+
+docker-tag: ## Tag Docker image for registry
+	docker tag $(DOCKER_IMAGE_NAME):$(DOCKER_TAG) $(DOCKER_REGISTRY)/$(DOCKER_IMAGE_NAME):$(DOCKER_TAG)
+
+docker-push: docker-tag ## Push Docker image to registry
+	docker push $(DOCKER_REGISTRY)/$(DOCKER_IMAGE_NAME):$(DOCKER_TAG)
+
+docker-clean: ## Remove Docker image
+	docker rmi $(DOCKER_IMAGE_NAME):$(DOCKER_TAG) || true
+
+docker-prune: ## Clean up Docker build cache
+	docker system prune -f
+	docker builder prune -f
+
+##@ Cleanup
+
+clean: ## Remove build artifacts and cache files
+	find . -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true
+	find . -type d -name "*.egg-info" -exec rm -rf {} + 2>/dev/null || true
+	find . -type d -name ".pytest_cache" -exec rm -rf {} + 2>/dev/null || true
+	find . -type d -name ".ruff_cache" -exec rm -rf {} + 2>/dev/null || true
+	find . -type f -name "*.pyc" -delete
+	find . -type f -name "*.pyo" -delete
+	find . -type f -name ".coverage" -delete
+	rm -rf htmlcov/
+	rm -rf dist/
+	rm -rf build/
+
+clean-outputs: ## Remove output files (masks, results CSVs)
+	rm -rf output/*
+	@echo "Output directory cleaned"
+
+clean-all: clean docker-clean ## Remove all build artifacts, cache, and Docker images
+
+##@ Model Management
+
+download-models: ## Download required models from HuggingFace
+	@echo "Models will be downloaded automatically on first run"
+	$(PYTHON) -c "from mosaic.gradio_app import download_and_process_models; download_and_process_models()"
+
+##@ Documentation
+
+docs-requirements: ## Show what needs to be documented
+	@echo "Documentation TODO:"
+	@echo "  - API documentation"
+	@echo "  - Model architecture details"
+	@echo "  - CLI usage examples"
+	@echo "  - Docker deployment guide"
+
+##@ CI/CD
+
+ci-test: install-dev test-coverage format-check ## Run all CI checks (no lint to save time)
+	@echo "All CI checks passed!"
+
+ci-test-strict: install-dev test-coverage format-check lint ## Run all CI checks including pylint
+	@echo "All strict CI checks passed!"
+
+ci-docker: docker-build ## Build Docker image for CI
+	@echo "Docker image built successfully"
+
+##@ Development Utilities
+
+shell: ## Open Python shell with project in path
+	$(PYTHON)
+
+ipython: ## Open IPython shell with project in path
+	uv run ipython
+
+notebook: ## Start Jupyter notebook server
+	uv run jupyter notebook
+
+check-deps: ## Check for outdated dependencies
+	uv pip list --outdated
+
+update-deps: ## Update dependencies (be careful!)
+	uv sync --upgrade
+
+lock: ## Update lock file
+	uv lock
+
+##@ Git Hooks
+
+pre-commit-install: ## Install pre-commit hooks
+	@echo "Setting up pre-commit hooks..."
+	@echo "#!/bin/sh" > .git/hooks/pre-commit
+	@echo "make format-check test-fast" >> .git/hooks/pre-commit
+	@chmod +x .git/hooks/pre-commit
+	@echo "Pre-commit hooks installed (format-check + test-fast)"
+
+pre-commit-uninstall: ## Uninstall pre-commit hooks
+	rm -f .git/hooks/pre-commit
+	@echo "Pre-commit hooks uninstalled"
+
+##@ Information
+
+info: ## Display project information
+	@echo "Mosaic - H&E Whole Slide Image Analysis"
+	@echo "========================================"
+	@echo ""
+	@echo "Python version:"
+	@$(PYTHON) --version
+	@echo ""
+	@echo "UV version:"
+	@uv --version
+	@echo ""
+	@echo "Project structure:"
+	@echo "  src/mosaic/          - Main application code"
+	@echo "  tests/               - Test suite"
+	@echo "  data/                - Input data directory"
+	@echo "  output/              - Analysis results"
+	@echo ""
+	@echo "Key commands:"
+	@echo "  make install-dev     - Setup development environment"
+	@echo "  make test            - Run test suite"
+	@echo "  make run-ui          - Launch web interface"
+	@echo "  make docker-build    - Build Docker image"
+
+version: ## Show version information
+	@$(PYTHON) -c "import mosaic; print(f'Mosaic version: {mosaic.__version__}')" 2>/dev/null || echo "Version info not available"
+
+tree: ## Show project directory tree (requires tree command)
+	@tree -L 3 -I '__pycache__|*.pyc|*.egg-info|.pytest_cache|.ruff_cache|htmlcov|.venv' . || echo "tree command not found. Install with: apt-get install tree"
+
+##@ Performance
+
+profile: ## Profile a single slide analysis (usage: make profile SLIDE=path/to/slide.svs)
+	$(PYTHON) -m cProfile -o profile.stats -m mosaic.gradio_app --slide-path $(SLIDE) --output-dir profile_output
+	$(PYTHON) -c "import pstats; p = pstats.Stats('profile.stats'); p.sort_stats('cumulative'); p.print_stats(20)"
+
+benchmark: ## Run performance benchmarks
+	@echo "Running benchmark suite..."
+	@echo "This will process the test slide and measure performance"
+	time $(PYTHON) -m mosaic.gradio_app --slide-path tests/testdata/948176.svs --output-dir benchmark_output

src/mosaic/analysis.py

@@ -26,8 +26,10 @@ except ImportError:
                 return lambda f: f
             return fn
 
+
 # Detect T4 hardware by checking actual GPU
 import torch
+
 IS_T4_GPU = False
 GPU_NAME = "Unknown"
 if not IS_ZEROGPU and torch.cuda.is_available():
@@ -64,18 +66,21 @@ from mosaic.inference import run_aeon, run_paladin
 from mosaic.data_directory import get_data_directory
 
 # Log hardware detection at module load
-logger.info(f"Hardware: {GPU_TYPE} | batch_size={DEFAULT_BATCH_SIZE}, num_workers={DEFAULT_NUM_WORKERS}")
+logger.info(
+    f"Hardware: {GPU_TYPE} | batch_size={DEFAULT_BATCH_SIZE}, num_workers={DEFAULT_NUM_WORKERS}"
+)
+
 
+def _extract_ctranspath_features(coords, slide_path, attrs, num_workers, model):
+    """Extract CTransPath features on GPU using pre-loaded model.
 
-def _extract_ctranspath_features(coords, slide_path, attrs, num_workers):
-    """Extract CTransPath features on GPU.
-    
     Args:
         coords: Tissue tile coordinates
         slide_path: Path to the whole slide image file
         attrs: Slide attributes
         num_workers: Number of worker processes
-        
+        model: Pre-loaded CTransPath model from ModelCache
+
     Returns:
         tuple: (ctranspath_features, coords)
     """
@@ -86,87 +91,92 @@ def _extract_ctranspath_features(coords, slide_path, attrs, num_workers):
     elif IS_T4_GPU:
         num_workers = DEFAULT_NUM_WORKERS
         batch_size = DEFAULT_BATCH_SIZE
-        logger.info(f"Running CTransPath on T4: processing {len(coords)} tiles with batch_size={batch_size}")
+        logger.info(
+            f"Running CTransPath on T4: processing {len(coords)} tiles with batch_size={batch_size}"
+        )
     else:
         num_workers = max(num_workers, 8)
         batch_size = 64
         logger.info(f"Running CTransPath with {num_workers} workers")
-    
+
     start_time = pd.Timestamp.now()
 
-    data_dir = get_data_directory()
     ctranspath_features, _ = get_features(
         coords,
         slide_path,
         attrs,
-        model_type=ModelType.CTRANSPATH,
-        model_path=str(data_dir / "ctranspath.pth"),
+        model=model,
         num_workers=num_workers,
         batch_size=batch_size,
         use_gpu=True,
     )
-    
+
     end_time = pd.Timestamp.now()
     logger.info(f"CTransPath extraction took {end_time - start_time}")
-    
+
     return ctranspath_features, coords
 
 
-def _extract_optimus_features(filtered_coords, slide_path, attrs, num_workers):
-    """Extract Optimus features on GPU.
-    
+def _extract_optimus_features(filtered_coords, slide_path, attrs, num_workers, model):
+    """Extract Optimus features on GPU using pre-loaded model.
+
     Args:
         filtered_coords: Filtered tissue tile coordinates
         slide_path: Path to the whole slide image file
         attrs: Slide attributes
         num_workers: Number of worker processes
-        
+        model: Pre-loaded Optimus model from ModelCache
+
     Returns:
         Optimus features
     """
     if IS_ZEROGPU:
         num_workers = 0
         batch_size = 128
-        logger.info(f"Running Optimus on ZeroGPU: processing {len(filtered_coords)} tiles")
+        logger.info(
+            f"Running Optimus on ZeroGPU: processing {len(filtered_coords)} tiles"
+        )
     elif IS_T4_GPU:
         num_workers = DEFAULT_NUM_WORKERS
         batch_size = DEFAULT_BATCH_SIZE
-        logger.info(f"Running Optimus on T4: processing {len(filtered_coords)} tiles with batch_size={batch_size}")
+        logger.info(
+            f"Running Optimus on T4: processing {len(filtered_coords)} tiles with batch_size={batch_size}"
+        )
     else:
         num_workers = max(num_workers, 8)
         batch_size = 64
         logger.info(f"Running Optimus with {num_workers} workers")
-    
+
     start_time = pd.Timestamp.now()
 
-    data_dir = get_data_directory()
     features, _ = get_features(
         filtered_coords,
         slide_path,
         attrs,
-        model_type=ModelType.OPTIMUS,
-        model_path=str(data_dir / "optimus.pkl"),
+        model=model,
         num_workers=num_workers,
         batch_size=batch_size,
         use_gpu=True,
     )
-    
+
     end_time = pd.Timestamp.now()
     logger.info(f"Optimus extraction took {end_time - start_time}")
-    
+
     return features
 
 
-def _run_aeon_inference(features, site_type, num_workers, sex=None, tissue_site_idx=None):
+def _run_aeon_inference(
+    features, site_type, num_workers, sex=None, tissue_site_idx=None
+):
     """Run Aeon cancer subtype inference on GPU.
-    
+
     Args:
         features: Optimus features
         site_type: Site type ("Primary" or "Metastatic")
         num_workers: Number of worker processes
         sex: Patient sex (0=Male, 1=Female), optional
         tissue_site_idx: Tissue site index (0-56), optional
-        
+
     Returns:
         Aeon results DataFrame
     """
@@ -179,7 +189,7 @@ def _run_aeon_inference(features, site_type, num_workers, sex=None, tissue_site_
     else:
         num_workers = max(num_workers, 8)
         logger.info(f"Running Aeon with num_workers={num_workers}")
-    
+
     start_time = pd.Timestamp.now()
     logger.info("Running Aeon for cancer subtype inference")
     data_dir = get_data_directory()
@@ -194,7 +204,7 @@ def _run_aeon_inference(features, site_type, num_workers, sex=None, tissue_site_
         use_cpu=False,
     )
     end_time = pd.Timestamp.now()
-    
+
     # Log memory stats if CUDA is available
     if torch.cuda.is_available():
         try:
@@ -207,19 +217,19 @@ def _run_aeon_inference(features, site_type, num_workers, sex=None, tissue_site_
             logger.info(f"Aeon inference took {end_time - start_time}")
     else:
         logger.info(f"Aeon inference took {end_time - start_time}")
-    
+
     return aeon_results
 
 
 def _run_paladin_inference(features, aeon_results, site_type, num_workers):
     """Run Paladin biomarker inference on GPU.
-    
+
     Args:
         features: Optimus features
         aeon_results: Aeon results DataFrame
         site_type: Site type ("Primary" or "Metastatic")
         num_workers: Number of worker processes
-        
+
     Returns:
         Paladin results DataFrame
     """
@@ -232,7 +242,7 @@ def _run_paladin_inference(features, aeon_results, site_type, num_workers):
     else:
         num_workers = max(num_workers, 8)
         logger.info(f"Running Paladin with num_workers={num_workers}")
-    
+
     start_time = pd.Timestamp.now()
     logger.info("Running Paladin for biomarker inference")
     data_dir = get_data_directory()
@@ -246,7 +256,7 @@ def _run_paladin_inference(features, aeon_results, site_type, num_workers):
         use_cpu=False,
     )
     end_time = pd.Timestamp.now()
-    
+
     # Log memory stats if CUDA is available
     if torch.cuda.is_available():
         try:
@@ -259,7 +269,7 @@ def _run_paladin_inference(features, aeon_results, site_type, num_workers):
             logger.info(f"Paladin inference took {end_time - start_time}")
     else:
         logger.info(f"Paladin inference took {end_time - start_time}")
-    
+
     return paladin_results
 
 
@@ -278,8 +288,16 @@ def _run_inference_pipeline_free(
 ):
     """Run inference pipeline with 60s GPU limit (for free users)."""
     return _run_inference_pipeline_impl(
-        coords, slide_path, attrs, site_type, sex, tissue_site_idx,
-        cancer_subtype, cancer_subtype_name_map, num_workers, progress
+        coords,
+        slide_path,
+        attrs,
+        site_type,
+        sex,
+        tissue_site_idx,
+        cancer_subtype,
+        cancer_subtype_name_map,
+        num_workers,
+        progress,
     )
 
 
@@ -298,8 +316,16 @@ def _run_inference_pipeline_pro(
 ):
     """Run inference pipeline with 300s GPU limit (for PRO users)."""
     return _run_inference_pipeline_impl(
-        coords, slide_path, attrs, site_type, sex, tissue_site_idx,
-        cancer_subtype, cancer_subtype_name_map, num_workers, progress
+        coords,
+        slide_path,
+        attrs,
+        site_type,
+        sex,
+        tissue_site_idx,
+        cancer_subtype,
+        cancer_subtype_name_map,
+        num_workers,
+        progress,
     )
 
 
@@ -315,11 +341,10 @@ def _run_inference_pipeline_impl(
     num_workers,
     progress,
 ):
-    """Run complete inference pipeline with separate GPU calls.
+    """Run complete inference pipeline using model cache.
 
-    This function orchestrates the GPU operations by calling separate functions
-    for each GPU-intensive task, allowing HF Spaces to allocate GPU resources
-    independently for each operation.
+    This function loads models once and reuses them throughout the pipeline,
+    orchestrating GPU operations for feature extraction and inference.
 
     Args:
         coords: Tissue tile coordinates
@@ -336,59 +361,84 @@ def _run_inference_pipeline_impl(
             - aeon_results: DataFrame with cancer subtype predictions and confidence scores
             - paladin_results: DataFrame with biomarker predictions
     """
-    # Step 2: Extract CTransPath features
-    progress(0.3, desc="Extracting CTransPath features")
-    ctranspath_features, coords = _extract_ctranspath_features(
-        coords, slide_path, attrs, num_workers
-    )
+    # Load all models once for the entire pipeline
+    from mosaic.model_manager import load_all_models
 
-    # Step 3: Filter features using marker classifier (CPU operation)
-    start_time = pd.Timestamp.now()
-    data_dir = get_data_directory()
-    marker_classifier = pickle.load(open(data_dir / "marker_classifier.pkl", "rb"))
-    progress(0.35, desc="Filtering features with marker classifier")
-    logger.info("Filtering features with marker classifier")
-    _, filtered_coords = filter_features(
-        ctranspath_features,
-        coords,
-        marker_classifier,
-        threshold=0.25,
-    )
-    end_time = pd.Timestamp.now()
-    logger.info(f"Feature filtering took {end_time - start_time}")
-    logger.info(
-        f"Filtered from {len(coords)} to {len(filtered_coords)} tiles using marker classifier"
-    )
+    progress(0.1, desc="Loading models")
+    logger.info("Loading models for inference pipeline")
+    model_cache = load_all_models(use_gpu=True)
 
-    # Step 4: Extract Optimus features
-    progress(0.4, desc="Extracting Optimus features")
-    features = _extract_optimus_features(filtered_coords, slide_path, attrs, num_workers)
+    try:
+        # Step 2: Extract CTransPath features using cached model
+        progress(0.3, desc="Extracting CTransPath features")
+        ctranspath_features, coords = _extract_ctranspath_features(
+            coords, slide_path, attrs, num_workers, model=model_cache.ctranspath_model
+        )
 
-    # Step 5: Run Aeon to predict histology if not supplied
-    if cancer_subtype == "Unknown":
-        progress(0.9, desc="Running Aeon for cancer subtype inference")
-        aeon_results = _run_aeon_inference(features, site_type, num_workers, sex, tissue_site_idx)
-    else:
-        cancer_subtype_code = cancer_subtype_name_map.get(cancer_subtype)
-        aeon_results = pd.DataFrame(
-            {
-                "Cancer Subtype": [cancer_subtype_code],
-                "Confidence": [1.0],
-            }
+        # Step 3: Filter features using cached marker classifier
+        start_time = pd.Timestamp.now()
+        progress(0.35, desc="Filtering features with marker classifier")
+        logger.info("Filtering features with marker classifier")
+        _, filtered_coords = filter_features(
+            ctranspath_features,
+            coords,
+            model_cache.marker_classifier,
+            threshold=0.25,
+        )
+        end_time = pd.Timestamp.now()
+        logger.info(f"Feature filtering took {end_time - start_time}")
+        logger.info(
+            f"Filtered from {len(coords)} to {len(filtered_coords)} tiles using marker classifier"
         )
-        logger.info(f"Using user-supplied cancer subtype: {cancer_subtype}")
 
-    # Step 6: Run Paladin to predict biomarkers
-    if len(aeon_results) == 0:
-        logger.warning("No Aeon results, skipping Paladin inference")
-        return None, None
-    
-    progress(0.95, desc="Running Paladin for biomarker inference")
-    paladin_results = _run_paladin_inference(features, aeon_results, site_type, num_workers)
+        # Step 4: Extract Optimus features using cached model
+        progress(0.4, desc="Extracting Optimus features")
+        features = _extract_optimus_features(
+            filtered_coords,
+            slide_path,
+            attrs,
+            num_workers,
+            model=model_cache.optimus_model,
+        )
 
-    aeon_results.set_index("Cancer Subtype", inplace=True)
+        # Step 5: Run Aeon to predict histology if not supplied
+        if cancer_subtype == "Unknown":
+            progress(0.9, desc="Running Aeon for cancer subtype inference")
+            aeon_results = _run_aeon_inference_with_model(
+                features,
+                model_cache.aeon_model,
+                model_cache.device,
+                site_type,
+                num_workers,
+                sex,
+                tissue_site_idx,
+            )
+        else:
+            cancer_subtype_code = cancer_subtype_name_map.get(cancer_subtype)
+            aeon_results = pd.DataFrame(
+                {
+                    "Cancer Subtype": [cancer_subtype_code],
+                    "Confidence": [1.0],
+                }
+            )
+            logger.info(f"Using user-supplied cancer subtype: {cancer_subtype}")
 
-    return aeon_results, paladin_results
+        # Step 6: Run Paladin to predict biomarkers
+        if len(aeon_results) == 0:
+            logger.warning("No Aeon results, skipping Paladin inference")
+            return None, None
+
+        progress(0.95, desc="Running Paladin for biomarker inference")
+        paladin_results = _run_paladin_inference_with_models(
+            features, aeon_results, site_type, model_cache, num_workers
+        )
+
+        aeon_results.set_index("Cancer Subtype", inplace=True)
+
+        return aeon_results, paladin_results
+    finally:
+        # Clean up models to free GPU memory
+        model_cache.cleanup()
 
 
 # ============================================================================
@@ -531,11 +581,10 @@ def _run_inference_pipeline_with_models(
     Returns:
         Tuple of (aeon_results, paladin_results)
     """
-    # Step 1: Extract CTransPath features (still uses mussel's get_features)
-    # Note: Feature extraction optimization can be added later if needed
+    # Step 1: Extract CTransPath features with PRE-LOADED model
     progress(0.3, desc="Extracting CTransPath features")
     ctranspath_features, coords = _extract_ctranspath_features(
-        coords, slide_path, attrs, num_workers
+        coords, slide_path, attrs, num_workers, model=model_cache.ctranspath_model
     )
 
     # Step 2: Filter features using pre-loaded marker classifier
@@ -554,9 +603,11 @@ def _run_inference_pipeline_with_models(
         f"Filtered from {len(coords)} to {len(filtered_coords)} tiles using marker classifier"
     )
 
-    # Step 3: Extract Optimus features (still uses mussel's get_features)
+    # Step 3: Extract Optimus features with PRE-LOADED model
     progress(0.5, desc="Extracting Optimus features")
-    features = _extract_optimus_features(filtered_coords, slide_path, attrs, num_workers)
+    features = _extract_optimus_features(
+        filtered_coords, slide_path, attrs, num_workers, model=model_cache.optimus_model
+    )
 
     # Step 4: Run Aeon inference with pre-loaded model (if cancer subtype unknown)
     aeon_results = None
@@ -564,7 +615,9 @@ def _run_inference_pipeline_with_models(
 
     # Check if cancer subtype is unknown
     if cancer_subtype in ["Unknown", None]:
-        logger.info("Running Aeon inference with PRE-LOADED model (cancer subtype unknown)")
+        logger.info(
+            "Running Aeon inference with PRE-LOADED model (cancer subtype unknown)"
+        )
         aeon_results = _run_aeon_inference_with_model(
             features,
             model_cache.aeon_model,  # Use pre-loaded Aeon model
@@ -593,116 +646,7 @@ def _run_inference_pipeline_with_models(
     return aeon_results, paladin_results
 
 
-def analyze_slide_with_models(
-    slide_path,
-    seg_config,
-    site_type,
-    sex,
-    tissue_site,
-    cancer_subtype,
-    cancer_subtype_name_map,
-    model_cache,
-    ihc_subtype="",
-    num_workers=4,
-    progress=None,
-):
-    """Analyze a slide using pre-loaded models (batch-optimized version).
-
-    This function is optimized for batch processing where models are loaded once
-    in a ModelCache and reused across multiple slides.
-
-    Args:
-        slide_path: Path to the slide file
-        seg_config: Segmentation configuration ("Biopsy", "Resection", or "TCGA")
-        site_type: "Primary" or "Metastatic"
-        sex: Patient sex ("Unknown", "Male", "Female")
-        tissue_site: Tissue site name
-        cancer_subtype: Known cancer subtype or "Unknown"
-        cancer_subtype_name_map: Dict mapping display names to OncoTree codes
-        model_cache: ModelCache instance with pre-loaded models
-        ihc_subtype: IHC subtype for breast cancer (optional)
-        num_workers: Number of workers for data loading
-        progress: Gradio progress tracker
-
-    Returns:
-        Tuple of (slide_mask, aeon_results, paladin_results)
-    """
-    from mosaic.inference.data import encode_sex, encode_tissue_site
-
-    if progress is None:
-        progress = lambda frac, desc: None  # No-op progress function
-
-    # Encode sex and tissue site
-    sex_idx = encode_sex(sex) if sex else None
-    tissue_site_idx = encode_tissue_site(tissue_site) if tissue_site else None
-
-    # Step 1: Convert seg_config string to config object
-    if isinstance(seg_config, str):
-        if seg_config == "Biopsy":
-            seg_config = BiopsySegConfig()
-        elif seg_config == "Resection":
-            seg_config = ResectionSegConfig()
-        elif seg_config == "TCGA":
-            seg_config = TcgaSegConfig()
-        else:
-            raise ValueError(f"Unknown segmentation configuration: {seg_config}")
-
-    # Step 2: Tissue segmentation (CPU operation, not affected by model caching)
-    progress(0.0, desc="Segmenting tissue")
-    logger.info(f"Segmenting tissue for slide: {slide_path}")
-    start_time = pd.Timestamp.now()
-
-    if values := segment_tissue(
-        slide_path=slide_path,
-        patch_size=224,
-        mpp=0.5,
-        seg_level=-1,
-        segment_threshold=seg_config.segment_threshold,
-        median_blur_ksize=seg_config.median_blur_ksize,
-        morphology_ex_kernel=seg_config.morphology_ex_kernel,
-        tissue_area_threshold=seg_config.tissue_area_threshold,
-        hole_area_threshold=seg_config.hole_area_threshold,
-        max_num_holes=seg_config.max_num_holes,
-    ):
-        polygon, _, coords, attrs = values
-    else:
-        logger.warning("No tissue detected in slide")
-        return None, None, None
-
-    end_time = pd.Timestamp.now()
-    logger.info(f"Tissue segmentation took {end_time - start_time}")
-    logger.info(f"Found {len(coords)} tissue tiles")
-
-    if len(coords) == 0:
-        logger.warning("No tissue tiles found in slide")
-        return None, None, None
-
-    # Step 2: Create slide mask visualization (CPU operation)
-    progress(0.2, desc="Creating slide mask")
-    logger.info("Drawing slide mask")
-    slide_mask = draw_slide_mask(
-        slide_path, polygon, outline="black", fill=(255, 0, 0, 80), vis_level=-1
-    )
-    logger.info("Slide mask drawn")
-
-    # Step 3: Run inference pipeline with pre-loaded models
-    aeon_results, paladin_results = _run_inference_pipeline_with_models(
-        coords,
-        slide_path,
-        attrs,
-        site_type,
-        sex_idx,
-        tissue_site_idx,
-        cancer_subtype,
-        cancer_subtype_name_map,
-        model_cache,
-        num_workers,
-        progress,
-    )
-
-    progress(1.0, desc="Analysis complete")
-
-    return slide_mask, aeon_results, paladin_results
+# Removed: analyze_slide_with_models merged into analyze_slide below
 
 
 def analyze_slide(
@@ -717,26 +661,27 @@ def analyze_slide(
     num_workers=4,
     progress=gr.Progress(track_tqdm=True),
     request: gr.Request = None,
+    model_cache=None,
 ):
     """Analyze a whole slide image for cancer subtype and biomarker prediction.
 
-    This function performs a complete analysis pipeline including:
-    1. Tissue segmentation (CPU-only, no GPU required)
-    2. GPU-intensive feature extraction and model inference
-
-    The GPU-intensive operations are handled by a separate function decorated
-    with @spaces.GPU to efficiently manage GPU resources on Hugging Face Spaces.
-    Tissue segmentation runs on CPU and is not included in the GPU allocation.
+    This function works in two modes:
+    1. **Single-slide mode** (model_cache=None): Loads models, analyzes one slide, cleans up
+    2. **Batch mode** (model_cache provided): Uses pre-loaded models for efficiency
 
     Args:
         slide_path: Path to the whole slide image file
         seg_config: Segmentation configuration, one of "Biopsy", "Resection", or "TCGA"
         site_type: Site type, either "Primary" or "Metastatic"
+        sex: Patient sex ("Unknown", "Male", "Female")
+        tissue_site: Tissue site name
         cancer_subtype: Cancer subtype (OncoTree code or "Unknown" for inference)
         cancer_subtype_name_map: Dictionary mapping cancer subtype names to codes
         ihc_subtype: IHC subtype for breast cancer (optional)
         num_workers: Number of worker processes for feature extraction
         progress: Gradio progress tracker for UI updates
+        request: Gradio request object (for HF Spaces authentication)
+        model_cache: Optional ModelCache with pre-loaded models (for batch processing)
 
     Returns:
         tuple: (slide_mask, aeon_results, paladin_results)
@@ -795,51 +740,6 @@ def analyze_slide(
     )
     logger.info("Slide mask drawn")
 
-    # Step 2-6: Run inference pipeline with GPU
-    # Check if user is logged in for longer GPU duration
-    is_logged_in = False
-    username = "anonymous"
-    if request is not None:
-        try:
-            # Check if user is logged in via JWT token in referer
-            # HF Spaces doesn't populate request.username but includes JWT in URL
-            if hasattr(request, 'headers'):
-                referer = request.headers.get('referer', '')
-                if '__sign=' in referer:
-                    # Extract and decode JWT token
-                    import re
-                    import json
-                    import base64
-                    
-                    match = re.search(r'__sign=([^&]+)', referer)
-                    if match:
-                        token = match.group(1)
-                        try:
-                            # JWT format: header.payload.signature
-                            # We only need the payload (middle part)
-                            parts = token.split('.')
-                            if len(parts) == 3:
-                                # Decode base64 payload (add padding if needed)
-                                payload = parts[1]
-                                payload += '=' * (4 - len(payload) % 4)
-                                decoded = base64.urlsafe_b64decode(payload)
-                                token_data = json.loads(decoded)
-                                
-                                # Check if user is in token
-                                if 'onBehalfOf' in token_data and 'user' in token_data['onBehalfOf']:
-                                    username = token_data['onBehalfOf']['user']
-                                    is_logged_in = True
-                                    logger.info(f"Found user in JWT token: {username}")
-                        except Exception as e:
-                            logger.warning(f"Failed to decode JWT: {e}")
-            
-            if IS_ZEROGPU:
-                logger.info(f"User: {username} | Logged in: {is_logged_in}")
-        except Exception as e:
-            logger.warning(f"Failed to detect user: {e}")
-            import traceback
-            logger.warning(traceback.format_exc())
-
     # Convert sex and tissue_site to indices for Aeon model
     from mosaic.inference.data import encode_sex, encode_tissue_site
 
@@ -851,10 +751,11 @@ def analyze_slide(
     if tissue_site is not None:
         tissue_site_idx = encode_tissue_site(tissue_site)
 
-    if is_logged_in:
-        if IS_ZEROGPU:
-            logger.info("Using 300s GPU allocation (logged-in user)")
-        aeon_results, paladin_results = _run_inference_pipeline_pro(
+    # Run inference pipeline - two modes based on model_cache
+    if model_cache is not None:
+        # Batch mode: use pre-loaded models
+        logger.info("Using pre-loaded models from ModelCache (batch mode)")
+        aeon_results, paladin_results = _run_inference_pipeline_with_models(
             coords,
             slide_path,
             attrs,
@@ -863,23 +764,91 @@ def analyze_slide(
             tissue_site_idx,
             cancer_subtype,
             cancer_subtype_name_map,
+            model_cache,
             num_workers,
             progress,
         )
     else:
-        if IS_ZEROGPU:
-            logger.info("Using 60s GPU allocation (anonymous user)")
-        aeon_results, paladin_results = _run_inference_pipeline_free(
-            coords,
-            slide_path,
-            attrs,
-            site_type,
-            sex_idx,
-            tissue_site_idx,
-            cancer_subtype,
-            cancer_subtype_name_map,
-            num_workers,
-            progress,
-        )
+        # Single-slide mode: load models on-demand
+        # Check if user is logged in for longer GPU duration (HF Spaces only)
+        is_logged_in = False
+        username = "anonymous"
+        if request is not None:
+            try:
+                # Check if user is logged in via JWT token in referer
+                # HF Spaces doesn't populate request.username but includes JWT in URL
+                if hasattr(request, "headers"):
+                    referer = request.headers.get("referer", "")
+                    if "__sign=" in referer:
+                        # Extract and decode JWT token
+                        import re
+                        import json
+                        import base64
+
+                        match = re.search(r"__sign=([^&]+)", referer)
+                        if match:
+                            token = match.group(1)
+                            try:
+                                # JWT format: header.payload.signature
+                                # We only need the payload (middle part)
+                                parts = token.split(".")
+                                if len(parts) == 3:
+                                    # Decode base64 payload (add padding if needed)
+                                    payload = parts[1]
+                                    payload += "=" * (4 - len(payload) % 4)
+                                    decoded = base64.urlsafe_b64decode(payload)
+                                    token_data = json.loads(decoded)
+
+                                    # Check if user is in token
+                                    if (
+                                        "onBehalfOf" in token_data
+                                        and "user" in token_data["onBehalfOf"]
+                                    ):
+                                        username = token_data["onBehalfOf"]["user"]
+                                        is_logged_in = True
+                                        logger.info(
+                                            f"Found user in JWT token: {username}"
+                                        )
+                            except Exception as e:
+                                logger.warning(f"Failed to decode JWT: {e}")
+
+                if IS_ZEROGPU:
+                    logger.info(f"User: {username} | Logged in: {is_logged_in}")
+            except Exception as e:
+                logger.warning(f"Failed to detect user: {e}")
+                import traceback
+
+                logger.warning(traceback.format_exc())
+
+        if is_logged_in:
+            if IS_ZEROGPU:
+                logger.info("Using 300s GPU allocation (logged-in user)")
+            aeon_results, paladin_results = _run_inference_pipeline_pro(
+                coords,
+                slide_path,
+                attrs,
+                site_type,
+                sex_idx,
+                tissue_site_idx,
+                cancer_subtype,
+                cancer_subtype_name_map,
+                num_workers,
+                progress,
+            )
+        else:
+            if IS_ZEROGPU:
+                logger.info("Using 60s GPU allocation (anonymous user)")
+            aeon_results, paladin_results = _run_inference_pipeline_free(
+                coords,
+                slide_path,
+                attrs,
+                site_type,
+                sex_idx,
+                tissue_site_idx,
+                cancer_subtype,
+                cancer_subtype_name_map,
+                num_workers,
+                progress,
+            )
 
     return slide_mask, aeon_results, paladin_results

src/mosaic/batch_analysis.py

@@ -1,238 +0,0 @@
-"""Batch processing coordinator for multi-slide analysis.
-
-This module provides optimized batch processing functionality that loads
-models once and reuses them across multiple slides, significantly reducing
-overhead compared to processing slides individually.
-"""
-
-from typing import Dict, List, Optional, Tuple
-import pandas as pd
-import time
-from loguru import logger
-
-from mosaic.model_manager import load_all_models
-from mosaic.analysis import analyze_slide_with_models
-
-
-def analyze_slides_batch(
-    slides: List[str],
-    settings_df: pd.DataFrame,
-    cancer_subtype_name_map: Dict[str, str],
-    num_workers: int = 4,
-    aggressive_memory_mgmt: Optional[bool] = None,
-    progress=None,
-) -> Tuple[List[Tuple], List[pd.DataFrame], List[pd.DataFrame]]:
-    """Analyze multiple slides with models loaded once for batch processing.
-
-    This function provides significant performance improvements over sequential
-    processing by loading all models once at the start, processing all slides
-    with the pre-loaded models, and cleaning up at the end.
-
-    Performance Benefits:
-    - ~90% reduction in model loading operations
-    - 25-45% overall speedup depending on model loading overhead
-    - Memory-efficient: same peak memory as single-slide processing
-
-    Args:
-        slides: List of slide file paths
-        settings_df: DataFrame with columns matching SETTINGS_COLUMNS from ui/utils.py
-        cancer_subtype_name_map: Dict mapping cancer subtype display names to OncoTree codes
-        num_workers: Number of CPU workers for data loading (default: 4)
-        aggressive_memory_mgmt: Memory management strategy:
-            - None: Auto-detect based on GPU type (T4 = True, A100 = False)
-            - True: T4-style aggressive cleanup (load/delete Paladin models per slide)
-            - False: Cache Paladin models across slides (requires >40GB GPU memory)
-        progress: Optional Gradio progress tracker
-
-    Returns:
-        Tuple of (all_slide_masks, all_aeon_results, all_paladin_results):
-            - all_slide_masks: List of (slide_mask_image, slide_name) tuples
-            - all_aeon_results: List of DataFrames with Aeon cancer subtype predictions
-            - all_paladin_results: List of DataFrames with Paladin biomarker predictions
-
-    Example:
-        ```python
-        slides = ["slide1.svs", "slide2.svs", "slide3.svs"]
-        settings_df = pd.DataFrame({
-            "Slide": ["slide1.svs", "slide2.svs", "slide3.svs"],
-            "Site Type": ["Primary", "Primary", "Metastatic"],
-            "Sex": ["Male", "Female", "Unknown"],
-            "Tissue Site": ["Lung", "Breast", "Unknown"],
-            "Cancer Subtype": ["Unknown", "Unknown", "LUAD"],
-            "IHC Subtype": ["", "HR+/HER2-", ""],
-            "Segmentation Config": ["Biopsy", "Resection", "Biopsy"],
-        })
-
-        masks, aeon, paladin = analyze_slides_batch(
-            slides, settings_df, cancer_subtype_name_map
-        )
-        ```
-
-    Notes:
-        - GPU memory requirements: ~9-15GB for typical batches
-        - T4 GPUs (16GB): Uses aggressive memory management automatically
-        - A100 GPUs (80GB): Can cache Paladin models for better performance
-        - Maintains backward compatibility: single slides can still use analyze_slide()
-    """
-    if progress is None:
-        progress = lambda frac, desc: None  # No-op progress function
-
-    num_slides = len(slides)
-    batch_start_time = time.time()
-
-    logger.info("=" * 80)
-    logger.info(f"BATCH PROCESSING: Starting analysis of {num_slides} slides")
-    logger.info("=" * 80)
-
-    # Step 1: Load all models once
-    progress(0.0, desc="Loading models for batch processing")
-    model_load_start = time.time()
-
-    try:
-        model_cache = load_all_models(
-            use_gpu=True,
-            aggressive_memory_mgmt=aggressive_memory_mgmt,
-        )
-
-        model_load_time = time.time() - model_load_start
-        logger.info(f"Model loading completed in {model_load_time:.2f}s")
-        logger.info("")
-
-        # Log memory strategy
-        if model_cache.aggressive_memory_mgmt:
-            logger.info(
-                "Memory strategy: AGGRESSIVE (T4-style) - "
-                "Paladin models loaded/freed per slide"
-            )
-        else:
-            logger.info(
-                "Using caching strategy (A100-style): "
-                "Paladin models will be cached across slides"
-            )
-
-    except Exception as e:
-        logger.error(f"Failed to load models: {e}")
-        raise
-
-    # Step 2: Process each slide with pre-loaded models
-    all_slide_masks = []
-    all_aeon_results = []
-    all_paladin_results = []
-    slide_times = []
-
-    logger.info("=" * 80)
-    logger.info("Processing slides with PRE-LOADED models (no model reloading!)")
-    logger.info("=" * 80)
-
-    try:
-        for idx, (slide_path, (_, row)) in enumerate(zip(slides, settings_df.iterrows())):
-            slide_name = slide_path.split("/")[-1] if "/" in slide_path else slide_path
-
-            # Update progress
-            progress_frac = (idx + 0.1) / num_slides
-            progress(progress_frac, desc=f"Analyzing slide {idx + 1}/{num_slides}: {slide_name}")
-
-            logger.info("")
-            logger.info(f"[{idx + 1}/{num_slides}] Processing: {slide_name}")
-            logger.info(f"         Using pre-loaded models (no disk I/O for core models)")
-            slide_start_time = time.time()
-
-            try:
-                # Use batch-optimized analysis with pre-loaded models
-                slide_mask, aeon_results, paladin_results = analyze_slide_with_models(
-                    slide_path=slide_path,
-                    seg_config=row["Segmentation Config"],
-                    site_type=row["Site Type"],
-                    sex=row.get("Sex", "Unknown"),
-                    tissue_site=row.get("Tissue Site", "Unknown"),
-                    cancer_subtype=row["Cancer Subtype"],
-                    cancer_subtype_name_map=cancer_subtype_name_map,
-                    model_cache=model_cache,
-                    ihc_subtype=row.get("IHC Subtype", ""),
-                    num_workers=num_workers,
-                    progress=progress,
-                )
-
-                slide_time = time.time() - slide_start_time
-                slide_times.append(slide_time)
-
-                # Collect results
-                if slide_mask is not None:
-                    all_slide_masks.append((slide_mask, slide_name))
-
-                if aeon_results is not None:
-                    # Add slide name to results for multi-slide batches
-                    if num_slides > 1:
-                        aeon_results.columns = [f"{slide_name}"]
-                    all_aeon_results.append(aeon_results)
-
-                if paladin_results is not None:
-                    # Add slide name column
-                    paladin_results.insert(
-                        0, "Slide", pd.Series([slide_name] * len(paladin_results))
-                    )
-                    all_paladin_results.append(paladin_results)
-
-                logger.info(f"[{idx + 1}/{num_slides}] ✓ Completed in {slide_time:.2f}s")
-
-            except Exception as e:
-                slide_time = time.time() - slide_start_time
-                slide_times.append(slide_time)
-                logger.exception(f"[{idx + 1}/{num_slides}] ✗ Failed after {slide_time:.2f}s: {e}")
-                # Continue with next slide instead of failing entire batch
-                continue
-
-    finally:
-        # Step 3: Always cleanup models (even if there were errors)
-        logger.info("")
-        logger.info("=" * 80)
-        logger.info("Cleaning up models...")
-        progress(0.99, desc="Cleaning up models")
-        model_cache.cleanup()
-        logger.info("✓ Model cleanup complete")
-
-    # Calculate batch statistics
-    batch_total_time = time.time() - batch_start_time
-    num_successful = len(all_slide_masks)
-    num_failed = num_slides - num_successful
-
-    # Log comprehensive summary
-    logger.info("=" * 80)
-    logger.info("BATCH PROCESSING SUMMARY")
-    logger.info("=" * 80)
-    logger.info(f"Total slides:        {num_slides}")
-    logger.info(f"Successfully processed: {num_successful}")
-    logger.info(f"Failed:              {num_failed}")
-    logger.info("")
-    logger.info(f"Model loading time:  {model_load_time:.2f}s (done ONCE for entire batch)")
-    logger.info(f"Total batch time:    {batch_total_time:.2f}s")
-
-    if slide_times:
-        avg_slide_time = sum(slide_times) / len(slide_times)
-        min_slide_time = min(slide_times)
-        max_slide_time = max(slide_times)
-        total_slide_time = sum(slide_times)
-
-        logger.info("")
-        logger.info("Per-slide processing times:")
-        logger.info(f"  Average: {avg_slide_time:.2f}s")
-        logger.info(f"  Min:     {min_slide_time:.2f}s")
-        logger.info(f"  Max:     {max_slide_time:.2f}s")
-        logger.info(f"  Total:   {total_slide_time:.2f}s")
-
-        # Calculate efficiency
-        overhead_time = batch_total_time - total_slide_time
-        logger.info("")
-        logger.info(f"Batch overhead:      {overhead_time:.2f}s ({overhead_time/batch_total_time*100:.1f}%)")
-        logger.info(f"Slide processing:    {total_slide_time:.2f}s ({total_slide_time/batch_total_time*100:.1f}%)")
-
-    logger.info("")
-    logger.info("✓ Batch processing optimization benefits:")
-    logger.info("  - Models loaded ONCE (not once per slide)")
-    logger.info("  - Reduced disk I/O for model loading")
-    logger.info(f"  - Processed {num_slides} slides with shared model cache")
-    logger.info("=" * 80)
-
-    progress(1.0, desc=f"Batch analysis complete ({num_successful}/{num_slides} successful)")
-
-    return all_slide_masks, all_aeon_results, all_paladin_results

src/mosaic/gradio_app.py

@@ -25,15 +25,15 @@ from mosaic.ui.utils import (
     SEX_OPTIONS,
 )
 from mosaic.analysis import analyze_slide
-from mosaic.batch_analysis import analyze_slides_batch
+from mosaic.model_manager import load_all_models
 
 
 def download_and_process_models():
-    """Download models from HuggingFace and initialize cancer subtype mappings.
+    """Download essential models from HuggingFace and initialize cancer subtype mappings.
 
-    Downloads the Paladin and Aeon models from the PDM-Group HuggingFace repository
-    to the HuggingFace cache directory and creates mappings between cancer subtype
-    names and OncoTree codes.
+    Downloads only the core models (CTransPath, Optimus, Aeon, marker classifier) and
+    metadata files from the PDM-Group HuggingFace repository. Paladin models are
+    downloaded on-demand when needed for inference.
 
     Returns:
         tuple: (cancer_subtype_name_map, reversed_cancer_subtype_name_map, cancer_subtypes)
@@ -41,47 +41,69 @@ def download_and_process_models():
             - reversed_cancer_subtype_name_map: Dict mapping OncoTree codes to display names
             - cancer_subtypes: List of all supported cancer subtype codes
     """
-    # Download to HF cache directory (not local_dir)
-    # This returns the path to the cached snapshot
-    logger.info("Downloading models from HuggingFace Hub to cache directory...")
+    # Download only essential files to HF cache directory
+    # Paladin models will be downloaded on-demand
+    logger.info(
+        "Downloading essential models from HuggingFace Hub (Paladin models loaded on-demand)..."
+    )
     cache_dir = snapshot_download(
         repo_id="PDM-Group/paladin-aeon-models",
+        allow_patterns=[
+            "*.csv",  # Model maps and metadata
+            "ctranspath.pth",  # CTransPath model
+            "aeon_model.pkl",  # Aeon model
+            "marker_classifier.pkl",  # Marker classifier
+            "tissue_site_*",  # Tissue site mappings
+        ],
         # No local_dir - use HF cache
     )
-    logger.info(f"Models downloaded to: {cache_dir}")
+    logger.info(f"Essential models downloaded to: {cache_dir}")
 
     # Set the data directory for other modules to use
     set_data_directory(cache_dir)
 
+    # Pre-download Optimus model from bioptimus/H-optimus-0
+    # This ensures it's cached at startup since it's needed for every slide
+    logger.info("Pre-downloading Optimus model from bioptimus/H-optimus-0...")
+    from mussel.models import ModelType, get_model_factory
+
+    optimus_factory = get_model_factory(ModelType.OPTIMUS)
+    # This will trigger the download and cache the model
+    _ = optimus_factory.get_model(
+        model_path="hf-hub:bioptimus/H-optimus-0",
+        use_gpu=False,  # Just download, don't load to GPU yet
+        gpu_device_id=None,
+    )
+    logger.info("✓ Optimus model cached")
+
     model_map = pd.read_csv(
         Path(cache_dir) / "paladin_model_map.csv",
     )
     cancer_subtypes = model_map["cancer_subtype"].unique().tolist()
     cancer_subtype_name_map = {"Unknown": "UNK"}
-    cancer_subtype_name_map.update({
-        f"{get_oncotree_code_name(code)} ({code})": code for code in cancer_subtypes
-    })
+    cancer_subtype_name_map.update(
+        {f"{get_oncotree_code_name(code)} ({code})": code for code in cancer_subtypes}
+    )
     reversed_cancer_subtype_name_map = {
         value: key for key, value in cancer_subtype_name_map.items()
     }
-    
-    # Set the global maps in the UI module
-    set_cancer_subtype_maps(cancer_subtype_name_map, reversed_cancer_subtype_name_map, cancer_subtypes)
-    
-    return cancer_subtype_name_map, reversed_cancer_subtype_name_map, cancer_subtypes
-
 
+    # Set the global maps in the UI module
+    set_cancer_subtype_maps(
+        cancer_subtype_name_map, reversed_cancer_subtype_name_map, cancer_subtypes
+    )
 
+    return cancer_subtype_name_map, reversed_cancer_subtype_name_map, cancer_subtypes
 
 
 def main():
     """Main entry point for the Mosaic application.
-    
+
     Parses command-line arguments and routes to the appropriate mode:
     - Single slide processing (--slide-path)
     - Batch processing (--slide-csv)
     - Web interface (default, no slide arguments)
-    
+
     Command-line arguments control analysis parameters like site type,
     cancer subtype, segmentation configuration, and output directory.
     """
@@ -160,7 +182,9 @@ def main():
         logger.add("debug.log", level="DEBUG")
         logger.debug("Debug logging enabled")
 
-    cancer_subtype_name_map, reversed_cancer_subtype_name_map, cancer_subtypes = download_and_process_models()
+    cancer_subtype_name_map, reversed_cancer_subtype_name_map, cancer_subtypes = (
+        download_and_process_models()
+    )
 
     if args.slide_path and not args.slide_csv:
         # Single slide processing mode
@@ -180,7 +204,12 @@ def main():
             ],
             columns=SETTINGS_COLUMNS,
         )
-        settings_df = validate_settings(settings_df, cancer_subtype_name_map, cancer_subtypes, reversed_cancer_subtype_name_map)
+        settings_df = validate_settings(
+            settings_df,
+            cancer_subtype_name_map,
+            cancer_subtypes,
+            reversed_cancer_subtype_name_map,
+        )
         slide_mask, aeon_results, paladin_results = analyze_slide(
             args.slide_path,
             args.segmentation_config,
@@ -218,24 +247,62 @@ def main():
         # Load and validate settings
         settings_df = load_settings(args.slide_csv)
         settings_df = validate_settings(
-            settings_df, cancer_subtype_name_map, cancer_subtypes, reversed_cancer_subtype_name_map
+            settings_df,
+            cancer_subtype_name_map,
+            cancer_subtypes,
+            reversed_cancer_subtype_name_map,
         )
 
         # Extract slide paths
         slides = settings_df["Slide"].tolist()
 
-        logger.info(f"Processing {len(slides)} slides in batch mode with models loaded once")
-
-        # Use batch processing (models loaded once)
-        all_slide_masks, all_aeon_results, all_paladin_results = analyze_slides_batch(
-            slides=slides,
-            settings_df=settings_df,
-            cancer_subtype_name_map=cancer_subtype_name_map,
-            num_workers=args.num_workers,
-            aggressive_memory_mgmt=None,  # Auto-detect GPU type
-            progress=None,
+        logger.info(
+            f"Processing {len(slides)} slides in batch mode with models loaded once"
         )
 
+        # Load models once for batch processing
+        model_cache = load_all_models(use_gpu=True, aggressive_memory_mgmt=None)
+
+        all_slide_masks = []
+        all_aeon_results = []
+        all_paladin_results = []
+
+        try:
+            # Process each slide with pre-loaded models
+            for idx, slide_path in enumerate(slides):
+                row = settings_df.iloc[idx]
+                slide_name = row["Slide"]
+
+                logger.info(f"[{idx + 1}/{len(slides)}] Processing: {slide_name}")
+
+                slide_mask, aeon_results, paladin_results = analyze_slide(
+                    slide_path=slide_path,
+                    seg_config=row["Segmentation Config"],
+                    site_type=row["Site Type"],
+                    sex=row.get("Sex", "Unknown"),
+                    tissue_site=row.get("Tissue Site", "Unknown"),
+                    cancer_subtype=row["Cancer Subtype"],
+                    cancer_subtype_name_map=cancer_subtype_name_map,
+                    ihc_subtype=row.get("IHC Subtype", ""),
+                    num_workers=args.num_workers,
+                    progress=lambda frac, desc: None,  # No-op progress for CLI
+                    request=None,
+                    model_cache=model_cache,
+                )
+
+                if slide_mask is not None:
+                    all_slide_masks.append((slide_mask, slide_name))
+                if aeon_results is not None:
+                    all_aeon_results.append(aeon_results)
+                if paladin_results is not None:
+                    paladin_results.insert(
+                        0, "Slide", pd.Series([slide_name] * len(paladin_results))
+                    )
+                    all_paladin_results.append(paladin_results)
+        finally:
+            logger.info("Cleaning up model cache")
+            model_cache.cleanup()
+
         # Save individual slide results
         for idx, (slide_mask, slide_name) in enumerate(all_slide_masks):
             mask_path = output_dir / f"{slide_name}_mask.png"
@@ -252,7 +319,9 @@ def main():
         if all_paladin_results:
             combined_paladin = pd.concat(all_paladin_results, ignore_index=True)
             for slide_name in combined_paladin["Slide"].unique():
-                slide_paladin = combined_paladin[combined_paladin["Slide"] == slide_name]
+                slide_paladin = combined_paladin[
+                    combined_paladin["Slide"] == slide_name
+                ]
                 paladin_output_path = output_dir / f"{slide_name}_paladin_results.csv"
                 slide_paladin.to_csv(paladin_output_path, index=False)
                 logger.info(f"Saved Paladin results to {paladin_output_path}")

src/mosaic/inference/aeon.py

@@ -80,8 +80,12 @@ def run_with_model(
         target_dict = json.loads(target_dict_str)
 
     histologies = target_dict["histologies"]
-    INT_TO_CANCER_TYPE_MAP_LOCAL = {i: histology for i, histology in enumerate(histologies)}
-    CANCER_TYPE_TO_INT_MAP_LOCAL = {v: k for k, v in INT_TO_CANCER_TYPE_MAP_LOCAL.items()}
+    INT_TO_CANCER_TYPE_MAP_LOCAL = {
+        i: histology for i, histology in enumerate(histologies)
+    }
+    CANCER_TYPE_TO_INT_MAP_LOCAL = {
+        v: k for k, v in INT_TO_CANCER_TYPE_MAP_LOCAL.items()
+    }
 
     # Calculate col_indices_to_drop using local mapping
     col_indices_to_drop_local = [
@@ -100,7 +104,9 @@ def run_with_model(
         tissue_site_idx=tissue_site_idx,
         n_max_tiles=20000,
     )
-    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+    dataloader = DataLoader(
+        dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers
+    )
 
     results = []
     batch = next(iter(dataloader))
@@ -140,8 +146,14 @@ def run_with_model(
 
 
 def run(
-    features, model_path, metastatic=False, batch_size=8, num_workers=8, use_cpu=False,
-    sex=None, tissue_site_idx=None
+    features,
+    model_path,
+    metastatic=False,
+    batch_size=8,
+    num_workers=8,
+    use_cpu=False,
+    sex=None,
+    tissue_site_idx=None,
 ):
     """Run Aeon model inference for cancer subtype prediction.
 
@@ -176,12 +188,20 @@ def run(
         target_dict_str = f.read().strip().replace("'", '"')
         target_dict = json.loads(target_dict_str)
 
-    histologies = target_dict['histologies']
-    INT_TO_CANCER_TYPE_MAP_LOCAL = {i: histology for i, histology in enumerate(histologies)}
-    CANCER_TYPE_TO_INT_MAP_LOCAL = {v: k for k, v in INT_TO_CANCER_TYPE_MAP_LOCAL.items()}
+    histologies = target_dict["histologies"]
+    INT_TO_CANCER_TYPE_MAP_LOCAL = {
+        i: histology for i, histology in enumerate(histologies)
+    }
+    CANCER_TYPE_TO_INT_MAP_LOCAL = {
+        v: k for k, v in INT_TO_CANCER_TYPE_MAP_LOCAL.items()
+    }
 
     # Calculate col_indices_to_drop using local mapping
-    col_indices_to_drop_local = [CANCER_TYPE_TO_INT_MAP_LOCAL[x] for x in CANCER_TYPES_TO_DROP if x in CANCER_TYPE_TO_INT_MAP_LOCAL]
+    col_indices_to_drop_local = [
+        CANCER_TYPE_TO_INT_MAP_LOCAL[x]
+        for x in CANCER_TYPES_TO_DROP
+        if x in CANCER_TYPE_TO_INT_MAP_LOCAL
+    ]
 
     site_type = SiteType.METASTASIS if metastatic else SiteType.PRIMARY
 
@@ -306,7 +326,9 @@ def main():
     tissue_site_idx = None
     if opt.tissue_site:
         tissue_site_idx = encode_tissue_site(opt.tissue_site)
-        logger.info(f"Using tissue site: {opt.tissue_site} (encoded as {tissue_site_idx})")
+        logger.info(
+            f"Using tissue site: {opt.tissue_site} (encoded as {tissue_site_idx})"
+        )
 
     results_df, part_embedding = run(
         features=features,

src/mosaic/inference/data.py

@@ -212,10 +212,10 @@ DEFAULT_TISSUE_SITE_IDX = 8
 
 def get_tissue_site_map():
     """Load tissue site name → index mapping from CSV.
-    
+
     Returns:
         dict: Mapping of tissue site names to indices (0-56)
-    
+
     Raises:
         FileNotFoundError: If the tissue site CSV file is not found
     """
@@ -232,17 +232,17 @@ def get_tissue_site_map():
                 f"Tissue site mapping file not found at {csv_path}. "
                 f"Please ensure the data directory contains 'tissue_site_original_to_idx.csv'."
             ) from e
-        
+
         _TISSUE_SITE_MAP = {}
         for _, row in df.iterrows():
-            _TISSUE_SITE_MAP[row['TISSUE_SITE']] = int(row['idx'])
-    
+            _TISSUE_SITE_MAP[row["TISSUE_SITE"]] = int(row["idx"])
+
     return _TISSUE_SITE_MAP
 
 
 def get_tissue_site_options():
     """Get sorted unique tissue site names for UI dropdowns.
-    
+
     Returns:
         list: Sorted list of unique tissue site names
     """
@@ -258,7 +258,7 @@ def get_sex_map():
 
     Returns:
         dict: Mapping of sex values to indices (0-2)
-    
+
     Raises:
         FileNotFoundError: If the sex mapping CSV file is not found
     """
@@ -278,7 +278,7 @@ def get_sex_map():
 
         _SEX_MAP = {}
         for _, row in df.iterrows():
-            _SEX_MAP[row['SEX']] = int(row['idx'])
+            _SEX_MAP[row["SEX"]] = int(row["idx"])
 
     return _SEX_MAP
 
@@ -299,10 +299,10 @@ def encode_sex(sex):
 
 def encode_tissue_site(site_name):
     """Convert tissue site name to index (0-56).
-    
+
     Args:
         site_name: Tissue site name from CSV
-        
+
     Returns:
         int: Tissue site index, defaults to DEFAULT_TISSUE_SITE_IDX ("Not Applicable")
     """
@@ -312,11 +312,11 @@ def encode_tissue_site(site_name):
 
 def tissue_site_to_one_hot(site_idx, num_classes=57):
     """Convert tissue site index to one-hot vector.
-    
+
     Args:
         site_idx: Index value (0-56 for tissue site, 0-2 for sex)
         num_classes: Number of classes (57 for tissue site, 3 for sex)
-        
+
     Returns:
         list: One-hot encoded vector
     """
@@ -395,22 +395,18 @@ class TileFeatureTensorDataset(Dataset):
         Returns:
             dict: the item
         """
-        result = {
-            "site": self.site_type.value,
-            "tile_tensor": self.features
-        }
-        
+        result = {"site": self.site_type.value, "tile_tensor": self.features}
+
         # Add sex and tissue_site if provided (for Aeon)
         if self.sex is not None:
             result["SEX"] = torch.tensor(
-                tissue_site_to_one_hot(self.sex, num_classes=3),
-                dtype=torch.float32
+                tissue_site_to_one_hot(self.sex, num_classes=3), dtype=torch.float32
             )
-        
+
         if self.tissue_site_idx is not None:
             result["TISSUE_SITE"] = torch.tensor(
                 tissue_site_to_one_hot(self.tissue_site_idx, num_classes=57),
-                dtype=torch.float32
+                dtype=torch.float32,
             )
-        
+
         return result

src/mosaic/inference/paladin.py

@@ -38,10 +38,10 @@ def load_model_map(model_map_path: str) -> dict[Any, Any]:
 
     A dict is returned, mapping each cancer_subtype to a table mapping a
     target to the pathname for the model that predicts it.
-    
+
     Args:
         model_map_path: Path to the CSV file containing the model map
-        
+
     Returns:
         Dictionary mapping cancer subtypes to their target-specific models
     """
@@ -58,10 +58,10 @@ def load_model_map(model_map_path: str) -> dict[Any, Any]:
 
 def load_aeon_scores(df: pd.DataFrame) -> dict[str, float]:
     """Load Aeon output table with cancer subtypes and confidence values.
-    
+
     Args:
         df: DataFrame with columns 'Cancer Subtype' and 'Confidence'
-    
+
     Returns:
         Dictionary mapping cancer subtypes to their confidence scores
     """
@@ -75,11 +75,11 @@ def load_aeon_scores(df: pd.DataFrame) -> dict[str, float]:
 
 def select_cancer_subtypes(aeon_scores: dict[str, float], k=1) -> list[str]:
     """Select the top k cancer subtypes based on Aeon confidence scores.
-    
+
     Args:
         aeon_scores: Dictionary mapping cancer subtypes to confidence scores
         k: Number of top subtypes to select (default: 1)
-        
+
     Returns:
         List of cancer subtype codes sorted by confidence (highest first)
     """
@@ -91,11 +91,11 @@ def select_cancer_subtypes(aeon_scores: dict[str, float], k=1) -> list[str]:
 
 def select_models(cancer_subtypes: list[str], model_map: dict[Any, Any]) -> list[Any]:
     """Select Paladin models for the given cancer subtypes.
-    
+
     Args:
         cancer_subtypes: List of cancer subtype codes
         model_map: Dictionary mapping cancer subtypes to their models
-        
+
     Returns:
         List of tuples (cancer_subtype, target, model_path)
     """
@@ -188,13 +188,13 @@ def run_model(device, dataset, model_path: str, num_workers, batch_size) -> floa
 
 def logits_to_point_estimates(logits):
     """Convert model logits to point estimates for beta-binomial distribution.
-    
+
     The logits tensor contains alpha and beta parameters interleaved.
     This function computes the mean of the beta-binomial distribution: alpha/(alpha+beta).
-    
+
     Args:
         logits: Tensor of shape (batch_size, 2*(n_tasks)) with alpha/beta parameters
-        
+
     Returns:
         Tensor of shape (batch_size, n_tasks) with point estimates
     """
@@ -215,10 +215,10 @@ def run(
     use_cpu: bool = False,
 ):
     """Run Paladin inference for biomarker prediction on a single slide.
-    
+
     Uses either Aeon predictions or user-provided cancer subtype codes to select
     the appropriate Paladin models for biomarker prediction.
-    
+
     Args:
         features: NumPy array of tile features extracted from the WSI
         aeon_results: DataFrame with Aeon predictions (Cancer Subtype, Confidence)
@@ -229,10 +229,10 @@ def run(
         batch_size: Batch size for inference
         num_workers: Number of workers for data loading
         use_cpu: Force CPU usage instead of GPU
-        
+
     Returns:
         DataFrame with columns: Cancer Subtype, Target, Score
-        
+
     Note:
         Either aeon_results or cancer_subtype_codes must be provided, but not both.
         Either model_map_path or model_path must be provided, but not both.

src/mosaic/model_manager.py

@@ -13,6 +13,7 @@ import torch
 from loguru import logger
 
 from mosaic.data_directory import get_data_directory
+from mussel.models import ModelType, get_model_factory
 
 
 class ModelCache:
@@ -50,7 +51,9 @@ class ModelCache:
         self.paladin_models: Dict[tuple, torch.nn.Module] = {}
         self.is_t4_gpu = is_t4_gpu
         self.aggressive_memory_mgmt = aggressive_memory_mgmt
-        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = device or torch.device(
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
 
     def cleanup_paladin(self):
         """Aggressively free all Paladin models from memory.
@@ -78,15 +81,18 @@ class ModelCache:
         self.cleanup_paladin()
 
         # Clean up core models
-        del self.ctranspath_model
-        del self.optimus_model
-        del self.marker_classifier
-        del self.aeon_model
-
-        self.ctranspath_model = None
-        self.optimus_model = None
-        self.marker_classifier = None
-        self.aeon_model = None
+        if self.ctranspath_model is not None:
+            del self.ctranspath_model
+            self.ctranspath_model = None
+        if self.optimus_model is not None:
+            del self.optimus_model
+            self.optimus_model = None
+        if self.marker_classifier is not None:
+            del self.marker_classifier
+            self.marker_classifier = None
+        if self.aeon_model is not None:
+            del self.aeon_model
+            self.aeon_model = None
 
         # Force garbage collection and GPU cache clearing
         gc.collect()
@@ -147,7 +153,9 @@ def load_all_models(
             if is_t4_gpu:
                 logger.info("  → Paladin models will be loaded and freed per slide")
             else:
-                logger.info("  → Paladin models will be cached and reused across slides")
+                logger.info(
+                    "  → Paladin models will be cached and reused across slides"
+                )
     elif use_gpu and not torch.cuda.is_available():
         logger.warning("GPU requested but CUDA not available, falling back to CPU")
         use_gpu = False
@@ -165,24 +173,37 @@ def load_all_models(
     if not ctranspath_path.exists():
         raise FileNotFoundError(f"CTransPath model not found at {ctranspath_path}")
 
-    # Note: CTransPath loading is handled by mussel, so we just store the path for now
-    # We'll integrate with mussel's model factory in the feature extraction wrappers
-    ctranspath_model = ctranspath_path
+    ctranspath_factory = get_model_factory(ModelType.CTRANSPATH)
+    ctranspath_model = ctranspath_factory.get_model(
+        str(ctranspath_path), use_gpu=use_gpu, gpu_device_id=0 if use_gpu else None
+    )
+    logger.info("✓ CTransPath model loaded")
 
-    # Load Optimus model
-    logger.info("Loading Optimus model...")
-    optimus_path = data_dir / "optimus.pkl"
-    if not optimus_path.exists():
-        raise FileNotFoundError(f"Optimus model not found at {optimus_path}")
+    if use_gpu and torch.cuda.is_available():
+        mem = torch.cuda.memory_allocated() / (1024**3)
+        logger.info(f"  GPU memory: {mem:.2f} GB")
 
-    # Note: Same as CTransPath, Optimus loading is handled by mussel
-    optimus_model = optimus_path
+    # Load Optimus model from Hugging Face Hub
+    logger.info("Loading Optimus model from bioptimus/H-optimus-0...")
+    optimus_factory = get_model_factory(ModelType.OPTIMUS)
+    optimus_model = optimus_factory.get_model(
+        model_path="hf-hub:bioptimus/H-optimus-0",
+        use_gpu=use_gpu,
+        gpu_device_id=0 if use_gpu else None,
+    )
+    logger.info("✓ Optimus model loaded")
+
+    if use_gpu and torch.cuda.is_available():
+        mem = torch.cuda.memory_allocated() / (1024**3)
+        logger.info(f"  GPU memory: {mem:.2f} GB")
 
     # Load Marker Classifier
     logger.info("Loading Marker Classifier...")
     marker_classifier_path = data_dir / "marker_classifier.pkl"
     if not marker_classifier_path.exists():
-        raise FileNotFoundError(f"Marker classifier not found at {marker_classifier_path}")
+        raise FileNotFoundError(
+            f"Marker classifier not found at {marker_classifier_path}"
+        )
 
     with open(marker_classifier_path, "rb") as f:
         marker_classifier = pickle.load(f)  # nosec
@@ -238,12 +259,14 @@ def load_paladin_model_for_inference(
     cache: ModelCache,
     model_path: Path,
 ) -> torch.nn.Module:
-    """Load a single Paladin model for inference.
+    """Load a single Paladin model for inference, downloading on-demand if needed.
 
     Implements adaptive loading strategy:
     - T4 GPU (aggressive mode): Load model fresh, caller must delete after use
     - A100 GPU (caching mode): Check cache, load if needed, return cached model
 
+    If the model file doesn't exist locally, downloads it from HuggingFace Hub.
+
     Args:
         cache: ModelCache instance managing loaded models
         model_path: Path to the Paladin model file
@@ -255,6 +278,8 @@ def load_paladin_model_for_inference(
         On T4 GPUs, caller MUST delete the model and call torch.cuda.empty_cache()
         after inference to avoid OOM errors.
     """
+    from huggingface_hub import hf_hub_download
+
     model_key = str(model_path)
 
     # Check cache first (only used in non-aggressive mode)
@@ -262,11 +287,32 @@ def load_paladin_model_for_inference(
         logger.info(f"  ✓ Using CACHED Paladin model: {model_path.name} (no disk I/O!)")
         return cache.paladin_models[model_key]
 
+    # Download model from HF Hub if it doesn't exist locally
+    if not model_path.exists():
+        logger.info(
+            f"  ⬇ Downloading Paladin model from HuggingFace Hub: {model_path.name}"
+        )
+        # Extract the relative path from the data directory
+        data_dir = get_data_directory()
+        relative_path = model_path.relative_to(data_dir)
+
+        downloaded_path = hf_hub_download(
+            repo_id="PDM-Group/paladin-aeon-models",
+            filename=str(relative_path),
+            cache_dir=data_dir.parent.parent,  # Use HF cache directory
+        )
+        model_path = Path(downloaded_path)
+        logger.info(f"  ✓ Downloaded to: {model_path}")
+
     # Load model from disk
     if cache.aggressive_memory_mgmt:
-        logger.info(f"  → Loading Paladin model: {model_path.name} (will free after use)")
+        logger.info(
+            f"  → Loading Paladin model: {model_path.name} (will free after use)"
+        )
     else:
-        logger.info(f"  → Loading Paladin model: {model_path.name} (will cache for reuse)")
+        logger.info(
+            f"  → Loading Paladin model: {model_path.name} (will cache for reuse)"
+        )
 
     with open(model_path, "rb") as f:
         model = pickle.load(f)  # nosec

src/mosaic/ui/app.py

@@ -24,7 +24,7 @@ from mosaic.ui.utils import (
     SETTINGS_COLUMNS,
 )
 from mosaic.analysis import analyze_slide
-from mosaic.batch_analysis import analyze_slides_batch
+from mosaic.model_manager import load_all_models
 
 current_dir = Path(__file__).parent.parent
 
@@ -45,6 +45,12 @@ def set_cancer_subtype_maps(csn_map, rcsn_map, cs):
 def analyze_slides(
     slides,
     settings_input,
+    site_type,
+    sex,
+    tissue_site,
+    cancer_subtype,
+    ihc_subtype,
+    seg_config,
     user_dir,
     progress=gr.Progress(track_tqdm=True),
     request: gr.Request = None,
@@ -52,61 +58,112 @@ def analyze_slides(
     if slides is None or len(slides) == 0:
         raise gr.Error("Please upload at least one slide.")
     if user_dir is None:
-        user_dir = create_user_directory(None, gr.Request())
+        if request is not None:
+            user_dir = create_user_directory(None, request)
+        if user_dir is None:
+            # Fallback to temp directory if session hash not available
+            import tempfile
+
+            user_dir = Path(tempfile.mkdtemp(prefix="mosaic_"))
+
+    # Handle empty settings_input (e.g., when dataframe is hidden for single slide)
+    # Regenerate settings from dropdowns if settings_input is empty
+    if settings_input is None or len(settings_input) == 0:
+        logger.info("Settings dataframe is empty, regenerating from dropdown values")
+        settings = []
+        for file in slides:
+            filename = file.name if hasattr(file, "name") else file
+            slide_name = filename.split("/")[-1]
+            settings.append(
+                [
+                    slide_name,
+                    site_type,
+                    sex,
+                    tissue_site,
+                    cancer_subtype,
+                    ihc_subtype,
+                    seg_config,
+                ]
+            )
+        settings_input = pd.DataFrame(settings, columns=SETTINGS_COLUMNS)
+
     settings_input = validate_settings(
-        settings_input, cancer_subtype_name_map, cancer_subtypes, reversed_cancer_subtype_name_map
+        settings_input,
+        cancer_subtype_name_map,
+        cancer_subtypes,
+        reversed_cancer_subtype_name_map,
     )
     if len(slides) != len(settings_input):
         raise gr.Error("Missing settings for uploaded slides")
 
-    # Use batch processing for multiple slides (models loaded once)
-    # Use single-slide processing for 1 slide (maintains exact same behavior)
+    all_slide_masks = []
+    all_aeon_results = []
+    all_paladin_results = []
+
+    # Load models once (for batch) or per-slide (for single)
+    model_cache = None
     if len(slides) > 1:
-        logger.info(f"Using batch processing for {len(slides)} slides")
-        progress(0.0, desc=f"Starting batch analysis ({len(slides)} slides)")
-
-        all_slide_masks, all_aeon_results, all_paladin_results = analyze_slides_batch(
-            slides=slides,
-            settings_df=settings_input,
-            cancer_subtype_name_map=cancer_subtype_name_map,
-            num_workers=4,
-            aggressive_memory_mgmt=None,  # Auto-detect GPU type
-            progress=progress,
-        )
+        logger.info(f"Batch mode: Loading models once for {len(slides)} slides")
+        progress(0.0, desc=f"Loading models for batch processing")
+        model_cache = load_all_models(use_gpu=True, aggressive_memory_mgmt=None)
     else:
-        # Single slide: use existing analyze_slide() for backward compatibility
-        logger.info("Using single-slide processing (1 slide)")
-        progress(0.0, desc="Starting single-slide analysis")
-
-        all_slide_masks = []
-        all_aeon_results = []
-        all_paladin_results = []
-
-        row = settings_input.iloc[0]
-        slide_name = row["Slide"]
-
-        slide_mask, aeon_results, paladin_results = analyze_slide(
-            slides[0],
-            row["Segmentation Config"],
-            row["Site Type"],
-            row["Sex"],
-            row["Tissue Site"],
-            row["Cancer Subtype"],
-            cancer_subtype_name_map,
-            row["IHC Subtype"],
-            progress=progress,
-            request=request,
-        )
+        logger.info("Single-slide mode: models loaded within analyze_slide")
+
+    try:
+        # Process all slides with unified analyze_slide function
+        for idx, slide_path in enumerate(slides):
+            row = settings_input.iloc[idx]
+            slide_name = row["Slide"]
+
+            logger.info(f"[{idx + 1}/{len(slides)}] Processing: {slide_name}")
+            slide_progress = idx / len(slides)
+            progress(slide_progress, desc=f"Analyzing slide {idx + 1}/{len(slides)}")
+
+            slide_mask, aeon_results, paladin_results = analyze_slide(
+                slide_path=slide_path,
+                seg_config=row["Segmentation Config"],
+                site_type=row["Site Type"],
+                sex=row.get("Sex", "Unknown"),
+                tissue_site=row.get("Tissue Site", "Unknown"),
+                cancer_subtype=row["Cancer Subtype"],
+                cancer_subtype_name_map=cancer_subtype_name_map,
+                ihc_subtype=row.get("IHC Subtype", ""),
+                num_workers=4,
+                progress=progress,
+                request=request,
+                model_cache=model_cache,  # Pre-loaded for batch, None for single
+            )
 
-        if slide_mask is not None:
-            all_slide_masks.append((slide_mask, slide_name))
-        if aeon_results is not None:
-            all_aeon_results.append(aeon_results)
-        if paladin_results is not None:
-            paladin_results.insert(
-                0, "Slide", pd.Series([slide_name] * len(paladin_results))
+            if slide_mask is not None:
+                all_slide_masks.append((slide_mask, slide_name))
+            if aeon_results is not None:
+                all_aeon_results.append(aeon_results)
+            if paladin_results is not None:
+                paladin_results.insert(
+                    0, "Slide", pd.Series([slide_name] * len(paladin_results))
+                )
+                all_paladin_results.append(paladin_results)
+
+            # Yield intermediate update to show slide masks as they're generated
+            # This allows the UI to update incrementally during processing
+            yield (
+                all_slide_masks.copy(),  # Current slide masks
+                gr.DataFrame(visible=False),  # aeon_output_table (not ready yet)
+                gr.DownloadButton(
+                    visible=False
+                ),  # aeon_download_button (not ready yet)
+                None,  # paladin_output_table (not ready yet)
+                gr.DownloadButton(
+                    visible=False
+                ),  # paladin_download_button (not ready yet)
+                user_dir,  # user_dir_state
             )
-            all_paladin_results.append(paladin_results)
+
+    finally:
+        # Clean up model cache if it was loaded for batch processing
+        if model_cache is not None:
+            logger.info("Cleaning up model cache")
+            model_cache.cleanup()
 
     progress(0.99, desc="Analysis complete, wrapping up results")
 
@@ -155,7 +212,8 @@ def analyze_slides(
 
     progress(1.0, desc="All done!")
 
-    return (
+    # Final yield with complete results
+    yield (
         all_slide_masks,
         combined_aeon_results,
         aeon_output,
@@ -273,17 +331,20 @@ def launch_gradio(server_name, server_port, share):
         )
         def clear_fn():
             return (
-                None,
-                None,
-                None,
-                None,
-                gr.Dataframe(visible=False),
-                gr.DownloadButton(visible=False),
-                gr.Dataframe(visible=False),
-                gr.File(visible=False),
+                None,  # input_slides
+                None,  # slide_masks
+                None,  # paladin_output_table
+                gr.DownloadButton(visible=False),  # paladin_download_button
+                gr.Dataframe(visible=False),  # aeon_output_table
+                gr.DownloadButton(visible=False),  # aeon_download_button
+                gr.Dataframe(visible=False),  # settings_input
+                gr.File(visible=False),  # settings_csv
             )
 
-        def get_settings(files, site_type, sex, tissue_site, cancer_subtype, ihc_subtype, seg_config):
+        def get_settings(
+            files, site_type, sex, tissue_site, cancer_subtype, ihc_subtype, seg_config
+        ):
+            """Generate initial settings DataFrame from uploaded files and dropdown values."""
             if files is None:
                 return pd.DataFrame()
             settings = []
@@ -291,22 +352,30 @@ def launch_gradio(server_name, server_port, share):
                 filename = file.name if hasattr(file, "name") else file
                 slide_name = filename.split("/")[-1]
                 settings.append(
-                    [slide_name, site_type, sex, tissue_site, cancer_subtype, ihc_subtype, seg_config]
+                    [
+                        slide_name,
+                        site_type,
+                        sex,
+                        tissue_site,
+                        cancer_subtype,
+                        ihc_subtype,
+                        seg_config,
+                    ]
                 )
             df = pd.DataFrame(settings, columns=SETTINGS_COLUMNS)
             return df
 
-        # Only display settings table and upload button if multiple slides are uploaded
-        @gr.on(
-            [
-                input_slides.change,
-                site_dropdown.change,
-                sex_dropdown.change,
-                tissue_site_dropdown.change,
-                cancer_subtype_dropdown.change,
-                ihc_subtype_dropdown.change,
-                seg_config_dropdown.change,
-            ],
+        def update_settings_column(settings_df, column_name, new_value):
+            """Update a specific column in the settings DataFrame."""
+            if settings_df is None or len(settings_df) == 0:
+                return settings_df
+            # Create a copy to avoid modifying the original
+            updated_df = settings_df.copy()
+            updated_df[column_name] = new_value
+            return updated_df
+
+        # Handle file uploads - regenerate entire settings table
+        @input_slides.change(
             inputs=[
                 input_slides,
                 site_dropdown,
@@ -318,22 +387,103 @@ def launch_gradio(server_name, server_port, share):
             ],
             outputs=[settings_input, settings_csv, ihc_subtype_dropdown],
         )
-        def update_settings(files, site_type, sex, tissue_site, cancer_subtype, ihc_subtype, seg_config):
+        def update_files(
+            files, site_type, sex, tissue_site, cancer_subtype, ihc_subtype, seg_config
+        ):
+            """Handle file upload - regenerate settings table from scratch."""
             has_ihc = "Breast" in cancer_subtype
             if not files:
                 return None, None, gr.Dropdown(visible=has_ihc)
             settings_df = get_settings(
-                files, site_type, sex, tissue_site, cancer_subtype, ihc_subtype, seg_config
+                files,
+                site_type,
+                sex,
+                tissue_site,
+                cancer_subtype,
+                ihc_subtype,
+                seg_config,
             )
             if settings_df is not None:
                 has_ihc = any("Breast" in cs for cs in settings_df["Cancer Subtype"])
             visible = files and len(files) > 1
             return (
-                gr.Dataframe(settings_df, visible=visible),
+                gr.Dataframe(value=settings_df, visible=visible),
                 gr.File(visible=visible),
                 gr.Dropdown(visible=has_ihc),
             )
 
+        # Handle individual dropdown changes - only update the relevant column
+        @site_dropdown.change(
+            inputs=[settings_input, site_dropdown],
+            outputs=[settings_input],
+        )
+        def update_site_type(settings_df, site_type):
+            """Update Site Type column when dropdown changes."""
+            if settings_df is None or len(settings_df) == 0:
+                return settings_df
+            updated_df = update_settings_column(settings_df, "Site Type", site_type)
+            return gr.Dataframe(value=updated_df)
+
+        @sex_dropdown.change(
+            inputs=[settings_input, sex_dropdown],
+            outputs=[settings_input],
+        )
+        def update_sex(settings_df, sex):
+            """Update Sex column when dropdown changes."""
+            if settings_df is None or len(settings_df) == 0:
+                return settings_df
+            updated_df = update_settings_column(settings_df, "Sex", sex)
+            return gr.Dataframe(value=updated_df)
+
+        @tissue_site_dropdown.change(
+            inputs=[settings_input, tissue_site_dropdown],
+            outputs=[settings_input],
+        )
+        def update_tissue_site(settings_df, tissue_site):
+            """Update Tissue Site column when dropdown changes."""
+            if settings_df is None or len(settings_df) == 0:
+                return settings_df
+            updated_df = update_settings_column(settings_df, "Tissue Site", tissue_site)
+            return gr.Dataframe(value=updated_df)
+
+        @cancer_subtype_dropdown.change(
+            inputs=[settings_input, cancer_subtype_dropdown],
+            outputs=[settings_input, ihc_subtype_dropdown],
+        )
+        def update_cancer_subtype(settings_df, cancer_subtype):
+            """Update Cancer Subtype column when dropdown changes."""
+            has_ihc = "Breast" in cancer_subtype
+            if settings_df is None or len(settings_df) == 0:
+                return settings_df, gr.Dropdown(visible=has_ihc)
+            updated_df = update_settings_column(
+                settings_df, "Cancer Subtype", cancer_subtype
+            )
+            return gr.Dataframe(value=updated_df), gr.Dropdown(visible=has_ihc)
+
+        @ihc_subtype_dropdown.change(
+            inputs=[settings_input, ihc_subtype_dropdown],
+            outputs=[settings_input],
+        )
+        def update_ihc_subtype(settings_df, ihc_subtype):
+            """Update IHC Subtype column when dropdown changes."""
+            if settings_df is None or len(settings_df) == 0:
+                return settings_df
+            updated_df = update_settings_column(settings_df, "IHC Subtype", ihc_subtype)
+            return gr.Dataframe(value=updated_df)
+
+        @seg_config_dropdown.change(
+            inputs=[settings_input, seg_config_dropdown],
+            outputs=[settings_input],
+        )
+        def update_seg_config(settings_df, seg_config):
+            """Update Segmentation Config column when dropdown changes."""
+            if settings_df is None or len(settings_df) == 0:
+                return settings_df
+            updated_df = update_settings_column(
+                settings_df, "Segmentation Config", seg_config
+            )
+            return gr.Dataframe(value=updated_df)
+
         @settings_csv.upload(
             inputs=[settings_csv],
             outputs=[settings_input],
@@ -349,6 +499,12 @@ def launch_gradio(server_name, server_port, share):
             inputs=[
                 input_slides,
                 settings_input,
+                site_dropdown,
+                sex_dropdown,
+                tissue_site_dropdown,
+                cancer_subtype_dropdown,
+                ihc_subtype_dropdown,
+                seg_config_dropdown,
                 user_dir_state,
             ],
             outputs=[
@@ -363,9 +519,14 @@ def launch_gradio(server_name, server_port, share):
             show_progress_on=paladin_output_table,
         )
         settings_input.change(
-            lambda df: validate_settings(df, cancer_subtype_name_map, cancer_subtypes, reversed_cancer_subtype_name_map),
+            lambda df: validate_settings(
+                df,
+                cancer_subtype_name_map,
+                cancer_subtypes,
+                reversed_cancer_subtype_name_map,
+            ),
             inputs=[settings_input],
-            outputs=[settings_input]
+            outputs=[settings_input],
         )
         demo.load(
             create_user_directory,

src/mosaic/ui/utils.py

@@ -61,13 +61,13 @@ def get_tissue_sites():
 
 def get_oncotree_code_name(code):
     """Retrieve the human-readable name for an OncoTree code.
-    
+
     Queries the OncoTree API to get the cancer subtype name corresponding
     to the given code. Results are cached to avoid repeated API calls.
-    
+
     Args:
         code: OncoTree code (e.g., "LUAD", "BRCA")
-        
+
     Returns:
         Human-readable cancer subtype name, or "Unknown" if not found
     """
@@ -108,16 +108,16 @@ def create_user_directory(state, request: gr.Request):
 
 def load_settings(slide_csv_path):
     """Load slide analysis settings from CSV file.
-    
+
     Loads the CSV and ensures all required columns are present, adding defaults
     for optional columns if they are missing.
-    
+
     Args:
         slide_csv_path: Path to the CSV file containing slide settings
-        
+
     Returns:
         DataFrame with columns: Slide, Site Type, Cancer Subtype, IHC Subtype, Segmentation Config
-        
+
     Raises:
         ValueError: If required columns are missing from the CSV
     """
@@ -138,21 +138,26 @@ def load_settings(slide_csv_path):
     return settings_df
 
 
-def validate_settings(settings_df, cancer_subtype_name_map, cancer_subtypes, reversed_cancer_subtype_name_map):
+def validate_settings(
+    settings_df,
+    cancer_subtype_name_map,
+    cancer_subtypes,
+    reversed_cancer_subtype_name_map,
+):
     """Validate and normalize slide analysis settings.
-    
+
     Checks each row for valid values and normalizes cancer subtype names.
     Generates warnings for invalid entries and replaces them with defaults.
-    
+
     Args:
         settings_df: DataFrame with slide settings to validate
         cancer_subtype_name_map: Dict mapping subtype display names to codes
         cancer_subtypes: List of valid cancer subtype codes
         reversed_cancer_subtype_name_map: Dict mapping codes to display names
-        
+
     Returns:
         Validated DataFrame with normalized values
-        
+
     Note:
         Invalid entries are replaced with defaults and warnings are displayed
         to the user via Gradio warnings.
@@ -215,13 +220,13 @@ def validate_settings(settings_df, cancer_subtype_name_map, cancer_subtypes, rev
 
 def export_to_csv(df):
     """Export a DataFrame to CSV file for download.
-    
+
     Args:
         df: DataFrame to export
-        
+
     Returns:
         Path to the exported CSV file
-        
+
     Raises:
         gr.Error: If the DataFrame is None or empty
     """

tests/benchmark_batch_performance.py

@@ -21,7 +21,9 @@ from mosaic.batch_analysis import analyze_slides_batch
 from mosaic.ui.utils import load_settings, validate_settings
 
 
-def benchmark_sequential_processing(slides, settings_df, cancer_subtype_name_map, num_workers):
+def benchmark_sequential_processing(
+    slides, settings_df, cancer_subtype_name_map, num_workers
+):
     """Benchmark traditional sequential processing (models loaded per slide)."""
     logger.info("=" * 80)
     logger.info("BENCHMARKING: Sequential Processing (OLD METHOD)")
@@ -51,13 +53,15 @@ def benchmark_sequential_processing(slides, settings_df, cancer_subtype_name_map
         slide_time = time.time() - slide_start
         logger.info(f"Slide {idx + 1} completed in {slide_time:.2f}s")
 
-        results.append({
-            "slide": slide_path,
-            "time": slide_time,
-            "has_mask": slide_mask is not None,
-            "has_aeon": aeon_results is not None,
-            "has_paladin": paladin_results is not None,
-        })
+        results.append(
+            {
+                "slide": slide_path,
+                "time": slide_time,
+                "has_mask": slide_mask is not None,
+                "has_aeon": aeon_results is not None,
+                "has_paladin": paladin_results is not None,
+            }
+        )
 
     total_time = time.time() - start_time
     peak_memory = torch.cuda.max_memory_allocated() if torch.cuda.is_available() else 0
@@ -79,7 +83,9 @@ def benchmark_sequential_processing(slides, settings_df, cancer_subtype_name_map
     }
 
 
-def benchmark_batch_processing(slides, settings_df, cancer_subtype_name_map, num_workers):
+def benchmark_batch_processing(
+    slides, settings_df, cancer_subtype_name_map, num_workers
+):
     """Benchmark optimized batch processing (models loaded once)."""
     logger.info("=" * 80)
     logger.info("BENCHMARKING: Batch Processing (NEW METHOD)")
@@ -128,7 +134,9 @@ def compare_results(sequential_stats, batch_stats):
 
     speedup = sequential_stats["total_time"] / batch_stats["total_time"]
     time_saved = sequential_stats["total_time"] - batch_stats["total_time"]
-    percent_faster = (1 - (batch_stats["total_time"] / sequential_stats["total_time"])) * 100
+    percent_faster = (
+        1 - (batch_stats["total_time"] / sequential_stats["total_time"])
+    ) * 100
 
     logger.info(f"Number of slides: {sequential_stats['num_slides']}")
     logger.info(f"")
@@ -141,9 +149,11 @@ def compare_results(sequential_stats, batch_stats):
 
     if torch.cuda.is_available():
         logger.info(f"")
-        logger.info(f"Sequential peak memory: {sequential_stats['peak_memory_gb']:.2f} GB")
+        logger.info(
+            f"Sequential peak memory: {sequential_stats['peak_memory_gb']:.2f} GB"
+        )
         logger.info(f"Batch peak memory:      {batch_stats['peak_memory_gb']:.2f} GB")
-        memory_diff = batch_stats['peak_memory_gb'] - sequential_stats['peak_memory_gb']
+        memory_diff = batch_stats["peak_memory_gb"] - sequential_stats["peak_memory_gb"]
         logger.info(f"Memory difference:      {memory_diff:+.2f} GB")
 
     logger.info("=" * 80)
@@ -161,31 +171,20 @@ def main():
     parser = argparse.ArgumentParser(
         description="Benchmark batch processing performance"
     )
+    parser.add_argument("--slides", nargs="+", help="List of slide paths to process")
     parser.add_argument(
-        "--slides",
-        nargs="+",
-        help="List of slide paths to process"
-    )
-    parser.add_argument(
-        "--slide-csv",
-        type=str,
-        help="CSV file with slide paths and settings"
+        "--slide-csv", type=str, help="CSV file with slide paths and settings"
     )
     parser.add_argument(
-        "--num-workers",
-        type=int,
-        default=4,
-        help="Number of workers for data loading"
+        "--num-workers", type=int, default=4, help="Number of workers for data loading"
     )
     parser.add_argument(
         "--skip-sequential",
         action="store_true",
-        help="Skip sequential benchmark (faster, only test batch mode)"
+        help="Skip sequential benchmark (faster, only test batch mode)",
     )
     parser.add_argument(
-        "--output",
-        type=str,
-        help="Save benchmark results to JSON file"
+        "--output", type=str, help="Save benchmark results to JSON file"
     )
 
     args = parser.parse_args()
@@ -195,27 +194,35 @@ def main():
 
     # Load cancer subtype mappings
     from mosaic.gradio_app import download_and_process_models
-    cancer_subtype_name_map, cancer_subtypes, reversed_cancer_subtype_name_map = download_and_process_models()
+
+    cancer_subtype_name_map, cancer_subtypes, reversed_cancer_subtype_name_map = (
+        download_and_process_models()
+    )
 
     # Prepare slides and settings
     if args.slide_csv:
         settings_df = load_settings(args.slide_csv)
         settings_df = validate_settings(
-            settings_df, cancer_subtype_name_map, cancer_subtypes, reversed_cancer_subtype_name_map
+            settings_df,
+            cancer_subtype_name_map,
+            cancer_subtypes,
+            reversed_cancer_subtype_name_map,
         )
         slides = settings_df["Slide"].tolist()
     else:
         slides = args.slides
         # Create default settings
-        settings_df = pd.DataFrame({
-            "Slide": slides,
-            "Site Type": ["Primary"] * len(slides),
-            "Sex": ["Unknown"] * len(slides),
-            "Tissue Site": ["Unknown"] * len(slides),
-            "Cancer Subtype": ["Unknown"] * len(slides),
-            "IHC Subtype": [""] * len(slides),
-            "Segmentation Config": ["Biopsy"] * len(slides),
-        })
+        settings_df = pd.DataFrame(
+            {
+                "Slide": slides,
+                "Site Type": ["Primary"] * len(slides),
+                "Sex": ["Unknown"] * len(slides),
+                "Tissue Site": ["Unknown"] * len(slides),
+                "Cancer Subtype": ["Unknown"] * len(slides),
+                "IHC Subtype": [""] * len(slides),
+                "Segmentation Config": ["Biopsy"] * len(slides),
+            }
+        )
 
     logger.info(f"Benchmarking with {len(slides)} slides")
     logger.info(f"GPU available: {torch.cuda.is_available()}")
@@ -239,8 +246,9 @@ def main():
         # Save results if requested
         if args.output:
             import json
+
             output_path = Path(args.output)
-            with open(output_path, 'w') as f:
+            with open(output_path, "w") as f:
                 json.dump(comparison, f, indent=2, default=str)
             logger.info(f"Benchmark results saved to {output_path}")
 

tests/conftest.py

@@ -3,22 +3,28 @@
 import sys
 from unittest.mock import MagicMock
 
+
 # Create mock for gradio with Error class
 class GradioMock(MagicMock):
     """Mock for gradio that supports Error and Warning classes."""
+
     Error = Exception
     Warning = lambda msg: None
     Request = MagicMock
     Progress = MagicMock
-    
+
+
 # Mock heavy dependencies before any imports
 # This is necessary to allow tests to run without full environment setup
-sys.modules['mussel'] = MagicMock()
-sys.modules['mussel.models'] = MagicMock()
-sys.modules['mussel.utils'] = MagicMock()
-sys.modules['mussel.utils.segment'] = MagicMock()
-sys.modules['mussel.cli'] = MagicMock()
-sys.modules['mussel.cli.tessellate'] = MagicMock()
-sys.modules['gradio'] = GradioMock()
-sys.modules['huggingface_hub'] = MagicMock()
-sys.modules['loguru'] = MagicMock()
+sys.modules["mussel"] = MagicMock()
+sys.modules["mussel.models"] = MagicMock()
+sys.modules["mussel.utils"] = MagicMock()
+sys.modules["mussel.utils.segment"] = MagicMock()
+sys.modules["mussel.cli"] = MagicMock()
+sys.modules["mussel.cli.tessellate"] = MagicMock()
+sys.modules["gradio"] = GradioMock()
+sys.modules["huggingface_hub"] = MagicMock()
+sys.modules["loguru"] = MagicMock()
+
+# Import fixtures from test_fixtures.py to make them available to all tests
+pytest_plugins = ["tests.test_fixtures"]

tests/test_batch_analysis.py

@@ -1,279 +0,0 @@
-"""Integration tests for batch_analysis module.
-
-Tests the batch processing coordinator and end-to-end batch workflow.
-"""
-
-import pytest
-import pandas as pd
-from pathlib import Path
-from unittest.mock import Mock, patch, MagicMock
-import numpy as np
-
-from mosaic.batch_analysis import analyze_slides_batch
-
-
-class TestAnalyzeSlidesBatch:
-    """Test analyze_slides_batch function."""
-
-    @pytest.fixture
-    def sample_settings_df(self):
-        """Create sample settings DataFrame for testing."""
-        return pd.DataFrame({
-            "Slide": ["slide1.svs", "slide2.svs", "slide3.svs"],
-            "Site Type": ["Primary", "Primary", "Metastatic"],
-            "Sex": ["Male", "Female", "Unknown"],
-            "Tissue Site": ["Lung", "Breast", "Unknown"],
-            "Cancer Subtype": ["Unknown", "Unknown", "LUAD"],
-            "IHC Subtype": ["", "HR+/HER2-", ""],
-            "Segmentation Config": ["Biopsy", "Resection", "Biopsy"],
-        })
-
-    @pytest.fixture
-    def cancer_subtype_name_map(self):
-        """Sample cancer subtype name mapping."""
-        return {
-            "Unknown": "Unknown",
-            "Lung Adenocarcinoma": "LUAD",
-            "Breast Invasive Ductal Carcinoma": "IDC",
-        }
-
-    @patch('mosaic.batch_analysis.load_all_models')
-    @patch('mosaic.batch_analysis.analyze_slide_with_models')
-    def test_batch_analysis_basic(
-        self, mock_analyze_slide, mock_load_models, sample_settings_df, cancer_subtype_name_map
-    ):
-        """Test basic batch analysis workflow."""
-        # Mock model cache
-        mock_cache = Mock()
-        mock_cache.cleanup = Mock()
-        mock_load_models.return_value = mock_cache
-
-        # Mock analyze_slide_with_models to return NEW DataFrames each time
-        def mock_analyze_side_effect(*args, **kwargs):
-            mock_mask = Mock()
-            # Aeon results should have Cancer Subtype as index, not a column
-            mock_aeon = pd.DataFrame({"Confidence": [0.95]}, index=pd.Index(["LUAD"], name="Cancer Subtype"))
-            mock_paladin = pd.DataFrame({
-                "Cancer Subtype": ["LUAD"],
-                "Biomarker": ["EGFR"],
-                "Score": [0.85]
-            })
-            return (mock_mask, mock_aeon, mock_paladin)
-
-        mock_analyze_slide.side_effect = mock_analyze_side_effect
-
-        slides = ["slide1.svs", "slide2.svs", "slide3.svs"]
-
-        # Run batch analysis
-        masks, aeon_results, paladin_results = analyze_slides_batch(
-            slides=slides,
-            settings_df=sample_settings_df,
-            cancer_subtype_name_map=cancer_subtype_name_map,
-            num_workers=4,
-        )
-
-        # Verify models were loaded once
-        mock_load_models.assert_called_once()
-
-        # Verify analyze_slide_with_models was called for each slide
-        assert mock_analyze_slide.call_count == 3
-
-        # Verify cleanup was called
-        mock_cache.cleanup.assert_called_once()
-
-        # Verify results structure
-        assert len(masks) == 3
-        assert len(aeon_results) == 3
-        assert len(paladin_results) == 3
-
-    @patch('mosaic.batch_analysis.load_all_models')
-    @patch('mosaic.batch_analysis.analyze_slide_with_models')
-    def test_batch_analysis_with_failures(
-        self, mock_analyze_slide, mock_load_models, sample_settings_df, cancer_subtype_name_map
-    ):
-        """Test batch analysis continues when individual slides fail."""
-        mock_cache = Mock()
-        mock_cache.cleanup = Mock()
-        mock_load_models.return_value = mock_cache
-
-        # First slide succeeds, second fails, third succeeds
-        def mock_analyze_side_effect(*args, **kwargs):
-            # Get the slide_path to determine which call this is
-            call_count = mock_analyze_slide.call_count
-            if call_count == 2:  # Second call (index 1)
-                raise RuntimeError("Slide processing failed")
-
-            mock_mask = Mock()
-            # Aeon results should have Cancer Subtype as index, not a column
-            mock_aeon = pd.DataFrame({"Confidence": [0.95]}, index=pd.Index(["LUAD"], name="Cancer Subtype"))
-            mock_paladin = pd.DataFrame({
-                "Cancer Subtype": ["LUAD"],
-                "Biomarker": ["EGFR"],
-                "Score": [0.85]
-            })
-            return (mock_mask, mock_aeon, mock_paladin)
-
-        mock_analyze_slide.side_effect = mock_analyze_side_effect
-
-        slides = ["slide1.svs", "slide2.svs", "slide3.svs"]
-
-        # Should not raise exception
-        masks, aeon_results, paladin_results = analyze_slides_batch(
-            slides=slides,
-            settings_df=sample_settings_df,
-            cancer_subtype_name_map=cancer_subtype_name_map,
-        )
-
-        # Should have results for 2 out of 3 slides
-        assert len(masks) == 2
-        assert len(aeon_results) == 2
-        assert len(paladin_results) == 2
-
-        # Cleanup should still be called
-        mock_cache.cleanup.assert_called_once()
-
-    @patch('mosaic.batch_analysis.load_all_models')
-    def test_batch_analysis_cleanup_on_error(
-        self, mock_load_models, sample_settings_df, cancer_subtype_name_map
-    ):
-        """Test cleanup is called even when load_all_models fails."""
-        mock_load_models.side_effect = RuntimeError("Failed to load models")
-
-        slides = ["slide1.svs"]
-
-        with pytest.raises(RuntimeError, match="Failed to load models"):
-            analyze_slides_batch(
-                slides=slides,
-                settings_df=sample_settings_df,
-                cancer_subtype_name_map=cancer_subtype_name_map,
-            )
-
-    @patch('mosaic.batch_analysis.load_all_models')
-    @patch('mosaic.batch_analysis.analyze_slide_with_models')
-    def test_batch_analysis_empty_results(
-        self, mock_analyze_slide, mock_load_models, sample_settings_df, cancer_subtype_name_map
-    ):
-        """Test batch analysis with slides that have no tissue."""
-        mock_cache = Mock()
-        mock_cache.cleanup = Mock()
-        mock_load_models.return_value = mock_cache
-
-        # All slides return None (no tissue found)
-        mock_analyze_slide.return_value = (None, None, None)
-
-        slides = ["slide1.svs", "slide2.svs"]
-
-        masks, aeon_results, paladin_results = analyze_slides_batch(
-            slides=slides,
-            settings_df=sample_settings_df[:2],
-            cancer_subtype_name_map=cancer_subtype_name_map,
-        )
-
-        # Should have empty results
-        assert len(masks) == 0
-        assert len(aeon_results) == 0
-        assert len(paladin_results) == 0
-
-        # Cleanup should still be called
-        mock_cache.cleanup.assert_called_once()
-
-    @patch('mosaic.batch_analysis.load_all_models')
-    @patch('mosaic.batch_analysis.analyze_slide_with_models')
-    def test_batch_analysis_aggressive_memory_management(
-        self, mock_analyze_slide, mock_load_models, sample_settings_df, cancer_subtype_name_map
-    ):
-        """Test batch analysis with explicit aggressive memory management."""
-        mock_cache = Mock()
-        mock_cache.cleanup = Mock()
-        mock_cache.aggressive_memory_mgmt = True
-        mock_load_models.return_value = mock_cache
-
-        mock_analyze_slide.return_value = (Mock(), Mock(), Mock())
-
-        slides = ["slide1.svs"]
-
-        analyze_slides_batch(
-            slides=slides,
-            settings_df=sample_settings_df[:1],
-            cancer_subtype_name_map=cancer_subtype_name_map,
-            aggressive_memory_mgmt=True,
-        )
-
-        # Verify aggressive_memory_mgmt was passed to load_all_models
-        mock_load_models.assert_called_once_with(
-            use_gpu=True,
-            aggressive_memory_mgmt=True,
-        )
-
-    @patch('mosaic.batch_analysis.load_all_models')
-    @patch('mosaic.batch_analysis.analyze_slide_with_models')
-    def test_batch_analysis_progress_tracking(
-        self, mock_analyze_slide, mock_load_models, sample_settings_df, cancer_subtype_name_map
-    ):
-        """Test batch analysis updates progress correctly."""
-        mock_cache = Mock()
-        mock_cache.cleanup = Mock()
-        mock_load_models.return_value = mock_cache
-
-        mock_analyze_slide.return_value = (Mock(), Mock(), Mock())
-
-        mock_progress = Mock()
-        slides = ["slide1.svs", "slide2.svs", "slide3.svs"]
-
-        analyze_slides_batch(
-            slides=slides,
-            settings_df=sample_settings_df,
-            cancer_subtype_name_map=cancer_subtype_name_map,
-            progress=mock_progress,
-        )
-
-        # Verify progress was called
-        assert mock_progress.call_count > 0
-
-        # Verify final progress call
-        final_call = mock_progress.call_args_list[-1]
-        assert final_call[0][0] == 1.0  # Should be 100% at end
-
-    @patch('mosaic.batch_analysis.load_all_models')
-    @patch('mosaic.batch_analysis.analyze_slide_with_models')
-    def test_batch_analysis_multi_slide_naming(
-        self, mock_analyze_slide, mock_load_models, sample_settings_df, cancer_subtype_name_map
-    ):
-        """Test that multi-slide results include slide names."""
-        mock_cache = Mock()
-        mock_cache.cleanup = Mock()
-        mock_load_models.return_value = mock_cache
-
-        # Return new DataFrames each time
-        def mock_analyze_side_effect(*args, **kwargs):
-            mock_mask = Mock()
-            # Aeon results should have Cancer Subtype as index, not a column
-            mock_aeon = pd.DataFrame({"Confidence": [0.95]}, index=pd.Index(["LUAD"], name="Cancer Subtype"))
-            mock_paladin = pd.DataFrame({
-                "Cancer Subtype": ["LUAD"],
-                "Biomarker": ["EGFR"],
-                "Score": [0.85]
-            })
-            return (mock_mask, mock_aeon, mock_paladin)
-
-        mock_analyze_slide.side_effect = mock_analyze_side_effect
-
-        slides = ["slide1.svs", "slide2.svs"]
-
-        masks, aeon_results, paladin_results = analyze_slides_batch(
-            slides=slides,
-            settings_df=sample_settings_df[:2],
-            cancer_subtype_name_map=cancer_subtype_name_map,
-        )
-
-        # Verify slide names are in results
-        assert len(masks) == 2
-        assert masks[0][1] == "slide1.svs"
-        assert masks[1][1] == "slide2.svs"
-
-        # Paladin results should have Slide column
-        assert "Slide" in paladin_results[0].columns
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])

tests/test_cli.py

@@ -0,0 +1,298 @@
+"""Tests for CLI execution modes and argument handling.
+
+This module tests the Mosaic CLI, including:
+- Argument parsing and routing
+- Single-slide processing mode
+- Batch CSV processing mode
+- Model download behavior
+- Output file generation
+"""
+
+import pytest
+from unittest.mock import Mock, patch, MagicMock, call
+from pathlib import Path
+import pandas as pd
+
+
+class TestArgumentParsing:
+    """Test CLI argument parsing and mode routing."""
+
+    @patch("mosaic.gradio_app.launch_gradio")
+    @patch("mosaic.gradio_app.download_and_process_models")
+    @patch("sys.argv", ["mosaic"])
+    def test_no_arguments_launches_web_interface(self, mock_download, mock_launch):
+        """Test no arguments routes to web interface mode."""
+        mock_download.return_value = ({}, {}, [])
+
+        from mosaic.gradio_app import main
+
+        main()
+
+        # Should call launch_gradio
+        assert mock_launch.called
+        assert mock_launch.call_count == 1
+
+    @patch("mosaic.gradio_app.analyze_slide")
+    @patch("mosaic.gradio_app.download_and_process_models")
+    @patch("sys.argv", ["mosaic", "--slide-path", "test.svs", "--output-dir", "out"])
+    def test_slide_path_routes_to_single_mode(self, mock_download, mock_analyze):
+        """Test --slide-path routes to single-slide mode."""
+        mock_download.return_value = ({"Unknown": "UNK"}, {"UNK": "Unknown"}, [])
+        mock_analyze.return_value = (None, None, None)
+
+        from mosaic.gradio_app import main
+
+        with patch("mosaic.gradio_app.Path.mkdir"):
+            main()
+
+        # Should call analyze_slide
+        assert mock_analyze.called
+
+    @patch("mosaic.gradio_app.load_all_models")
+    @patch("mosaic.gradio_app.load_settings")
+    @patch("mosaic.gradio_app.validate_settings")
+    @patch("mosaic.gradio_app.analyze_slide")
+    @patch("mosaic.gradio_app.download_and_process_models")
+    @patch("sys.argv", ["mosaic", "--slide-csv", "test.csv", "--output-dir", "out"])
+    def test_slide_csv_routes_to_batch_mode(
+        self,
+        mock_download,
+        mock_analyze,
+        mock_validate,
+        mock_load_settings,
+        mock_load_models,
+    ):
+        """Test --slide-csv routes to batch mode."""
+        mock_download.return_value = ({"Unknown": "UNK"}, {"UNK": "Unknown"}, [])
+        mock_load_settings.return_value = pd.DataFrame(
+            {
+                "Slide": ["test.svs"],
+                "Site Type": ["Primary"],
+                "Sex": ["Unknown"],
+                "Tissue Site": ["Unknown"],
+                "Cancer Subtype": ["Unknown"],
+                "IHC Subtype": [""],
+                "Segmentation Config": ["Biopsy"],
+            }
+        )
+        mock_validate.return_value = mock_load_settings.return_value
+        mock_analyze.return_value = (None, None, None)
+
+        mock_cache = Mock()
+        mock_cache.cleanup = Mock()
+        mock_load_models.return_value = mock_cache
+
+        from mosaic.gradio_app import main
+
+        with patch("mosaic.gradio_app.Path.mkdir"):
+            main()
+
+        # Should call load_all_models (batch mode)
+        assert mock_load_models.called
+
+
+class TestSingleSlideMode:
+    """Test single-slide processing mode."""
+
+    @patch("mosaic.gradio_app.Path.mkdir")
+    @patch("mosaic.gradio_app.analyze_slide")
+    @patch("mosaic.gradio_app.download_and_process_models")
+    def test_analyze_slide_called_with_correct_params(
+        self, mock_download, mock_analyze, mock_mkdir, cli_args_single
+    ):
+        """Test analyze_slide called with correct parameters in single mode."""
+        mock_download.return_value = ({"Unknown": "UNK"}, {"UNK": "Unknown"}, [])
+        mock_analyze.return_value = (None, None, None)
+
+        # Patch ArgumentParser to return our test args
+        with patch(
+            "mosaic.gradio_app.ArgumentParser.parse_args", return_value=cli_args_single
+        ):
+            from mosaic.gradio_app import main
+
+            main()
+
+        # Verify analyze_slide was called
+        assert mock_analyze.called
+        call_args = mock_analyze.call_args[0]  # Positional args
+
+        # Check key parameters (analyze_slide uses positional args)
+        assert call_args[0] == cli_args_single.slide_path  # slide_path
+        assert call_args[1] == cli_args_single.segmentation_config  # seg_config
+        assert call_args[2] == cli_args_single.site_type  # site_type
+
+    @patch("PIL.Image.Image.save")
+    @patch("mosaic.gradio_app.Path.mkdir")
+    @patch("mosaic.gradio_app.analyze_slide")
+    @patch("mosaic.gradio_app.download_and_process_models")
+    def test_output_files_saved_correctly(
+        self,
+        mock_download,
+        mock_analyze,
+        mock_mkdir,
+        mock_save,
+        cli_args_single,
+        mock_analyze_slide_results,
+    ):
+        """Test output files are saved with correct names."""
+        from PIL import Image
+
+        mock_download.return_value = ({"Unknown": "UNK"}, {"UNK": "Unknown"}, [])
+
+        # Mock analyze_slide to return results
+        mask, aeon_results, paladin_results = mock_analyze_slide_results
+        mock_analyze.return_value = (mask, aeon_results, paladin_results)
+
+        # Patch ArgumentParser
+        with patch(
+            "mosaic.gradio_app.ArgumentParser.parse_args", return_value=cli_args_single
+        ):
+            # Patch DataFrame.to_csv to avoid actual file writes
+            with patch("pandas.DataFrame.to_csv"):
+                from mosaic.gradio_app import main
+
+                main()
+
+        # Verify save was called for mask
+        assert mock_save.called
+
+
+class TestBatchCsvMode:
+    """Test batch CSV processing mode."""
+
+    @patch("mosaic.gradio_app.Path.mkdir")
+    @patch("mosaic.gradio_app.load_all_models")
+    @patch("mosaic.gradio_app.analyze_slide")
+    @patch("mosaic.gradio_app.validate_settings")
+    @patch("mosaic.gradio_app.load_settings")
+    @patch("mosaic.gradio_app.download_and_process_models")
+    def test_load_all_models_called_once(
+        self,
+        mock_download,
+        mock_load_settings,
+        mock_validate,
+        mock_analyze,
+        mock_load_models,
+        mock_mkdir,
+        cli_args_batch,
+        sample_settings_df,
+        mock_analyze_slide_results,
+    ):
+        """Test load_all_models called once in batch mode."""
+        from PIL import Image
+
+        mock_download.return_value = ({"Unknown": "UNK"}, {"UNK": "Unknown"}, [])
+        mock_load_settings.return_value = sample_settings_df
+        mock_validate.return_value = sample_settings_df
+
+        # Return fresh DataFrames on each call to avoid mutation
+        def mock_analyze_side_effect(*args, **kwargs):
+            mask = Image.new("RGB", (100, 100), color="red")
+            aeon_results = pd.DataFrame(
+                {"Cancer Subtype": ["LUAD"], "Confidence": [0.95]}
+            )
+            paladin_results = pd.DataFrame(
+                {
+                    "Cancer Subtype": ["LUAD", "LUAD", "LUAD"],
+                    "Biomarker": ["TP53", "KRAS", "EGFR"],
+                    "Score": [0.85, 0.72, 0.63],
+                }
+            )
+            return (mask, aeon_results, paladin_results)
+
+        mock_analyze.side_effect = mock_analyze_side_effect
+
+        mock_cache = Mock()
+        mock_cache.cleanup = Mock()
+        mock_load_models.return_value = mock_cache
+
+        with patch(
+            "mosaic.gradio_app.ArgumentParser.parse_args", return_value=cli_args_batch
+        ):
+            with patch("pandas.DataFrame.to_csv"):
+                with patch("PIL.Image.Image.save"):
+                    from mosaic.gradio_app import main
+
+                    main()
+
+        # load_all_models should be called exactly once
+        assert mock_load_models.call_count == 1
+
+        # analyze_slide should be called for each slide (3 times)
+        assert mock_analyze.call_count == 3
+
+        # All analyze_slide calls should receive the model_cache
+        for call in mock_analyze.call_args_list:
+            assert call[1]["model_cache"] == mock_cache
+
+        # cleanup should be called
+        assert mock_cache.cleanup.called
+
+    @patch("mosaic.gradio_app.Path.mkdir")
+    @patch("mosaic.gradio_app.load_all_models")
+    @patch("mosaic.gradio_app.analyze_slide")
+    @patch("mosaic.gradio_app.validate_settings")
+    @patch("mosaic.gradio_app.load_settings")
+    @patch("mosaic.gradio_app.download_and_process_models")
+    def test_combined_outputs_generated(
+        self,
+        mock_download,
+        mock_load_settings,
+        mock_validate,
+        mock_analyze,
+        mock_load_models,
+        mock_mkdir,
+        cli_args_batch,
+        sample_settings_df,
+        mock_analyze_slide_results,
+    ):
+        """Test combined output files are generated in batch mode."""
+        from PIL import Image
+
+        mock_download.return_value = (
+            {"Unknown": "UNK", "Lung Adenocarcinoma (LUAD)": "LUAD"},
+            {"UNK": "Unknown", "LUAD": "Lung Adenocarcinoma (LUAD)"},
+            ["LUAD"],
+        )
+        mock_load_settings.return_value = sample_settings_df
+        mock_validate.return_value = sample_settings_df
+
+        # Return fresh DataFrames on each call
+        def mock_analyze_side_effect(*args, **kwargs):
+            mask = Image.new("RGB", (100, 100), color="red")
+            aeon_results = pd.DataFrame(
+                {"Cancer Subtype": ["LUAD"], "Confidence": [0.95]}
+            )
+            paladin_results = pd.DataFrame(
+                {
+                    "Cancer Subtype": ["LUAD", "LUAD", "LUAD"],
+                    "Biomarker": ["TP53", "KRAS", "EGFR"],
+                    "Score": [0.85, 0.72, 0.63],
+                }
+            )
+            return (mask, aeon_results, paladin_results)
+
+        mock_analyze.side_effect = mock_analyze_side_effect
+
+        mock_cache = Mock()
+        mock_cache.cleanup = Mock()
+        mock_load_models.return_value = mock_cache
+
+        csv_calls = []
+
+        def track_csv_write(path, *args, **kwargs):
+            """Track CSV file writes."""
+            csv_calls.append(str(path))
+
+        with patch(
+            "mosaic.gradio_app.ArgumentParser.parse_args", return_value=cli_args_batch
+        ):
+            with patch("pandas.DataFrame.to_csv", side_effect=track_csv_write):
+                with patch("PIL.Image.Image.save"):
+                    from mosaic.gradio_app import main
+
+                    main()
+
+        # Should have combined files
+        combined_files = [c for c in csv_calls if "combined" in c]
+        assert len(combined_files) >= 2  # combined_aeon and combined_paladin

tests/test_fixtures.py

@@ -0,0 +1,377 @@
+"""Shared fixtures and utilities for UI and CLI tests.
+
+This module provides reusable fixtures for testing the Mosaic Gradio UI and CLI,
+including mock file objects, settings DataFrames, cancer subtype mappings, and
+utility functions for test setup/teardown.
+"""
+
+import tempfile
+from pathlib import Path
+from unittest.mock import Mock
+import pandas as pd
+import numpy as np
+import pytest
+from PIL import Image
+
+
+# ============================================================================
+# File and Path Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def test_slide_path():
+    """Path to actual test slide for integration tests."""
+    return Path("tests/testdata/948176.svs")
+
+
+@pytest.fixture
+def temp_output_dir():
+    """Temporary directory for test outputs."""
+    with tempfile.TemporaryDirectory(prefix="mosaic_test_") as tmpdir:
+        yield Path(tmpdir)
+
+
+@pytest.fixture
+def mock_user_dir(temp_output_dir):
+    """Mock user directory (same as temp_output_dir for simplicity)."""
+    return temp_output_dir
+
+
+# ============================================================================
+# Mock File Upload Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def sample_files_single():
+    """Mock single file upload."""
+    mock_file = Mock()
+    mock_file.name = "test_slide_1.svs"
+    return [mock_file]
+
+
+@pytest.fixture
+def sample_files_multiple():
+    """Mock multiple file uploads (3 files)."""
+    files = []
+    for i in range(1, 4):
+        mock_file = Mock()
+        mock_file.name = f"test_slide_{i}.svs"
+        files.append(mock_file)
+    return files
+
+
+def create_mock_file(filename):
+    """Create a mock file object with specified filename.
+
+    Args:
+        filename: Name for the mock file
+
+    Returns:
+        Mock object with .name attribute
+    """
+    mock_file = Mock()
+    mock_file.name = filename
+    return mock_file
+
+
+# ============================================================================
+# Settings DataFrame Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def sample_settings_df():
+    """Sample settings DataFrame with 3 slides."""
+    return pd.DataFrame(
+        {
+            "Slide": ["slide1.svs", "slide2.svs", "slide3.svs"],
+            "Site Type": ["Primary", "Metastatic", "Primary"],
+            "Sex": ["Unknown", "Female", "Male"],
+            "Tissue Site": ["Lung", "Liver", "Unknown"],
+            "Cancer Subtype": ["Unknown", "Lung Adenocarcinoma (LUAD)", "Unknown"],
+            "IHC Subtype": ["", "", ""],
+            "Segmentation Config": ["Biopsy", "Resection", "TCGA"],
+        }
+    )
+
+
+def create_settings_df(n_rows, **kwargs):
+    """Generate a test settings DataFrame with specified number of rows.
+
+    Args:
+        n_rows: Number of rows to generate
+        **kwargs: Column overrides (e.g., site_type="Metastatic")
+
+    Returns:
+        DataFrame with SETTINGS_COLUMNS
+    """
+    defaults = {
+        "Slide": [f"slide_{i}.svs" for i in range(1, n_rows + 1)],
+        "Site Type": ["Primary"] * n_rows,
+        "Sex": ["Unknown"] * n_rows,
+        "Tissue Site": ["Unknown"] * n_rows,
+        "Cancer Subtype": ["Unknown"] * n_rows,
+        "IHC Subtype": [""] * n_rows,
+        "Segmentation Config": ["Biopsy"] * n_rows,
+    }
+
+    # Override with any provided kwargs
+    for key, value in kwargs.items():
+        column_name = key.replace("_", " ").title()
+        if column_name in defaults:
+            if isinstance(value, list):
+                defaults[column_name] = value
+            else:
+                defaults[column_name] = [value] * n_rows
+
+    return pd.DataFrame(defaults)
+
+
+# ============================================================================
+# CSV File Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def sample_csv_valid():
+    """Temporary CSV file with valid settings."""
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
+        f.write(
+            "Slide,Site Type,Sex,Tissue Site,Cancer Subtype,IHC Subtype,Segmentation Config\n"
+        )
+        f.write("slide1.svs,Primary,Unknown,Lung,Unknown,,Biopsy\n")
+        f.write(
+            "slide2.svs,Metastatic,Female,Liver,Lung Adenocarcinoma (LUAD),,Resection\n"
+        )
+        f.write("slide3.svs,Primary,Male,Unknown,Unknown,,TCGA\n")
+        f.flush()
+        yield f.name
+    Path(f.name).unlink(missing_ok=True)
+
+
+@pytest.fixture
+def sample_csv_invalid():
+    """Temporary CSV file with invalid values (for validation testing)."""
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
+        f.write(
+            "Slide,Site Type,Sex,Tissue Site,Cancer Subtype,IHC Subtype,Segmentation Config\n"
+        )
+        f.write(
+            "slide1.svs,InvalidSite,InvalidSex,InvalidTissue,InvalidSubtype,InvalidIHC,InvalidConfig\n"
+        )
+        f.write(
+            "slide2.svs,Primary,Unknown,Lung,BRCA,HR+/HER2+,Biopsy\n"
+        )  # Valid breast cancer
+        f.flush()
+        yield f.name
+    Path(f.name).unlink(missing_ok=True)
+
+
+@pytest.fixture
+def sample_csv_minimal():
+    """Temporary CSV file with only required columns (missing optional columns)."""
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
+        f.write("Slide,Site Type,Cancer Subtype\n")
+        f.write("slide1.svs,Primary,Unknown\n")
+        f.write("slide2.svs,Metastatic,LUAD\n")
+        f.flush()
+        yield f.name
+    Path(f.name).unlink(missing_ok=True)
+
+
+# ============================================================================
+# Cancer Subtype Mapping Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def mock_cancer_subtype_maps():
+    """Mock cancer subtype mappings for testing."""
+    cancer_subtype_name_map = {
+        "Unknown": "UNK",
+        "Lung Adenocarcinoma (LUAD)": "LUAD",
+        "Breast Invasive Carcinoma (BRCA)": "BRCA",
+        "Colorectal Adenocarcinoma (COAD)": "COAD",
+        "Prostate Adenocarcinoma (PRAD)": "PRAD",
+    }
+
+    reversed_cancer_subtype_name_map = {
+        "UNK": "Unknown",
+        "LUAD": "Lung Adenocarcinoma (LUAD)",
+        "BRCA": "Breast Invasive Carcinoma (BRCA)",
+        "COAD": "Colorectal Adenocarcinoma (COAD)",
+        "PRAD": "Prostate Adenocarcinoma (PRAD)",
+    }
+
+    cancer_subtypes = ["LUAD", "BRCA", "COAD", "PRAD"]
+
+    return cancer_subtype_name_map, reversed_cancer_subtype_name_map, cancer_subtypes
+
+
+# ============================================================================
+# Mock Analysis Results Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def mock_analyze_slide_results():
+    """Mock results from analyze_slide function."""
+    # Create a simple test mask image
+    mask = Image.new("RGB", (100, 100), color="red")
+
+    # Create Aeon results DataFrame
+    aeon_results = pd.DataFrame(
+        {
+            "Cancer Subtype": ["LUAD"],
+            "Confidence": [0.95],
+        }
+    )
+
+    # Create Paladin results DataFrame (NOTE: No "Slide" column - that gets added by CLI/UI)
+    paladin_results = pd.DataFrame(
+        {
+            "Cancer Subtype": ["LUAD", "LUAD", "LUAD"],
+            "Biomarker": ["TP53", "KRAS", "EGFR"],
+            "Score": [0.85, 0.72, 0.63],
+        }
+    )
+
+    return (mask, aeon_results, paladin_results)
+
+
+@pytest.fixture
+def mock_model_cache():
+    """Mock ModelCache with test models."""
+    from unittest.mock import Mock
+
+    cache = Mock()
+    cache.ctranspath_model = Mock()
+    cache.optimus_model = Mock()
+    cache.marker_classifier = Mock()
+    cache.aeon_model = Mock()
+    cache.paladin_models = {}
+    cache.device = Mock()
+    cache.cleanup = Mock()
+
+    return cache
+
+
+# ============================================================================
+# CLI Argument Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def cli_args_single():
+    """Mock argparse Namespace for single-slide mode."""
+    from argparse import Namespace
+
+    return Namespace(
+        debug=False,
+        server_name="0.0.0.0",
+        server_port=None,
+        share=False,
+        slide_path="tests/testdata/948176.svs",
+        slide_csv=None,
+        output_dir="test_output",
+        site_type="Primary",
+        sex="Unknown",
+        tissue_site="Unknown",
+        cancer_subtype="Unknown",
+        ihc_subtype="",
+        segmentation_config="Biopsy",
+        num_workers=4,
+    )
+
+
+@pytest.fixture
+def cli_args_batch(sample_csv_valid):
+    """Mock argparse Namespace for batch mode."""
+    from argparse import Namespace
+
+    return Namespace(
+        debug=False,
+        server_name="0.0.0.0",
+        server_port=None,
+        share=False,
+        slide_path=None,
+        slide_csv=sample_csv_valid,
+        output_dir="test_output",
+        site_type="Primary",
+        sex="Unknown",
+        tissue_site="Unknown",
+        cancer_subtype="Unknown",
+        ihc_subtype="",
+        segmentation_config="Biopsy",
+        num_workers=4,
+    )
+
+
+# ============================================================================
+# Utility Functions
+# ============================================================================
+
+
+def verify_csv_output(path, expected_columns):
+    """Validate CSV file structure.
+
+    Args:
+        path: Path to CSV file
+        expected_columns: List of expected column names
+
+    Returns:
+        DataFrame loaded from CSV
+
+    Raises:
+        AssertionError: If CSV is invalid or missing columns
+    """
+    assert Path(path).exists(), f"CSV file not found: {path}"
+
+    df = pd.read_csv(path)
+    assert not df.empty, f"CSV file is empty: {path}"
+
+    missing_cols = set(expected_columns) - set(df.columns)
+    assert not missing_cols, f"Missing columns in CSV: {missing_cols}"
+
+    return df
+
+
+def mock_gradio_components():
+    """Context manager to mock Gradio component classes.
+
+    Usage:
+        with mock_gradio_components() as mocks:
+            # Gradio components are mocked
+            result = function_that_returns_gr_components()
+            # Verify mocks
+            assert mocks['Dataframe'].called
+    """
+    from unittest.mock import patch, Mock
+
+    mocks = {
+        "Dataframe": Mock(return_value=Mock()),
+        "File": Mock(return_value=Mock()),
+        "DownloadButton": Mock(return_value=Mock()),
+        "Dropdown": Mock(return_value=Mock()),
+        "Gallery": Mock(return_value=Mock()),
+        "Error": Exception,  # gr.Error is an exception
+        "Warning": Mock(),
+    }
+
+    patches = []
+    for name, mock_obj in mocks.items():
+        patch_obj = patch(f"mosaic.ui.app.gr.{name}", mock_obj)
+        patches.append(patch_obj)
+
+    # Start all patches
+    for p in patches:
+        p.start()
+
+    try:
+        yield mocks
+    finally:
+        # Stop all patches
+        for p in patches:
+            p.stop()

tests/test_gradio_app.py

@@ -71,14 +71,16 @@ class TestLoadSettings:
         reversed_cancer_subtype_name_map = {
             value: key for key, value in cancer_subtype_name_map.items()
         }
-        return cancer_subtype_name_map, cancer_subtypes, reversed_cancer_subtype_name_map
+        return (
+            cancer_subtype_name_map,
+            cancer_subtypes,
+            reversed_cancer_subtype_name_map,
+        )
 
     @pytest.fixture
     def temp_settings_csv(self):
         """Create a temporary settings CSV file with all columns."""
-        with tempfile.NamedTemporaryFile(
-            mode="w", delete=False, suffix=".csv"
-        ) as f:
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".csv") as f:
             f.write("Slide,Site Type,Cancer Subtype,IHC Subtype,Segmentation Config\n")
             f.write("slide1.svs,Primary,Unknown,,Biopsy\n")
             f.write("slide2.svs,Metastatic,Unknown,,Resection\n")
@@ -89,9 +91,7 @@ class TestLoadSettings:
     @pytest.fixture
     def temp_minimal_settings_csv(self):
         """Create a temporary settings CSV file with minimal columns."""
-        with tempfile.NamedTemporaryFile(
-            mode="w", delete=False, suffix=".csv"
-        ) as f:
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".csv") as f:
             f.write("Slide,Site Type\n")
             f.write("slide1.svs,Primary\n")
             f.write("slide2.svs,Metastatic\n")
@@ -129,9 +129,7 @@ class TestLoadSettings:
 
     def test_load_settings_missing_required_column_raises_error(self):
         """Test that missing required column raises ValueError."""
-        with tempfile.NamedTemporaryFile(
-            mode="w", delete=False, suffix=".csv"
-        ) as f:
+        with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".csv") as f:
             f.write("RandomColumn\n")
             f.write("value\n")
             temp_path = f.name

tests/test_model_manager.py

@@ -10,7 +10,11 @@ from unittest.mock import Mock, patch, MagicMock
 import pickle
 import gc
 
-from mosaic.model_manager import ModelCache, load_all_models, load_paladin_model_for_inference
+from mosaic.model_manager import (
+    ModelCache,
+    load_all_models,
+    load_paladin_model_for_inference,
+)
 
 
 class TestModelCache:
@@ -73,9 +77,11 @@ class TestModelCache:
 
         assert cache.paladin_models == {}
 
-    @patch('torch.cuda.is_available', return_value=True)
-    @patch('torch.cuda.empty_cache')
-    def test_cleanup_paladin_clears_cuda_cache(self, mock_empty_cache, mock_cuda_available):
+    @patch("torch.cuda.is_available", return_value=True)
+    @patch("torch.cuda.empty_cache")
+    def test_cleanup_paladin_clears_cuda_cache(
+        self, mock_empty_cache, mock_cuda_available
+    ):
         """Test cleanup_paladin calls torch.cuda.empty_cache()."""
         cache = ModelCache()
         cache.paladin_models = {"model1": Mock()}
@@ -107,52 +113,52 @@ class TestModelCache:
 class TestLoadAllModels:
     """Test load_all_models function."""
 
-    @patch('torch.cuda.is_available', return_value=False)
+    @patch("torch.cuda.is_available", return_value=False)
     def test_load_models_cpu_only(self, mock_cuda_available):
         """Test loading models when CUDA is not available."""
-        with patch('builtins.open', create=True) as mock_open:
-            with patch('pickle.load') as mock_pickle:
+        with patch("builtins.open", create=True) as mock_open:
+            with patch("pickle.load") as mock_pickle:
                 # Mock the pickle loads
                 mock_pickle.return_value = Mock()
 
                 # Mock file exists checks
-                with patch.object(Path, 'exists', return_value=True):
+                with patch.object(Path, "exists", return_value=True):
                     cache = load_all_models(use_gpu=False)
 
         assert cache is not None
         assert cache.device == torch.device("cpu")
         assert cache.aggressive_memory_mgmt is False
 
-    @patch('torch.cuda.is_available', return_value=True)
-    @patch('torch.cuda.get_device_name', return_value="NVIDIA A100")
+    @patch("torch.cuda.is_available", return_value=True)
+    @patch("torch.cuda.get_device_name", return_value="NVIDIA A100")
     def test_load_models_a100_gpu(self, mock_get_device, mock_cuda_available):
         """Test loading models on A100 GPU (high memory)."""
-        with patch('builtins.open', create=True):
-            with patch('pickle.load') as mock_pickle:
+        with patch("builtins.open", create=True):
+            with patch("pickle.load") as mock_pickle:
                 mock_model = Mock()
                 mock_model.to = Mock(return_value=mock_model)
                 mock_model.eval = Mock()
                 mock_pickle.return_value = mock_model
 
-                with patch.object(Path, 'exists', return_value=True):
+                with patch.object(Path, "exists", return_value=True):
                     cache = load_all_models(use_gpu=True, aggressive_memory_mgmt=None)
 
         assert cache.device == torch.device("cuda")
         assert cache.is_t4_gpu is False
         assert cache.aggressive_memory_mgmt is False  # A100 should use caching
 
-    @patch('torch.cuda.is_available', return_value=True)
-    @patch('torch.cuda.get_device_name', return_value="Tesla T4")
+    @patch("torch.cuda.is_available", return_value=True)
+    @patch("torch.cuda.get_device_name", return_value="Tesla T4")
     def test_load_models_t4_gpu(self, mock_get_device, mock_cuda_available):
         """Test loading models on T4 GPU (low memory)."""
-        with patch('builtins.open', create=True):
-            with patch('pickle.load') as mock_pickle:
+        with patch("builtins.open", create=True):
+            with patch("pickle.load") as mock_pickle:
                 mock_model = Mock()
                 mock_model.to = Mock(return_value=mock_model)
                 mock_model.eval = Mock()
                 mock_pickle.return_value = mock_model
 
-                with patch.object(Path, 'exists', return_value=True):
+                with patch.object(Path, "exists", return_value=True):
                     cache = load_all_models(use_gpu=True, aggressive_memory_mgmt=None)
 
         assert cache.device == torch.device("cuda")
@@ -161,33 +167,36 @@ class TestLoadAllModels:
 
     def test_load_models_missing_aeon_file(self):
         """Test load_all_models raises error when Aeon model file is missing."""
+
         def exists_side_effect(self):
             # Return True for marker_classifier and optimus, False for aeon
             filename = str(self)
-            if 'aeon_model.pkl' in filename:
+            if "aeon_model.pkl" in filename:
                 return False
             return True
 
-        with patch.object(Path, 'exists', exists_side_effect):
+        with patch.object(Path, "exists", exists_side_effect):
             with pytest.raises(FileNotFoundError, match="Aeon model not found"):
-                with patch('builtins.open', create=True):
-                    with patch('pickle.load'):
+                with patch("builtins.open", create=True):
+                    with patch("pickle.load"):
                         load_all_models(use_gpu=False)
 
-    @patch('torch.cuda.is_available', return_value=True)
+    @patch("torch.cuda.is_available", return_value=True)
     def test_load_models_explicit_aggressive_mode(self, mock_cuda_available):
         """Test explicit aggressive memory management setting."""
-        with patch('torch.cuda.get_device_name', return_value="NVIDIA A100"):
-            with patch('builtins.open', create=True):
-                with patch('pickle.load') as mock_pickle:
+        with patch("torch.cuda.get_device_name", return_value="NVIDIA A100"):
+            with patch("builtins.open", create=True):
+                with patch("pickle.load") as mock_pickle:
                     mock_model = Mock()
                     mock_model.to = Mock(return_value=mock_model)
                     mock_model.eval = Mock()
                     mock_pickle.return_value = mock_model
 
-                    with patch.object(Path, 'exists', return_value=True):
+                    with patch.object(Path, "exists", return_value=True):
                         # Force aggressive mode even on A100
-                        cache = load_all_models(use_gpu=True, aggressive_memory_mgmt=True)
+                        cache = load_all_models(
+                            use_gpu=True, aggressive_memory_mgmt=True
+                        )
 
         assert cache.aggressive_memory_mgmt is True  # Should respect explicit setting
 
@@ -200,8 +209,8 @@ class TestLoadPaladinModelForInference:
         cache = ModelCache(aggressive_memory_mgmt=True, device=torch.device("cpu"))
         model_path = Path("data/paladin/test_model.pkl")
 
-        with patch('builtins.open', create=True):
-            with patch('pickle.load') as mock_pickle:
+        with patch("builtins.open", create=True):
+            with patch("pickle.load") as mock_pickle:
                 mock_model = Mock()
                 mock_model.to = Mock(return_value=mock_model)
                 mock_model.eval = Mock()
@@ -220,8 +229,8 @@ class TestLoadPaladinModelForInference:
         cache = ModelCache(aggressive_memory_mgmt=False, device=torch.device("cpu"))
         model_path = Path("data/paladin/test_model.pkl")
 
-        with patch('builtins.open', create=True):
-            with patch('pickle.load') as mock_pickle:
+        with patch("builtins.open", create=True):
+            with patch("pickle.load") as mock_pickle:
                 mock_model = Mock()
                 mock_model.to = Mock(return_value=mock_model)
                 mock_model.eval = Mock()
@@ -243,7 +252,7 @@ class TestLoadPaladinModelForInference:
         cache.paladin_models[str(model_path)] = cached_model
 
         # Load model - should return cached version without pickle.load
-        with patch('pickle.load') as mock_pickle:
+        with patch("pickle.load") as mock_pickle:
             model = load_paladin_model_for_inference(cache, model_path)
 
         assert model == cached_model

tests/test_regression_single_slide.py

@@ -30,13 +30,14 @@ class TestSingleSlideRegression:
             "Lung Adenocarcinoma": "LUAD",
         }
 
-    @patch('mosaic.analysis.segment_tissue')
-    @patch('mosaic.analysis.draw_slide_mask')
-    @patch('mosaic.analysis._extract_ctranspath_features')
-    @patch('mosaic.analysis.filter_features')
-    @patch('mosaic.analysis._extract_optimus_features')
-    @patch('mosaic.analysis._run_aeon_inference')
-    @patch('mosaic.analysis._run_paladin_inference')
+    @patch("mosaic.analysis.segment_tissue")
+    @patch("mosaic.analysis.draw_slide_mask")
+    @patch("mosaic.model_manager.load_all_models")
+    @patch("mosaic.analysis._extract_ctranspath_features")
+    @patch("mosaic.analysis.filter_features")
+    @patch("mosaic.analysis._extract_optimus_features")
+    @patch("mosaic.analysis._run_aeon_inference_with_model")
+    @patch("mosaic.analysis._run_paladin_inference_with_models")
     def test_single_slide_analyze_slide_unchanged(
         self,
         mock_paladin,
@@ -44,6 +45,7 @@ class TestSingleSlideRegression:
         mock_optimus,
         mock_filter,
         mock_ctranspath,
+        mock_load_models,
         mock_mask,
         mock_segment,
         mock_slide_path,
@@ -60,6 +62,16 @@ class TestSingleSlideRegression:
         mock_mask_image = Mock()
         mock_mask.return_value = mock_mask_image
 
+        # Mock ModelCache with required attributes
+        mock_model_cache = Mock()
+        mock_model_cache.ctranspath_model = Mock()
+        mock_model_cache.optimus_model = Mock()
+        mock_model_cache.marker_classifier = Mock()
+        mock_model_cache.aeon_model = Mock()
+        mock_model_cache.device = Mock()
+        mock_model_cache.cleanup = Mock()
+        mock_load_models.return_value = mock_model_cache
+
         mock_features = np.random.rand(100, 768)
         mock_ctranspath.return_value = (mock_features, mock_coords)
 
@@ -69,17 +81,14 @@ class TestSingleSlideRegression:
         mock_optimus_features = np.random.rand(50, 1536)
         mock_optimus.return_value = mock_optimus_features
 
-        mock_aeon_results = pd.DataFrame({
-            "Cancer Subtype": ["LUAD", "LUSC"],
-            "Confidence": [0.85, 0.15]
-        })
+        mock_aeon_results = pd.DataFrame(
+            {"Cancer Subtype": ["LUAD", "LUSC"], "Confidence": [0.85, 0.15]}
+        )
         mock_aeon.return_value = mock_aeon_results
 
-        mock_paladin_results = pd.DataFrame({
-            "Cancer Subtype": ["LUAD"],
-            "Biomarker": ["EGFR"],
-            "Score": [0.75]
-        })
+        mock_paladin_results = pd.DataFrame(
+            {"Cancer Subtype": ["LUAD"], "Biomarker": ["EGFR"], "Score": [0.75]}
+        )
         mock_paladin.return_value = mock_paladin_results
 
         # Run analyze_slide
@@ -107,10 +116,10 @@ class TestSingleSlideRegression:
         assert isinstance(aeon_results, pd.DataFrame)
         assert isinstance(paladin_results, pd.DataFrame)
 
-    @patch('mosaic.ui.app.analyze_slide')
-    @patch('mosaic.ui.app.create_user_directory')
-    @patch('mosaic.ui.app.validate_settings')
-    @patch('pandas.DataFrame.to_csv')  # Mock CSV writing to avoid directory issues
+    @patch("mosaic.ui.app.analyze_slide")
+    @patch("mosaic.ui.app.create_user_directory")
+    @patch("mosaic.ui.app.validate_settings")
+    @patch("pandas.DataFrame.to_csv")  # Mock CSV writing to avoid directory issues
     def test_gradio_single_slide_uses_analyze_slide(
         self,
         mock_to_csv,
@@ -121,40 +130,53 @@ class TestSingleSlideRegression:
         """Test that Gradio UI uses analyze_slide for single slide (not batch mode)."""
         # Setup
         import tempfile
+
         with tempfile.TemporaryDirectory() as tmpdir:
             mock_dir = Path(tmpdir) / "test_user"
             mock_dir.mkdir()
             mock_create_dir.return_value = mock_dir
 
-            settings_df = pd.DataFrame({
-                "Slide": ["test.svs"],
-                "Site Type": ["Primary"],
-                "Sex": ["Male"],
-                "Tissue Site": ["Lung"],
-                "Cancer Subtype": ["Unknown"],
-                "IHC Subtype": [""],
-                "Segmentation Config": ["Biopsy"],
-            })
+            settings_df = pd.DataFrame(
+                {
+                    "Slide": ["test.svs"],
+                    "Site Type": ["Primary"],
+                    "Sex": ["Male"],
+                    "Tissue Site": ["Lung"],
+                    "Cancer Subtype": ["Unknown"],
+                    "IHC Subtype": [""],
+                    "Segmentation Config": ["Biopsy"],
+                }
+            )
             mock_validate.return_value = settings_df
 
             mock_mask = Mock()
             mock_aeon = pd.DataFrame({"Cancer Subtype": ["LUAD"], "Confidence": [0.9]})
-            mock_paladin = pd.DataFrame({
-                "Cancer Subtype": ["LUAD"],
-                "Biomarker": ["EGFR"],
-                "Score": [0.8]
-            })
+            mock_paladin = pd.DataFrame(
+                {"Cancer Subtype": ["LUAD"], "Biomarker": ["EGFR"], "Score": [0.8]}
+            )
             mock_analyze_slide.return_value = (mock_mask, mock_aeon, mock_paladin)
 
             from mosaic.ui.app import cancer_subtype_name_map
 
-            # Call analyze_slides with a single slide
-            with patch('mosaic.ui.app.get_oncotree_code_name', return_value="Lung Adenocarcinoma"):
-                masks, aeon, aeon_btn, paladin, paladin_btn, user_dir = analyze_slides(
+            # Call analyze_slides with a single slide (generator function)
+            with patch(
+                "mosaic.ui.app.get_oncotree_code_name",
+                return_value="Lung Adenocarcinoma",
+            ):
+                gen = analyze_slides(
                     slides=["test.svs"],
                     settings_input=settings_df,
+                    site_type="Primary",
+                    sex="Male",
+                    tissue_site="Lung",
+                    cancer_subtype="Unknown",
+                    ihc_subtype="",
+                    seg_config="Biopsy",
                     user_dir=mock_dir,
                 )
+                # Consume generator to get final result
+                results = list(gen)
+                masks, aeon, aeon_btn, paladin, paladin_btn, user_dir = results[-1]
 
             # Verify analyze_slide was called (not analyze_slides_batch)
             mock_analyze_slide.assert_called_once()
@@ -162,10 +184,11 @@ class TestSingleSlideRegression:
             # Verify results
             assert len(masks) == 1
 
-
-    @patch('mosaic.analysis.segment_tissue')
-    @patch('mosaic.analysis.gr.Warning')
-    def test_single_slide_no_tissue_found(self, mock_warning, mock_segment, mock_slide_path, cancer_subtype_name_map):
+    @patch("mosaic.analysis.segment_tissue")
+    @patch("mosaic.analysis.gr.Warning")
+    def test_single_slide_no_tissue_found(
+        self, mock_warning, mock_segment, mock_slide_path, cancer_subtype_name_map
+    ):
         """Test single-slide analysis when no tissue is found."""
         # No tissue tiles found
         mock_segment.return_value = None  # segment_tissue returns None when no tissue
@@ -187,18 +210,20 @@ class TestSingleSlideRegression:
         # Verify warning was raised
         mock_warning.assert_called_once()
 
-    @patch('mosaic.analysis.segment_tissue')
-    @patch('mosaic.analysis.draw_slide_mask')
-    @patch('mosaic.analysis._extract_ctranspath_features')
-    @patch('mosaic.analysis.filter_features')
-    @patch('mosaic.analysis._extract_optimus_features')
-    @patch('mosaic.analysis._run_paladin_inference')
+    @patch("mosaic.analysis.segment_tissue")
+    @patch("mosaic.analysis.draw_slide_mask")
+    @patch("mosaic.model_manager.load_all_models")
+    @patch("mosaic.analysis._extract_ctranspath_features")
+    @patch("mosaic.analysis.filter_features")
+    @patch("mosaic.analysis._extract_optimus_features")
+    @patch("mosaic.analysis._run_paladin_inference_with_models")
     def test_single_slide_known_cancer_subtype_skips_aeon(
         self,
         mock_paladin,
         mock_optimus,
         mock_filter,
         mock_ctranspath,
+        mock_load_models,
         mock_mask,
         mock_segment,
         mock_slide_path,
@@ -211,16 +236,25 @@ class TestSingleSlideRegression:
         mock_attrs = {}
         mock_segment.return_value = (mock_polygon, None, mock_coords, mock_attrs)
         mock_mask.return_value = Mock()
+
+        # Mock ModelCache
+        mock_model_cache = Mock()
+        mock_model_cache.ctranspath_model = Mock()
+        mock_model_cache.optimus_model = Mock()
+        mock_model_cache.marker_classifier = Mock()
+        mock_model_cache.aeon_model = Mock()
+        mock_model_cache.device = Mock()
+        mock_model_cache.cleanup = Mock()
+        mock_load_models.return_value = mock_model_cache
+
         mock_ctranspath.return_value = (np.random.rand(10, 768), np.array([[0, 0]]))
         mock_filter.return_value = (None, np.array([[0, 0]]))
         mock_optimus.return_value = np.random.rand(10, 1536)
-        mock_paladin.return_value = pd.DataFrame({
-            "Cancer Subtype": ["LUAD"],
-            "Biomarker": ["EGFR"],
-            "Score": [0.8]
-        })
+        mock_paladin.return_value = pd.DataFrame(
+            {"Cancer Subtype": ["LUAD"], "Biomarker": ["EGFR"], "Score": [0.8]}
+        )
 
-        with patch('mosaic.analysis._run_aeon_inference') as mock_aeon:
+        with patch("mosaic.analysis._run_aeon_inference_with_model") as mock_aeon:
             slide_mask, aeon_results, paladin_results = analyze_slide(
                 slide_path=mock_slide_path,
                 seg_config="Biopsy",
@@ -244,6 +278,7 @@ class TestBackwardCompatibility:
     def test_analyze_slide_signature_unchanged(self):
         """Test that analyze_slide function signature is unchanged."""
         from inspect import signature
+
         sig = signature(analyze_slide)
 
         # Verify required parameters exist
@@ -261,8 +296,8 @@ class TestBackwardCompatibility:
 
     def test_analyze_slide_return_type_unchanged(self):
         """Test that analyze_slide returns the same tuple structure."""
-        with patch('mosaic.analysis.segment_tissue', return_value=None):  # No tissue
-            with patch('mosaic.analysis.gr.Warning'):  # Mock the warning
+        with patch("mosaic.analysis.segment_tissue", return_value=None):  # No tissue
+            with patch("mosaic.analysis.gr.Warning"):  # Mock the warning
                 result = analyze_slide(
                     slide_path="test.svs",
                     seg_config="Biopsy",

tests/test_ui_components.py

@@ -0,0 +1,302 @@
+"""Tests for Gradio UI components and their interactions.
+
+This module tests the Mosaic Gradio UI components, including:
+- Settings validation
+- Analysis workflow
+"""
+
+import pytest
+import pandas as pd
+from unittest.mock import Mock, patch, MagicMock
+from pathlib import Path
+
+# Import after mocking (mocks are set up in conftest.py)
+from mosaic.ui.app import (
+    analyze_slides,
+    set_cancer_subtype_maps,
+)
+from mosaic.ui.utils import SETTINGS_COLUMNS
+
+
+class TestSettingsValidation:
+    """Test settings validation logic."""
+
+    @patch("mosaic.ui.utils.gr.Warning")
+    def test_invalid_cancer_subtype_defaults_to_unknown(
+        self, mock_warning, mock_cancer_subtype_maps
+    ):
+        """Test invalid cancer subtype generates warning and defaults to Unknown."""
+        from mosaic.ui.utils import validate_settings
+
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+
+        # Create DataFrame with invalid cancer subtype
+        df = pd.DataFrame(
+            {
+                "Slide": ["test.svs"],
+                "Site Type": ["Primary"],
+                "Sex": ["Unknown"],
+                "Tissue Site": ["Unknown"],
+                "Cancer Subtype": ["InvalidSubtype"],
+                "IHC Subtype": [""],
+                "Segmentation Config": ["Biopsy"],
+            }
+        )
+
+        result = validate_settings(
+            df, cancer_subtype_name_map, cancer_subtypes, reversed_map
+        )
+
+        # Should default to Unknown
+        assert result.iloc[0]["Cancer Subtype"] == "Unknown"
+        # Warning should be called
+        assert mock_warning.called
+
+    @patch("mosaic.ui.utils.gr.Warning")
+    def test_invalid_site_type_defaults_to_primary(
+        self, mock_warning, mock_cancer_subtype_maps
+    ):
+        """Test invalid site type generates warning and defaults to Primary."""
+        from mosaic.ui.utils import validate_settings
+
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+
+        df = pd.DataFrame(
+            {
+                "Slide": ["test.svs"],
+                "Site Type": ["InvalidSite"],
+                "Sex": ["Unknown"],
+                "Tissue Site": ["Unknown"],
+                "Cancer Subtype": ["Unknown"],
+                "IHC Subtype": [""],
+                "Segmentation Config": ["Biopsy"],
+            }
+        )
+
+        result = validate_settings(
+            df, cancer_subtype_name_map, cancer_subtypes, reversed_map
+        )
+
+        assert result.iloc[0]["Site Type"] == "Primary"
+        assert mock_warning.called
+
+
+class TestAnalysisWorkflow:
+    """Test analysis workflow with mocked analyze_slide."""
+
+    @patch("mosaic.ui.app.analyze_slide")
+    @patch("mosaic.ui.app.create_user_directory")
+    def test_single_slide_analysis_no_model_cache(
+        self,
+        mock_create_dir,
+        mock_analyze,
+        sample_files_single,
+        mock_analyze_slide_results,
+        mock_cancer_subtype_maps,
+        temp_output_dir,
+    ):
+        """Test single slide analysis doesn't load model cache."""
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+        set_cancer_subtype_maps(cancer_subtype_name_map, reversed_map, cancer_subtypes)
+
+        # Setup mocks
+        mock_create_dir.return_value = temp_output_dir
+        mock_analyze.return_value = mock_analyze_slide_results
+
+        # Generate settings DataFrame manually
+        settings_df = pd.DataFrame(
+            {
+                "Slide": ["test_slide_1.svs"],
+                "Site Type": ["Primary"],
+                "Sex": ["Unknown"],
+                "Tissue Site": ["Unknown"],
+                "Cancer Subtype": ["Unknown"],
+                "IHC Subtype": [""],
+                "Segmentation Config": ["Biopsy"],
+            }
+        )
+
+        # Call analyze_slides (generator)
+        gen = analyze_slides(
+            sample_files_single,
+            settings_df,
+            "Primary",
+            "Unknown",
+            "Unknown",
+            "Unknown",
+            "",
+            "Biopsy",
+            temp_output_dir,
+        )
+
+        # Consume generator
+        results = list(gen)
+
+        # Should yield at least once (intermediate + final)
+        assert len(results) >= 1
+
+        # analyze_slide should be called once
+        assert mock_analyze.call_count == 1
+
+        # Should be called with model_cache=None (single-slide mode)
+        call_kwargs = mock_analyze.call_args[1]
+        assert call_kwargs["model_cache"] is None
+
+    @patch("mosaic.ui.app.load_all_models")
+    @patch("mosaic.ui.app.analyze_slide")
+    @patch("mosaic.ui.app.create_user_directory")
+    def test_batch_analysis_loads_model_cache_once(
+        self,
+        mock_create_dir,
+        mock_analyze,
+        mock_load_models,
+        sample_files_multiple,
+        mock_analyze_slide_results,
+        mock_model_cache,
+        mock_cancer_subtype_maps,
+        temp_output_dir,
+    ):
+        """Test batch analysis loads models once and reuses cache."""
+        from PIL import Image
+
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+        set_cancer_subtype_maps(cancer_subtype_name_map, reversed_map, cancer_subtypes)
+
+        # Setup mocks - return new DataFrames on each call to avoid mutation issues
+        def mock_analyze_side_effect(*args, **kwargs):
+            mask = Image.new("RGB", (100, 100), color="red")
+            aeon_results = pd.DataFrame(
+                {"Cancer Subtype": ["LUAD"], "Confidence": [0.95]}
+            )
+            paladin_results = pd.DataFrame(
+                {
+                    "Cancer Subtype": ["LUAD", "LUAD", "LUAD"],
+                    "Biomarker": ["TP53", "KRAS", "EGFR"],
+                    "Score": [0.85, 0.72, 0.63],
+                }
+            )
+            return (mask, aeon_results, paladin_results)
+
+        mock_create_dir.return_value = temp_output_dir
+        mock_load_models.return_value = mock_model_cache
+        mock_analyze.side_effect = mock_analyze_side_effect
+
+        # Generate settings DataFrame manually for 3 files
+        settings_df = pd.DataFrame(
+            {
+                "Slide": ["test_slide_1.svs", "test_slide_2.svs", "test_slide_3.svs"],
+                "Site Type": ["Primary", "Primary", "Primary"],
+                "Sex": ["Unknown", "Unknown", "Unknown"],
+                "Tissue Site": ["Unknown", "Unknown", "Unknown"],
+                "Cancer Subtype": ["Unknown", "Unknown", "Unknown"],
+                "IHC Subtype": ["", "", ""],
+                "Segmentation Config": ["Biopsy", "Biopsy", "Biopsy"],
+            }
+        )
+
+        # Call analyze_slides
+        gen = analyze_slides(
+            sample_files_multiple,
+            settings_df,
+            "Primary",
+            "Unknown",
+            "Unknown",
+            "Unknown",
+            "",
+            "Biopsy",
+            temp_output_dir,
+        )
+
+        # Consume generator
+        results = list(gen)
+
+        # load_all_models should be called once
+        assert mock_load_models.call_count == 1
+
+        # analyze_slide should be called 3 times (once per file)
+        assert mock_analyze.call_count == 3
+
+        # All calls should use the same model_cache
+        for call in mock_analyze.call_args_list:
+            assert call[1]["model_cache"] == mock_model_cache
+
+        # cleanup should be called
+        assert mock_model_cache.cleanup.called
+
+    @patch("mosaic.ui.app.create_user_directory")
+    def test_no_slides_raises_error(
+        self, mock_create_dir, mock_cancer_subtype_maps, temp_output_dir
+    ):
+        """Test that no slides uploaded raises gr.Error."""
+        import gradio as gr
+
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+        set_cancer_subtype_maps(cancer_subtype_name_map, reversed_map, cancer_subtypes)
+
+        mock_create_dir.return_value = temp_output_dir
+
+        # Call with no slides
+        gen = analyze_slides(
+            None,
+            None,
+            "Primary",
+            "Unknown",
+            "Unknown",
+            "Unknown",
+            "",
+            "Biopsy",
+            temp_output_dir,
+        )
+
+        # Should raise gr.Error
+        with pytest.raises(gr.Error):
+            next(gen)
+
+    @patch("mosaic.ui.app.create_user_directory")
+    def test_settings_mismatch_raises_error(
+        self,
+        mock_create_dir,
+        sample_files_multiple,
+        sample_settings_df,
+        mock_cancer_subtype_maps,
+        temp_output_dir,
+    ):
+        """Test that settings count mismatch raises gr.Error."""
+        import gradio as gr
+
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+        set_cancer_subtype_maps(cancer_subtype_name_map, reversed_map, cancer_subtypes)
+
+        mock_create_dir.return_value = temp_output_dir
+
+        # sample_files_multiple has 3 files, sample_settings_df has 3 rows
+        # Manually create mismatch by using only 2 files
+        two_files = sample_files_multiple[:2]
+
+        gen = analyze_slides(
+            two_files,
+            sample_settings_df,
+            "Primary",
+            "Unknown",
+            "Unknown",
+            "Unknown",
+            "",
+            "Biopsy",
+            temp_output_dir,
+        )
+
+        # Should raise gr.Error about mismatch
+        with pytest.raises(gr.Error):
+            next(gen)

tests/test_ui_events.py

@@ -0,0 +1,349 @@
+"""Tests for UI event handlers and state management.
+
+This module tests complex event interactions in the Mosaic Gradio UI, including:
+- Settings state management across events
+- Generator behavior and incremental updates
+- Error and warning display
+"""
+
+import pytest
+import pandas as pd
+from unittest.mock import Mock, patch, MagicMock
+from pathlib import Path
+import inspect
+
+from mosaic.ui.app import (
+    analyze_slides,
+    set_cancer_subtype_maps,
+)
+from mosaic.ui.utils import SETTINGS_COLUMNS, validate_settings, load_settings
+
+
+class TestSettingsStateManagement:
+    """Test settings state management across multiple events."""
+
+    def test_csv_upload_replaces_settings(
+        self, sample_csv_valid, mock_cancer_subtype_maps
+    ):
+        """Test CSV upload replaces existing settings."""
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+
+        # Load CSV
+        loaded_df = load_settings(sample_csv_valid)
+        validated_df = validate_settings(
+            loaded_df, cancer_subtype_name_map, cancer_subtypes, reversed_map
+        )
+
+        # Verify new settings loaded
+        assert len(validated_df) == 3
+        assert validated_df.iloc[0]["Slide"] == "slide1.svs"
+        assert validated_df.iloc[1]["Slide"] == "slide2.svs"
+
+
+class TestGeneratorBehavior:
+    """Test generator behavior for incremental updates."""
+
+    @patch("mosaic.ui.app.analyze_slide")
+    @patch("mosaic.ui.app.create_user_directory")
+    def test_analyze_slides_is_generator(
+        self,
+        mock_create_dir,
+        mock_analyze,
+        sample_files_single,
+        mock_analyze_slide_results,
+        mock_cancer_subtype_maps,
+        temp_output_dir,
+    ):
+        """Test analyze_slides returns a generator."""
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+        set_cancer_subtype_maps(cancer_subtype_name_map, reversed_map, cancer_subtypes)
+
+        mock_create_dir.return_value = temp_output_dir
+        mock_analyze.return_value = mock_analyze_slide_results
+
+        settings_df = pd.DataFrame(
+            {
+                "Slide": ["test_slide_1.svs"],
+                "Site Type": ["Primary"],
+                "Sex": ["Unknown"],
+                "Tissue Site": ["Unknown"],
+                "Cancer Subtype": ["Unknown"],
+                "IHC Subtype": [""],
+                "Segmentation Config": ["Biopsy"],
+            }
+        )
+
+        result = analyze_slides(
+            sample_files_single,
+            settings_df,
+            "Primary",
+            "Unknown",
+            "Unknown",
+            "Unknown",
+            "",
+            "Biopsy",
+            temp_output_dir,
+        )
+
+        # Verify it's a generator
+        assert inspect.isgenerator(result)
+
+    @patch("mosaic.ui.app.load_all_models")
+    @patch("mosaic.ui.app.analyze_slide")
+    @patch("mosaic.ui.app.create_user_directory")
+    def test_intermediate_yields_update_masks_only(
+        self,
+        mock_create_dir,
+        mock_analyze,
+        mock_load_models,
+        sample_files_multiple,
+        mock_analyze_slide_results,
+        mock_model_cache,
+        mock_cancer_subtype_maps,
+        temp_output_dir,
+    ):
+        """Test intermediate yields show only slide masks."""
+        from PIL import Image
+
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+        set_cancer_subtype_maps(cancer_subtype_name_map, reversed_map, cancer_subtypes)
+
+        mock_create_dir.return_value = temp_output_dir
+        mock_load_models.return_value = mock_model_cache
+
+        # Return fresh DataFrames on each call
+        def mock_analyze_side_effect(*args, **kwargs):
+            mask = Image.new("RGB", (100, 100), color="red")
+            aeon_results = pd.DataFrame(
+                {"Cancer Subtype": ["LUAD"], "Confidence": [0.95]}
+            )
+            paladin_results = pd.DataFrame(
+                {
+                    "Cancer Subtype": ["LUAD", "LUAD", "LUAD"],
+                    "Biomarker": ["TP53", "KRAS", "EGFR"],
+                    "Score": [0.85, 0.72, 0.63],
+                }
+            )
+            return (mask, aeon_results, paladin_results)
+
+        mock_analyze.side_effect = mock_analyze_side_effect
+
+        settings_df = pd.DataFrame(
+            {
+                "Slide": ["test_slide_1.svs", "test_slide_2.svs", "test_slide_3.svs"],
+                "Site Type": ["Primary", "Primary", "Primary"],
+                "Sex": ["Unknown", "Unknown", "Unknown"],
+                "Tissue Site": ["Unknown", "Unknown", "Unknown"],
+                "Cancer Subtype": ["Unknown", "Unknown", "Unknown"],
+                "IHC Subtype": ["", "", ""],
+                "Segmentation Config": ["Biopsy", "Biopsy", "Biopsy"],
+            }
+        )
+
+        gen = analyze_slides(
+            sample_files_multiple,
+            settings_df,
+            "Primary",
+            "Unknown",
+            "Unknown",
+            "Unknown",
+            "",
+            "Biopsy",
+            temp_output_dir,
+        )
+
+        # Get first intermediate yield (after first slide)
+        first_yield = next(gen)
+
+        # Should be tuple with 6 elements
+        assert len(first_yield) == 6
+
+        # First element is slide_masks (should have 1 entry)
+        slide_masks = first_yield[0]
+        assert len(slide_masks) == 1
+
+    @patch("mosaic.ui.app.load_all_models")
+    @patch("mosaic.ui.app.analyze_slide")
+    @patch("mosaic.ui.app.create_user_directory")
+    def test_final_yield_has_complete_results(
+        self,
+        mock_create_dir,
+        mock_analyze,
+        mock_load_models,
+        sample_files_multiple,
+        mock_analyze_slide_results,
+        mock_model_cache,
+        mock_cancer_subtype_maps,
+        temp_output_dir,
+    ):
+        """Test final yield contains complete results."""
+        from PIL import Image
+
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+        set_cancer_subtype_maps(cancer_subtype_name_map, reversed_map, cancer_subtypes)
+
+        mock_create_dir.return_value = temp_output_dir
+        mock_load_models.return_value = mock_model_cache
+
+        # Return fresh DataFrames on each call
+        def mock_analyze_side_effect(*args, **kwargs):
+            mask = Image.new("RGB", (100, 100), color="red")
+            aeon_results = pd.DataFrame(
+                {"Cancer Subtype": ["LUAD"], "Confidence": [0.95]}
+            )
+            paladin_results = pd.DataFrame(
+                {
+                    "Cancer Subtype": ["LUAD", "LUAD", "LUAD"],
+                    "Biomarker": ["TP53", "KRAS", "EGFR"],
+                    "Score": [0.85, 0.72, 0.63],
+                }
+            )
+            return (mask, aeon_results, paladin_results)
+
+        mock_analyze.side_effect = mock_analyze_side_effect
+
+        settings_df = pd.DataFrame(
+            {
+                "Slide": ["test_slide_1.svs", "test_slide_2.svs", "test_slide_3.svs"],
+                "Site Type": ["Primary", "Primary", "Primary"],
+                "Sex": ["Unknown", "Unknown", "Unknown"],
+                "Tissue Site": ["Unknown", "Unknown", "Unknown"],
+                "Cancer Subtype": ["Unknown", "Unknown", "Unknown"],
+                "IHC Subtype": ["", "", ""],
+                "Segmentation Config": ["Biopsy", "Biopsy", "Biopsy"],
+            }
+        )
+
+        gen = analyze_slides(
+            sample_files_multiple,
+            settings_df,
+            "Primary",
+            "Unknown",
+            "Unknown",
+            "Unknown",
+            "",
+            "Biopsy",
+            temp_output_dir,
+        )
+
+        # Consume generator to get final yield
+        results = list(gen)
+        final_yield = results[-1]
+
+        # Final yield should have all results
+        assert len(final_yield) == 6
+        slide_masks = final_yield[0]
+        assert len(slide_masks) == 3  # All 3 slides
+
+
+class TestErrorDisplay:
+    """Test error and warning display behavior."""
+
+    @patch("mosaic.ui.app.create_user_directory")
+    def test_no_slides_raises_gr_error(
+        self, mock_create_dir, mock_cancer_subtype_maps, temp_output_dir
+    ):
+        """Test that no slides raises gr.Error."""
+        import gradio as gr
+
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+        set_cancer_subtype_maps(cancer_subtype_name_map, reversed_map, cancer_subtypes)
+
+        mock_create_dir.return_value = temp_output_dir
+
+        gen = analyze_slides(
+            None,
+            None,
+            "Primary",
+            "Unknown",
+            "Unknown",
+            "Unknown",
+            "",
+            "Biopsy",
+            temp_output_dir,
+        )
+
+        # Should raise gr.Error
+        with pytest.raises(gr.Error):
+            next(gen)
+
+    @patch("mosaic.ui.utils.gr.Warning")
+    def test_validation_warnings_shown(self, mock_warning, mock_cancer_subtype_maps):
+        """Test validation warnings are displayed."""
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+
+        # Create DataFrame with multiple invalid values
+        df = pd.DataFrame(
+            {
+                "Slide": ["test1.svs", "test2.svs"],
+                "Site Type": ["InvalidSite", "Primary"],
+                "Sex": ["Unknown", "InvalidSex"],
+                "Tissue Site": ["Unknown", "Unknown"],
+                "Cancer Subtype": ["InvalidSubtype", "Unknown"],
+                "IHC Subtype": ["", ""],
+                "Segmentation Config": ["Biopsy", "InvalidConfig"],
+            }
+        )
+
+        result = validate_settings(
+            df, cancer_subtype_name_map, cancer_subtypes, reversed_map
+        )
+
+        # Should have warning calls (at least 1 for the multiple invalid values)
+        assert mock_warning.call_count >= 1
+
+        # Verify defaults applied
+        assert result.iloc[0]["Site Type"] == "Primary"  # Invalid → Primary
+        assert result.iloc[0]["Cancer Subtype"] == "Unknown"  # Invalid → Unknown
+        assert result.iloc[1]["Sex"] == "Unknown"  # Invalid → Unknown
+        assert result.iloc[1]["Segmentation Config"] == "Biopsy"  # Invalid → Biopsy
+
+    @patch("mosaic.ui.app.create_user_directory")
+    def test_settings_mismatch_raises_gr_error(
+        self,
+        mock_create_dir,
+        sample_files_multiple,
+        sample_settings_df,
+        mock_cancer_subtype_maps,
+        temp_output_dir,
+    ):
+        """Test settings/files count mismatch raises gr.Error."""
+        import gradio as gr
+
+        cancer_subtype_name_map, reversed_map, cancer_subtypes = (
+            mock_cancer_subtype_maps
+        )
+        set_cancer_subtype_maps(cancer_subtype_name_map, reversed_map, cancer_subtypes)
+
+        mock_create_dir.return_value = temp_output_dir
+
+        # Create mismatch: 2 files but 3 settings rows
+        two_files = sample_files_multiple[:2]
+
+        gen = analyze_slides(
+            two_files,
+            sample_settings_df,
+            "Primary",
+            "Unknown",
+            "Unknown",
+            "Unknown",
+            "",
+            "Biopsy",
+            temp_output_dir,
+        )
+
+        # Should raise gr.Error about mismatch
+        with pytest.raises(gr.Error):
+            next(gen)