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) <noreply@anthropic.com>
This commit is contained in:
Ben Stull
@@ -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()
|
||||
Reference in New Issue
Block a user