forge/backend/app/services/image_upscaler.py

"""Image Upscaler Service - Topaz Labs API

Available Models:
- proteus: General enhancement with fine-tuning parameters (default)
- artemis: Detail enhancement and noise reduction
- gaia: Specialized for HD/4K upscaling
- iris: Noise and compression artifact reduction
- nyx: Low light and high ISO recovery
- rhea: Detail recovery for older/degraded images
- theia: High-fidelity upscaling

Output Options:
- Scale: 2x, 4x, 6x, 8x (up to 16K)
- Output formats: png, jpg, tiff
- Face enhancement: auto-detect and enhance faces
- Noise reduction: 0-100
- Sharpening: 0-100
- Grain recovery: preserve film grain
"""
import httpx
import os
from uuid import uuid4
from datetime import datetime
import asyncio
from typing import Optional, Dict, Any

from app.database import SessionLocal
from app.models.job import Job
from app.models.asset import Asset
from app.config import settings


# Topaz enhancement models with their specialties
TOPAZ_MODELS = {
    "proteus": {
        "name": "Proteus",
        "description": "General enhancement with fine control over noise, blur, and compression",
        "parameters": ["noise_reduction", "sharpening", "compression_recovery", "detail_enhancement"],
        "best_for": "General purpose, low to medium quality footage"
    },
    "artemis": {
        "name": "Artemis",
        "description": "Detail enhancement with noise reduction",
        "parameters": ["noise_reduction", "detail_recovery"],
        "best_for": "Details in low-noise footage"
    },
    "gaia": {
        "name": "Gaia",
        "description": "Specialized for upscaling HD to 4K/8K",
        "parameters": ["detail_level", "anti_aliasing"],
        "best_for": "High-resolution upscaling from HD source"
    },
    "iris": {
        "name": "Iris",
        "description": "Noise and compression artifact reduction",
        "parameters": ["noise_reduction", "compression_recovery", "debanding"],
        "best_for": "Heavily compressed or noisy images"
    },
    "nyx": {
        "name": "Nyx",
        "description": "Low light and high ISO recovery",
        "parameters": ["noise_reduction", "shadow_recovery", "highlight_recovery"],
        "best_for": "Dark or high-ISO images"
    },
    "rhea": {
        "name": "Rhea",
        "description": "Detail recovery for older/degraded images",
        "parameters": ["detail_recovery", "texture_enhancement"],
        "best_for": "Scanned photos, old digital images"
    },
    "theia": {
        "name": "Theia",
        "description": "High-fidelity detail enhancement",
        "parameters": ["detail_level", "texture_preservation"],
        "best_for": "Maximum detail retention"
    },
    "auto": {
        "name": "Auto",
        "description": "Automatically select best model for input",
        "parameters": [],
        "best_for": "When unsure which model to use"
    }
}


