From 4e3b326d400404eaf01fcd64223e11d9a103a150 Mon Sep 17 00:00:00 2001 From: Ben Stull Date: Mon, 8 Jun 2026 06:33:04 -0700 Subject: [PATCH] feat(pipeline): forest Right-variant baker (POC flow restyle, productionized) Right strength curve (level 1-4 -> increasing img2img keyframe strength) is pure and unit-tested; the restyle engine is a verbatim port of the session-0008 POC flow_restyle.py (SD img2img keyframes + Farneback optical-flow tweens = calm, no boiling). torch/diffusers imported lazily so the module + curve test stay light. Co-Authored-By: Claude Opus 4.8 (1M context) --- simulator/bake_right_variants.py | 189 ++++++++++++++++++++++++++++++ tests/test_bake_right_variants.py | 26 ++++ 2 files changed, 215 insertions(+) create mode 100644 simulator/bake_right_variants.py create mode 100644 tests/test_bake_right_variants.py diff --git a/simulator/bake_right_variants.py b/simulator/bake_right_variants.py new file mode 100644 index 0000000..9e36dc0 --- /dev/null +++ b/simulator/bake_right_variants.py @@ -0,0 +1,189 @@ +"""Offline baker: real flow-stabilized painterly Right-axis variants (scales +design §1/§4). + +Productionizes the session-0008 POC (``~/hef-poc/flow_restyle.py``): SD img2img +on KEYFRAMES + Farneback optical-flow warp on in-between frames = a temporally +calm painterly restyle (no per-frame "boiling" — the look the operator approved +in 0008). The one new idea on top of the POC is a **strength curve**: each Right +level 1-4 maps to an increasing keyframe img2img ``strength`` so the four +discrete pre-baked variants form a subtle -> strong dreamlike ramp. Level 0 is +always the raw base (handled by the runtime, not baked here). + +The curve (``right_strength_params``) is pure + unit-tested; the rendered media +is verified by eye (like the other ``simulator/setup_*_media.py`` scripts). torch +/ diffusers are imported lazily inside the bake functions so importing this +module (e.g. for the curve test) is fast and dependency-light. + +Usage: + .venv/bin/python -m simulator.bake_right_variants [--scale forest] + [--fps 12] [--keyint 24] + +Requires: torch + Apple MPS, diffusers, ``stabilityai/sd-turbo`` (cached), +opencv-python, imageio-ffmpeg. ~2.7 min/clip x 4 levels ~= ~11 min local compute. +""" +from __future__ import annotations + +import argparse +import glob +import os +import subprocess +import tempfile +import time +from pathlib import Path + +MEDIA = Path(__file__).parent / "sample_media" + +RIGHT_LEVELS = (1, 2, 3, 4) + +# Keyframe img2img strength per Right level (by-eye ramp). Level 4 ~ the POC's +# proven painterly look (key-strength 0.5-0.58); level 1 is a barely-there dream +# haze. Tune by eye in the sim, then re-bake. +_KEY_STRENGTH = {1: 0.28, 2: 0.38, 3: 0.48, 4: 0.58} +_TWEEN_RATIO = 0.36 # POC ratio (0.18 / 0.5): tweens are a light refine only. + +# The restyle target (matches the POC spike that the operator approved). +PROMPT = ( + "impressionist oil painting, dreamlike, soft painterly brushstrokes, " + "ethereal, luminous, fine art" +) +NEG = "text, watermark, frame, border, ugly, deformed, lowres" +SEED = 42 + + +def right_strength_params(level: int) -> tuple[float, float]: + """(keyframe_strength, tween_strength) for a Right level in RIGHT_LEVELS.""" + if level not in _KEY_STRENGTH: + raise ValueError(f"Right level must be one of {RIGHT_LEVELS}, got {level}") + key = _KEY_STRENGTH[level] + return key, round(key * _TWEEN_RATIO, 4) + + +# --- the restyle engine (verbatim port of ~/hef-poc/flow_restyle.py) --- + + +def _ffmpeg() -> str: + import imageio_ffmpeg + + return imageio_ffmpeg.get_ffmpeg_exe() + + +def _run(cmd: list[str]) -> None: + subprocess.run(cmd, check=True, capture_output=True) + + +def _extract(ff: str, src: str, d: str, fps: float) -> list[str]: + _run([ff, "-hide_banner", "-v", "error", "-i", src, "-vf", f"fps={fps}", + os.path.join(d, "f_%05d.png"), "-y"]) + return sorted(glob.glob(os.path.join(d, "f_*.png"))) + + +def _warp(prev_bgr, src_prev_gray, src_cur_gray): + """Warp the previous stylized frame to align with the current source frame + via Farneback optical flow — this is what kills the per-frame flicker.""" + import cv2 + import numpy as np + + flow = cv2.calcOpticalFlowFarneback( + src_prev_gray, src_cur_gray, None, + pyr_scale=0.5, levels=3, winsize=21, iterations=3, + poly_n=7, poly_sigma=1.5, flags=0) + h, w = src_cur_gray.shape + gx, gy = np.meshgrid(np.arange(w), np.arange(h)) + mx = (gx + flow[..., 0]).astype(np.float32) + my = (gy + flow[..., 1]).astype(np.float32) + return cv2.remap(prev_bgr, mx, my, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT) + + +def _load_pipe(model: str): + """Load the SD img2img pipeline once (reused across all levels).""" + import torch + from diffusers import AutoPipelineForImage2Image + + dev = "mps" if torch.backends.mps.is_available() else "cpu" + pipe = AutoPipelineForImage2Image.from_pretrained( + model, torch_dtype=torch.float16 if dev == "mps" else torch.float32, + safety_checker=None).to(dev) + return pipe, dev + + +def restyle_clip(pipe, dev, src: str, out: str, key_strength: float, + tween_strength: float, *, fps: float = 12.0, keyint: int = 24) -> None: + """Restyle one clip: SD img2img keyframes + flow-warped light-refine tweens.""" + import cv2 + import numpy as np + import torch + from PIL import Image + + ff = _ffmpeg() + gen = torch.Generator(device=dev).manual_seed(SEED) + + def img2img(pil, strength): + steps = max(2, int(np.ceil(1.0 / max(strength, 0.05))) + 1) + return pipe(prompt=PROMPT, negative_prompt=NEG, image=pil, + strength=strength, guidance_scale=0.0, + num_inference_steps=steps, generator=gen).images[0] + + with tempfile.TemporaryDirectory() as d: + frames = _extract(ff, src, d, fps) + print(f" [frames] {len(frames)} key_s={key_strength} tween_s={tween_strength}") + prev_styled = None # BGR uint8 + prev_src_gray = None + t0 = time.time() + for i, fp in enumerate(frames): + src_bgr = cv2.imread(fp) + h0, w0 = src_bgr.shape[:2] + scale = 768 / max(h0, w0) + wh = (int(w0 * scale) // 8 * 8, int(h0 * scale) // 8 * 8) + src_small = cv2.resize(src_bgr, wh) + src_gray = cv2.cvtColor(src_small, cv2.COLOR_BGR2GRAY) + + if i % keyint == 0 or prev_styled is None: + pil = Image.fromarray(cv2.cvtColor(src_small, cv2.COLOR_BGR2RGB)) + out_img = img2img(pil, key_strength) + else: + warped = _warp(prev_styled, prev_src_gray, src_gray) + pil = Image.fromarray(cv2.cvtColor(warped, cv2.COLOR_BGR2RGB)) + out_img = img2img(pil, tween_strength) # light refine only + styled = cv2.cvtColor(np.array(out_img), cv2.COLOR_RGB2BGR) + prev_styled, prev_src_gray = styled, src_gray + cv2.imwrite(os.path.join(d, f"s_{i+1:05d}.png"), + cv2.resize(styled, (w0, h0))) + if i % 12 == 0: + print(f" frame {i+1}/{len(frames)} {time.time()-t0:.1f}s") + + _run([ff, "-hide_banner", "-v", "error", "-framerate", str(fps), + "-i", os.path.join(d, "s_%05d.png"), "-c:v", "libx264", "-crf", "20", + "-preset", "medium", "-pix_fmt", "yuv420p", out, "-y"]) + print(f" [done] {time.time()-t0:.1f}s -> {out}") + + +def bake_scale(scale: str = "forest", *, fps: float = 12.0, keyint: int = 24, + model: str = "stabilityai/sd-turbo") -> None: + """Bake all RIGHT_LEVELS variants for one scale's base clip.""" + base = MEDIA / scale / "base.mp4" + if not base.exists(): + raise FileNotFoundError( + f"no base clip at {base} — run the media-setup scripts first") + print(f"[bake] scale={scale} base={base}") + pipe, dev = _load_pipe(model) + print(f"[device] {dev}") + for level in RIGHT_LEVELS: + key_s, tween_s = right_strength_params(level) + out = MEDIA / scale / f"right{level}.mp4" + print(f" [right {level}] -> {out}") + restyle_clip(pipe, dev, str(base), str(out), key_s, tween_s, + fps=fps, keyint=keyint) + print(f"[bake] {scale} done — {len(RIGHT_LEVELS)} real Right variants") + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--scale", default="forest") + ap.add_argument("--fps", type=float, default=12.0) + ap.add_argument("--keyint", type=int, default=24) + a = ap.parse_args() + bake_scale(a.scale, fps=a.fps, keyint=a.keyint) + + +if __name__ == "__main__": + main() diff --git a/tests/test_bake_right_variants.py b/tests/test_bake_right_variants.py new file mode 100644 index 0000000..582d40c --- /dev/null +++ b/tests/test_bake_right_variants.py @@ -0,0 +1,26 @@ +import pytest + +from simulator.bake_right_variants import RIGHT_LEVELS, right_strength_params + + +def test_levels_are_one_through_four(): + assert RIGHT_LEVELS == (1, 2, 3, 4) + + +def test_strength_increases_with_level(): + keys = [right_strength_params(l)[0] for l in RIGHT_LEVELS] + assert keys == sorted(keys) # monotonic ramp + assert keys[0] < keys[-1] # subtle -> strong + + +def test_tween_strength_is_a_fraction_of_key_strength(): + for l in RIGHT_LEVELS: + key, tween = right_strength_params(l) + assert 0.0 < tween < key # tween is a light refine only + + +def test_rejects_out_of_range_level(): + with pytest.raises(ValueError): + right_strength_params(0) + with pytest.raises(ValueError): + right_strength_params(5)