"""Assert-based self-check for Gemini frame JSON parsing + timestamp list.""" from __future__ import annotations from src.extract_frames import clamp_timestamp, extract_frames # noqa: F401 from src.gemini_analyze import parse_scene_analysis, timestamps_from_analysis def main() -> None: raw = """ { "overall_quality": "360p phone UGC, soft and compressed", "frames": [ {"timestamp_sec": 12.4, "label": "hook", "description": "person talking", "quality_cues": "soft"}, {"timestamp_sec": 12.6, "label": "dup", "description": "almost same", "quality_cues": "soft"}, {"timestamp_sec": 90, "label": "demo", "description": "product closeup", "quality_cues": "blocky"}, {"timestamp_sec": 200.5, "label": "cta", "description": "end screen", "quality_cues": "noisy"} ] } """ analysis = parse_scene_analysis(raw) # auto: no user max — keep all after dedupe assert analysis.overall_quality.startswith("360p"), analysis.overall_quality # 12.4 and 12.6 are < 0.5s apart → dedupe keeps one assert len(analysis.frames) == 3, [f.timestamp_sec for f in analysis.frames] ts = timestamps_from_analysis(analysis) assert ts == [12.4, 90.0, 200.5], ts capped = parse_scene_analysis(raw, max_frames=2) assert len(capped.frames) == 2 assert timestamps_from_analysis(capped) == [12.4, 90.0] from src.gemini_analyze import ( SCRIPT_MIN_FRAMES, _build_prompt, _build_script_only_prompt, _frames_cap, _script_frames_cap, ) prompt = _build_prompt(duration_sec=30.0, script=None) assert "DISTINCT" in prompt assert "fixed count" in prompt.lower() or "YOURSELF" in prompt assert _frames_cap(None) == 32 assert _frames_cap(5) == 5 assert SCRIPT_MIN_FRAMES == 30 assert _script_frames_cap(None) >= SCRIPT_MIN_FRAMES assert _script_frames_cap(10) >= SCRIPT_MIN_FRAMES script_prompt = _build_prompt(duration_sec=30.0, script="Hook.\nBody.\nCTA.") assert "SCRIPT decides" in script_prompt or "story beats" in script_prompt assert "DISTINCT shots / images" not in script_prompt assert "Hook." in script_prompt assert str(SCRIPT_MIN_FRAMES) in script_prompt assert "AT LEAST" in script_prompt only = _build_script_only_prompt("Hook.\nBody.\nCTA.") assert "No video" in only or "script-only" in only.lower() assert "Hook." in only and "CTA." in only assert "Files API" not in only assert "DISTINCT shots / images" not in only assert "timestamp_sec" in only assert str(SCRIPT_MIN_FRAMES) in only assert "AT LEAST" in only # Script mode may keep close timestamps (different beats, same source moment) close_beats = parse_scene_analysis( { "overall_quality": "x", "frames": [ { "timestamp_sec": 5.0, "label": "a", "description": "d", "quality_cues": "q", "script_line": "beat one", "regen_brief": "r1", }, { "timestamp_sec": 5.2, "label": "b", "description": "d", "quality_cues": "q", "script_line": "beat two", "regen_brief": "r2", }, ], }, dedupe_close=False, ) assert len(close_beats.frames) == 2 from src.regenerate import ( ASPECT_RATIOS, DEFAULT_ASPECT_RATIO, GROK_MODEL, _caption_remove_prompt, _near_copy_prompt, _script_t2i_prompt, _normalize_aspect_ratio, _scrub, ) assert DEFAULT_ASPECT_RATIO == "auto" assert "9:16" in ASPECT_RATIOS and "16:9" in ASPECT_RATIOS assert _normalize_aspect_ratio("9:16") == "9:16" assert _normalize_aspect_ratio("nope") == "auto" _check_story_order() _check_cast_routing() _check_text_cleanup() _check_retry_and_rejection() _check_references() _check_crop_to_ratio() _check_fal_payload() _check_animate_payload() _check_clip_join() fenced = parse_scene_analysis( '```json\n{"overall_quality":"x","frames":[{"timestamp_sec":1,"label":"a","description":"d","quality_cues":"q"}]}\n```', max_frames=5, ) assert timestamps_from_analysis(fenced) == [1.0] # Gemini sometimes returns times past EOF (e.g. 856s on a ~746s video) past_eof = parse_scene_analysis( { "overall_quality": "x", "frames": [ {"timestamp_sec": 10, "label": "a", "description": "d", "quality_cues": "q"}, {"timestamp_sec": 856, "label": "b", "description": "d", "quality_cues": "q"}, ], }, max_frames=8, duration_sec=746.173, ) assert timestamps_from_analysis(past_eof) == [10.0, clamp_timestamp(856, 746.173)] assert past_eof.frames[-1].timestamp_sec < 746.173 scripted = parse_scene_analysis( { "overall_quality": "soft phone", "frames": [ { "timestamp_sec": 5, "label": "hook", "description": "person talking to camera", "quality_cues": "soft", "script_line": "Stop wasting money on bad ads.", "regen_brief": "Creator pointing at camera, serious expression", } ], }, max_frames=4, duration_sec=100, ) assert scripted.frames[0].script_line.startswith("Stop wasting") assert "pointing" in scripted.frames[0].regen_brief from src.gemini_analyze import FrameSpec plain = FrameSpec(timestamp_sec=0, description="house with trucks", quality_cues="soft") caption_prompt = _near_copy_prompt(plain) assert "caption" in caption_prompt.lower() or "timestamp" in caption_prompt.lower() assert "police" not in caption_prompt.lower() assert "dead" not in caption_prompt.lower() assert GROK_MODEL == "xai/grok-imagine-image" assert "Erase every letter" in _caption_remove_prompt(0) assert "Aggressively wipe" in _caption_remove_prompt(1) near = _near_copy_prompt(scripted.frames[0], "soft phone") assert "Stop wasting money" in near assert "police" not in _scrub("Two police officers stand on the porch. Soft light.") assert "Soft light" in _scrub("Two police officers stand on the porch. Soft light.") t2i = _script_t2i_prompt(scripted.frames[0], "soft phone") assert "Stop wasting money" in t2i assert "pointing" in t2i.lower() or "Creator" in t2i assert "Edit this phone-video" not in t2i assert "No captions" in t2i or "no captions" in t2i.lower() # pipeline accepts script-only (no video) at the API boundary import inspect from src.pipeline import run_pipeline sig = inspect.signature(run_pipeline) assert "video_path" in sig.parameters assert sig.parameters["video_path"].default is None _check_parallel_regen() _check_edit_approve() _check_reel() _check_script_leads_prompt() _check_short_script_still_runs() _check_voiceover() print("selfcheck ok") def _check_script_leads_prompt() -> None: """The beat outranks the source frame, or the edit endpoint returns a near-copy.""" from src.gemini_analyze import FrameSpec from src.regenerate import _script_match_prompt frame = FrameSpec( timestamp_sec=3, description="woman on a sofa", script_line="I saved £240 in ten minutes.", regen_brief="grinning at her phone in a supermarket car park", ) prompt = _script_match_prompt(frame) assert "car park" in prompt and "£240" in prompt # what the old wording promised the model, and why every frame came back as the source assert "Adjust expression/pose only as needed" not in prompt, prompt assert "Keep the same people, framing" not in prompt, prompt assert "do not return a copy of the input" in prompt, prompt # the source frame is explicitly demoted to people + look, and the beat decides the rest assert "The beat decides what is in the picture" in prompt, prompt assert "the still only supplies the people" in prompt, prompt def _check_short_script_still_runs() -> None: """A script with fewer beats than the floor generates them, it does not abort the run.""" import json from src import gemini_analyze as ga def frames(n: int) -> str: return json.dumps( { "overall_quality": "soft phone", "cast": [{"id": "c1", "name": "mum", "description": "30s, brown hair"}], "frames": [ { "timestamp_sec": float(i), "label": f"beat {i}", "description": "d", "quality_cues": "q", "script_line": f"line {i}", "regen_brief": f"brief {i}", "character_ids": ["c1"], } for i in range(n) ], } ) calls: list[int] = [] real_client = ga._gemini_client real_json = ga._generate_json ga._gemini_client = lambda **kw: (object(), "fake-model") ga._generate_json = lambda client, model_id, contents: ( calls.append(len(contents)) or frames(6) ) try: analysis = ga.analyze_script("Hook. Body. CTA.") finally: ga._gemini_client = real_client ga._generate_json = real_json # one corrective retry was still attempted, then the short answer was accepted assert len(calls) == 2, calls assert len(analysis.frames) == 6, len(analysis.frames) assert [f.script_line for f in analysis.frames] == [f"line {i}" for i in range(6)] def _check_voiceover() -> None: """TTS payload, and that the mux keeps the video's length and drops its own audio.""" import json import subprocess import tempfile from pathlib import Path from src import voice from src.voice import ( MAX_CHARS, TTS_MODEL, clamp_speed, estimate_cost, normalize_voice, script_text_for, ) assert TTS_MODEL == "fal-ai/elevenlabs/tts/multilingual-v2" assert clamp_speed(3.0) == 1.2 and clamp_speed(0.1) == 0.7 # endpoint rejects the rest assert normalize_voice("") == "Rachel" assert round(estimate_cost("x" * 1000), 4) == 0.10 class _Spec: def __init__(self, line: str) -> None: self.script_line = line class _Res: def __init__(self, line: str) -> None: self.spec = _Spec(line) # one script line spread over several beats must be read once, not three times assert script_text_for([_Res("a"), _Res("a"), _Res("b")]) == "a\nb" assert script_text_for(" pasted script ") == "pasted script" calls: list[tuple[str, dict]] = [] real_sub = voice.fal_client.subscribe real_get = voice.requests.get voice.fal_client.subscribe = lambda m, arguments=None, **kw: ( calls.append((m, arguments)) or {"audio": {"url": "https://f/vo.mp3"}} ) voice.requests.get = lambda url, **kw: type( "R", (), {"content": b"mp3", "raise_for_status": lambda self: None} )() import os real_key = os.environ.get("FAL_KEY") os.environ["FAL_KEY"] = "test" try: with tempfile.TemporaryDirectory() as tmp: out = voice.generate_voiceover( " Save money today. " + "x" * MAX_CHARS, Path(tmp) / "voiceover.mp3", voice="George", speed=2.0, ) assert out.read_bytes() == b"mp3" model, args = calls[-1] assert model == TTS_MODEL, model assert args["voice"] == "George" and args["speed"] == 1.2, args assert len(args["text"]) == MAX_CHARS and args["text"].startswith("Save money"), len( args["text"] ) assert sorted(args) == ["similarity_boost", "speed", "stability", "text", "voice"], args finally: voice.fal_client.subscribe = real_sub voice.requests.get = real_get if real_key is None: os.environ.pop("FAL_KEY", None) else: os.environ["FAL_KEY"] = real_key import shutil if not (shutil.which("ffmpeg") and shutil.which("ffprobe")): return from src.assemble_video import attach_audio, has_audio with tempfile.TemporaryDirectory() as tmp: tmpdir = Path(tmp) silent = tmpdir / "silent.mp4" # 3s video, no audio subprocess.run( ["ffmpeg", "-nostdin", "-y", "-loglevel", "error", "-f", "lavfi", "-i", "color=c=black:s=64x64:d=3", "-pix_fmt", "yuv420p", str(silent)], check=True, ) long_audio = tmpdir / "vo.mp3" # 10s of tone — longer than the video on purpose subprocess.run( ["ffmpeg", "-nostdin", "-y", "-loglevel", "error", "-f", "lavfi", "-i", "sine=frequency=440:duration=10", str(long_audio)], check=True, ) voiced = attach_audio(silent, long_audio, tmpdir / "voiced.mp4") assert has_audio(voiced) probe = subprocess.run( ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "json", str(voiced)], capture_output=True, text=True, ) seconds = float(json.loads(probe.stdout)["format"]["duration"]) # the video's length wins: a 10s voiceover on a 3s video must not stretch the video assert 2.8 < seconds < 3.6, seconds def _check_reel() -> None: """Concat list is exact; the encode itself runs only when ffmpeg is installed.""" import json import shutil import subprocess import tempfile from pathlib import Path from src.assemble_video import DEFAULT_SECONDS_EACH, _filter_complex, build_reel assert DEFAULT_SECONDS_EACH == 4.0 fc = _filter_complex(3, 240, 426) assert fc.count("scale=240:426") == 3 # one chain per image, no more, no less assert fc.endswith("[v0][v1][v2]concat=n=3:v=1:a=0[out]"), fc assert "force_original_aspect_ratio=decrease" in fc and "setsar=1" in fc if not (shutil.which("ffmpeg") and shutil.which("ffprobe")): print("(skipped reel encode — ffmpeg not on PATH)") return with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) imgs = [] # deliberately mismatched sizes — the reel must letterbox, not fail for i, size in enumerate(("240x426", "200x400", "240x426")): p = root / f"img_{i}.jpg" subprocess.run( ["ffmpeg", "-nostdin", "-y", "-loglevel", "error", "-f", "lavfi", "-i", f"color=c=blue:s={size}", "-frames:v", "1", str(p)], capture_output=True, ) assert p.is_file(), p imgs.append(p) out = build_reel(imgs, root / "reel.mp4", seconds_each=2.0) assert out.is_file() and out.stat().st_size > 0 probe = subprocess.run( ["ffprobe", "-v", "error", "-show_entries", "format=duration:stream=width,height", "-of", "json", str(out)], capture_output=True, text=True, ) data = json.loads(probe.stdout) dur = float(data["format"]["duration"]) # exact, not approximate: every image must get its full slice, including the last assert abs(dur - 6.0) < 0.1, f"3 images x 2s should be 6s, got {dur}" stream = data["streams"][0] assert (stream["width"], stream["height"]) == (240, 426), stream assert stream["width"] % 2 == 0 and stream["height"] % 2 == 0 # yuv420p needs even # None / missing entries are skipped rather than crashing the mux out2 = build_reel([None, imgs[0], root / "gone.jpg"], root / "r2.mp4", seconds_each=1.0) assert out2.is_file() # Mixed orientations must land on a real image's shape, never a square mash-up of # max(width) x max(height). mixed = [] for name, size in (("wide.jpg", "640x360"), ("tall.jpg", "270x480")): p = root / name subprocess.run( ["ffmpeg", "-nostdin", "-y", "-loglevel", "error", "-f", "lavfi", "-i", f"color=c=green:s={size}", "-frames:v", "1", str(p)], capture_output=True, ) mixed.append(p) out3 = build_reel(mixed, root / "r3.mp4", seconds_each=1.0) probe3 = subprocess.run( ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=width,height", "-of", "json", str(out3)], capture_output=True, text=True, ) s3 = json.loads(probe3.stdout)["streams"][0] assert (s3["width"], s3["height"]) == (640, 360), s3 # the larger-area image wins try: build_reel([], root / "empty.mp4") except ValueError: pass else: raise AssertionError("build_reel must reject an empty image list") # build_run_reel must never raise — a failed mux cannot sink a paid-for run from src.gemini_analyze import FrameSpec, SceneAnalysis from src.pipeline import FrameResult, RunResult, build_run_reel def result_with(paths): return RunResult( run_id="r", run_dir=root, analysis=SceneAnalysis(), frames=[FrameResult(spec=FrameSpec(timestamp_sec=0), regenerated=p) for p in paths], seconds_each=1.0, ) assert build_run_reel(result_with([])) is None # no frames at all assert build_run_reel(result_with([root / "nope.jpg"])) is None # frames all missing logs: list[str] = [] ok = result_with(imgs) reel = build_run_reel(ok, seconds_each=1.0, log=logs.append) assert reel is not None and reel.name == "reel.mp4" assert ok.reel == reel and ok.seconds_each == 1.0 assert any("3 × 1s" in m for m in logs), logs def _check_cast_routing() -> None: """Each shot gets only its own characters' references — the core of the fix.""" import tempfile from pathlib import Path from src import pipeline from src.gemini_analyze import CharacterSpec, FrameSpec, parse_scene_analysis from src.regenerate import _cast_clause, _identity_lock # --- Gemini side: cast survives parsing, unknown ids are dropped --- parsed = parse_scene_analysis( { "overall_quality": "soft", "cast": [ {"id": "c1", "name": "young mum", "description": "30s, brown bob, red coat"}, {"id": "c2", "name": "driver", "description": "50s man, grey beard, blue polo"}, {"id": "c1", "name": "dupe", "description": "ignored duplicate"}, ], "frames": [ {"timestamp_sec": 0, "label": "a", "description": "d", "quality_cues": "q", "script_line": "one", "regen_brief": "b", "character_ids": ["c1"]}, {"timestamp_sec": 1, "label": "b", "description": "d", "quality_cues": "q", "script_line": "two", "regen_brief": "b", "character_ids": ["c2"]}, {"timestamp_sec": 2, "label": "c", "description": "d", "quality_cues": "q", "script_line": "three", "regen_brief": "b", "character_ids": ["c1", "c2"]}, {"timestamp_sec": 3, "label": "pack", "description": "d", "quality_cues": "q", "script_line": "four", "regen_brief": "b", "character_ids": []}, {"timestamp_sec": 4, "label": "ghost", "description": "d", "quality_cues": "q", "script_line": "five", "regen_brief": "b", "character_ids": ["c9"]}, ], }, sort_by_time=False, dedupe_close=False, ) assert [c.id for c in parsed.cast] == ["c1", "c2"], "duplicate ids must collapse" assert parsed.frames[3].character_ids == [], "a product shot keeps nobody" assert parsed.frames[4].character_ids == [], "an undefined id must be dropped" assert [c.label for c in parsed.characters_in(parsed.frames[2])] == ["young mum", "driver"] # --- prompt side: who is in the shot, and which reference is whom --- mum, driver = parsed.cast clause = _cast_clause([mum]) assert "young mum" in clause and "red coat" in clause assert "driver" not in clause, "a solo shot must not describe the other character" lock = _identity_lock([mum, driver], 2) assert "reference image 1 is young mum" in lock assert "reference image 2 is driver" in lock assert "not swap" in lock or "Do not swap" in lock assert _identity_lock([], 0) == "", "no cast and no refs means no identity clause" # --- routing side: the reference actually handed to each frame --- with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) sheets = {} for cid in ("c1", "c2"): p = root / "cast" / f"{cid}.jpg" p.parent.mkdir(parents=True, exist_ok=True) p.write_bytes(b"x") sheets[cid] = p uploaded = root / "uploaded.jpg" uploaded.write_bytes(b"u") seen: dict[int, tuple[tuple[str, ...], tuple[str, ...]]] = {} def fake_regen(original, frame, *, out_path, references=None, characters=None, **kw): idx = int(Path(out_path).stem.split("_")[1]) seen[idx] = ( tuple(Path(r).name for r in (references or [])), tuple(c.id for c in (characters or [])), ) Path(out_path).parent.mkdir(parents=True, exist_ok=True) Path(out_path).write_bytes(b"i") return Path(out_path) real, pipeline.regenerate_frame = pipeline.regenerate_frame, fake_regen try: pipeline._regen_all( parsed.frames, [None] * len(parsed.frames), overall_quality="", dest=root, aspect_ratio="auto", cleanup_pass=False, workers=4, log=lambda m: None, references=[uploaded], anchor_first=True, cast=parsed.by_id(), sheets=sheets, ) assert seen[0] == (("c1.jpg",), ("c1",)), seen[0] # the whole bug: c1's face must never reach c2's shot assert seen[1] == (("c2.jpg",), ("c2",)), seen[1] assert seen[2] == (("c1.jpg", "c2.jpg"), ("c1", "c2")), seen[2] assert seen[3] == ((), ()), "product shot must get no face at all" assert seen[4] == ((), ()), "dropped id must not fall back to someone else" assert all("uploaded.jpg" not in refs for refs, _ in seen.values()) # with no cast at all, the uploaded reference still applies everywhere seen.clear() plain = [FrameSpec(timestamp_sec=0, script_line="x") for _ in range(3)] pipeline._regen_all( plain, [None] * 3, overall_quality="", dest=root, aspect_ratio="auto", cleanup_pass=False, workers=3, log=lambda m: None, references=[uploaded], ) assert all(refs == ("uploaded.jpg",) for refs, _ in seen.values()), seen finally: pipeline.regenerate_frame = real def _check_text_cleanup() -> None: """Cleanup passes are driven by the text check, not run blindly.""" import os import tempfile from pathlib import Path from src import regenerate from src.gemini_analyze import FrameSpec from src.regenerate import _text_check_enabled, _text_fix_attempts, regenerate_frame assert _text_fix_attempts() == 1 # one extra pass by default, as before os.environ["MAX_TEXT_FIX_PASSES"] = "2" assert _text_fix_attempts() == 2 os.environ["MAX_TEXT_FIX_PASSES"] = "99" assert _text_fix_attempts() == 3, "must stay bounded — each pass costs money" del os.environ["MAX_TEXT_FIX_PASSES"] assert _text_check_enabled() os.environ["TEXT_CHECK"] = "0" assert not _text_check_enabled() del os.environ["TEXT_CHECK"] verdicts: list[bool | None] = [] calls: list[str] = [] real_run = regenerate._run_fal_resilient real_detect = regenerate.detect_readable_text real_write = regenerate._write_image def fake_run(src, prompt, *, aspect_ratio=None, references=None): # the video-only path's FIRST prompt is itself a caption remover, so count calls # rather than classifying them calls.append("cleanup" if "Aggressively wipe" in prompt else "generate") return b"img" regenerate._run_fal_resilient = fake_run regenerate.detect_readable_text = lambda p, **kw: verdicts.pop(0) if verdicts else False regenerate._write_image = lambda data, dest, ar: (dest.write_bytes(data), dest)[1] os.environ["FAL_KEY"] = "test" try: with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) src = root / "src.jpg" src.write_bytes(b"seed") plain = FrameSpec(timestamp_sec=0, description="a house") # clean first result → no second call. This is the saving versus the old code, # which always paid for two passes. calls.clear(); verdicts[:] = [False] regenerate_frame(src, plain, out_path=root / "a.jpg") assert len(calls) == 1, calls # text seen → one cleanup pass, then verified clean calls.clear(); verdicts[:] = [True, False] notes: list[str] = [] regenerate_frame(src, plain, out_path=root / "b.jpg", on_note=notes.append) assert len(calls) == 2 and calls[1] == "cleanup", calls assert notes == [], notes # still dirty after the allowed passes → flagged for the user, not retried forever calls.clear(); verdicts[:] = [True, True] notes = [] regenerate_frame(src, plain, out_path=root / "c.jpg", on_note=notes.append) assert len(calls) == 2, calls assert notes and "text" in notes[0], notes # checker unavailable (None) → behave like the old always-clean-up default calls.clear(); verdicts[:] = [None, None] regenerate_frame(src, plain, out_path=root / "d.jpg") assert len(calls) == 2, calls # fast mode → single call, no check at all calls.clear(); verdicts[:] = [True, True] regenerate_frame(src, plain, out_path=root / "e.jpg", cleanup_pass=False) assert len(calls) == 1, calls # script-only frames are now checked too — they never were before calls.clear(); verdicts[:] = [True, False] beat = FrameSpec(timestamp_sec=0, script_line="Buy now", regen_brief="a kitchen") regenerate_frame(None, beat, out_path=root / "f.jpg") assert len(calls) == 2 and calls[0] == "generate", calls # no leftover temp files from the cleanup loop assert not list(root.glob("*_txt*.jpg")), list(root.glob("*_txt*.jpg")) finally: regenerate._run_fal_resilient = real_run regenerate.detect_readable_text = real_detect regenerate._write_image = real_write def _check_story_order() -> None: """Script mode must keep Gemini's story order, not re-sort beats by timestamp.""" beats = [ {"timestamp_sec": t, "label": f"b{n}", "description": "d", "quality_cues": "q", "script_line": f"{n}. line", "regen_brief": f"r{n}"} for n, t in enumerate([2.0, 30.0, 8.0, 45.0, 5.0], start=1) ] data = {"overall_quality": "x", "frames": beats} story = [f"{n}. line" for n in range(1, 6)] kept = parse_scene_analysis(data, dedupe_close=False, sort_by_time=False, duration_sec=60) assert [f.script_line for f in kept.frames] == story, [f.script_line for f in kept.frames] # timestamps still travel with their beat, they are just no longer the sort key assert [f.timestamp_sec for f in kept.frames] == [2.0, 30.0, 8.0, 45.0, 5.0] # no-script mode still orders by time, where the timeline *is* the running order timed = parse_scene_analysis(data, duration_sec=60) assert [f.timestamp_sec for f in timed.frames] == [2.0, 5.0, 8.0, 30.0, 45.0] import inspect from src.gemini_analyze import analyze_script, analyze_video # both script paths must opt out of sorting assert "sort_by_time=False" in inspect.getsource(analyze_script) assert "sort_by_time=not has_script" in inspect.getsource(analyze_video) def _check_retry_and_rejection() -> None: """Transient errors retry; refused prompts fall back to a reworded one.""" from src import regenerate from src.regenerate import _is_transient, is_content_rejection, with_retries class HTTPish(Exception): def __init__(self, status): super().__init__(f"http {status}") self.status_code = status class Timeouty(Exception): pass Timeouty.__name__ = "FalClientTimeoutError" assert _is_transient(HTTPish(429)) and _is_transient(HTTPish(503)) assert not _is_transient(HTTPish(400)) and not _is_transient(HTTPish(422)) assert _is_transient(Timeouty()) # a 429 is a rate limit, never a moderation refusal assert not is_content_rejection(HTTPish(429)) assert is_content_rejection(HTTPish(422)) assert is_content_rejection(Exception("Request blocked by content moderation")) assert not is_content_rejection(Exception("some unrelated failure")) real_sleep = regenerate.time.sleep regenerate.time.sleep = lambda s: None try: tries = {"n": 0} def flaky(): tries["n"] += 1 if tries["n"] < 3: raise HTTPish(429) return "ok" assert with_retries(flaky) == "ok" and tries["n"] == 3 # permanent failures must not be retried — that would just burn time hard = {"n": 0} def refused(): hard["n"] += 1 raise HTTPish(422) try: with_retries(refused) except Exception: pass assert hard["n"] == 1, hard # and retries are bounded always = {"n": 0} def down(): always["n"] += 1 raise HTTPish(503) try: with_retries(down, attempts=3) except Exception: pass assert always["n"] == 3, always finally: regenerate.time.sleep = real_sleep # the reworded fallback strips the sensitive sentence that triggered the refusal safe = regenerate._safe_variant("A cop holds a gun. A woman smiles at the camera.") assert "gun" not in safe and "smiles" in safe assert "brand-safe" in safe def _check_references() -> None: """Identity references ride along in image_urls and change the prompt.""" from pathlib import Path from src import regenerate from src.gemini_analyze import FrameSpec from src.regenerate import ( MAX_IMAGE_INPUTS, MAX_REFERENCES, GROK_MODEL, GROK_MODEL_EDIT, _run_fal, _scene_from_reference_prompt, ) assert MAX_IMAGE_INPUTS == 3 and MAX_REFERENCES == 2 frame = FrameSpec(timestamp_sec=0, script_line="Buy now", regen_brief="woman in a kitchen") prompt = _scene_from_reference_prompt(frame, "", [], 1) assert "woman in a kitchen" in prompt and "Buy now" in prompt assert "same face" in prompt # the identity lock # The scene must come before the reference is mentioned, and the reference must be # identity-only: leading with "keep the same person" returned the portrait with a new # background and lost the beat. assert prompt.index("woman in a kitchen") < prompt.index("reference image"), prompt assert "ONLY the person's identity" in prompt, prompt assert "Ignore its background, framing, pose and expression" in prompt, prompt here = Path(__file__) calls: list[tuple[str, dict]] = [] real_sub = regenerate.fal_client.subscribe real_up = regenerate.upload_image real_dl = regenerate._download_output regenerate.fal_client.subscribe = lambda m, arguments=None, **kw: ( calls.append((m, arguments)) or {"images": [{"url": "https://f/o.jpg"}]} ) regenerate.upload_image = lambda p, **kw: f"https://f/{Path(p).name}" regenerate._download_output = lambda out: b"img" try: # no source frame + a reference → edit that reference instead of text-to-image _run_fal(None, "p", aspect_ratio="9:16", references=[here]) model, args = calls[-1] assert model == GROK_MODEL_EDIT, model assert args["image_urls"] == [f"https://f/{here.name}"], args # no source, no reference → still plain text-to-image _run_fal(None, "p", aspect_ratio="9:16") assert calls[-1][0] == GROK_MODEL, calls[-1][0] assert "image_urls" not in calls[-1][1] # source frame + references → frame first, references after, capped at 3 total _run_fal(here, "p", aspect_ratio="auto", references=[here, here, here, here]) _, args = calls[-1] assert len(args["image_urls"]) == MAX_IMAGE_INPUTS, args["image_urls"] # a reference path that does not exist is ignored, not sent as a broken URL _run_fal(here, "p", aspect_ratio="auto", references=[Path("/nope/x.jpg")]) assert len(calls[-1][1]["image_urls"]) == 1 finally: regenerate.fal_client.subscribe = real_sub regenerate.upload_image = real_up regenerate._download_output = real_dl def _check_animate_payload() -> None: """Arguments sent to the image-to-video endpoint, plus the cost estimate.""" import os import tempfile from pathlib import Path from src import animate from src.animate import ( DEFAULT_CLIP_SECONDS, GROK_VIDEO_MODEL, _extract_video_url, animate_frame, clamp_seconds, estimate_cost, motion_prompt, normalize_video_aspect_ratio, ) from src.gemini_analyze import FrameSpec assert DEFAULT_CLIP_SECONDS == 4 # fal takes an integer 1..15 assert clamp_seconds(4) == 4 and clamp_seconds(4.4) == 4 assert clamp_seconds(0) == 1 and clamp_seconds(99) == 15 assert clamp_seconds(None) == 4 and clamp_seconds("x") == 4 # the video enum is narrower than the image one — 2:1 is valid for stills, not clips assert normalize_video_aspect_ratio("9:16") == "9:16" assert normalize_video_aspect_ratio("2:1") == "auto" assert normalize_video_aspect_ratio(None) == "auto" # 720p is ~1.75x 480p; a 4s 720p clip is ~$0.56 assert abs(estimate_cost(1, 4, "720p") - 0.56) < 0.01 assert abs(estimate_cost(10, 4, "480p") - 3.20) < 0.01 assert estimate_cost(0, 4, "720p") == 0 prompt = motion_prompt(FrameSpec(timestamp_sec=0, regen_brief="woman pointing at camera")) assert "woman pointing" in prompt assert "No new text" in prompt and "No scene cuts" in prompt # sensitive wording is scrubbed here too, same as for stills assert "police" not in motion_prompt(FrameSpec(timestamp_sec=0, regen_brief="a police car.")) assert "zoom slowly" in motion_prompt(None, "zoom slowly") assert _extract_video_url({"video": {"url": "https://x/v.mp4"}}) == "https://x/v.mp4" assert _extract_video_url({"video": None}) is None assert _extract_video_url({}) is None assert _extract_video_url(None) is None calls: list[tuple[str, dict]] = [] real_sub = animate.fal_client.subscribe real_up = animate.upload_image # animate imports it by name, so patch it there real_get = animate.requests.get class FakeResp: content = b"fakemp4bytes" def raise_for_status(self): return None with tempfile.TemporaryDirectory() as tmp: img = Path(tmp) / "f.jpg" img.write_bytes(b"jpeg") animate.fal_client.subscribe = lambda m, arguments=None, **kw: ( calls.append((m, arguments)) or {"video": {"url": "https://fal.media/v.mp4"}} ) animate.upload_image = lambda p, **kw: "https://fal.media/in.jpg" animate.requests.get = lambda url, **kw: FakeResp() os.environ["FAL_KEY"] = "test-key" try: out = animate_frame(img, Path(tmp) / "c.mp4", seconds=4, resolution="480p", aspect_ratio="9:16") assert out.is_file() and out.read_bytes() == b"fakemp4bytes" model, args = calls[-1] assert model == GROK_VIDEO_MODEL, model assert args["duration"] == 4 and isinstance(args["duration"], int) assert args["resolution"] == "480p" assert args["aspect_ratio"] == "9:16" assert args["image_url"] == "https://fal.media/in.jpg" assert isinstance(args["prompt"], str) and args["prompt"] # unsupported values must be corrected before they reach fal animate_frame(img, Path(tmp) / "c2.mp4", seconds=99, resolution="4k", aspect_ratio="2:1") _, args = calls[-1] assert args["duration"] == 15 assert args["resolution"] == "720p" assert args["aspect_ratio"] == "auto" finally: animate.fal_client.subscribe = real_sub animate.upload_image = real_up animate.requests.get = real_get def _check_clip_join() -> None: """Joining real clips: order, total duration, and audio handling.""" import json import shutil import subprocess import tempfile from pathlib import Path from src.assemble_video import build_clip_reel, has_audio if not (shutil.which("ffmpeg") and shutil.which("ffprobe")): print("(skipped clip join — ffmpeg not on PATH)") return def make_clip(path: Path, seconds: int, size: str, audio: bool) -> Path: cmd = ["ffmpeg", "-nostdin", "-y", "-loglevel", "error", "-f", "lavfi", "-i", f"testsrc=s={size}:rate=30:d={seconds}"] if audio: cmd += ["-f", "lavfi", "-i", f"sine=frequency=440:duration={seconds}", "-c:a", "aac", "-shortest"] cmd += ["-c:v", "libx264", "-pix_fmt", "yuv420p", "-t", str(seconds), str(path)] subprocess.run(cmd, capture_output=True) assert path.is_file() and path.stat().st_size > 0, path return path def probe(path: Path) -> dict: out = subprocess.run( ["ffprobe", "-v", "error", "-show_entries", "format=duration:stream=width,height,codec_type", "-of", "json", str(path)], capture_output=True, text=True, ).stdout return json.loads(out) with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) # three 4s clips, all with audio, one at an odd size clips = [ make_clip(root / "a.mp4", 4, "480x854", True), make_clip(root / "b.mp4", 4, "360x640", True), make_clip(root / "c.mp4", 4, "480x854", True), ] assert all(has_audio(c) for c in clips) out = build_clip_reel(clips, root / "final.mp4") info = probe(out) dur = float(info["format"]["duration"]) assert abs(dur - 12.0) < 0.3, f"3 x 4s should be ~12s, got {dur}" kinds = {s["codec_type"] for s in info["streams"]} assert "audio" in kinds, "audio must survive when every clip has it" vid = next(s for s in info["streams"] if s["codec_type"] == "video") assert (vid["width"], vid["height"]) == (480, 854), vid # largest clip's shape # one silent clip in the mix → video-only output rather than desynced audio mixed = [clips[0], make_clip(root / "d.mp4", 4, "480x854", False)] out2 = build_clip_reel(mixed, root / "mixed.mp4") info2 = probe(out2) assert {s["codec_type"] for s in info2["streams"]} == {"video"} assert abs(float(info2["format"]["duration"]) - 8.0) < 0.3 try: build_clip_reel([], root / "none.mp4") except ValueError: pass else: raise AssertionError("build_clip_reel must reject an empty list") def _check_fal_payload() -> None: """The exact arguments sent to fal — the API itself is never called in tests. Guards the two things the OpenAPI schema pins down: the text-to-image endpoint has no "auto" in its aspect_ratio enum (so the key must be omitted), and the edit endpoint takes `image_urls` as a list, not a single `image`. """ from pathlib import Path from src import regenerate from src.gemini_analyze import FrameSpec from src.regenerate import GROK_MODEL, GROK_MODEL_EDIT, _run_fal, _extract_url calls: list[tuple[str, dict]] = [] def fake_subscribe(model, arguments=None, **kw): calls.append((model, arguments)) return {"images": [{"url": "https://fal.media/out.jpg"}], "revised_prompt": "x"} real_sub = regenerate.fal_client.subscribe real_up = regenerate.fal_client.upload_file real_dl = regenerate._download_output regenerate.fal_client.subscribe = fake_subscribe regenerate.fal_client.upload_file = lambda p, **kw: f"https://fal.media/in/{Path(p).name}" regenerate._download_output = lambda out: b"\xff\xd8jpegbytes" try: # text-to-image: no image_urls, and "auto" is dropped (not a valid enum value there) _run_fal(None, "a prompt", aspect_ratio="auto") model, args = calls[-1] assert model == GROK_MODEL, model assert "image_urls" not in args assert "aspect_ratio" not in args, args assert args["prompt"] == "a prompt" and args["num_images"] == 1 assert args["output_format"] == "jpeg" _run_fal(None, "p", aspect_ratio="9:16") assert calls[-1][1]["aspect_ratio"] == "9:16" # edit: image_urls is a list, and "auto" IS valid (means keep the input's ratio) _run_fal(Path(__file__), "edit me", aspect_ratio="auto") model, args = calls[-1] assert model == GROK_MODEL_EDIT, model assert isinstance(args["image_urls"], list) and len(args["image_urls"]) == 1 assert args["image_urls"][0].startswith("https://fal.media/in/") assert args["aspect_ratio"] == "auto" _run_fal(Path(__file__), "edit me", aspect_ratio="1:1") assert calls[-1][1]["aspect_ratio"] == "1:1" # an unsupported ratio must degrade to auto, never reach fal and 422 _run_fal(Path(__file__), "edit me", aspect_ratio="7:3") assert calls[-1][1]["aspect_ratio"] == "auto" finally: regenerate.fal_client.subscribe = real_sub regenerate.fal_client.upload_file = real_up regenerate._download_output = real_dl # fal response shape: {"images": [{"url": ...}]} assert _extract_url({"images": [{"url": "https://x/y.jpg"}]}) == "https://x/y.jpg" assert _extract_url({"images": []}) is None assert _extract_url({}) is None assert _extract_url(None) is None assert _extract_url("https://x/y.jpg") == "https://x/y.jpg" _check_regenerate_paths() def _check_regenerate_paths() -> None: """All three regenerate_frame branches still produce a file, with fal faked out.""" import os import shutil import subprocess import tempfile from pathlib import Path from src import regenerate from src.assemble_video import probe_image_size from src.gemini_analyze import FrameSpec from src.regenerate import regenerate_frame if not (shutil.which("ffmpeg") and shutil.which("ffprobe")): print("(skipped regenerate path check — ffmpeg not on PATH)") return with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) src_frame = root / "src.jpg" subprocess.run( ["ffmpeg", "-nostdin", "-y", "-loglevel", "error", "-f", "lavfi", "-i", "testsrc=s=720x1280", "-frames:v", "1", str(src_frame)], capture_output=True, ) # whatever fal "returns" is this 9:16 image, so a 1:1 request must be corrected locally payload = src_frame.read_bytes() calls: list[str] = [] real_sub = regenerate.fal_client.subscribe real_up = regenerate.fal_client.upload_file real_dl = regenerate._download_output regenerate.fal_client.subscribe = lambda m, arguments=None, **kw: calls.append(m) or { "images": [{"url": "https://fal.media/out.jpg"}] } regenerate.fal_client.upload_file = lambda p, **kw: "https://fal.media/in.jpg" regenerate._download_output = lambda out: payload # pretend every result is text-free; the text loop itself is covered separately real_detect = regenerate.detect_readable_text regenerate.detect_readable_text = lambda p, **kw: False os.environ["FAL_KEY"] = "test-key" try: spec_plain = FrameSpec(timestamp_sec=0, description="a house") out = regenerate_frame(src_frame, spec_plain, out_path=root / "a.jpg", aspect_ratio="auto") # verified clean → one call, where the old code always paid for two assert out.is_file() and len(calls) == 1, calls assert not list(root.glob("a_txt*.jpg")), "cleanup temps must be removed" calls.clear() out = regenerate_frame(src_frame, spec_plain, out_path=root / "b.jpg", aspect_ratio="auto", cleanup_pass=False) assert out.is_file() and len(calls) == 1, calls # fast mode = single pass calls.clear() spec_script = FrameSpec(timestamp_sec=0, script_line="Buy now", regen_brief="smile") out = regenerate_frame(src_frame, spec_script, out_path=root / "c.jpg", aspect_ratio="auto") assert out.is_file() and len(calls) == 1, calls assert not (root / "c_tmp.jpg").exists(), "scripted temp must be cleaned up" calls.clear() out = regenerate_frame(None, spec_script, out_path=root / "d.jpg", aspect_ratio="auto") assert out.is_file() and len(calls) == 1, calls # script-only = one T2I call # the safety net: fal handed back 9:16 while 1:1 was requested out = regenerate_frame(src_frame, spec_plain, out_path=root / "e.jpg", aspect_ratio="1:1", cleanup_pass=False) assert probe_image_size(out) == (720, 720), probe_image_size(out) finally: regenerate.fal_client.subscribe = real_sub regenerate.fal_client.upload_file = real_up regenerate._download_output = real_dl regenerate.detect_readable_text = real_detect def _check_crop_to_ratio() -> None: """The requested aspect ratio must reach the model as the input image's real shape.""" import shutil import subprocess import tempfile from pathlib import Path from src.assemble_video import probe_image_size from src.regenerate import _crop_to_ratio if not (shutil.which("ffmpeg") and shutil.which("ffprobe")): print("(skipped crop check — ffmpeg not on PATH)") return with tempfile.TemporaryDirectory() as tmp: tall = Path(tmp) / "tall.jpg" # 9:16, like a real UGC upload subprocess.run( ["ffmpeg", "-nostdin", "-y", "-loglevel", "error", "-f", "lavfi", "-i", "testsrc=s=720x1280", "-frames:v", "1", str(tall)], capture_output=True, ) assert probe_image_size(tall) == (720, 1280) square = _crop_to_ratio(tall, "1:1") assert square is not None assert probe_image_size(square) == (720, 720), probe_image_size(square) square.unlink(missing_ok=True) wide = _crop_to_ratio(tall, "16:9") assert wide is not None w, h = probe_image_size(wide) assert w == 720 and abs(w / h - 16 / 9) < 0.02, (w, h) wide.unlink(missing_ok=True) # already the requested shape → no wasted re-encode assert _crop_to_ratio(tall, "9:16") is None # garbage ratio must not blow up a frame assert _crop_to_ratio(tall, "auto") is None assert _crop_to_ratio(Path(tmp) / "missing.jpg", "1:1") is None def _check_edit_approve() -> None: """Approving an edit replaces the image in place and keeps the replaced version.""" import tempfile from pathlib import Path from src.pipeline import apply_edit, edit_candidate_path from src.regenerate import _edit_prompt prompt = _edit_prompt(" make it daytime ") assert "make it daytime" in prompt assert "ONLY the requested change" in prompt assert "No captions" in prompt # user wording is passed through verbatim, unlike Gemini's scrubbed descriptions assert "police" in _edit_prompt("add a police car") with tempfile.TemporaryDirectory() as tmp: run = Path(tmp) regen = run / "regen" / "frame_03_regen.jpg" regen.parent.mkdir(parents=True) regen.write_bytes(b"v1") c1 = edit_candidate_path(run, 3) assert c1.name == "frame_03_edit1.jpg" and c1.parent.name == "edits" c1.parent.mkdir(parents=True, exist_ok=True) c1.write_bytes(b"v2") # a second unapproved attempt must not clobber the first c2 = edit_candidate_path(run, 3) assert c2.name == "frame_03_edit2.jpg", c2 final = apply_edit(regen, c1) assert final == regen assert regen.read_bytes() == b"v2" # replaced in place assert not c1.exists() # candidate moved, not copied assert regen.with_name("frame_03_regen_prev1.jpg").read_bytes() == b"v1" # approving again keeps a second backup rather than overwriting the first c2.write_bytes(b"v3") apply_edit(regen, c2) assert regen.read_bytes() == b"v3" assert regen.with_name("frame_03_regen_prev1.jpg").read_bytes() == b"v1" assert regen.with_name("frame_03_regen_prev2.jpg").read_bytes() == b"v2" # earlier generation failed → candidate is adopted as-is, nothing to back up c3 = edit_candidate_path(run, 7) c3.write_bytes(b"new") assert apply_edit(None, c3) == c3 try: apply_edit(regen, run / "edits" / "missing.jpg") except FileNotFoundError: pass else: raise AssertionError("apply_edit must reject a missing candidate") def _check_parallel_regen() -> None: """Concurrent regen keeps frame order and isolates per-frame failures.""" import tempfile from pathlib import Path from src import pipeline from src.gemini_analyze import FrameSpec specs = [FrameSpec(timestamp_sec=float(i), label=f"b{i}") for i in range(9)] calls: list[int] = [] def fake_regen(original, frame, *, out_path, **kwargs): idx = int(Path(out_path).stem.split("_")[1]) calls.append(idx) if idx == 4: raise RuntimeError("boom") assert kwargs["cleanup_pass"] is False return Path(out_path) real, pipeline.regenerate_frame = pipeline.regenerate_frame, fake_regen logs: list[str] = [] try: with tempfile.TemporaryDirectory() as tmp: results = pipeline._regen_all( specs, [None] * len(specs), overall_quality="soft", dest=Path(tmp), aspect_ratio="auto", cleanup_pass=False, workers=4, log=logs.append, ) finally: pipeline.regenerate_frame = real assert len(results) == len(specs) # results stay in submission order even though completion order is arbitrary assert [r.spec.timestamp_sec for r in results] == [float(i) for i in range(9)] assert sorted(calls) == list(range(9)), calls assert results[4].regenerated is None and "boom" in (results[4].error or "") for i, r in enumerate(results): if i != 4: assert r.error is None and r.regenerated is not None assert r.regenerated.name == f"frame_{i:02d}_regen.jpg" assert any("9/9" in m for m in logs), logs _check_anchor_first() def _check_anchor_first() -> None: """Script-only runs anchor on their own first frame so the cast stays consistent.""" import tempfile from pathlib import Path from src import pipeline from src.gemini_analyze import FrameSpec specs = [FrameSpec(timestamp_sec=0, script_line=f"beat {i}") for i in range(5)] seen: list[tuple[int, tuple[str, ...]]] = [] def fake_regen(original, frame, *, out_path, references=None, **kwargs): idx = int(Path(out_path).stem.split("_")[1]) seen.append((idx, tuple(Path(r).name for r in (references or [])))) Path(out_path).parent.mkdir(parents=True, exist_ok=True) Path(out_path).write_bytes(b"x") return Path(out_path) real, pipeline.regenerate_frame = pipeline.regenerate_frame, fake_regen try: with tempfile.TemporaryDirectory() as tmp: logs: list[str] = [] results = pipeline._regen_all( specs, [None] * 5, overall_quality="", dest=Path(tmp), aspect_ratio="auto", cleanup_pass=False, workers=4, log=logs.append, anchor_first=True, ) assert all(r.regenerated for r in results) by_idx = dict(seen) assert seen[0][0] == 0, "the anchor frame must be generated first" assert by_idx[0] == (), "the anchor itself has no reference" anchor = "frame_00_regen.jpg" for i in range(1, 5): assert by_idx[i] == (anchor,), (i, by_idx[i]) assert any("anchor" in m for m in logs), logs # an explicit reference wins: no serial anchor pass, every frame uses the reference seen.clear() with tempfile.TemporaryDirectory() as tmp: ref = Path(tmp) / "hero.jpg" ref.write_bytes(b"ref") pipeline._regen_all( specs, [None] * 5, overall_quality="", dest=Path(tmp), aspect_ratio="auto", cleanup_pass=False, workers=4, log=lambda m: None, references=[ref], anchor_first=True, ) assert all(refs == ("hero.jpg",) for _, refs in seen), seen # a failed anchor must not abort the run seen.clear() def fail_first(original, frame, *, out_path, references=None, **kwargs): idx = int(Path(out_path).stem.split("_")[1]) if idx == 0: raise RuntimeError("anchor boom") return fake_regen(original, frame, out_path=out_path, references=references) pipeline.regenerate_frame = fail_first with tempfile.TemporaryDirectory() as tmp: logs = [] results = pipeline._regen_all( specs, [None] * 5, overall_quality="", dest=Path(tmp), aspect_ratio="auto", cleanup_pass=False, workers=4, log=logs.append, anchor_first=True, ) assert results[0].regenerated is None and "boom" in (results[0].error or "") assert all(r.regenerated for r in results[1:]), "other frames must still generate" assert all(refs == () for _, refs in seen), seen finally: pipeline.regenerate_frame = real if __name__ == "__main__": main()