Spaces:
Running on Zero
Running on Zero
| """Flora-TempSeg demo: fine-grained video temporal segmentation on Qwen3.5-9B. | |
| Reproduces the released model's native inference path (nvidia/Flora-TempSeg): | |
| 64 ordered temporal slots @ 4 Hz per 16-second window (12 s emit core, 2 s guards), | |
| 20 Hz local ticks, trained pool prompt with the canonical window envelope, | |
| thinking disabled, greedy decoding, <=1024 new tokens, FP32 Gated-DeltaNet | |
| parameters and FP32 LoRA adapter on a BF16 base. | |
| """ | |
| import spaces # ZeroGPU: must be imported before torch | |
| import json | |
| import math | |
| import re | |
| import time | |
| from bisect import bisect_left | |
| from dataclasses import dataclass | |
| from fractions import Fraction | |
| from pathlib import Path | |
| import av | |
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| # --------------------------------------------------------------------------- | |
| # Model identity (pinned revisions from the release manifest models.json) | |
| # --------------------------------------------------------------------------- | |
| TEMPSEG_REPO = "nvidia/Flora-TempSeg" | |
| TEMPSEG_REVISION = "114a1076f70eb53f4a58ac045505db3087e79347" | |
| BASE_REPO = "Qwen/Qwen3.5-9B" | |
| BASE_REVISION = "c202236235762e1c871ad0ccb60c8ee5ba337b9a" | |
| TICK_US = 50_000 # 20 Hz local ticks | |
| SLOT_SPACING_US = 250_000 # 4 Hz slots | |
| SLOT_COUNT = 64 # 16 s view | |
| SLOT_FPS = 4 | |
| LEFT_GUARD_US = RIGHT_GUARD_US = 2_000_000 | |
| CORE_US = 12_000_000 | |
| VIEW_US = LEFT_GUARD_US + CORE_US + RIGHT_GUARD_US | |
| MAX_NEW_TOKENS = 1024 | |
| CONTEXT_BUDGET = 32_768 | |
| MAX_SEGMENTS = 32 | |
| MAX_DURATION_S = 120 | |
| MAX_UPLOAD_BYTES = 256 * 1024 * 1024 | |
| MAX_DIMENSION = 4096 | |
| MAX_FRAME_COUNT = 7200 | |
| # Trained pool prompts (verbatim from the release TEMPSEG_PROMPTS.json). | |
| POOL_PROMPTS = { | |
| "human": ( | |
| "<video>\n" | |
| "The video contains 64 ordered temporal slots sampled at 4 Hz for one 16-second local view.\n" | |
| "The emit core is 12 seconds with a 2-second guard on each side.\n" | |
| "Edge padding repeats the nearest real frame; ignore slots outside the valid tick envelope.\n" | |
| "Times use 20 Hz local ticks in the half-open range [0,320).\n" | |
| "Canonical window envelope:\n" | |
| '{"valid_start_tick":0,"valid_end_tick":320,"core_start_tick":40,"core_end_tick":280}\n' | |
| "Emit only subtask intervals with positive half-open intersection with [core_start_tick,core_end_tick).\n" | |
| "A quantized boundary alias may instead end exactly at core_start_tick or start exactly at core_end_tick.\n" | |
| "lo=true requires s==valid_start_tick; ro=true requires e==valid_end_tick.\n" | |
| "Return exactly one JSON object and no prose or Markdown.\n" | |
| "Use only the top-level key segments; each row's keys must appear exactly as s,e,lo,ro.\n" | |
| "Preserve concurrent or duplicate intervals as separate rows and sort rows by (s,e,lo,ro).\n" | |
| "An empty list is allowed only when no interval has emit-core support.\n" | |
| ), | |
| "yam": ( | |
| "<video>\n" | |
| "The video contains 64 ordered temporal slots sampled at 4 Hz for one 16-second local view.\n" | |
| "The emit core is 12 seconds with a 2-second guard on each side.\n" | |
| "Segment meaningful local operations. Acquisition and delivery are distinct when observable; retain distinct directional handovers. Ordinary contacts, adjustments and low-motion pauses within the same ongoing goal do not create new events. Keep clear failure/retry or execution-quality transitions as separate intervals even when their intended operation is unchanged. Represent genuine pure rest as an interval, not a gap.\n" | |
| "Edge padding repeats the nearest real sampled frame; ignore slots outside the valid tick envelope.\n" | |
| "Times use 20 Hz local ticks in the half-open range [0,320).\n" | |
| "Canonical window envelope:\n" | |
| '{"core_end_tick":280,"core_start_tick":40,"valid_end_tick":320,"valid_start_tick":0}\n' | |
| "Emit only intervals with positive half-open intersection with [core_start_tick,core_end_tick). Clip intervals at valid video limits; lo/ro indicate continuation beyond the visible limit. lo=true requires s==valid_start_tick; ro=true requires e==valid_end_tick.\n" | |
| "Return exactly one JSON object and no prose or Markdown, with only top-level key segments. Each row has keys exactly in order s,e,lo,ro. Sort rows by (s,e,lo,ro)." | |
| ), | |
| } | |
| POOL_SHA256 = { | |
| "human": "21521d5e3b854488f022f402cbb3c31137e2a1a8a114837b16e353c631bc350a", | |
| "yam": "a50644d790524e1d4c80a5ad154e940d9bb512b0816eb9d2703cf5721a2f1943", | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Media probing (production limits, PyAV) | |
| # --------------------------------------------------------------------------- | |
| class InputError(ValueError): | |
| """Raised when an uploaded video violates the model's input contract.""" | |
| class SourceVideo: | |
| """Frame index + presentation timestamps (µs, first frame = 0) for one upload.""" | |
| path: Path | |
| rate: Fraction | |
| pts_us: tuple | |
| duration_us: int | |
| size: tuple | |
| def duration_s(self) -> float: | |
| return self.duration_us / 1_000_000 | |
| def probe(path) -> SourceVideo: | |
| """Index every source frame's presentation timestamp (first frame = 0 µs).""" | |
| path = Path(path) | |
| if path.stat().st_size > MAX_UPLOAD_BYTES: | |
| raise InputError("Video exceeds 256 MiB.") | |
| with av.open(str(path), options={"err_detect": "explode"}) as container: | |
| if len(container.streams.video) != 1: | |
| raise InputError("Use a file containing exactly one video stream.") | |
| stream = container.streams.video[0] | |
| stream.codec_context.thread_count = 1 | |
| rate = stream.average_rate | |
| if not rate or rate < 4: | |
| raise InputError("The input needs at least 4 source frames per second.") | |
| size = (stream.codec_context.width, stream.codec_context.height) | |
| if min(size) <= 0 or max(size) > MAX_DIMENSION: | |
| raise InputError("Video dimensions must be at most 4096 pixels per side.") | |
| times, zero, last_length = [], None, None | |
| for frame in container.decode(stream): | |
| if frame.pts is None: | |
| raise InputError("Every source frame needs a presentation timestamp.") | |
| zero = frame.pts if zero is None else zero | |
| t = Fraction(frame.pts - zero) * stream.time_base * 1_000_000 | |
| if times and t <= times[-1]: | |
| raise InputError("Source timestamps must be strictly increasing.") | |
| if len(times) >= MAX_FRAME_COUNT: | |
| raise InputError(f"Video exceeds {MAX_FRAME_COUNT} frames.") | |
| if t > MAX_DURATION_S * 1_000_000: | |
| raise InputError(f"Video exceeds {MAX_DURATION_S:g} seconds.") | |
| times.append(t) | |
| last_length = Fraction(frame.duration) * stream.time_base * 1_000_000 if frame.duration else None | |
| if len(times) < 2: | |
| raise InputError("The video contains too few source frames.") | |
| tail = last_length or Fraction(1_000_000, rate) | |
| duration_us = int(round(times[-1] + tail)) | |
| if duration_us > MAX_DURATION_S * 1_000_000: | |
| raise InputError(f"Video exceeds {MAX_DURATION_S:g} seconds.") | |
| return SourceVideo(path, rate, tuple(times), duration_us, size) | |
| def decode_selected(video: SourceVideo, wanted: set) -> dict: | |
| """Decode only the requested source frame indices as RGB PIL images.""" | |
| wanted = set(wanted) | |
| result = {} | |
| if not wanted: | |
| return result | |
| if min(wanted) < 0 or max(wanted) >= len(video.pts_us): | |
| raise InputError("A requested source frame is outside the video.") | |
| with av.open(str(video.path), options={"err_detect": "explode"}) as container: | |
| stream = container.streams.video[0] | |
| stream.codec_context.thread_count = 1 | |
| zero = None | |
| for i, frame in enumerate(container.decode(stream)): | |
| zero = frame.pts if zero is None else zero | |
| if i in wanted: | |
| if Fraction(frame.pts - zero) * stream.time_base * 1_000_000 != video.pts_us[i]: | |
| raise InputError("Source frame timestamps changed while decoding.") | |
| result[i] = frame.to_image().convert("RGB") | |
| if i >= max(wanted): | |
| break | |
| if set(result) != wanted: | |
| raise InputError("A requested source frame could not be decoded.") | |
| return result | |
| # --------------------------------------------------------------------------- | |
| # Window scheduling (production scope scheduler, no coarse priors) | |
| # --------------------------------------------------------------------------- | |
| class Window: | |
| """One fine window over the source video.""" | |
| window_index: int | |
| requested_start_us: int | |
| requested_end_us: int | |
| valid_start_tick: int | |
| valid_end_tick: int | |
| core_start_tick: int | |
| core_end_tick: int | |
| global_tick_offset: int # requested_start_us // TICK_US | |
| def _floor_div(a, b): | |
| return a // b | |
| def _ceil_div(a, b): | |
| return -(-a // b) | |
| def schedule_windows(duration_us: int) -> list: | |
| """Emit cores advance by 12 s; guards are clipped at the video edges.""" | |
| windows = [] | |
| index = 0 | |
| while True: | |
| core_start_us = index * CORE_US | |
| if core_start_us >= duration_us: | |
| break | |
| core_end_us = min(core_start_us + CORE_US, duration_us) | |
| requested_start_us = core_start_us - LEFT_GUARD_US | |
| requested_end_us = requested_start_us + VIEW_US | |
| valid_start_us = max(0, requested_start_us) | |
| valid_end_us = min(duration_us, requested_end_us) | |
| windows.append( | |
| Window( | |
| window_index=index, | |
| requested_start_us=requested_start_us, | |
| requested_end_us=requested_end_us, | |
| valid_start_tick=_floor_div(valid_start_us - requested_start_us, TICK_US), | |
| valid_end_tick=_ceil_div(valid_end_us - requested_start_us, TICK_US), | |
| core_start_tick=_floor_div(core_start_us - requested_start_us, TICK_US), | |
| core_end_tick=_ceil_div(core_end_us - requested_start_us, TICK_US), | |
| global_tick_offset=requested_start_us // TICK_US, | |
| ) | |
| ) | |
| index += 1 | |
| if not windows: | |
| raise InputError("Video is too short to segment.") | |
| return windows | |
| # --------------------------------------------------------------------------- | |
| # Frame plan: 64 slots @ 4 Hz, nearest source frame per slot, edge padding | |
| # --------------------------------------------------------------------------- | |
| def _nearest_index(requested_us: int, times, max_selection_error_us: int) -> int: | |
| """Index of the source frame nearest requested_us (within half a slot).""" | |
| requested = Fraction(requested_us, 1) | |
| position = bisect_left(times, requested) | |
| candidates = [i for i in (position - 1, position) if 0 <= i < len(times)] | |
| if not candidates: | |
| raise InputError("No decoded frame is available for an in-video slot.") | |
| selected = min(candidates, key=lambda i: (abs(times[i] - requested), times[i], i)) | |
| if abs(times[selected] - requested) > max_selection_error_us: | |
| raise InputError("Nearest decoded frame is farther than half a profile slot.") | |
| return selected | |
| def frame_plan(window: Window, video: SourceVideo) -> list: | |
| """Slot -> source frame index map (64 slots, nearest frame, edge padding).""" | |
| pts = video.pts_us | |
| if not pts or pts[0] != 0 or any(a >= b for a, b in zip(pts, pts[1:])): | |
| raise InputError("Source timestamps must start at zero and strictly increase.") | |
| if pts[-1] >= video.duration_us: | |
| raise InputError("The last frame lies outside the video duration.") | |
| requested = [window.requested_start_us + i * SLOT_SPACING_US for i in range(SLOT_COUNT)] | |
| pad = [not 0 <= t < video.duration_us for t in requested] | |
| indices = [None if padding else _nearest_index(t, pts, SLOT_SPACING_US // 2) for t, padding in zip(requested, pad)] | |
| real = [i for i, v in enumerate(indices) if v is not None] | |
| if not real: | |
| raise InputError("No decoded frame is available for this window.") | |
| first, last = real[0], real[-1] | |
| indices[:first] = [indices[first]] * first | |
| indices[last + 1:] = [indices[last]] * (SLOT_COUNT - last - 1) | |
| observed = [v for v, p in zip(indices, pad) if not p] | |
| if len(set(observed)) != len(observed): | |
| raise InputError("Real temporal slots must not duplicate evidence.") | |
| return indices | |
| # --------------------------------------------------------------------------- | |
| # Window prompt, instruction suffix, validation, projection | |
| # --------------------------------------------------------------------------- | |
| _ENVELOPE_RE = re.compile(r"(Canonical window envelope:\s*)(\{[^{}]+\})") | |
| def window_prompt(window: Window, pool: str) -> str: | |
| """Trained pool prompt with this window's envelope JSON substituted.""" | |
| template = POOL_PROMPTS[pool] | |
| matches = list(_ENVELOPE_RE.finditer(template)) | |
| if len(matches) != 1: | |
| raise ValueError("pool template must contain exactly one envelope line") | |
| match = matches[0] | |
| keys = list(json.loads(match[2])) | |
| if set(keys) != {"valid_start_tick", "valid_end_tick", "core_start_tick", "core_end_tick"}: | |
| raise ValueError("envelope keys are invalid") | |
| envelope = json.dumps( | |
| {k: getattr(window, k) for k in keys}, | |
| separators=(",", ":"), | |
| ) | |
| return template[: match.start(2)] + envelope + template[match.end(2):] | |
| def instruction_suffix(instruction: str) -> str: | |
| """Optional episode instruction, appended as sorted JSON.""" | |
| value = (instruction or "").strip() | |
| if not value: | |
| return "" | |
| if len(value) > 2000: | |
| raise InputError("Keep the optional task instruction below 2000 characters.") | |
| return "\n" + json.dumps( | |
| {"episode_instruction": value}, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False | |
| ) | |
| def _object_without_duplicate_keys(pairs): | |
| result = {} | |
| for key, value in pairs: | |
| if key in result: | |
| raise ValueError(f"duplicate JSON key is forbidden: {key}") | |
| result[key] = value | |
| return result | |
| def parse_prediction(text: str, window: Window) -> list: | |
| """Strict production validation of the model's rows against the envelope.""" | |
| if not isinstance(text, str) or not text: | |
| raise ValueError("The model did not return a complete output for this input.") | |
| try: | |
| value = json.loads(text, object_pairs_hook=_object_without_duplicate_keys) | |
| except json.JSONDecodeError as exc: | |
| raise ValueError(f"Model output is not valid JSON: {exc}") from exc | |
| if not isinstance(value, dict) or list(value) != ["segments"]: | |
| raise ValueError("Output must be exactly one object with the single key segments.") | |
| rows = value["segments"] | |
| if not isinstance(rows, list): | |
| raise ValueError("segments must be a list.") | |
| if len(rows) > MAX_SEGMENTS: | |
| raise ValueError(f"prediction contains more than {MAX_SEGMENTS} segments") | |
| parsed, previous = [], None | |
| for index, row in enumerate(rows): | |
| if not isinstance(row, dict) or list(row) != ["s", "e", "lo", "ro"]: | |
| raise ValueError(f"segment {index}: keys must appear exactly as s,e,lo,ro") | |
| s, e, lo, ro = row["s"], row["e"], row["lo"], row["ro"] | |
| if isinstance(s, bool) or isinstance(e, bool) or not isinstance(s, int) or not isinstance(e, int): | |
| raise ValueError(f"segment {index}: ticks must be integers") | |
| if type(lo) is not bool or type(ro) is not bool: | |
| raise ValueError(f"segment {index}: open flags must be JSON booleans") | |
| if not window.valid_start_tick <= s < e <= window.valid_end_tick: | |
| raise ValueError(f"segment {index}: lies outside the valid tick envelope") | |
| if lo and s != window.valid_start_tick: | |
| raise ValueError(f"segment {index}: lo=true requires s==valid_start_tick") | |
| if ro and e != window.valid_end_tick: | |
| raise ValueError(f"segment {index}: ro=true requires e==valid_end_tick") | |
| key = (s, e, lo, ro) | |
| if previous is not None and key < previous: | |
| raise ValueError("segments are not in canonical order") | |
| previous = key | |
| core_intersects = s < window.core_end_tick and e > window.core_start_tick | |
| core_touches = e == window.core_start_tick or s == window.core_end_tick | |
| if not (core_intersects or core_touches): | |
| continue # production drops guard-only fragments | |
| parsed.append(key) | |
| return parsed | |
| def project_rows(rows, window: Window, duration_us: int) -> list: | |
| """Local ticks -> source-video microseconds (production projection).""" | |
| events = [] | |
| for s, e, lo, ro in rows: | |
| start_us = max(0, (window.global_tick_offset + s) * TICK_US) | |
| end_us = min(duration_us, (window.global_tick_offset + e) * TICK_US) | |
| if end_us <= start_us: | |
| continue | |
| completeness = { | |
| (False, False): "complete", | |
| (True, False): "left_censored", | |
| (False, True): "right_censored", | |
| (True, True): "both_censored", | |
| }[(lo, ro)] | |
| events.append( | |
| { | |
| "start_s": round(start_us / 1_000_000, 3), | |
| "end_s": round(end_us / 1_000_000, 3), | |
| "completeness": completeness, | |
| "window": window.window_index, | |
| } | |
| ) | |
| return events | |
| def stitch_events(events: list) -> list: | |
| """Merge window events into video-level intervals (simplified vs production). | |
| Censored events that touch across a 12 s core seam are joined; the | |
| production stitcher additionally resolves overlapping components within a | |
| 2-tick tolerance, which needs captioner support that is out of scope here. | |
| """ | |
| ordered = sorted(events, key=lambda e: (e["start_s"], e["end_s"])) | |
| merged = [] | |
| for event in ordered: | |
| if merged: | |
| prev = merged[-1] | |
| touches = event["start_s"] <= prev["end_s"] + 1e-9 | |
| if touches and ( | |
| prev["completeness"] in ("right_censored", "both_censored") | |
| or event["completeness"] in ("left_censored", "both_censored") | |
| ): | |
| prev["end_s"] = max(prev["end_s"], event["end_s"]) | |
| prev["completeness"] = "stitched" | |
| prev["windows"].append(event["window"]) | |
| continue | |
| merged.append({**event, "windows": [event["window"]]}) | |
| return merged | |
| # --------------------------------------------------------------------------- | |
| # Preprocessing helpers | |
| # --------------------------------------------------------------------------- | |
| def regularized_resize_dims(height, width, min_pixels=256, max_pixels=40000): | |
| """Aspect-preserving pixel-area regularization (production helper).""" | |
| if min(height, width, min_pixels, max_pixels) <= 0 or min_pixels > max_pixels: | |
| raise ValueError("regularized resize arguments are invalid") | |
| area = height * width | |
| if area > max_pixels: | |
| scale = math.sqrt(max_pixels / area) | |
| return int(height * scale), int(width * scale) | |
| if area < min_pixels: | |
| scale = math.sqrt(min_pixels / area) | |
| return int(height * scale), int(width * scale) | |
| return height, width | |
| def human_transport(image: Image.Image) -> Image.Image: | |
| """JPEG q98 4:4:4 round-trip the trained human pool applies to every frame.""" | |
| import io | |
| stream = io.BytesIO() | |
| image.save(stream, format="JPEG", quality=98, subsampling=0, optimize=False) | |
| stream.seek(0) | |
| with Image.open(stream) as decoded: | |
| return decoded.convert("RGB") | |
| # --------------------------------------------------------------------------- | |
| # Model loading (module scope, ZeroGPU eager) | |
| # --------------------------------------------------------------------------- | |
| print("Loading Flora-TempSeg: Qwen3.5-9B base + FP32 GDN + FP32 LoRA adapter ...", flush=True) | |
| _LOAD_STARTED = time.time() | |
| from huggingface_hub import hf_hub_download | |
| from peft import PeftConfig, get_peft_model, set_peft_model_state_dict | |
| from safetensors.torch import load_file | |
| from transformers import AutoModelForImageTextToText, AutoProcessor | |
| from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5PreTrainedModel | |
| from transformers.models.qwen3_vl.video_processing_qwen3_vl import Qwen3VLVideoProcessor | |
| from transformers.video_utils import VideoMetadata | |
| # Keep Gated DeltaNet parameters FP32 (native precision contract). | |
| Qwen3_5PreTrainedModel._keep_in_fp32_modules_strict = sorted( | |
| set(getattr(Qwen3_5PreTrainedModel, "_keep_in_fp32_modules_strict", None) or ()) | {"A_log", "dt_bias"} | |
| ) | |
| _BASE = AutoModelForImageTextToText.from_pretrained( | |
| BASE_REPO, | |
| revision=BASE_REVISION, | |
| dtype=torch.bfloat16, | |
| attn_implementation="sdpa", | |
| ) | |
| _ADAPTER_FILE = hf_hub_download(TEMPSEG_REPO, "adapter_model.safetensors", revision=TEMPSEG_REVISION) | |
| MODEL = get_peft_model( | |
| _BASE, | |
| PeftConfig.from_pretrained(TEMPSEG_REPO, revision=TEMPSEG_REVISION), | |
| adapter_name="candidate", | |
| autocast_adapter_dtype=True, | |
| ) | |
| set_peft_model_state_dict(MODEL, load_file(str(_ADAPTER_FILE), device="cpu"), adapter_name="candidate") | |
| MODEL.to("cuda").eval() | |
| # BF16 temporal compute autocast around generate (native contract). | |
| _ORIGINAL_GENERATE = MODEL.generate | |
| def _generate_bf16(*args, **kwargs): | |
| with torch.autocast(device_type="cuda", dtype=torch.bfloat16): | |
| return _ORIGINAL_GENERATE(*args, **kwargs) | |
| MODEL.generate = _generate_bf16 | |
| PROCESSOR = AutoProcessor.from_pretrained( | |
| TEMPSEG_REPO, | |
| revision=TEMPSEG_REVISION, | |
| use_fast=True, | |
| min_pixels=1024, | |
| max_pixels=589824, | |
| video_processor=Qwen3VLVideoProcessor.from_pretrained( | |
| TEMPSEG_REPO, revision=TEMPSEG_REVISION, fps=SLOT_FPS, max_frames=128 | |
| ), | |
| ) | |
| PROCESSOR.chat_template = Path( | |
| hf_hub_download(TEMPSEG_REPO, "chat_template.jinja", revision=TEMPSEG_REVISION) | |
| ).read_text() | |
| print(f"Model ready in {time.time() - _LOAD_STARTED:.1f}s", flush=True) | |
| def build_inputs(frames, window: Window, pool: str, instruction: str): | |
| """Processor inputs for one window (do_sample_frames=False, 64-slot metadata).""" | |
| regularized = [] | |
| for image in frames: | |
| image = image.convert("RGB") | |
| height, width = regularized_resize_dims(height=image.height, width=image.width) | |
| if image.size != (width, height): | |
| image = image.resize((width, height), Image.Resampling.BICUBIC) | |
| regularized.append(image) | |
| prompt = window_prompt(window, pool) + instruction_suffix(instruction) | |
| text = PROCESSOR.apply_chat_template( | |
| [{"role": "user", "content": prompt.replace("<video>", PROCESSOR.video_token)}], | |
| tokenize=False, | |
| add_generation_prompt=True, | |
| enable_thinking=False, | |
| ) | |
| inputs = PROCESSOR( | |
| text=[text], | |
| videos=[regularized], | |
| video_metadata=[VideoMetadata(total_num_frames=SLOT_COUNT, fps=SLOT_FPS, frames_indices=list(range(SLOT_COUNT)))], | |
| do_sample_frames=False, | |
| return_tensors="pt", | |
| ) | |
| inputs.pop("video_metadata", None) | |
| return inputs, text | |
| def run_window(window: Window, video: SourceVideo, pool: str, instruction: str): | |
| """Inference for one window: plan -> decode -> prompt -> greedy generate.""" | |
| indices = frame_plan(window, video) | |
| images = decode_selected(video, set(indices)) | |
| if pool == "human": | |
| images = {k: human_transport(v) for k, v in images.items()} | |
| frames = [images[i] for i in indices] | |
| inputs, prompt = build_inputs(frames, window, pool, instruction) | |
| n_tokens = inputs["input_ids"].shape[1] | |
| if n_tokens + MAX_NEW_TOKENS > CONTEXT_BUDGET: | |
| raise InputError("The selected window exceeds the native model context budget.") | |
| args = {k: v.to("cuda") if torch.is_tensor(v) else v for k, v in inputs.items()} | |
| eos = PROCESSOR.tokenizer.eos_token_id | |
| started = time.perf_counter() | |
| with torch.inference_mode(): | |
| output = MODEL.generate( | |
| **args, | |
| do_sample=False, | |
| num_beams=1, | |
| max_new_tokens=MAX_NEW_TOKENS, | |
| eos_token_id=eos, | |
| pad_token_id=PROCESSOR.tokenizer.pad_token_id, | |
| use_cache=True, | |
| max_time=300, | |
| ) | |
| elapsed = time.perf_counter() - started | |
| ids = output[0, n_tokens:].cpu().tolist() | |
| stop = bool(ids) and ids[-1] in ([eos] if isinstance(eos, int) else eos) | |
| text = PROCESSOR.tokenizer.decode(ids, skip_special_tokens=True).strip() | |
| if not stop or not text: | |
| raise InputError("The model did not return a complete output for this input.") | |
| return text, elapsed | |
| # --------------------------------------------------------------------------- | |
| # Timeline visualization | |
| # --------------------------------------------------------------------------- | |
| def timeline_html(events: list, duration_s: float) -> str: | |
| """SVG timeline with one colored bar per detected subtask interval.""" | |
| if duration_s <= 0: | |
| return "<p>No timeline.</p>" | |
| width, track = 940, 46 | |
| lanes = [] | |
| for event in events: | |
| placed = None | |
| for lane in lanes: | |
| if event["start_s"] >= lane["end"]: | |
| placed = lane | |
| break | |
| if placed is None: | |
| placed = {"end": 0, "bars": []} | |
| lanes.append(placed) | |
| placed["end"] = max(placed["end"], event["end_s"]) | |
| placed["bars"].append(event) | |
| n_lanes = max(1, min(len(lanes), 3)) | |
| height = 12 + n_lanes * track + 22 | |
| scale = width / duration_s | |
| palette = ["#5b6bd9", "#8f6bcb", "#4f9d8a", "#c96f9e", "#c98a5b", "#6bb1d9"] | |
| parts = [ | |
| f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" ' | |
| f'style="width:100%;background:#f8f8fb;border-radius:10px">' | |
| ] | |
| for li, lane in enumerate(lanes[:3]): | |
| y = 12 + li * track | |
| parts.append( | |
| f'<line x1="0" y1="{y + track // 2}" x2="{width}" y2="{y + track // 2}" stroke="#e2e2ee" stroke-width="1"/>' | |
| ) | |
| for bi, bar in enumerate(lane["bars"]): | |
| x = bar["start_s"] * scale | |
| w = max(3.0, (bar["end_s"] - bar["start_s"]) * scale) | |
| color = palette[(li + bi) % len(palette)] | |
| title = f'{bar["start_s"]:.2f}s – {bar["end_s"]:.2f}s ({bar["completeness"]})' | |
| parts.append( | |
| f'<g><rect x="{x:.1f}" y="{y + 8}" width="{w:.1f}" height="{track - 16}" rx="6" fill="{color}" opacity="0.85">' | |
| f'<title>{title}</title></rect>' | |
| f'<text x="{x + w / 2:.1f}" y="{y + track / 2 + 4}" font-size="11" fill="#fff" text-anchor="middle">{bar["start_s"]:.1f}</text></g>' | |
| ) | |
| step = max(1, int(duration_s // 8 or 1)) | |
| for t in range(0, int(math.floor(duration_s)) + 1, step): | |
| x = t * scale | |
| parts.append( | |
| f'<line x1="{x:.1f}" y1="0" x2="{x:.1f}" y2="{12 + n_lanes * track}" stroke="#d5d5e4" stroke-width="1"/>' | |
| ) | |
| parts.append( | |
| f'<text x="{x:.1f}" y="{height - 8}" font-size="12" fill="#666" text-anchor="middle">{t}s</text>' | |
| ) | |
| parts.append("</svg>") | |
| return "".join(parts) | |
| # --------------------------------------------------------------------------- | |
| # Gradio interface | |
| # --------------------------------------------------------------------------- | |
| CSS = """ | |
| .gradio-container .prose {font-size: 15px;} | |
| #timeline svg {max-width: 100%; height: auto;} | |
| """ | |
| def _estimate_duration(video_path, pool="human", instruction="", *args, **kwargs): | |
| """Reserve GPU time for one request, sized from measured per-window cost. | |
| Measured on ZeroGPU: 1.9-5.0 s of generation per 16-second window on | |
| 10 fps clips; 1.4x margin plus an allowance for worst-case dense | |
| predictions (32 rows) and heavy source decode. | |
| """ | |
| try: | |
| with av.open(str(video_path)) as container: | |
| stream = container.streams.video[0] | |
| duration_s = float(stream.duration * stream.time_base) if stream.duration else 0.0 | |
| n_windows = max(1, math.ceil(duration_s / (CORE_US / 1_000_000))) | |
| except Exception: | |
| n_windows = 10 # 120 s worst case | |
| return min(240, 20 + 15 * n_windows) | |
| def segment_video(video_path, pool="human", instruction="", progress=gr.Progress(track_tqdm=True)): | |
| """Segment an uploaded video into subtask intervals and draw them on its timeline. | |
| Args: | |
| video_path: uploaded video file (≤120 s, ≤256 MiB, ≥4 fps, one stream). | |
| pool: trained prompt pool — "human" for egocentric human-action video, | |
| "yam" for single-view robot manipulation. | |
| instruction: optional free-text episode instruction (≤2000 characters). | |
| progress: Gradio progress bar. | |
| Returns: | |
| A dict of merged intervals, a markdown summary, a timeline SVG, and a | |
| per-window audit dict (envelope, raw model JSON, projected rows). | |
| """ | |
| if not video_path: | |
| raise gr.Error("Upload a video first.") | |
| pool = "human" if pool in (None, "", "human", "Human") else "yam" | |
| video = probe(video_path) | |
| windows = schedule_windows(video.duration_us) | |
| progress(0.05, desc=f"{len(windows)} window(s) planned") | |
| raw_windows, events, timings = [], [], [] | |
| for i, window in enumerate(windows): | |
| text, elapsed = run_window(window, video, pool, instruction) | |
| rows = parse_prediction(text, window) | |
| window_events = project_rows(rows, window, video.duration_us) | |
| events.extend(window_events) | |
| timings.append(elapsed) | |
| raw_windows.append( | |
| { | |
| "window": window.window_index, | |
| "requested_view_start_s": window.requested_start_us / 1_000_000, | |
| "envelope": { | |
| "valid_start_tick": window.valid_start_tick, | |
| "valid_end_tick": window.valid_end_tick, | |
| "core_start_tick": window.core_start_tick, | |
| "core_end_tick": window.core_end_tick, | |
| }, | |
| "model_json": json.loads(text), | |
| "projected": window_events, | |
| "generation_s": round(elapsed, 2), | |
| } | |
| ) | |
| progress((i + 1) / (len(windows) + 1), desc=f"Window {i + 1}/{len(windows)} done") | |
| merged = stitch_events(events) | |
| result = { | |
| "video_duration_s": round(video.duration_s, 3), | |
| "pool": pool, | |
| "num_windows": len(windows), | |
| "raw_windows": raw_windows, | |
| "segments": [ | |
| { | |
| "start_s": s["start_s"], | |
| "end_s": s["end_s"], | |
| "completeness": s["completeness"], | |
| "windows": s["windows"], | |
| } | |
| for s in merged | |
| ], | |
| "total_generation_s": round(sum(timings), 2), | |
| } | |
| timeline = timeline_html(merged, video.duration_s) | |
| summary = ( | |
| f"**{len(merged)}** interval(s) across {len(windows)} window(s) · " | |
| f"video {video.duration_s:.1f}s · pool `{pool}` · generation {sum(timings):.1f}s" | |
| ) | |
| return result, summary, timeline, raw_windows | |
| with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="Flora-TempSeg · video temporal segmentation") as demo: | |
| gr.HTML( | |
| """ | |
| <div style="display:flex;align-items:baseline;gap:12px;flex-wrap:wrap"> | |
| <h1 style="margin:0">Flora-TempSeg <span style="font-size:.5em;color:#888">· fine-grained video temporal segmentation</span></h1> | |
| </div> | |
| <p style="margin-top:6px;max-width:900px"> | |
| A Qwen3.5-9B LoRA (hybrid linear attention) that emits <b>subtask interval boundaries</b> for long egocentric / | |
| robot-manipulation video. Each 16-second window is sampled as 64 ordered slots @ 4 Hz with a 12-second emit | |
| core; boundaries come back as 20 Hz local ticks, projected to video seconds and stitched across windows. | |
| The model emits <i>boundaries only</i> — it does not name actions (that is the separate Flora captioner). | |
| </p> | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| video_in = gr.Video(label="Video (≤120 s · ≤256 MiB · ≥4 fps · one stream)", sources=["upload"]) | |
| pool_in = gr.Radio( | |
| ["human", "yam"], | |
| value="human", | |
| label="Pool prompt", | |
| info="human: egocentric human-action videos · yam: single-view robot manipulation", | |
| ) | |
| instruction_in = gr.Textbox( | |
| label="Optional episode instruction (≤2000 chars)", | |
| placeholder="e.g. put the desktop items into the drawer and close the drawer", | |
| ) | |
| run = gr.Button("Segment video", variant="primary") | |
| with gr.Column(scale=2): | |
| summary_md = gr.Markdown("Upload a video and press **Segment video**.") | |
| timeline = gr.HTML(elem_id="timeline", value="<div></div>") | |
| segments_json = gr.JSON(label="Segmented intervals (projected to video seconds)") | |
| with gr.Accordion("Per-window raw output (local ticks, envelope, model JSON)", open=False): | |
| windows_json = gr.JSON(label="Per-window audit") | |
| gr.Examples( | |
| examples=[ | |
| [ | |
| "assets/droid_drawer.mp4", | |
| "yam", | |
| "put the desktop items into the drawer and close the drawer", | |
| ], | |
| [ | |
| "assets/galaxea_basket.mp4", | |
| "yam", | |
| "place the food in the right basket with the right gripper, and place the non food items in the left basket with the left gripper.", | |
| ], | |
| ], | |
| inputs=[video_in, pool_in, instruction_in], | |
| outputs=[segments_json, summary_md, timeline, windows_json], | |
| fn=segment_video, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Examples (WGO-Bench episodes, Apache-2.0 re-upload)", | |
| ) | |
| run.click( | |
| segment_video, | |
| inputs=[video_in, pool_in, instruction_in], | |
| outputs=[segments_json, summary_md, timeline, windows_json], | |
| api_name="segment", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(mcp_server=True) | |