models/style_transfer.py

"""Lazy TensorFlow Hub arbitrary style transfer wrapper."""

from __future__ import annotations

import os
import time
from functools import lru_cache

import cv2
import numpy as np

os.environ.setdefault("TFHUB_MODEL_LOAD_FORMAT", "COMPRESSED")


@lru_cache(maxsize=1)
def _load_model():
    import tensorflow_hub as hub

    return hub.load("https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2")


def _to_tensor(image: np.ndarray, max_size: int):
    import tensorflow as tf

    img = np.asarray(image).astype(np.float32) / 255.0
    h, w = img.shape[:2]
    scale = min(1.0, max_size / max(h, w))
    if scale < 1:
        img = cv2.resize(img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
    return tf.constant(img[None, ...])


def stylize(content: np.ndarray, style: np.ndarray, max_size: int = 512) -> tuple[np.ndarray, float, str]:
    try:
        start = time.perf_counter()
        model = _load_model()
        output = model(_to_tensor(content, max_size), _to_tensor(style, 256))[0]
        arr = np.clip(np.array(output[0]) * 255, 0, 255).astype(np.uint8)
        return arr, time.perf_counter() - start, "Style transfer complete."
    except Exception as exc:
        return np.asarray(content).astype(np.uint8), 0.0, f"Style transfer unavailable: {exc}"