augmenator/style_transfer.py

"""Neural style transfer: ONNX Model Zoo + PyTorch HF weights (CPU-friendly)."""

from __future__ import annotations

import random
from dataclasses import dataclass
from typing import Literal

import cv2
import numpy as np
import onnxruntime as ort
import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms

from augmenator.style_net import StyleNet
from augmenator.transform_net import TransformNet

StyleBackend = Literal["onnx", "transformnet", "stylenet", "classical"]


@dataclass(frozen=True)
class StyleSpec:
    tag: str
    label: str
    backend: StyleBackend
    repo: str = ""
    file: str = ""
    repo_type: str = "model"


STYLE_CATALOG: tuple[StyleSpec, ...] = (
    StyleSpec("style_candy", "candy", "onnx", "onnxmodelzoo/candy-9", "candy-9.onnx"),
    StyleSpec("style_mosaic", "mosaic", "onnx", "onnxmodelzoo/mosaic-9", "mosaic-9.onnx"),
    StyleSpec(
        "style_rain_princess",
        "rain-princess",
        "onnx",
        "onnxmodelzoo/rain-princess-9",
        "rain-princess-9.onnx",
    ),
    StyleSpec("style_udnie", "udnie", "onnx", "onnxmodelzoo/udnie-9", "udnie-9.onnx"),
    StyleSpec(
        "style_pointilism",
        "pointilism",
        "onnx",
        "onnxmodelzoo/pointilism-9",
        "pointilism-9.onnx",
    ),
    StyleSpec(
        "style_starry_night",
        "starry-night",
        "transformnet",
        "ebylmz/fast-neural-style-transfer",
        "models/starry_night_cw2.0_sw400000.0_tw2.0.pth",
        repo_type="space",
    ),
    StyleSpec(
        "style_sketch",
        "sketch",
        "stylenet",
        "Ateshh/mini-style-transfer",
        "sketch.pth",
    ),
    StyleSpec("style_random", "random", "classical"),
)

CONCRETE_STYLE_TAGS = frozenset(
    spec.tag for spec in STYLE_CATALOG if spec.tag != "style_random"
)
STYLE_TAGS = frozenset(spec.tag for spec in STYLE_CATALOG)
STYLE_MODELS = {spec.tag: {"label": spec.label, "repo": spec.repo, "file": spec.file} for spec in STYLE_CATALOG}

_STYLE_BY_TAG = {spec.tag: spec for spec in STYLE_CATALOG}

MAX_EDGE = 512
ONNX_MODEL_SIZE = 224

_onnx_sessions: dict[str, ort.InferenceSession] = {}
_torch_models: dict[str, torch.nn.Module] = {}
_sketch_weights_missing = False

_IMAGENET_NORMALIZE = transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225],
)


def _resize_for_cpu(image: Image.Image) -> tuple[Image.Image, Image.Image, tuple[int, int]]:
    original = image.convert("RGB")
    width, height = original.size
    scale = min(1.0, MAX_EDGE / max(width, height))
    if scale < 1.0:
        proc_w = max(1, int(width * scale))
        proc_h = max(1, int(height * scale))
        working = original.resize((proc_w, proc_h), Image.Resampling.LANCZOS)
    else:
        working = original
    return original, working, (width, height)


def _blend_strength(original: Image.Image, styled: Image.Image, strength: float) -> Image.Image:
    alpha = min(1.0, max(0.5, strength))
    if alpha < 1.0:
        return Image.blend(original, styled, alpha)
    return styled


def _get_onnx_session(spec: StyleSpec) -> ort.InferenceSession:
    if spec.tag not in _onnx_sessions:
        path = hf_hub_download(repo_id=spec.repo, filename=spec.file)
        _onnx_sessions[spec.tag] = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
    return _onnx_sessions[spec.tag]


def _load_transformnet(spec: StyleSpec) -> TransformNet:
    if spec.tag not in _torch_models:
        path = hf_hub_download(repo_id=spec.repo, filename=spec.file, repo_type=spec.repo_type)
        model = TransformNet()
        state_dict = torch.load(path, map_location="cpu", weights_only=True)
        model.load_state_dict(state_dict)
        model.eval()
        _torch_models[spec.tag] = model
    return _torch_models[spec.tag]  # type: ignore[return-value]


