| """Independent native Flora TempSeg → Captioner workflow on a CUDA Job. |
| |
| The default profile uses NVIDIA's reference decoder. The opt-in accelerated |
| profile batches independent requests and requires fused GDN kernels; both use |
| the same planner, processor, mixed precision loader and greedy decoding policy. |
| """ |
|
|
| import asyncio |
| import hashlib |
| import os |
| import sys |
| import tempfile |
| import threading |
| import time |
| from concurrent.futures import ThreadPoolExecutor |
| from fractions import Fraction |
| from itertools import islice |
| from pathlib import Path |
|
|
| FLORA = { |
| "model_id": "nvidia/Flora-Captioner", |
| "code_revision": "f736ca36925ae55d2f7915993ed1bc9f87288bb3", |
| "weights_revision": "fef0bbfe11829ca7f3539d58da7a6a2c2f109dc6", |
| "segmenter_id": "nvidia/Flora-TempSeg", |
| "segmenter_revision": "114a1076f70eb53f4a58ac045505db3087e79347", |
| "segmenter_sha256": "1d80fe544e4cd1a7ff6b0ea4be71e988972408dab81d4c8af72e3a628739d6c3", |
| "segmenter_base_id": "Qwen/Qwen3.5-9B", |
| "segmenter_base_revision": "c202236235762e1c871ad0ccb60c8ee5ba337b9a", |
| "temporal_profile": "yam", |
| "hardware": "h200", |
| "resident_models": True, |
| "license": "NVIDIA OneWay Noncommercial License", |
| "role": "TempSeg predicts and stitches intervals; Captioner writes their labels independently", |
| } |
|
|
|
|
| def reference_runtime(): |
| from huggingface_hub import snapshot_download |
|
|
| |
| |
| directory = str( |
| Path(os.environ.get("FLORA_CODE_CACHE", Path.home() / ".cache/flora/code")) / FLORA["code_revision"] |
| ) |
| snapshot_download( |
| FLORA["model_id"], |
| revision=FLORA["code_revision"], |
| allow_patterns=["batch_utils/**"], |
| local_dir=directory, |
| ) |
| sys.path.insert(0, directory) |
| from batch_utils.bootstrap import production |
| from batch_utils.config import MODELS |
| from batch_utils.models import Session, generate |
| from batch_utils.preprocessing import instruction_suffix |
| from flora_unified.unified_input import load_frames, prompt |
|
|
| expected = { |
| "tempseg": (FLORA["segmenter_id"], FLORA["segmenter_revision"]), |
| "tempseg_base": (FLORA["segmenter_base_id"], FLORA["segmenter_base_revision"]), |
| "captioner": (FLORA["model_id"], FLORA["weights_revision"]), |
| } |
| if any((MODELS[k]["repo"], MODELS[k]["revision"]) != value for k, value in expected.items()): |
| raise ValueError("Native Flora model revisions differ from the reviewed configuration") |
| if MODELS["tempseg"]["sha256"] != FLORA["segmenter_sha256"]: |
| raise ValueError("Native Flora adapter checksum differs from the reviewed configuration") |
|
|
| return production, Session, generate, instruction_suffix, load_frames, prompt |
|
|
|
|
| class TemporalPredictor: |
| """Adapt LeRobot frame access to NVIDIA's unchanged window/stitch engine.""" |
|
|
| def __init__(self, engine, source, binding, record, frames): |
| self.engine, self.source, self.binding = engine, source, binding |
| self.record, self.frames = record, frames |
|
|
| def predict(self, request): |
| from batch_utils.preprocessing import temporal_inputs |
| from flora_runtime.video_segmentation.fine_subtask import FINE_RESPONSE_SCHEMA |
|
|
| if ( |
| request["source"]["source_id"] != self.source.source_id |
| or request["model_binding"] != self.binding |
| ): |
| raise ValueError("Native Flora temporal request binding changed") |
| production = self.engine.reference[0] |
| pts_us = [Fraction(str(t)) * 1_000_000 for t in self.record.frame_timestamps] |
| predictions, payloads = [], iter(request["windows"]) |
| while batch := list(islice(payloads, self.engine.batch_size)): |
| windows, requests = [], [] |
| for payload in batch: |
| self.engine.check_deadline() |
| window = production.window_from_payload(payload, self.source) |
| plan = production.temporal_frame_plan(window, pts_us) |
| images = self.frames.frames_at( |
| self.record, [self.record.frame_timestamps[i] for i in plan["source_frame_indices"]] |
| ) |
| inputs, _ = temporal_inputs( |
| images, |
| window, |
| FLORA["temporal_profile"], |
| self.engine.processor, |
| ["main"], |
| self.record.episode_task or "", |
| ) |
| windows.append(window) |
| requests.append(inputs) |
| for window, prediction in zip(windows, self.engine.generate(requests, 1024), strict=True): |
| predictions.append({"scope_id": window.scope_id, "prediction_json": prediction}) |
| return { |
| "schema_version": FINE_RESPONSE_SCHEMA, |
| "source_id": self.source.source_id, |
| "model_io_contract": self.binding["model_io_contract"], |
| "model_binding": self.binding, |
| "predictions": predictions, |
| } |
|
|
|
|
| class FloraCaptioner: |
| def __init__( |
| self, |
| deadline=None, |
| cache_models=False, |
| execution="reference", |
| batch_size=4, |
| vision_frame_batching=False, |
| ): |
| if execution not in {"reference", "accelerated"}: |
| raise ValueError("Unknown Flora execution profile") |
| if not isinstance(batch_size, int) or isinstance(batch_size, bool) or not 1 <= batch_size <= 8: |
| raise ValueError("Flora batch size must be between 1 and 8") |
| self.lock = threading.RLock() |
| self.deadline = deadline |
| self.execution = execution |
| if type(vision_frame_batching) is not bool or vision_frame_batching and execution != "accelerated": |
| raise ValueError("Vision frame batching requires the accelerated Flora profile") |
| self.vision_frame_batching = vision_frame_batching |
| self.batch_size = batch_size if execution == "accelerated" else 1 |
| self.kernel_versions = {} |
| if execution == "accelerated": |
| from flora_acceleration import require_fast_kernels |
|
|
| self.kernel_versions = require_fast_kernels() |
| self.reference = reference_runtime() |
| self.cache_models = cache_models or execution == "accelerated" |
| self.session = None if self.cache_models else self.reference[1]() |
| self.sessions = {} |
| self.cached_models = {} |
| self.model = self.processor = None |
| self.loaded_kind = None |
| self.prefetch_pool = None |
| self.prefetch = None |
| if execution == "accelerated": |
| from batch_utils.models import snapshot, validate_manifest |
|
|
| def prepare_captioner(): |
| path = snapshot("captioner") |
| |
| |
| validate_manifest(path) |
| return path |
|
|
| self.prefetch_pool = ThreadPoolExecutor(max_workers=1) |
| self.prefetch = self.prefetch_pool.submit(prepare_captioner) |
|
|
| def close(self): |
| if self.prefetch_pool is not None: |
| self.prefetch_pool.shutdown(wait=True, cancel_futures=True) |
| self.model = self.processor = None |
| self.cached_models.clear() |
| for session in self.sessions.values(): |
| session.close() |
| self.sessions.clear() |
| if self.session is not None: |
| self.session.close() |
| self.loaded_kind = None |
|
|
| def check_resident_capacity(self): |
| import torch |
|
|
| |
| |
| if torch.cuda.get_device_properties(0).total_memory < 120 * 1024**3: |
| raise ValueError("Keeping both Flora models loaded requires an H200-class GPU with 120 GiB") |
|
|
| def check_deadline(self): |
| if self.deadline and time.time() >= self.deadline: |
| from provider import DeadlineExceededError |
|
|
| raise DeadlineExceededError("Flora Job time limit reached; completed episodes are saved.") |
|
|
| def annotate(self, record, frames, fps): |
| """Predict boundaries with TempSeg, then caption them without a provider call.""" |
| with self.lock: |
| spans = self.segment(record, frames, fps) |
| return self.relabel(record, spans, frames, fps) |
|
|
| def segment(self, record, frames, fps): |
| production = self.reference[0] |
| if fps < 4 or not record.frame_timestamps: |
| raise ValueError("Flora temporal segmentation needs at least 4 source FPS") |
| if any(abs(t - i / fps) > 1e-4 for i, t in enumerate(record.frame_timestamps)): |
| raise ValueError("The Studio Flora robot profile needs an episode-relative constant-rate clock") |
| self.load("tempseg") |
| source = production.AnnotationSource( |
| source_id=f"studio_episode_{record.episode_index}", |
| location=f"lerobot:episode/{record.episode_index}", |
| |
| |
| duration_s=round(len(record.frame_timestamps) / fps * 1_000_000) / 1_000_000, |
| fps=float(fps), |
| ) |
| pool = FLORA["temporal_profile"] |
| binding = { |
| "candidate_id": "flora_final", |
| "checkpoint_sha256": FLORA["segmenter_sha256"], |
| "model_io_contract": production.fine_window_profile(production.PROFILES[pool]).model_io_contract, |
| } |
| predictor = TemporalPredictor(self, source, binding, record, frames) |
| native = production.production_segmenter(predictor, binding, pool) |
| segments = asyncio.run(native.segment_async(source)) |
| if not segments: |
| raise ValueError("Flora TempSeg did not return any subtask intervals") |
| return [ |
| { |
| "start": segment.span.start_s, |
| "end": segment.span.end_s, |
| "metadata": dict(segment.metadata), |
| } |
| for segment in segments |
| ] |
|
|
| def samples(self, record, spans, frames, fps, directory): |
| """Native duration-adaptive BEFORE/TARGET/AFTER plans without reference text.""" |
| production, _, _, suffix, load_frames, prompt = self.reference |
| rate = Fraction(str(fps)) |
| plans, dispositions = production.caption_plans( |
| [ |
| {"start_s": span["start"], "end_s": span["end"], "metadata": span.get("metadata", {})} |
| for span in spans |
| ], |
| { |
| "fps_num": rate.numerator, |
| "fps_den": rate.denominator, |
| "frame_count": len(record.frame_timestamps), |
| }, |
| ) |
| wanted = sorted({i for plan in plans for block in plan["blocks"] for i in block["frame_indices"]}) |
| images = frames.frames_at(record, [record.frame_timestamps[i] for i in wanted]) |
| receipts = {} |
| for index, image in zip(wanted, images, strict=True): |
| path = directory / f"{index:012d}.png" |
| image.convert("RGB").save(path, format="PNG") |
| receipts[index] = { |
| "path": path.name, |
| "png_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), |
| "size": list(image.size), |
| } |
| by_event = {plan["event_id"]: plan for plan in plans} |
| for disposition in dispositions: |
| plan = by_event.get(disposition["event_id"]) |
| if plan is None: |
| yield {"status": disposition["status"], "index": disposition["event_index"]} |
| continue |
| sample = production.caption_sample(plan, receipts, ["main"]) |
| blocks, _ = load_frames(sample, directory) |
| content = prompt(sample, self.processor.video_token) + suffix(record.episode_task) |
| yield { |
| "status": "ready", |
| "index": plan["event_index"], |
| "plan": plan, |
| "sample": sample, |
| "frames": blocks, |
| "content": content, |
| } |
|
|
| def caption_inputs(self, sample): |
| from transformers.video_utils import VideoMetadata |
|
|
| self.check_deadline() |
| proc = self.processor |
| text = proc.apply_chat_template( |
| [{"role": "user", "content": sample["content"]}], |
| tokenize=False, |
| add_generation_prompt=True, |
| enable_thinking=False, |
| ) |
| metadata = [ |
| VideoMetadata( |
| total_num_frames=len(frames), |
| fps=sample["plan"]["source_fps_num"] / sample["plan"]["source_fps_den"], |
| frames_indices=list(block["block_local_timestamp_indices"]), |
| ) |
| for frames, block in zip(sample["frames"], sample["sample"]["blocks"], strict=True) |
| ] |
| inputs = proc( |
| text=[text], |
| videos=sample["frames"], |
| video_metadata=metadata, |
| do_sample_frames=False, |
| return_metadata=True, |
| return_tensors="pt", |
| ) |
| return inputs |
|
|
| def generate(self, requests, output_limit): |
| """Retain reference batch-one behavior unless acceleration is explicitly selected.""" |
| if self.execution == "accelerated": |
| from flora_acceleration import generate_many |
|
|
| return generate_many( |
| self.model, |
| self.processor, |
| requests, |
| output_limit, |
| batch_size=self.batch_size, |
| check_deadline=self.check_deadline, |
| ) |
| return [self.reference[2](self.model, self.processor, inputs, output_limit) for inputs in requests] |
|
|
| def infer(self, sample): |
| return self.generate([self.caption_inputs(sample)], 512)[0] |
|
|
| def load(self, kind="captioner"): |
| self.check_deadline() |
| if kind == "captioner" and self.prefetch is not None: |
| self.prefetch.result() |
| if self.cache_models: |
| if kind not in self.cached_models: |
| self.check_resident_capacity() |
| |
| |
| session = self.reference[1]() |
| model, processor = session.load(kind) |
| if self.vision_frame_batching: |
| from batch_utils.models import runtime_event |
| from flora_vision import configure_vision_batching |
|
|
| layers = configure_vision_batching(model) |
| runtime_event("vision_frames_batched", kind=kind, layers=layers, backend="sdpa") |
| self.cached_models[kind] = model, processor |
| self.sessions[kind] = session |
| self.model, self.processor = self.cached_models[kind] |
| self.loaded_kind = kind |
| return |
| if self.model is None or self.loaded_kind != kind: |
| |
| self.model = self.processor = None |
| self.model, self.processor = self.session.load(kind) |
| self.loaded_kind = kind |
|
|
| def relabel(self, record, spans, frames, fps): |
| with self.lock, tempfile.TemporaryDirectory(prefix="flora-caption-") as temporary: |
| self.load() |
| result = [dict(span) for span in spans] |
| samples = iter(self.samples(record, spans, frames, fps, Path(temporary))) |
| while batch := list(islice(samples, self.batch_size)): |
| if any(sample["status"] != "ready" for sample in batch): |
| raise ValueError("Flora cannot caption a source-limited target interval") |
| captions = self.generate([self.caption_inputs(sample) for sample in batch], 512) |
| for sample, caption in zip(batch, captions, strict=True): |
| if not caption.strip(): |
| raise ValueError("Flora returned an empty caption") |
| result[sample["index"]]["text"] = caption.strip() |
| return result |
|
|
