"""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)