| """Paired native-Flora/Qwen3.8 captioning on frozen ABC reference intervals. |
| |
| Run in a bounded GPU Job. Input contains interval boundaries and episode tasks, |
| never reference caption text. Each target is saved before continuing. |
| """ |
|
|
| import base64 |
| import io |
| import json |
| import os |
| import tempfile |
| import time |
| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| import av |
| import requests |
| from flora_captioner import FLORA, FloraCaptioner |
| from huggingface_hub import HfApi, hf_hub_download |
| from PIL import Image |
|
|
| QWEN = "Qwen/Qwen3.8-27B:novita" |
|
|
|
|
| class NativeFrames: |
| def __init__(self, entry, path): |
| self.entry, self.path, self.cache = entry, path, {} |
|
|
| def frames_at(self, record, timestamps): |
| fps = self.entry["fps"] |
| indices = [min(self.entry["length"] - 1, max(0, round(t * fps))) for t in timestamps] |
| with av.open(str(self.path)) as container: |
| stream = container.streams.video[0] |
| stream.codec_context.thread_count = 2 |
| for index in sorted(set(indices) - self.cache.keys()): |
| target = self.entry["video_offset"] + index / fps |
| container.seek(int(target / stream.time_base), stream=stream, backward=True) |
| closest, distance = None, float("inf") |
| for frame in container.decode(stream): |
| if frame.pts is None: |
| continue |
| stamp = float(frame.pts * stream.time_base) |
| if abs(stamp - target) < distance: |
| closest, distance = frame, abs(stamp - target) |
| if stamp >= target: |
| break |
| if closest is None or distance > 1 / fps + 1e-4: |
| raise ValueError("Source frame is missing or outside its expected clock") |
| self.cache[index] = closest.to_image().convert("RGB") |
| return [self.cache[index] for index in indices] |
|
|
|
|
| def qwen_messages(sample, record): |
| content = [ |
| { |
| "type": "text", |
| "text": ( |
| "Write one concise, self-contained English imperative for the local operation attempted in TARGET. " |
| "Identify concrete objects and the working hand. Use BEFORE/AFTER only to disambiguate TARGET; " |
| "do not include their operations. Do not invent hidden objects, goals or success. " |
| "Return only the instruction, on one line. All frames use one top camera. " |
| f"Episode instruction: {record.episode_task}" |
| ), |
| } |
| ] |
| for images, block in zip(sample["frames"], sample["sample"]["blocks"], strict=True): |
| content.append({"type": "text", "text": block["role"] + ":"}) |
| for frame, index in zip(images, block["original_descriptor"]["frame_indices"], strict=True): |
| image = (frame if isinstance(frame, Image.Image) else Image.fromarray(frame)).convert("RGB") |
| buffer = io.BytesIO() |
| image.save(buffer, format="JPEG", quality=84) |
| content.extend( |
| [ |
| {"type": "text", "text": f"Source frame {index}"}, |
| { |
| "type": "image_url", |
| "image_url": { |
| "url": "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode() |
| }, |
| }, |
| ] |
| ) |
| return [{"role": "user", "content": content}] |
|
|
|
|
| def qwen_caption(messages, state): |
| |
| reservation = (1_000_000 * 0.42 + 512 * 3.0) / 1e6 |
| if state["reserved_usd"] + reservation > 3: |
| raise RuntimeError("Qwen inference cap reached") |
| state["reserved_usd"] += reservation |
| response = requests.post( |
| "https://router.huggingface.co/v1/chat/completions", |
| headers={"Authorization": "Bearer " + os.environ["HF_TOKEN"]}, |
| json={ |
| "model": QWEN, |
| "messages": messages, |
| "temperature": 0, |
| "max_tokens": 512, |
| "chat_template_kwargs": {"enable_thinking": False}, |
| }, |
| timeout=(20, 300), |
| ) |
| if not response.ok: |
| raise RuntimeError(f"Qwen HTTP {response.status_code}; no automatic paid retry") |
| result = response.json() |
| usage = result.get("usage", {}) |
| if all(type(usage.get(key)) is int for key in ("prompt_tokens", "completion_tokens")): |
| cost = (usage["prompt_tokens"] * 0.42 + usage["completion_tokens"] * 3) / 1e6 |
| state["reserved_usd"] += cost - reservation |
| state["measured_usd"] += cost |
| choice = result["choices"][0] |
| if choice.get("finish_reason") == "length": |
| raise ValueError("Qwen response exceeded output limit") |
| caption = (choice["message"].get("content") or "").strip() |
| if not caption: |
| raise ValueError("Qwen returned an empty caption") |
| return caption |
|
|
|
|
| def main(): |
| api = HfApi() |
| repo = os.environ["BENCHMARK_REPO"] |
| revision = os.environ["BENCHMARK_REVISION"] |
| inputs = json.loads( |
| Path(hf_hub_download(repo, "inputs.json", repo_type="dataset", revision=revision)).read_text() |
| ) |
| work = Path("/workspace/results") |
| work.mkdir(parents=True, exist_ok=True) |
| state = { |
| "flora": FLORA, |
| "qwen": QWEN, |
| "input_revision": revision, |
| "reserved_usd": 0.0, |
| "measured_usd": 0.0, |
| "episodes": [], |
| "started_at": time.time(), |
| "status": "running", |
| } |
|
|
| def save(): |
| (work / "run.json").write_text(json.dumps(state, indent=2)) |
| api.upload_folder( |
| repo_id=repo, |
| repo_type="dataset", |
| folder_path=work, |
| path_in_repo="results", |
| commit_message="Checkpoint paired caption benchmark", |
| ) |
|
|
| engine = FloraCaptioner() |
| try: |
| engine.load() |
| for entry in inputs["episodes"]: |
| path = hf_hub_download( |
| inputs["repo_id"], entry["video_path"], repo_type="dataset", revision=inputs["revision"] |
| ) |
| record = SimpleNamespace( |
| episode_index=entry["episode_index"], |
| episode_task=entry["instruction"], |
| frame_timestamps=tuple(i / entry["fps"] for i in range(entry["length"])), |
| ) |
| frames = NativeFrames(entry, path) |
| result = {"episode_index": entry["episode_index"], "segments": [], "status": "running"} |
| state["episodes"].append(result) |
| with tempfile.TemporaryDirectory(prefix="flora-abc-") as temporary: |
| for sample in engine.samples(record, entry["spans"], frames, entry["fps"], Path(temporary)): |
| span = entry["spans"][sample["index"]] |
| output = {**span, "index": sample["index"], "status": sample["status"]} |
| result["segments"].append(output) |
| if sample["status"] == "ready": |
| output["input_blocks"] = sample["plan"]["blocks"] |
| started = time.monotonic() |
| output["flora_caption"] = engine.infer(sample).strip() |
| output["flora_seconds"] = time.monotonic() - started |
| started = time.monotonic() |
| output["qwen_caption"] = qwen_caption(qwen_messages(sample, record), state) |
| output["qwen_seconds"] = time.monotonic() - started |
| output["status"] = "complete" |
| save() |
| result["status"] = "complete" |
| frames.cache.clear() |
| state["status"] = "complete" |
| except Exception as exc: |
| state.update(status="failed", error_type=type(exc).__name__, error=str(exc)[:300]) |
| raise |
| finally: |
| state["finished_at"] = time.time() |
| save() |
| engine.close() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