async def upscale(job_id: str):
    """Upscale image using Topaz Labs API

    Input parameters:
    - scale: Upscale factor (2, 4, 6, 8)
    - model: Enhancement model (see TOPAZ_MODELS)
    - output_format: 'png', 'jpg', 'tiff' (default: png)
    - face_enhancement: Boolean to enable face detection and enhancement
    - noise_reduction: 0-100, amount of noise removal
    - sharpening: 0-100, output sharpening level
    - compression_recovery: 0-100, recover compression artifacts
    - detail_enhancement: 0-100, enhance fine details
    - preserve_grain: Boolean to preserve film grain
    - output_quality: 1-100 for jpg output (default: 95)
    """
    db = SessionLocal()
    try:
        job = db.query(Job).filter(Job.id == job_id).first()
        if not job:
            return

        input_data = job.input_data
        input_asset_ids = job.input_asset_ids

        if not input_asset_ids:
            raise ValueError("No input asset provided")

        # Get input asset
        input_asset = db.query(Asset).filter(Asset.id == input_asset_ids[0]).first()
        if not input_asset:
            raise ValueError("Input asset not found")

        # Extract parameters
        scale = input_data.get("scale", 2)
        model = input_data.get("model", "auto")
        output_format = input_data.get("output_format", "png")
        face_enhancement = input_data.get("face_enhancement", False)
        noise_reduction = input_data.get("noise_reduction")
        sharpening = input_data.get("sharpening")
        compression_recovery = input_data.get("compression_recovery")
        detail_enhancement = input_data.get("detail_enhancement")
        preserve_grain = input_data.get("preserve_grain", False)
        output_quality = input_data.get("output_quality", 95)

        job.progress = 10
        job.api_provider = "topaz"
        job.api_model = model
        db.commit()

        # Read input image
        with open(input_asset.file_path, "rb") as f:
            image_data = f.read()

        # Calculate output dimensions
        original_width = input_asset.width or 1920
        original_height = input_asset.height or 1080
        output_width = original_width * scale
        output_height = original_height * scale

        job.progress = 20
        db.commit()

        # Build enhancement parameters
        enhance_params: Dict[str, Any] = {
            "output_height": str(output_height),
            "output_width": str(output_width),
            "output_format": output_format,
            "model": model,
            "face_enhancement": "true" if face_enhancement else "false"
        }

        # Add model-specific parameters if provided
        if noise_reduction is not None:
            enhance_params["noise_reduction"] = str(min(100, max(0, noise_reduction)))
        if sharpening is not None:
            enhance_params["sharpening"] = str(min(100, max(0, sharpening)))
        if compression_recovery is not None:
            enhance_params["compression_recovery"] = str(min(100, max(0, compression_recovery)))
        if detail_enhancement is not None:
            enhance_params["detail_enhancement"] = str(min(100, max(0, detail_enhancement)))
        if preserve_grain:
            enhance_params["preserve_grain"] = "true"
        if output_format == "jpg":
            enhance_params["quality"] = str(output_quality)

        # Call Topaz API
        async with httpx.AsyncClient(timeout=600) as client:
            # Start async enhancement
            response = await client.post(
                "https://api.topazlabs.com/image/v1/enhance/async",
                headers={
                    "X-API-Key": settings.topaz_api_key,
                    "Accept": "application/json"
                },
                files={"image": (input_asset.original_filename, image_data, input_asset.mime_type)},
                data=enhance_params
            )
            response.raise_for_status()
            result = response.json()

            request_id = result.get("id") or result.get("requestId")

            job.progress = 40
            job.api_request_id = request_id
            db.commit()

            # Poll for completion
            output_url = None
            for i in range(180):  # Wait up to 6 minutes for large upscales
                await asyncio.sleep(2)

                status_response = await client.get(
                    f"https://api.topazlabs.com/image/v1/enhance/{request_id}/status",
                    headers={"X-API-Key": settings.topaz_api_key}
                )
                status_data = status_response.json()
                status = status_data.get("status", "")

                if status == "completed":
                    output_url = status_data.get("outputUrl") or status_data.get("output_url")
                    break
                elif status == "failed":
                    raise ValueError(f"Topaz enhancement failed: {status_data.get('error')}")

                job.progress = min(40 + (i * 0.28), 85)
                db.commit()

            if output_url:
                # Download result
                img_response = await client.get(output_url)
                upscaled_data = img_response.content

                job.progress = 90
                db.commit()

                # Determine output extension
                ext_map = {"png": ".png", "jpg": ".jpg", "jpeg": ".jpg", "tiff": ".tiff"}
                ext = ext_map.get(output_format, ".png")
                mime_map = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg", "tiff": "image/tiff"}
                mime = mime_map.get(output_format, "image/png")

                # Save output
                filename = f"upscaled_{scale}x_{model}_{uuid4()}{ext}"
                storage_path = os.path.join(settings.storage_path, "images")
                os.makedirs(storage_path, exist_ok=True)
                file_path = os.path.join(storage_path, filename)

                with open(file_path, "wb") as f:
                    f.write(upscaled_data)

                # Create output asset
                output_asset = Asset(
                    user_id=job.user_id,
                    project_id=job.project_id,
                    original_filename=filename,
                    stored_filename=filename,
                    file_path=file_path,
                    file_type="image",
                    mime_type=mime,
                    file_size_bytes=len(upscaled_data),
                    width=output_width,
                    height=output_height,
                    source_module="image_upscaler",
                    source_job_id=job.id,
                    parent_asset_id=input_asset.id,
                    asset_metadata={
                        "scale": scale,
                        "model": model,
                        "face_enhancement": face_enhancement,
                        "noise_reduction": noise_reduction,
                        "sharpening": sharpening,
                        "original_dimensions": f"{original_width}x{original_height}",
                        "output_dimensions": f"{output_width}x{output_height}"
                    }
                )
                db.add(output_asset)
                db.commit()
                db.refresh(output_asset)

                job.output_asset_ids = [output_asset.id]
                job.output_data = {"asset_id": str(output_asset.id), "file_path": file_path}

        job.progress = 100
        job.status = "completed"
        job.completed_at = datetime.utcnow()
        db.commit()

    except Exception as e:
        job.status = "failed"
        job.error_message = str(e)
        db.commit()
    finally:
        db.close()


def get_available_models() -> Dict[str, Any]:
    """Get all available Topaz upscaling models and their capabilities"""
    return TOPAZ_MODELS