Add ONNX neural style transfer with embedding-matched keywords

Five CPU-friendly fast neural style models (cartoon, mosaic, impressionist, abstract, pointillism) selected via the existing MiniLM embedding planner.

README.md

@@ -23,6 +23,13 @@ Upload an image and type an instruction. Transforms are selected by **semantic s
 
 ## Supported instructions
 
+**Artistic style transfer** (ONNX fast neural style, CPU-friendly — one style per request)
+- `cartoon` / `cartoon style` — candy/comic look
+- `mosaic` / `mosaic painting` — mosaic art
+- `impressionist` / `painterly` — soft painted look
+- `abstract art` / `cubist style` — bold abstract look
+- `pointillism` — dotted painting style
+
 **Color / classical** (examples — paraphrases work too)
 - `rotate`, `perspective`, `make it brighter`, `more contrast`, `vintage warm look`, `flip`, `blur softly`
 
@@ -67,5 +74,6 @@ First run downloads the embedding model, RapidOCR, and rembg u2netp weights.
 | `sentence-transformers/all-MiniLM-L6-v2` | Instruction → keyword matching | Apache 2.0 |
 | RapidOCR (ONNX) | Text region detection | Apache 2.0 |
 | rembg u2netp | Foreground segmentation | MIT |
+| ONNX fast neural style (onnxmodelzoo) | Artistic style transfer | See ONNX Model Zoo |
 
 See [backgrounds/README.md](backgrounds/README.md) for background prompt examples.

app.py

@@ -48,6 +48,9 @@ EXAMPLE_PROMPTS = [
     ["add cutout"],
     ["transparent cutout"],
     ["perspective transform"],
+    ["cartoon style"],
+    ["mosaic painting"],
+    ["impressionist look"],
     ["flip horizontally"],
     ["blur everything softly"],
 ]

pipeline/augment.py

@@ -5,6 +5,7 @@ import numpy as np
 from PIL import Image, ImageEnhance, ImageFilter, ImageOps
 
 from pipeline.spatial import apply_spatial_ops
+from pipeline.style_transfer import STYLE_MODELS, STYLE_TAGS, apply_style
 
 ALLOWED_TAGS = {
     "brighten",
@@ -27,7 +28,7 @@ ALLOWED_TAGS = {
     "tint_blue",
     "invert",
     "perspective",
-}
+} | set(STYLE_TAGS)
 
 
 def _clamp_strength(strength: float) -> float:
@@ -129,6 +130,9 @@ def _apply_tag(image: Image.Image, tag: str, strength: float) -> Image.Image:
         return cropped.resize((w, h), Image.Resampling.LANCZOS)
     if tag == "perspective":
         return _apply_perspective(image, s)
+    if tag in STYLE_TAGS:
+        styled = apply_style(image, tag, strength)
+        return styled, f"style({STYLE_MODELS[tag]['label']})"
 
     return image
 
@@ -185,7 +189,18 @@ def apply_augmentations(
         )
 
     for tag in tags:
-        if tag not in ALLOWED_TAGS:
+        if tag not in ALLOWED_TAGS or tag in STYLE_TAGS:
+            continue
+        out = _apply_tag(result, tag, strength)
+        if isinstance(out, tuple):
+            result, label = out
+            applied.append(label)
+        else:
+            result = out
+            applied.append(tag)
+
+    for tag in tags:
+        if tag not in STYLE_TAGS:
             continue
         out = _apply_tag(result, tag, strength)
         if isinstance(out, tuple):

pipeline/embedding_planner.py

@@ -16,6 +16,13 @@ MUTUAL_EXCLUSION_GROUPS: tuple[tuple[str, ...], ...] = (
     ("tint_red", "tint_green", "tint_blue"),
     ("cover_random", "cover_avoid_text", "cover_and_cutout", "cover_and_cutout_avoid_text"),
     ("add_cutout_transparent", "add_cutout_solid", "add_cutout", "add_cutout_avoid_text"),
+    (
+        "style_candy",
+        "style_mosaic",
+        "style_rain_princess",
+        "style_udnie",
+        "style_pointilism",
+    ),
 )
 
 _model: SentenceTransformer | None = None

pipeline/keyword_catalog.py

@@ -250,6 +250,72 @@ AUGMENT_KEYWORDS: tuple[AugmentKeyword, ...] = (
         kind="tag",
         phrases=("blur", "soften", "make blurry", "gaussian blur", "soft focus"),
     ),
