Curate a video dataset for generative video

git clone https://github.com/unionai/unionai-examples
cd unionai-examples/v2/tutorials/video_data_curation
python curate.py                          # all five films: 31 chunks, 56 minutes
python curate.py --smoke                  # 4 one-minute chunks from two films, no artifact
python curate.py --set min_aesthetic=5.0  # change a threshold; everything upstream is cached

SOURCES = {
    s.name: s
    for s in [
        Source(
            "big_buck_bunny", "Big Buck Bunny",
            f"{COMMONS}/c/c0/Big_Buck_Bunny_4K.webm/Big_Buck_Bunny_4K.webm.720p.vp9.webm",
            634.5, "(c) copyright 2008, Blender Foundation, peach.blender.org", "CC BY 3.0",
        ),
        Source(
            "sintel", "Sintel",
            f"{COMMONS}/f/f1/Sintel_movie_4K.webm/Sintel_movie_4K.webm.720p.vp9.webm",
            888.0, "(c) copyright Blender Foundation, durian.blender.org", "CC BY 3.0",
        ),
        Source(
            "tears_of_steel", "Tears of Steel",
            "https://download.blender.org/demo/movies/ToS/tears_of_steel_720p.mov",
            734.0, "(CC) Blender Foundation, mango.blender.org", "CC BY 3.0",
        ),
        Source(
            "elephants_dream", "Elephants Dream",
            "https://download.blender.org/ED/elephantsdream-720-h264-st-aac.mov",
            653.0, "(c) copyright 2006, Blender Foundation, Netherlands Media Art Institute, orange.blender.org", "CC BY 2.5",
        ),
        Source(
            "spring", "Spring",
            f"{COMMONS}/a/a5/Spring_-_Blender_Open_Movie.webm/Spring_-_Blender_Open_Movie.webm.720p.vp9.webm",
            464.0, "(CC) Blender Foundation, studio.blender.org", "CC BY 4.0",
        ),
    ]
}

def plan_chunks(names: list[str], chunk_s: int, head_s: int = 0, tail_s: int = 0) -> list[Chunk]:
    """Split each film into chunks of ``chunk_s`` seconds.

    With ``head_s`` and ``tail_s`` set, only the film's opening (title cards) and ending (credits)
    are used, which keeps a smoke run short and still gives the text filter work to do.
    """
    chunks = []
    for name in names:
        dur = int(SOURCES[name].duration_s)
        windows = [(0, dur)] if not head_s else [(0, min(head_s, dur - tail_s))] + ([(dur - tail_s, dur)] if tail_s else [])
        for a, b in windows:
            chunks += [Chunk(name, t, min(t + chunk_s, b)) for t in range(a, b, chunk_s) if min(t + chunk_s, b) - t >= 10]
    return chunks

def decode_window(url: str, start_s: float, duration_s: float, out: str) -> None:
    """Read ``duration_s`` seconds of a remote film into a 360p, 24 fps H.264 proxy.

    ``-ss`` before ``-i`` seeks with HTTP range requests, so only this window is downloaded.
    A keyframe every second keeps cutting clips out of the proxy fast and exact.
    """
    ffmpeg("-user_agent", USER_AGENT, "-ss", str(start_s), "-i", url, "-t", str(duration_s),
           "-vf", f"fps={FPS},scale=-2:{HEIGHT}", "-an", "-c:v", "libx264", "-preset", "veryfast",
           "-crf", "20", "-g", str(FPS), "-pix_fmt", "yuv420p", out)

def detect_shots(path: str) -> list[tuple[int, int]]:
    """Shot boundaries as (first frame, end frame) with PySceneDetect's adaptive content detector."""
    from scenedetect import AdaptiveDetector, SceneManager, open_video

    video = open_video(path)
    manager = SceneManager()
    manager.add_detector(AdaptiveDetector())
    manager.detect_scenes(video)
    return [(a.get_frames(), b.get_frames()) for a, b in manager.get_scene_list(start_in_scene=True)]

