Spaces:
Running on Zero
Running on Zero
| """Captioners: Florence-2 for images, Qwen3.5 vision-language models for whole videos.""" | |
| import re | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from torchcodec.decoders import VideoDecoder | |
| from transformers import AutoProcessor, Florence2ForConditionalGeneration, Qwen3_5ForConditionalGeneration | |
| from hub import model_path | |
| FLORENCE = {"dir": "florence-2-large", "repo": "florence-community/Florence-2-large"} | |
| CAPTION_TASKS = {"Brief": "<CAPTION>", "Detailed": "<DETAILED_CAPTION>", "More detailed": "<MORE_DETAILED_CAPTION>"} | |
| DEFAULT_CAPTION_DETAIL = "Brief" | |
| CAPTION_MAX_NEW_TOKENS = 256 | |
| # All Qwen3.5 vision-language models, loaded from the Hub cache (too large to vendor). | |
| # Marlin-2B is a Qwen3.5-2B fine-tune trained on one fixed prompt; it answers with a | |
| # Scene paragraph and <start - end> events, which format_marlin turns into Markdown. | |
| VIDEO_CAPTIONERS = { | |
| "Qwen3.5-4B": {"repo": "Qwen/Qwen3.5-4B", "marlin": False}, | |
| "Qwen3.5-9B": {"repo": "Qwen/Qwen3.5-9B", "marlin": False}, | |
| "Marlin-2B (scene + events)": {"repo": "NemoStation/Marlin-2B", "marlin": True}, | |
| } | |
| DEFAULT_VIDEO_CAPTIONERS = ["Qwen3.5-4B"] | |
| DEFAULT_VIDEO_PROMPT = "Write a one-sentence caption that describes this video." | |
| MARLIN_PROMPT = ( | |
| "Provide a spatial description of this clip followed by time-ranged events.\n" | |
| "For each event, give the time range as <start - end> and a short description." | |
| ) | |
| MARLIN_PIXELS_PER_FRAME = 200704 # ~448x448, the training-time cap | |
| VIDEO_CAPTION_MAX_NEW_TOKENS = {"qwen": 256, "marlin": 1024} | |
| VIDEO_MIN_FRAMES = 4 | |
| _CACHE: dict[str, tuple] = {} | |
| _THINK_BLOCK = re.compile(r"^\s*<think>.*?</think>\s*", re.DOTALL) | |
| _THINK_TAG = re.compile(r"^\s*<think>\s*") | |
| _MARLIN_EVENT = re.compile( | |
| r"^\s*<?\s*(\d+\.?\d*)\s*(?:seconds?|secs?|s)?\s*-\s*(\d+\.?\d*)\s*(?:seconds?|secs?|s)?\s*>?\s*[:\-]?\s*(.+?)\s*$" | |
| ) | |
| def needs_prompt(captioners: list[str]) -> bool: | |
| """True when any ticked captioner takes the user's prompt (Marlin has its own).""" | |
| return any(not VIDEO_CAPTIONERS[n]["marlin"] for n in captioners or []) | |
| def load_image_captioner(): | |
| """Load the Florence-2 (processor, model) once; kept on CPU until a GPU call.""" | |
| folder = FLORENCE["dir"] | |
| if folder not in _CACHE: | |
| path = model_path(FLORENCE) | |
| _CACHE[folder] = (AutoProcessor.from_pretrained(path), Florence2ForConditionalGeneration.from_pretrained(path).eval()) | |
| return _CACHE[folder] | |
| def load_video_captioner(repo: str): | |
| """Load a Qwen3.5 video captioner (processor, model) once in bf16; kept on CPU until a GPU call.""" | |
| if repo not in _CACHE: | |
| processor = AutoProcessor.from_pretrained(repo) | |
| model = Qwen3_5ForConditionalGeneration.from_pretrained(repo, dtype=torch.bfloat16).eval() | |
| _CACHE[repo] = (processor, model) | |
| return _CACHE[repo] | |
| def caption_image(processor, model, rgb: np.ndarray, detail: str, device: str) -> str: | |
| """Caption one RGB frame with Florence-2 at the given detail level.""" | |
| task = CAPTION_TASKS[detail] | |
| pil = Image.fromarray(rgb) | |
| inputs = processor(text=task, images=pil, return_tensors="pt").to(device) | |
| generated = model.generate(**inputs, max_new_tokens=CAPTION_MAX_NEW_TOKENS, num_beams=3) | |
| text = processor.batch_decode(generated, skip_special_tokens=False)[0] | |
| parsed = processor.post_process_generation(text, task=task, image_size=pil.size) | |
| return parsed[task].strip() | |
| def count_frames(path: str, fps: float, max_frames: int) -> int: | |
| """Frames to sample at `fps`, clamped to [VIDEO_MIN_FRAMES, max_frames] and the video length. | |
| The Qwen3.5 video processor ignores per-call min/max frame limits, so the clamp | |
| happens here and the result goes to the processor as `num_frames` with `fps` cleared. | |
| """ | |
| metadata = VideoDecoder(path).metadata | |
| wanted = int(metadata.duration_seconds * fps) | |
| return min(max(wanted, VIDEO_MIN_FRAMES), int(max_frames), metadata.num_frames) | |
| def caption_video(processor, model, path: str, prompt: str, num_frames: int, device: str, marlin: bool) -> str: | |
| """Caption one video with a Qwen3.5 model from `num_frames` evenly sampled frames.""" | |
| messages = [{"role": "user", "content": [{"type": "video", "path": path}, {"type": "text", "text": prompt}]}] | |
| processor_kwargs = {"num_frames": num_frames, "fps": None, "cap_pixels_per_frame": True} | |
| template_kwargs = {} | |
| if marlin: | |
| processor_kwargs["size"] = {"shortest_edge": 4096, "longest_edge": MARLIN_PIXELS_PER_FRAME * num_frames} | |
| else: | |
| template_kwargs["enable_thinking"] = False | |
| inputs = processor.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| tokenize=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| processor_kwargs=processor_kwargs, | |
| **template_kwargs, | |
| ).to(device) | |
| max_new_tokens = VIDEO_CAPTION_MAX_NEW_TOKENS["marlin" if marlin else "qwen"] | |
| output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False) | |
| text = processor.decode(output_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True) | |
| return strip_thinking(text) | |
| def strip_thinking(text: str) -> str: | |
| """Drop a leading `<think>...</think>` block, or a bare `<think>` prefix, as Marlin emits.""" | |
| return _THINK_TAG.sub("", _THINK_BLOCK.sub("", text)).strip() | |
| def format_marlin(text: str) -> str: | |
| """Turn Marlin's `Scene: ... Events: <a - b> ...` output into Markdown.""" | |
| lines = [] | |
| for raw in text.splitlines(): | |
| line = raw.strip() | |
| if not line: | |
| continue | |
| event = _MARLIN_EVENT.match(line) | |
| if event: | |
| start, end, desc = float(event.group(1)), float(event.group(2)), event.group(3) | |
| lines.append(f"- `{start:5.1f}s - {end:5.1f}s` {desc}") | |
| elif line.lower().startswith(("scene:", "events:")): | |
| head, _, rest = line.partition(":") | |
| lines.append(f"**{head.strip()}**" + (f" {rest.strip()}" if rest.strip() else "")) | |
| else: | |
| lines.append(line) | |
| return "\n".join(lines) | |