+    # Neural style transfer (ONNX fast neural style)
+    AugmentKeyword(
+        id="style_candy",
+        kind="tag",
+        phrases=(
+            "cartoon",
+            "cartoon style",
+            "comic style",
+            "comic book look",
+            "illustrated look",
+            "candy style",
+            "make it cartoon",
+            "like a cartoon",
+        ),
+    ),
+    AugmentKeyword(
+        id="style_mosaic",
+        kind="tag",
+        phrases=(
+            "mosaic",
+            "mosaic style",
+            "mosaic painting",
+            "stained glass look",
+            "tiled art style",
+            "mosaic art",
+        ),
+    ),
+    AugmentKeyword(
+        id="style_rain_princess",
+        kind="tag",
+        phrases=(
+            "impressionist",
+            "impressionist style",
+            "impressionist painting",
+            "painterly",
+            "soft painting",
+            "rain princess style",
+            "artistic painting",
+            "painted look",
+        ),
+    ),
+    AugmentKeyword(
+        id="style_udnie",
+        kind="tag",
+        phrases=(
+            "abstract art",
+            "abstract style",
+            "abstract painting",
+            "cubist style",
+            "expressive brushstrokes",
+            "udnie style",
+            "bold artistic style",
+        ),
+    ),
+    AugmentKeyword(
+        id="style_pointilism",
+        kind="tag",
+        phrases=(
+            "pointillism",
+            "pointillist",
+            "pointillist style",
+            "dotted painting",
+            "dot art",
+            "stippled look",
+        ),
+    ),
 )
 
 KEYWORD_BY_ID = {kw.id: kw for kw in AUGMENT_KEYWORDS}

pipeline/style_transfer.py

@@ -0,0 +1,81 @@
+"""Fast neural style transfer via ONNX Model Zoo models (CPU-friendly)."""
+
+import numpy as np
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+STYLE_MODELS: dict[str, dict[str, str]] = {
+    "style_candy": {"repo": "onnxmodelzoo/candy-9", "file": "candy-9.onnx", "label": "candy"},
+    "style_mosaic": {"repo": "onnxmodelzoo/mosaic-9", "file": "mosaic-9.onnx", "label": "mosaic"},
+    "style_rain_princess": {
+        "repo": "onnxmodelzoo/rain-princess-9",
+        "file": "rain-princess-9.onnx",
+        "label": "rain-princess",
+    },
+    "style_udnie": {"repo": "onnxmodelzoo/udnie-9", "file": "udnie-9.onnx", "label": "udnie"},
+    "style_pointilism": {
+        "repo": "onnxmodelzoo/pointilism-9",
+        "file": "pointilism-9.onnx",
+        "label": "pointilism",
+    },
+}
+
+STYLE_TAGS = frozenset(STYLE_MODELS)
+
+MAX_EDGE = 512
+MODEL_SIZE = 224
+
+_sessions: dict[str, ort.InferenceSession] = {}
+
+
+def _get_session(tag: str) -> ort.InferenceSession:
+    if tag not in _sessions:
+        meta = STYLE_MODELS[tag]
+        path = hf_hub_download(repo_id=meta["repo"], filename=meta["file"])
+        _sessions[tag] = ort.InferenceSession(path, providers=["CPUExecutionProvider"])
+    return _sessions[tag]
+
+
+def warmup(tag: str = "style_candy") -> None:
+    """Pre-download and cache one style model (optional; others load on demand)."""
+    if tag in STYLE_MODELS:
+        _get_session(tag)
+
+
+def apply_style(image: Image.Image, tag: str, strength: float = 1.0) -> Image.Image:
+    if tag not in STYLE_MODELS:
+        return image
+
+    session = _get_session(tag)
+    input_name = session.get_inputs()[0].name
+
+    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
+        proc_w, proc_h = width, height
+
+    model_input = working.resize((MODEL_SIZE, 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)
+    styled = Image.fromarray(styled_arr).resize((proc_w, proc_h), Image.Resampling.LANCZOS)
+
+    if (proc_w, proc_h) != (width, height):
+        styled = styled.resize((width, height), Image.Resampling.LANCZOS)
+
+    alpha = min(1.0, max(0.5, strength))
+    if alpha < 1.0:
+        styled = Image.blend(original, styled, alpha)
+
+    return styled

requirements.txt

@@ -8,4 +8,5 @@ numpy>=1.24.0
 rapidocr-onnxruntime>=1.3.0
 opencv-python-headless>=4.8.0
 rembg[cpu]>=2.0.50
+onnxruntime>=1.16.0
 requests>=2.28.0

scripts/validate_planner.py

@@ -18,6 +18,9 @@ planning_cases = [
     ("replace background with beach sunset", ["replace_background"]),
     ("put a table behind it", None),
     ("blur everything softly", ["blur"]),
+    ("cartoon style", ["style_candy"]),
+    ("mosaic painting", ["style_mosaic"]),
+    ("impressionist look", ["style_rain_princess"]),
     ("asdf qwerty nonsense", None),
     ("more contrast", ["contrast_up"]),
     ("transparent cutout", ["add_cutout_transparent"]),