def shots_to_clips(shots: list[tuple[int, int]]) -> tuple[list[tuple[int, int, int]], int]:
    """Trim each shot, drop the ones under 2 s, and split the ones over 10 s into equal pieces.

    Returns (start frame, end frame, shot index) per clip, and the number of shots dropped.
    """
    clips, short = [], 0
    for i, (a, b) in enumerate(shots):
        a, b = a + TRIM, b - TRIM
        n = b - a
        if n < MIN_S * FPS:
            short += 1
            continue
        k = math.ceil(n / (MAX_S * FPS))
        clips += [(a + round(j * n / k), a + round((j + 1) * n / k), i) for j in range(k)]
    return clips, short

def sample_frames(proxy: str, clips: list[tuple[int, int, int]], ids: list[str], out_dir: str) -> list[dict]:
    """One pass over the proxy: save each clip's frame pairs and compute its CPU statistics."""
    import cv2
    import numpy as np

    wanted: dict[int, list[tuple[int, int, str]]] = {}  # frame -> [(clip, sample point, "a"|"b")]
    for c, (a, b, _) in enumerate(clips):
        for k, p in enumerate(POSITIONS):
            f = a + int(p * (b - a - GAP))
            wanted.setdefault(f, []).append((c, k, "a"))
            wanted.setdefault(f + GAP, []).append((c, k, "b"))

    stats = [{"brightness": [], "sharpness": [], "frame_diff": []} for _ in clips]
    first: dict[tuple[int, int], np.ndarray] = {}
    cap = cv2.VideoCapture(proxy)
    f, last = 0, max(wanted)
    while f <= last:
        ok, frame = cap.read()
        if not ok:
            break
        for c, k, side in wanted.get(f, []):
            cv2.imwrite(f"{out_dir}/frames/{ids[c]}_{k}{side}.jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 92])
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            if side == "a":
                first[(c, k)] = gray
                stats[c]["brightness"].append(float(gray.mean()))
                stats[c]["sharpness"].append(float(cv2.Laplacian(gray, cv2.CV_64F).var()))
                if k == 1:  # a small thumbnail for the report and the card
                    h = 108
                    thumb = cv2.resize(frame, (round(frame.shape[1] * h / frame.shape[0]), h), interpolation=cv2.INTER_AREA)
                    cv2.imwrite(f"{out_dir}/thumbs/{ids[c]}.jpg", thumb, [cv2.IMWRITE_JPEG_QUALITY, 80])
            else:
                diff = np.abs(gray.astype(np.int16) - first.pop((c, k)).astype(np.int16)).mean()
                stats[c]["frame_diff"].append(float(diff) / GAP)
        f += 1
    cap.release()
    return [{k: round(sum(v) / max(len(v), 1), 3) for k, v in s.items()} for s in stats]

@cpu_env.task(cache=CACHE, retries=2)
async def cut_chunk(source: str, start_s: int, end_s: int) -> tuple[Dir, str]:
    """Decode one chunk of a film, cut it into clips, and sample each clip's frames."""
    src = SOURCES[source]
    out = tempfile.mkdtemp(prefix="chunk-")
    for sub in ("frames", "thumbs", "clips"):
        os.makedirs(f"{out}/{sub}")
    proxy = f"{tempfile.mkdtemp()}/proxy.mp4"
    began = t = time.time()
    seconds = {}

    shots.decode_window(src.url, start_s, end_s - start_s, proxy)
    seconds["decode"], t = time.time() - t, time.time()
    found = shots.detect_shots(proxy)
    clips, short = shots.shots_to_clips(found)
    seconds["shots"], t = time.time() - t, time.time()
    ids = [f"{source}-{start_s:04d}-{i:03d}" for i in range(len(clips))]
    stats = shots.sample_frames(proxy, clips, ids, out)
    shots.cut_clips(proxy, clips, ids, out)
    seconds["sample_and_cut"] = time.time() - t

    rows = [
        {"id": cid, "source": source, "chunk": f"{source}-{start_s:04d}", "shot": shot,
         "start_s": round(start_s + a / shots.FPS, 3), "end_s": round(start_s + b / shots.FPS, 3),
         "duration_s": round((b - a) / shots.FPS, 3), **s}
        for cid, (a, b, shot), s in zip(ids, clips, stats)
    ]
    shots.write_index(out, rows)
    summary = {"shots": len(found), "short_shots": short, "clips": len(rows), "footage_s": end_s - start_s,
               "seconds": {k: round(v, 1) for k, v in seconds.items()}, "replica": socket.gethostname(),
               "ran": [began, time.time()]}
    print(f"{source} {start_s}-{end_s}s: {len(found)} shots -> {len(rows)} clips in {sum(seconds.values()):.0f}s")
    return await Dir.from_local(out), json.dumps({"summary": summary, "clips": rows})

def scorers() -> dict:
    """CLIP ViT-L/14 with the LAION aesthetic head, RAFT optical flow, and the CRAFT text detector."""
    with _lock:
        if "clip" in _models:
            return _models
        import easyocr
        import torch
        from torchvision.models.optical_flow import Raft_Large_Weights, raft_large
        from transformers import CLIPModel, CLIPProcessor

        t0 = time.time()
        clip_id = "openai/clip-vit-large-patch14"
        head = torch.nn.Sequential(  # the improved-aesthetic-predictor MLP, on L2-normalized CLIP embeddings
            torch.nn.Linear(768, 1024), torch.nn.Dropout(0.2), torch.nn.Linear(1024, 128), torch.nn.Dropout(0.2),
            torch.nn.Linear(128, 64), torch.nn.Dropout(0.1), torch.nn.Linear(64, 16), torch.nn.Linear(16, 1),
        )
        state = torch.hub.load_state_dict_from_url(AESTHETIC_HEAD, map_location="cpu", file_name="aesthetic-l14.pth")
        head.load_state_dict({k.removeprefix("layers."): v for k, v in state.items()})
        weights = Raft_Large_Weights.DEFAULT
        _models.update(
            clip=CLIPModel.from_pretrained(clip_id, torch_dtype=torch.float16).cuda().eval(),
            clip_processor=CLIPProcessor.from_pretrained(clip_id),
            aesthetic=head.cuda().eval(),
            raft=raft_large(weights=weights).cuda().eval(),
            raft_transforms=weights.transforms(),
            # Detection only: we need where the text is, not what it says.
            ocr=easyocr.Reader(["en"], gpu=True, recognizer=False, verbose=False),
            load_s=round(time.time() - t0, 1),
        )
        print(f"loaded scorers on {replica()} in {_models['load_s']}s")
        return _models

def aesthetic(m: dict, images: list) -> list[float]:
    """The LAION aesthetic predictor (a 1 to 10 scale) on each image."""
    import torch

    with torch.inference_mode():
        pixels = m["clip_processor"](images=images, return_tensors="pt")["pixel_values"].cuda().half()
        emb = m["clip"].get_image_features(pixel_values=pixels).float()
        return m["aesthetic"](emb / emb.norm(dim=-1, keepdim=True)).squeeze(-1).tolist()

def motion(m: dict, firsts: list, seconds: list, height: int = 256) -> list[float]:
    """Mean RAFT optical-flow magnitude per pair, in pixels per frame at ``height`` px tall."""
    import numpy as np
    import torch
    import torch.nn.functional as F

    w = round(firsts[0].width * height / firsts[0].height / 8) * 8  # RAFT needs sides divisible by 8

    def batch(images):
        t = torch.from_numpy(np.stack([np.asarray(i) for i in images])).permute(0, 3, 1, 2).cuda()
        return F.interpolate(t.float(), size=(height, w), mode="bilinear", antialias=True).clamp(0, 255).to(torch.uint8)

    with torch.inference_mode():
        a, b = m["raft_transforms"](batch(firsts), batch(seconds))
        flow = m["raft"](a, b)[-1]  # the last of RAFT's iterative refinements
        return (flow.norm(dim=1).mean(dim=(1, 2)) / shots.GAP).tolist()

def text_area(m: dict, images: list) -> list[float]:
    """Share of each frame covered by detected text boxes."""
    import numpy as np

    out = []
    for img in images:
        arr = np.asarray(img)
        boxes, polys = m["ocr"].detect(arr)
        area = sum((x1 - x0) * (y1 - y0) for x0, x1, y0, y1 in boxes[0])
        for p in polys[0]:  # rotated boxes, by the shoelace formula
            xs, ys = [q[0] for q in p], [q[1] for q in p]
            area += abs(sum(xs[i] * ys[i - 1] - xs[i - 1] * ys[i] for i in range(len(p)))) / 2
        out.append(min(area / (arr.shape[0] * arr.shape[1]), 1.0))
    return out

def score_clips(chunk_dir: str, clip_ids: list[str]) -> tuple[list[dict], dict]:
    """Score every clip of one chunk. Returns per-clip scores and GPU seconds per scorer."""
    import torch

    m = scorers()
    seconds = {"aesthetic": 0.0, "motion": 0.0, "text": 0.0}
    rows = []
    for cid in clip_ids:
        firsts, pairs = _frames(chunk_dir, cid, "a"), _frames(chunk_dir, cid, "b")
        row = {"id": cid}
        for name, fn, agg in (
            ("aesthetic", lambda: aesthetic(m, firsts), lambda v: sum(v) / len(v)),
            ("motion", lambda: motion(m, firsts, pairs), lambda v: sum(v) / len(v)),
            ("text", lambda: text_area(m, firsts), max),  # a title card in any sampled frame counts
        ):
            t0 = time.time()
            row[name] = round(float(agg(fn())), 4)
            torch.cuda.synchronize()
            seconds[name] += time.time() - t0
        rows.append(row)
    return rows, {k: round(v, 2) for k, v in seconds.items()}

@score_env.task(cache=CACHE, retries=2)
async def score_chunk(chunk: Dir) -> str:
    """Aesthetic, motion and text scores for every clip of one chunk, on one GPU."""
    began = time.time()
    local = await chunk.download()
    with open(f"{local}/clips.json") as f:
        ids = [c["id"] for c in json.load(f)]
    cold = "clip" not in scoring._models
    t0 = time.time()
    rows, seconds = await asyncio.to_thread(scoring.score_clips, local, ids)
    print(f"scored {len(ids)} clips in {time.time() - t0:.1f}s on {scoring.replica()} {seconds}")
    return json.dumps({"scores": rows, "seconds": seconds, "replica": scoring.replica(),
                       "load_s": scoring._models["load_s"] if cold else 0.0, "ran": [began, time.time()]})

@dataclass
class Thresholds:
    """Every clip must pass all five filters. Tune them on the report's histograms."""

    min_brightness: float = 20.0  # mean luma over the 4 frames, 0-255: drops dark and mostly black shots
    max_brightness: float = 235.0  # drops white flashes
    min_sharpness: float = 20.0  # variance of the Laplacian at 360p: drops blurry frames
    max_text: float = 0.02  # share of the frame covered by text: drops titles and credits
    min_aesthetic: float = 4.5  # LAION aesthetic score, 1-10
    min_motion: float = 0.1  # optical flow, px per frame at 256 px tall: drops still shots
    max_motion: float = 20.0  # drops whip pans and flashes that flow can't follow

FILTERS = ("exposure", "blur", "text", "aesthetic", "motion")

def select(clip: dict, t: Thresholds) -> dict:
    """Which filters a clip passes, the first one it fails, and whether it is kept."""
    passed = {
        "exposure": t.min_brightness <= clip["brightness"] <= t.max_brightness,
        "blur": clip["sharpness"] >= t.min_sharpness,
        "text": clip["text"] <= t.max_text,
        "aesthetic": clip["aesthetic"] >= t.min_aesthetic,
        "motion": t.min_motion <= clip["motion"] <= t.max_motion,
    }
    first_fail = next((f for f in FILTERS if not passed[f]), "")
    return {**{f"pass_{f}": ok for f, ok in passed.items()}, "first_fail": first_fail, "keep": not first_fail}

def captioner(model_id: str) -> tuple[dict, float]:
    """The captioning model and its processor; also how long this call spent loading them (0 if warm)."""
    with _lock:
        if _models.get("vlm_id") == model_id:
            return _models, 0.0
        import torch
        from transformers import AutoModelForImageTextToText, AutoProcessor

        t0 = time.time()
        processor = AutoProcessor.from_pretrained(model_id)
        processor.image_processor.do_image_splitting = False  # one tile per frame is plenty at 360p
        model = AutoModelForImageTextToText.from_pretrained(model_id, torch_dtype=torch.bfloat16).cuda().eval()
        _models.update(vlm_id=model_id, vlm=model, vlm_processor=processor)
        load_s = round(time.time() - t0, 1)
        print(f"loaded {model_id} on {replica()} in {load_s}s")
        return _models, load_s

def caption(model_id: str, chunk_dir: str, clip_id: str, max_new_tokens: int = 72) -> tuple[str, float, float]:
    """Caption one clip from its four sampled frames. Returns the caption, GPU seconds and load seconds."""
    import torch

    m, load_s = captioner(model_id)
    images = _frames(chunk_dir, clip_id, "a")
    content = [{"type": "image"} for _ in images] + [{"type": "text", "text": CAPTION_PROMPT}]
    processor = m["vlm_processor"]
    prompt = processor.apply_chat_template([{"role": "user", "content": content}], add_generation_prompt=True)
    with _lock:  # one generation at a time on the GPU; downloads for other calls overlap with it
        t0 = time.time()
        inputs = processor(text=prompt, images=[images], return_tensors="pt").to("cuda", dtype=torch.bfloat16)
        with torch.inference_mode():
            out = m["vlm"].generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
        text = processor.batch_decode(out[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0]
        return " ".join(text.split()), time.time() - t0, load_s

@caption_env.task(cache=CACHE, retries=2)
async def caption_clip(chunk: Dir, clip_id: str, model: str = VLM) -> str:
    """Caption one kept clip from its four sampled frames. Cached per clip."""
    began = time.time()
    local = tempfile.mkdtemp()
    os.makedirs(f"{local}/frames")
    await asyncio.gather(*[
        File.from_existing_remote(f"{chunk.path}/frames/{clip_id}_{k}a.jpg").download(f"{local}/frames/{clip_id}_{k}a.jpg")
        for k in range(len(shots.POSITIONS))
    ])
    text, seconds, load_s = await asyncio.to_thread(scoring.caption, model, local, clip_id)
    return json.dumps({"caption": text, "seconds": round(seconds, 3), "replica": scoring.replica(), "load_s": load_s,
                       "ran": [began, time.time()]})

# unionai-reuse provides the bridge that reusable environments need.
cpu_image = (
    flyte.Image.from_debian_base(name="video-curation-cpu")
    .with_apt_packages("ffmpeg")
    .with_pip_packages(
        "scenedetect==0.6.7", "opencv-python-headless==4.11.0.86", "numpy<2.3", "pandas==2.3.3",
        "pyarrow==21.0.0", "Pillow==11.3.0", "unionai-reuse==0.2.1",
    )
)
# The default PyPI wheels of torch 2.9 carry CUDA 12.8.
gpu_image = (
    flyte.Image.from_debian_base(name="video-curation-gpu")
    .with_pip_packages("torch==2.9.1", "torchvision==0.24.1")
    .with_pip_packages(
        "transformers==4.57.1", "accelerate==1.10.1", "num2words==0.5.14", "easyocr==1.7.2",
        "opencv-python-headless==4.11.0.86", "numpy<2.3", "Pillow==11.3.0", "unionai-reuse==0.2.1",
    )
)

# Decoding and shot detection: a warm pool of CPU replicas that takes one chunk at a time each.
cpu_env = flyte.TaskEnvironment(
    name="video-curation-cut",
    image=cpu_image,
    resources=flyte.Resources(cpu=4, memory="8Gi", disk="20Gi"),
    reusable=flyte.ReusePolicy(replicas=(1, 5), concurrency=1, idle_ttl=600, scaledown_ttl=120),
)

# Scoring: CLIP, RAFT and the text detector load once per replica and stay on the GPU.
score_env = flyte.TaskEnvironment(
    name="video-curation-score",
    image=gpu_image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="L4:1"),
    reusable=flyte.ReusePolicy(replicas=2, concurrency=1, idle_ttl=600),
)

# Captioning: a different model with a different lifetime, so its own pool. Several calls
# share a replica: while one clip generates, the next one's frames download.
caption_env = flyte.TaskEnvironment(
    name="video-curation-caption",
    image=gpu_image,
    resources=flyte.Resources(cpu=4, memory="16Gi", gpu="L4:1"),
    reusable=flyte.ReusePolicy(replicas=(1, 2), concurrency=4, idle_ttl=600, scaledown_ttl=180),
)

driver_env = flyte.TaskEnvironment(
    name="video-curation-driver",
    image=cpu_image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    depends_on=[cpu_env, score_env, caption_env],
)

# "auto" cache keys hash each task's own source, not the helper modules it calls. Salting with
# a hash of those modules means editing shots.py or scoring.py invalidates the cached results.
HELPERS = ("sources.py", "shots.py", "scoring.py")
SALT = hashlib.sha256(b"".join((Path(__file__).parent / f).read_bytes() for f in HELPERS)).hexdigest()[:12]
CACHE = flyte.Cache("auto", salt=SALT)

@driver_env.task(report=True, produces_artifacts=True)
async def curate(
    sources: list[str] = list(SOURCES),
    chunk_s: int = 120,
    head_s: int = 0,
    tail_s: int = 0,
    thresholds: Thresholds = Thresholds(),
    vlm: str = VLM,
    sample_size: int = 24,
    publish: bool = True,
) -> Dir:
    """Turn whole films (or their first ``head_s`` and last ``tail_s`` seconds) into a captioned clip dataset."""
    started = time.time()
    chunks = plan_chunks(sources, chunk_s, head_s, tail_s)
    clips: list[dict] = []
    summaries: dict[str, dict] = {}
    dirs: dict[str, Dir] = {}
    calls: list[dict] = []  # one row per task call: stage, timing, replica, model load, cache hit
    lock = asyncio.Lock()

    async def call(stage: str, label: str, coro) -> tuple:
        """Await one task call and record it for the timeline and the throughput table."""
        called = time.time()
        result = await coro
        info = json.loads(result[1] if stage == "cut" else result)
        info = info.get("summary", info)
        began, ended = info["ran"]
        # A cache hit returns the outputs of an earlier run, with that run's timestamps.
        row = {"stage": stage, "label": label, "replica": info["replica"], "load_s": info.get("load_s", 0.0),
               "cached": began < called - 1, "start": began, "end": ended}
        calls.append(row)
        return result, row

    async def process(chunk) -> None:
        with flyte.group(chunk.id):
            (d, cut), row = await call("cut", chunk.id, cut_chunk(chunk.source, chunk.start_s, chunk.end_s))
            cut = json.loads(cut)
            row.update(clips=len(cut["clips"]), busy_s=sum(cut["summary"]["seconds"].values()))
            scored = {"scores": []}
            if cut["clips"]:
                scored, row = await call("score", chunk.id, score_chunk(d))
                scored = json.loads(scored)
                row.update(clips=len(cut["clips"]), busy_s=sum(scored["seconds"].values()), per_scorer=scored["seconds"])
            mine = [{**c, **s} for c, s in zip(cut["clips"], scored["scores"])]
            for c in mine:
                c.update(select(c, thresholds))

            async def cap(c: dict) -> None:
                r, row = await call("caption", c["id"], caption_clip(d, c["id"], vlm))
                r = json.loads(r)
                row.update(clips=1, busy_s=r["seconds"])
                c["caption"] = r["caption"]

            await asyncio.gather(*[cap(c) for c in mine if c["keep"]])
        async with lock:
            dirs[chunk.id], summaries[chunk.id] = d, cut["summary"]
            clips.extend(mine)
            await flyte.report.replace.aio(
                reports.live(clips, summaries, calls, len(chunks), thresholds, started, time.time()), do_flush=True)

    await asyncio.gather(*[process(c) for c in chunks])
    wall_s = time.time() - started
    clips.sort(key=lambda c: c["id"])
    for c in clips:
        c.setdefault("caption", "")
        c["cut_uri"] = f"{dirs[c['chunk']].path}/clips/{c['id']}.mp4"  # the cut stage's copy, every clip
        c["clip"] = f"clips/{c['id']}.mp4" if c["keep"] else ""  # inside the dataset, kept clips only

    # Thumbnails and frame strips for the report and the card, fetched only for the clips shown.
    shown = reports.pick_examples(clips, thresholds, sample_size)
    thumbs = dict(zip(shown["thumbs"], await asyncio.gather(*[
        read_remote(f"{dirs[c.rsplit('-', 1)[0]].path}/thumbs/{c}.jpg") for c in shown["thumbs"]])))
    frames = {cid: await asyncio.gather(*[read_remote(f"{dirs[cid.rsplit('-', 1)[0]].path}/frames/{cid}_{k}a.jpg")
                                          for k in range(len(shots.POSITIONS))]) for cid in shown["sample"]}

    out = build_dataset(clips, shown["sample"], frames, thresholds)
    # Copy every kept clip into the dataset, so it doesn't depend on intermediate task outputs.
    gate = asyncio.Semaphore(32)

    async def copy(c: dict) -> None:
        async with gate:
            await File.from_existing_remote(c["cut_uri"]).download(f"{out}/{c['clip']}")

    await asyncio.gather(*[copy(c) for c in clips if c["keep"]])

    ctx = dict(clips=clips, summaries=summaries, calls=calls, thresholds=thresholds, wall_s=wall_s, origin=started,
               thumbs=thumbs, frames=frames, sample=shown["sample"], vlm=vlm)
    await flyte.report.replace.aio(reports.overview(**ctx))
    flyte.report.get_tab("Filters").replace(reports.filters_tab(**ctx))
    flyte.report.get_tab("Captions").replace(reports.captions_tab(**ctx))
    await flyte.report.flush.aio()

    dataset = await Dir.from_local(out)
    if not publish:
        return dataset
    kept = [c for c in clips if c["keep"]]
    card = await artifacts.Card.create_from.aio(content=reports.card(**ctx), format="html", card_type="data")
    metadata = artifacts.Metadata(
        name="blender-open-movie-clips",
        kind="data",
        description=f"{len(kept)} captioned 2-10 s clips ({sum(c['duration_s'] for c in kept) / 60:.1f} min) "
                    f"curated from {len(clips)} candidates in {len(sources)} Blender open movies (CC BY)",
        attrs={"clips": str(len(kept)), "candidates": str(len(clips)), "captioner": vlm,
               "thresholds": json.dumps(asdict(thresholds)), "license": "CC BY (see ATTRIBUTION.txt)"},
        card=card,
    )
    return artifacts.new(dataset, metadata)

def build_dataset(clips: list[dict], sample: list[str], frames: dict, thresholds: Thresholds) -> str:
    """The manifest as parquet, frame strips for a sample of clips, and the attribution."""
    import pandas as pd

    out = tempfile.mkdtemp(prefix="dataset-")
    os.makedirs(f"{out}/sample")
    os.makedirs(f"{out}/clips")
    columns = ["id", "source", "chunk", "shot", "start_s", "end_s", "duration_s", "brightness", "sharpness",
               "frame_diff", "text", "aesthetic", "motion", *[f"pass_{f}" for f in FILTERS], "first_fail",
               "keep", "caption", "clip", "cut_uri"]
    df = pd.DataFrame(clips)[columns]
    df["title"] = df["source"].map(lambda s: SOURCES[s].title)
    df["license"] = df["source"].map(lambda s: SOURCES[s].license)
    df.to_parquet(f"{out}/manifest.parquet", index=False)
    for cid in sample:
        with open(f"{out}/sample/{cid}.jpg", "wb") as f:
            f.write(reports.strip_jpeg(frames[cid]))
    with open(f"{out}/thresholds.json", "w") as f:
        json.dump(asdict(thresholds), f, indent=2)
    with open(f"{out}/ATTRIBUTION.txt", "w") as f:
        f.write("Clips cut from Blender Foundation open movies, used under Creative Commons Attribution licenses.\n\n")
        f.writelines(f"{s.title}: {s.credit}. {s.license}. Source: {s.url}\n" for s in SOURCES.values())
    return out

python curate.py --set min_aesthetic=5.0