def _load_stylenet(spec: StyleSpec) -> StyleNet:
    global _sketch_weights_missing
    if spec.tag not in _torch_models:
        try:
            path = hf_hub_download(repo_id=spec.repo, filename=spec.file)
        except Exception as exc:
            _sketch_weights_missing = True
            raise RuntimeError(
                f"Sketch weights not found on Hugging Face ({spec.repo}/{spec.file}). "
                "Using classical pencil-sketch fallback."
            ) from exc
        model = StyleNet()
        model.load_state_dict(torch.load(path, map_location="cpu", weights_only=True))
        model.eval()
        _torch_models[spec.tag] = model
    return _torch_models[spec.tag]  # type: ignore[return-value]


def _apply_onnx_style(working: Image.Image, spec: StyleSpec) -> Image.Image:
    session = _get_onnx_session(spec)
    input_name = session.get_inputs()[0].name
    model_input = working.resize((ONNX_MODEL_SIZE, ONNX_MODEL_SIZE), Image.Resampling.LANCZOS)
    tensor = np.array(model_input).astype(np.float32)
    tensor = np.transpose(tensor, (2, 0, 1))
    tensor = np.expand_dims(tensor, axis=0)
    output = session.run(None, {input_name: tensor})[0]
    styled_arr = np.clip(output[0], 0, 255).transpose(1, 2, 0).astype(np.uint8)
    return Image.fromarray(styled_arr).resize(working.size, Image.Resampling.LANCZOS)


def _apply_transformnet_style(working: Image.Image, spec: StyleSpec) -> Image.Image:
    model = _load_transformnet(spec)
    transform = transforms.Compose([transforms.ToTensor(), _IMAGENET_NORMALIZE])
    tensor = transform(working).unsqueeze(0)
    with torch.no_grad():
        output = model(tensor).squeeze(0).cpu().numpy()
    mean = np.array([0.485, 0.456, 0.406]).reshape(3, 1, 1)
    std = np.array([0.229, 0.224, 0.225]).reshape(3, 1, 1)
    img = np.clip(output * std + mean, 0.0, 1.0)
    styled_arr = (img.transpose(1, 2, 0) * 255).astype(np.uint8)
    return Image.fromarray(styled_arr)


def _apply_stylenet_style(working: Image.Image, spec: StyleSpec) -> Image.Image:
    model = _load_stylenet(spec)
    transform = transforms.Compose([transforms.ToTensor(), _IMAGENET_NORMALIZE])
    tensor = transform(working).unsqueeze(0)
    with torch.no_grad():
        output = model(tensor).squeeze(0).clamp(0, 1)
    return transforms.ToPILImage()(output.cpu())


def _apply_classical_sketch(working: Image.Image) -> Image.Image:
    gray = cv2.cvtColor(np.array(working.convert("RGB")), cv2.COLOR_RGB2GRAY)
    inverted = 255 - gray
    blurred = cv2.GaussianBlur(inverted, (21, 21), 0)
    sketch = cv2.divide(gray, 255 - blurred, scale=256)
    rgb = cv2.cvtColor(sketch, cv2.COLOR_GRAY2RGB)
    return Image.fromarray(rgb)


def warmup(tag: str = "style_candy") -> None:
    """Pre-download and cache one style model (optional; others load on demand)."""
    if tag not in _STYLE_BY_TAG or tag == "style_random":
        return
    spec = _STYLE_BY_TAG[tag]
    if spec.backend == "onnx":
        _get_onnx_session(spec)
    elif spec.backend == "transformnet":
        _load_transformnet(spec)
    elif spec.backend == "stylenet":
        try:
            _load_stylenet(spec)
        except RuntimeError:
            pass


def apply_style(image: Image.Image, tag: str, strength: float = 1.0) -> Image.Image:
    if tag == "style_random":
        tag = random.choice(sorted(CONCRETE_STYLE_TAGS))

    if tag not in _STYLE_BY_TAG or tag == "style_random":
        return image

    spec = _STYLE_BY_TAG[tag]
    original, working, (width, height) = _resize_for_cpu(image)

    try:
        if spec.backend == "onnx":
            styled = _apply_onnx_style(working, spec)
        elif spec.backend == "transformnet":
            styled = _apply_transformnet_style(working, spec)
        elif spec.backend == "stylenet":
            try:
                styled = _apply_stylenet_style(working, spec)
            except RuntimeError:
                print(
                    "Note: Ateshh/mini-style-transfer sketch.pth is not on Hugging Face; "
                    "using classical pencil-sketch fallback."
                )
                styled = _apply_classical_sketch(working)
        else:
            return image
    except Exception as exc:
        if spec.tag == "style_sketch":
            print(f"Sketch style load failed ({exc}); using classical pencil-sketch fallback.")
            styled = _apply_classical_sketch(working)
        else:
            raise

    if working.size != (width, height):
        styled = styled.resize((width, height), Image.Resampling.LANCZOS)

    return _blend_strength(original, styled, strength)