Spaces:
Running on Zero
Running on Zero
| """Preencode unified omni samples with the released MiniMax-H3 ``ref2va`` front end. | |
| Every sample of the input JSONL is:: | |
| {"id": "...", "target": "clip.mp4", "text": "...", | |
| "observations": [{"media": "x.mp4|x.png", "frames": "all" | [0, 5], "degree": 0.999}, ...], | |
| "known_mask": "mask.mp4" | null, "source": "dataset-name"} | |
| ``frames`` says how an observation relates to the clip: ``"all"`` or a list of target latent frames | |
| for an aligned one, ``null`` for a free one. No field names a task or a condition type. | |
| The prompt embeddings, their token tags and the observation latents come from the released modular | |
| blocks (setup -> Qwen3-VL presentation -> reference VAE encode), so a sample is conditioned exactly | |
| as the released checkpoint expects. The target goes through the same reference encoder, on its own, | |
| to get latents in the same convention. One worker per GPU: ``--shard i --shards n``. | |
| Media and samples may be local paths or storage URLs (``<scheme>://...``, read and written through | |
| solarwm.storage); with URLs, node-local scratch only ever holds what a worker is busy with. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import contextlib | |
| import hashlib | |
| import json | |
| import math | |
| import os | |
| import subprocess | |
| import sys | |
| import threading | |
| import time | |
| import traceback | |
| import numpy as np | |
| import torch | |
| from solarwm import storage | |
| ROOT = os.environ.get("H3_REF_ROOT", "MiniMax-H3") | |
| # MiniMax-H3's one frame rate. Every video handed to it is put on this grid first and then cut to | |
| # the frames asked for, so a latent's place in time is a count of these frames from the clip's | |
| # start - not a fraction of however long the clip happened to be. | |
| FPS = 24.0 | |
| IMAGE_EXT = (".png", ".jpg", ".jpeg", ".webp", ".bmp") | |
| # The shortest video H3's reference blocks take (read at 2 fps, sampled frames merged in pairs). | |
| STILL_FRAMES = 13 | |
| SCRATCH = os.environ.get("OMNI_SCRATCH", "/tmp/omni_pre") | |
| class WriteRefused(RuntimeError): | |
| """The storage would not take a member this worker had already encoded.""" | |
| class FetchFailed(RuntimeError): | |
| """A row's media could not be fetched - most often because it is not there yet.""" | |
| def copy(source: str, destination: str) -> None: | |
| """Copy one file to or from storage (solarwm.storage retries a busy store a few times). | |
| Which way it went decides whose failure it is: a write the storage refused says nothing about | |
| the row, a fetch that failed says the row's media is missing - and that row is set aside.""" | |
| try: | |
| if storage.is_remote(destination): | |
| storage.put(source, destination) | |
| else: | |
| storage.get(source, destination) | |
| except Exception as exc: | |
| error = WriteRefused if storage.is_remote(destination) else FetchFailed | |
| raise error(f"storage copy {source} -> {destination}: {str(exc)[-300:]}") from exc | |
| def sweep(cache: str, keep_seconds: float = 900.0) -> None: | |
| """Drop cached media no worker has touched lately; a sibling may still be using a fresh one.""" | |
| now = time.time() | |
| for name in os.listdir(cache): | |
| with contextlib.suppress(OSError): | |
| if now - os.path.getatime(f"{cache}/{name}") > keep_seconds: | |
| os.remove(f"{cache}/{name}") | |
| class Sink: | |
| """Where encoded samples go: a remote object-storage prefix, or a local directory.""" | |
| def __init__(self, out: str, shard: int) -> None: | |
| self.out, self.remote = out.rstrip("/"), "://" in out | |
| # Scratch is per set, not just per shard. A node encodes one set after another, and an | |
| # index named only by its shard number would be found by the next set, appended to, and | |
| # published under that set's prefix - carrying the previous set's rows into it. | |
| self.stage = f"{SCRATCH}/out/{self.out.rsplit('/', 1)[-1]}" | |
| self.index_name = f"index.{shard:05d}.jsonl" | |
| self.index_path = f"{self.stage}/{self.index_name}" | |
| os.makedirs(f"{self.stage}/samples", exist_ok=True) | |
| if not self.remote: | |
| os.makedirs(f"{self.out}/samples", exist_ok=True) | |
| self.index_path = f"{self.out}/{self.index_name}" | |
| elif not os.path.exists(self.index_path): | |
| # A restarted shard picks up its own index and skips what it already wrote. | |
| with contextlib.suppress(Exception): # a first run has no index yet | |
| copy(f"{self.out}/{self.index_name}", self.index_path) | |
| def done(self) -> set[str]: | |
| if not os.path.exists(self.index_path): | |
| return set() | |
| with open(self.index_path) as lines: | |
| return {json.loads(line)["id"] for line in lines if line.strip()} | |
| def member(self, name: str) -> str: | |
| """Local path to write a member to; ``publish`` moves it where it belongs.""" | |
| return f"{self.stage if self.remote else self.out}/{name}" | |
| def publish(self, name: str) -> None: | |
| if self.remote: | |
| # Copied, then dropped here - not moved. A move that is tried again after it failed | |
| # has nothing left to move, because the first try took the file whether or not it | |
| # reported success, and the worker then throws away a member the storage already holds. | |
| copy(f"{self.stage}/{name}", f"{self.out}/{name}") | |
| with contextlib.suppress(OSError): | |
| os.remove(f"{self.stage}/{name}") | |
| def publish_index(self) -> None: | |
| if self.remote: | |
| copy(self.index_path, f"{self.out}/{self.index_name}") | |
| def fetch_fresh(path: str, destination: str) -> str: | |
| """The current version of a sample list, one copy per node. | |
| A list is rewritten under its own name, so a copy is kept per version - its size and time in | |
| the storage - and a new version is fetched as soon as there is one. The biggest lists run to | |
| gigabytes; every worker of every node fetching its own each pass would read terabytes | |
| for nothing. The first worker to get the lock fetches, the others wait and share its copy. | |
| ``destination`` is kept for the callers' sake; the copy lives in the node's list cache. | |
| """ | |
| import fcntl | |
| if "://" not in path: | |
| return path | |
| found = storage.stat(path) | |
| version = f"{found[0]}_{int(found[1] or 0)}" if found else "unknown" | |
| cache = os.environ.get("OMNI_LIST_CACHE", "/tmp/omni_lists") | |
| os.makedirs(cache, exist_ok=True) | |
| local = f"{cache}/{os.path.basename(path)}.{version}" | |
| if os.path.exists(local): | |
| return local | |
| with open(f"{local}.lock", "w") as lock: | |
| fcntl.flock(lock, fcntl.LOCK_EX) | |
| if not os.path.exists(local): | |
| copy(path, f"{local}.part") | |
| os.replace(f"{local}.part", local) | |
| # Older versions of the same list are no longer anyone's. | |
| for name in os.listdir(cache): | |
| if name.startswith(os.path.basename(path) + ".") and not name.startswith(os.path.basename(local)): | |
| try: | |
| os.remove(f"{cache}/{name}") | |
| except OSError: | |
| pass | |
| return local | |
| def localize(path: str) -> str: | |
| """A remote path is fetched to local scratch once; anything else is used as is.""" | |
| if "://" not in path: | |
| return path | |
| import hashlib | |
| local = f"{SCRATCH}/in/{hashlib.sha1(path.encode()).hexdigest()}{os.path.splitext(path)[1]}" | |
| if not os.path.exists(local): | |
| os.makedirs(f"{SCRATCH}/in", exist_ok=True) | |
| part = f"{local}.{os.getpid()}.part" # the workers of a node share this cache | |
| copy(path, part) | |
| os.replace(part, local) | |
| return local | |
| def open_media(path: str): | |
| """A PyAV container that tolerates undecodable metadata; newer PyAV releases dropped the option.""" | |
| import av | |
| try: | |
| return av.open(path, metadata_errors="ignore") | |
| except TypeError: | |
| return av.open(path) | |
| def read_video(path: str, seconds: float = 0.0) -> tuple[np.ndarray, float]: | |
| """Decode a clip - only its first ``seconds`` when given. | |
| H3 keeps 6.6 s of any clip, and OmniHuman alone holds 4K clips a thousand frames long: decoding | |
| all of one to use its start kept one worker busy for most of a pass while its seven siblings, | |
| finished, held idle cards waiting for it. | |
| """ | |
| import av | |
| path = localize(path) | |
| container = open_media(path) | |
| stream = container.streams.video[0] | |
| stream.thread_type = "AUTO" # frame threads: otherwise a worker decodes on a single core | |
| fps = float(stream.average_rate) | |
| limit = int(math.ceil(seconds * fps)) + 2 if seconds else 0 | |
| frames = [] | |
| for frame in container.decode(stream): | |
| frames.append(frame.to_ndarray(format="rgb24")) | |
| if limit and len(frames) >= limit: | |
| break | |
| container.close() | |
| return np.stack(frames), fps | |
| def prepared(frames: np.ndarray, fps: float, num_frames: int) -> np.ndarray: | |
| """A clip exactly as H3 would make it - on 24 fps, cut to ``num_frames``, LANCZOS-resized to the | |
| canvas its own aspect resolves to - with the resize spread over threads. | |
| H3 does this itself, one frame after another on one core, while the card waits: a quarter of | |
| every card's time went to it. A clip handed over already at 24 fps and at that size goes | |
| through H3 untouched, so doing the same steps here, in parallel, changes no pixel. | |
| """ | |
| from concurrent.futures import ThreadPoolExecutor | |
| from diffusers.modular_pipelines.minimax_h3.modular_pipeline import resolve_canvas_size | |
| from PIL import Image | |
| frames = on_grid(frames, fps)[:num_frames] | |
| # The released checkpoint's canvas rule: axes on multiples of 32, short edge 768, area 768x1344. | |
| height, width = resolve_canvas_size(frames.shape[2], frames.shape[1], 32, 768, 768 * 1344) | |
| if frames.shape[1:3] == (height, width): | |
| return np.ascontiguousarray(frames) | |
| def one(frame: np.ndarray) -> np.ndarray: | |
| return np.asarray(Image.fromarray(frame).resize((width, height), Image.Resampling.LANCZOS)) | |
| with ThreadPoolExecutor(16) as pool: | |
| return np.stack(list(pool.map(one, frames))) | |
| def frame_rate(path: str) -> float: | |
| """A clip's frame rate from its container.""" | |
| import av | |
| container = open_media(localize(path)) | |
| rate = float(container.streams.video[0].average_rate) | |
| container.close() | |
| return rate | |
| def label_mask(path: str, label: int, fps: float, start: float, span: float): | |
| """One object's region out of a per-frame label map (VPData's all_masks.npz: every frame of the | |
| clip, each pixel the number of the object on it), cut to the stretch that is needed, as frames | |
| the rest of the mask handling reads: white where the object is. It runs on the clip's clock - | |
| its frame i is the clip's frame i - so the clip's own rate is its rate.""" | |
| labels = np.load(localize(path))["arr_0"] | |
| first = int(round(start * fps)) | |
| last = int(math.ceil((start + span) * fps)) + 2 | |
| # Already cut at the start: the caller must not cut it again. | |
| return (labels[first:last] == label).astype(np.uint8)[..., None] * 255, fps | |
| def frame_count(path: str) -> int: | |
| """How many frames a clip holds, from its container rather than by decoding it.""" | |
| import av | |
| container = open_media(localize(path)) | |
| stream = container.streams.video[0] | |
| count = int(stream.frames or 0) | |
| if not count: | |
| seconds = float(stream.duration * stream.time_base) if stream.duration else (container.duration or 0) / 1e6 | |
| count = int(round(seconds * float(stream.average_rate))) | |
| container.close() | |
| return count | |
| def on_grid(frames: np.ndarray, fps: float) -> np.ndarray: | |
| """A clip put on 24 fps the way H3 puts a reference there: whole frames dropped or held.""" | |
| if fps == FPS: | |
| return frames | |
| scale = FPS / fps | |
| slots = np.floor(np.arange(frames.shape[0]) * scale + 0.5).astype(np.int64) | |
| end = math.floor(frames.shape[0] * scale + 0.5) | |
| return np.repeat(frames, np.diff(slots, append=end), axis=0) | |
| def crop_scale(size: tuple[int, int], canvas_shape: tuple[int, int]) -> tuple[float, float]: | |
| """What fraction of the height and width a centre-crop to the canvas aspect keeps.""" | |
| height, width = size | |
| wanted = canvas_shape[1] / canvas_shape[0] | |
| if abs(width / height - wanted) < 1e-3: | |
| return 1.0, 1.0 | |
| if width / height > wanted: | |
| return 1.0, round(height * wanted) / width | |
| return round(width / wanted) / height, 1.0 | |
| def fit(frames: np.ndarray, canvas_shape: tuple[int, int]) -> np.ndarray: | |
| """Centre-crop to the canvas's own aspect, because the encoder will not stretch to it. | |
| The reference blocks fit a frame inside the canvas and keep its aspect, so a clip that is not | |
| 1344:768 comes out on a smaller grid than the model's - 48x64 for four-thirds, 48x48 for a | |
| square render - and the reader will not admit it. Cropping keeps every pixel it keeps real; | |
| padding would teach the model to draw the bars. | |
| """ | |
| height, width = frames.shape[1:3] | |
| wanted = canvas_shape[1] / canvas_shape[0] | |
| if abs(width / height - wanted) < 1e-3: | |
| return frames | |
| if width / height > wanted: # too wide: take a full-height slice | |
| keep = int(round(height * wanted)) | |
| start = (width - keep) // 2 | |
| return frames[:, :, start : start + keep] | |
| keep = int(round(width / wanted)) # too tall: take a full-width slice | |
| start = (height - keep) // 2 | |
| return frames[:, start : start + keep, :] | |
| def region(frames: np.ndarray, box) -> np.ndarray: | |
| """The part of every frame a dataset's own crop box keeps, ``[x0, x1, y0, y1]`` in the clip's | |
| pixels. Some sets cut their clips down to a region of the frame - past a subtitle, a watermark, | |
| a letterbox, or onto the part that holds the subject - and name the box rather than re-encode | |
| millions of files; whatever is laid over the clip frame for frame is cut with it.""" | |
| if not box: | |
| return frames | |
| x0, x1, y0, y1 = (int(v) for v in box) | |
| height, width = frames.shape[1:3] | |
| if not (0 <= x0 < x1 <= width and 0 <= y0 < y1 <= height): | |
| raise ValueError(f"crop box {list(box)} does not fit the {width}x{height} frame") | |
| return frames[:, y0:y1, x0:x1] | |
| def within(image, limit: float = 4.0): | |
| """The centre of an image cut to the widest and tallest shapes H3 takes, 4:1 and 1:4.""" | |
| width, height = image.size | |
| if width > limit * height: | |
| keep = int(limit * height) | |
| return image.crop(((width - keep) // 2, 0, (width - keep) // 2 + keep, height)) | |
| if height > limit * width: | |
| keep = int(limit * width) | |
| return image.crop((0, (height - keep) // 2, width, (height - keep) // 2 + keep)) | |
| return image | |
| def aligned_image(image, canvas_shape: tuple[int, int], box=None): | |
| """An image laid over the target's canvas (a first frame, a still of an aligned stream): cut to the | |
| canvas's aspect and LANCZOS-resized to its size, as every frame of the target is. Handed to H3 as | |
| an image, a first frame became a 128x224 latent - off the target's 48x84 grid - and some seven | |
| thousand extra prompt tokens; ``reference`` hands it over as a one-frame video at this size.""" | |
| from PIL import Image | |
| pixels = fit(region(np.asarray(image)[None], box), canvas_shape)[0] | |
| height, width = canvas_shape | |
| if pixels.shape[:2] == (height, width): | |
| return Image.fromarray(pixels) | |
| return Image.fromarray(pixels).resize((width, height), Image.Resampling.LANCZOS) | |
| def reference( | |
| path: str, | |
| canvas_shape: tuple[int, int], | |
| video=None, | |
| crop: bool = True, | |
| start: float = 0.0, | |
| num_frames: int = 158, | |
| box=None, | |
| ): | |
| """One reference for H3. ``crop`` puts it on the canvas's shape; only what is laid over the | |
| target's canvas - the target itself, and an observation aligned with it frame for frame - has to | |
| be. A free reference (a portrait, an object, another shot) sits on its own spatial grid, as it | |
| does in H3 itself, so it keeps its own shape: cutting a full-length portrait to the canvas's | |
| 7:4 kept a band across its middle, a torso with no face and no feet.""" | |
| from diffusers.modular_pipelines.minimax_h3 import ( | |
| MiniMaxH3ImageReference, | |
| MiniMaxH3VideoReference, | |
| ) | |
| from PIL import Image | |
| if path.lower().endswith(IMAGE_EXT): | |
| image = Image.open(localize(path)).convert("RGB") | |
| if not crop: | |
| return MiniMaxH3ImageReference(image=within(image)) | |
| # A still laid over the target is one frame of it, so it is encoded the way the target's frames | |
| # are: as video, on the target's grid (as an image H3 brings it to its own still size). H3 | |
| # takes no video shorter than STILL_FRAMES, so the still is held that long; the causal VAE's | |
| # first latent sees the first frame alone, and that latent is all the sample keeps. | |
| frame = np.asarray(aligned_image(image, canvas_shape, box))[None] | |
| return MiniMaxH3VideoReference(frames=np.repeat(frame, STILL_FRAMES, axis=0), fps=FPS) | |
| frames, fps = video if video is not None else read_video(path, start + num_frames / FPS + 0.5) | |
| if start: | |
| frames = frames[int(round(start * fps)) :] | |
| # The target's own crop box, for a reference laid over it; a free one is never cut. | |
| frames = fit(region(frames, box), canvas_shape) if crop else frames | |
| return MiniMaxH3VideoReference(frames=prepared(frames, fps, num_frames), fps=FPS) | |
| def build_pipelines(device: str): | |
| from diffusers import ComponentsManager, ModularPipeline | |
| from diffusers.modular_pipelines import SequentialPipelineBlocks | |
| blocks = ModularPipeline.from_pretrained(ROOT).blocks | |
| manager = ComponentsManager() | |
| # Weights stay on the card. Offloading moved the text encoder over PCIe on every call, which | |
| # cost little per sample once warm but made the first call of every process several times | |
| # slower. The heaviest sample, 47 latents with four observations, fits on the card; | |
| # offloading remains, behind a switch, for a shape that would not fit. | |
| offload = os.environ.get("OMNI_CPU_OFFLOAD") == "1" | |
| if offload: | |
| manager.enable_auto_cpu_offload(device=device) | |
| def chain(workflow: str, names: tuple[str, ...]): | |
| sub = blocks.get_workflow(workflow).sub_blocks | |
| return SequentialPipelineBlocks.from_blocks_dict( | |
| {name: sub[name] for name in names} | |
| ).init_pipeline(ROOT, components_manager=manager) | |
| conditioned = chain("ref2va", ("before_encode", "text_encoder", "vae_encoder")) | |
| conditioned.load_components(dtype=torch.bfloat16) | |
| if not offload: | |
| for name in conditioned.pretrained_component_names: | |
| component = getattr(conditioned, name, None) | |
| if isinstance(component, torch.nn.Module): | |
| component.to(device) | |
| # One copy of every weight: the other chains borrow the components (the text encoder is 62 GB). | |
| loaded = {name: getattr(conditioned, name) for name in conditioned.pretrained_component_names} | |
| latents_only = chain("ref2va", ("before_encode", "vae_encoder")) | |
| text_only = chain("t2va", ("text_encoder",)) | |
| for pipe in (latents_only, text_only): | |
| pipe.update_components( | |
| **{k: v for k, v in loaded.items() if k in pipe.pretrained_component_names} | |
| ) | |
| return conditioned, latents_only, text_only | |
| def canvas(path: str, box=None) -> tuple[int, int]: | |
| """The canvas a clip is encoded on: landscape or portrait by the part of the frame that is kept.""" | |
| import av | |
| stream = open_media(localize(path)).streams.video[0] | |
| width, height = stream.width, stream.height | |
| if box: | |
| width, height = int(box[1]) - int(box[0]), int(box[3]) - int(box[2]) | |
| return (768, 1344) if width >= height else (1344, 768) | |
| def source_size(path: str) -> tuple[int, int]: | |
| import av | |
| stream = open_media(localize(path)).streams.video[0] | |
| return stream.height, stream.width | |
| def encode(sample: dict, pipes, num_frames: int) -> tuple[dict, dict]: | |
| conditioned, latents_only, text_only = pipes | |
| box = sample.get("crop") | |
| # A crop box moves the picture's origin and scales its field of view, and a camera track's | |
| # intrinsics would have to follow; nothing here does that, so the two are refused together | |
| # before any card time is spent. | |
| if box and sample.get("camera"): | |
| raise ValueError("a crop box and a camera track together are not supported") | |
| height, width = canvas(sample["target"], box) | |
| size = {"height": height, "width": width, "num_frames": num_frames} | |
| shape = (height, width) | |
| # Where in the clip the sample starts, in seconds. The encoder keeps 6.6 s of a clip; a sample | |
| # whose moment is later in a longer clip - a keyframe at 20 s of 22 - names where its stretch | |
| # begins, and the target, every observation on the clip's clock, the camera and the mask all | |
| # start there together. Free references are not on the clip's clock and are not moved. | |
| start = float(sample.get("start") or 0.0) | |
| # A start given as a frame number is turned into seconds by the clip's own rate: a list's | |
| # rounded "fps 23" for a 23.976 clip would put frame 736 1.3 s off. | |
| if sample.get("start_frame"): | |
| start = int(sample["start_frame"]) / frame_rate(sample["target"]) | |
| span = start + num_frames / FPS + 0.5 # what of the clip the latents can hold, and a margin | |
| video = read_video(sample["target"], span) | |
| video = (region(video[0], box), video[1]) | |
| # The clip as it was made, before anything is cut off: counted by the container, since only its | |
| # start was decoded. A camera track made one pose per frame of the whole clip needs it. | |
| whole_frames = max(frame_count(sample["target"]), int(video[0].shape[0])) | |
| if start: | |
| video = (video[0][int(round(start * video[1])) :], video[1]) | |
| if len(video[0]) < 2: | |
| raise ValueError(f"start {start:.2f} s is past the end of the clip") | |
| # Where the sample's stretch of the clip ends, as a frame of the clip, when a set curates a | |
| # stretch shorter than the encoder's window: reading on past it would take in what the set cut | |
| # away - often the next shot. Everything on the clip's clock is held to the same stretch, | |
| # because it is counted from the target's frames below. | |
| if sample.get("end_frame"): | |
| keep = int(sample["end_frame"]) - int(round(start * video[1])) | |
| if keep < 2: | |
| raise ValueError(f"end frame {sample['end_frame']} leaves {keep} frames after the start") | |
| video = (video[0][:keep], video[1]) | |
| clip_frames, clip_fps = int(video[0].shape[0]), float(video[1]) | |
| target_state = latents_only( | |
| references=[reference(sample["target"], shape, video, num_frames=num_frames)], | |
| prompt="", | |
| **size, | |
| ) | |
| del video | |
| target = target_state.get("condition_latents")[0][0] # [24, F, h, w] | |
| frames = int(target.shape[1]) | |
| # What of the clip the target latents hold: H3 puts the clip on 24 fps and keeps the first | |
| # num_frames, so a clip longer than that is encoded from its start and the rest is not there. | |
| # Anything else that runs on the clip's clock - a camera track, a mask that came with the | |
| # clip - has to be read over this same stretch, or it describes seconds the latents never saw. | |
| window = min(int(math.floor(clip_frames * FPS / clip_fps + 0.5)), num_frames) | |
| def seconds(latent: int) -> float: | |
| """Where a latent sits in the clip, in seconds from its start.""" | |
| return latent * (window - 1) / max(frames - 1, 1) / FPS | |
| # Two latents is five frames, a fifth of a second: a still with a timestamp. Nothing about | |
| # motion, continuation or timing can be learned from it, so it is refused by name here and | |
| # counted, rather than encoded and quietly drawn in training. | |
| if frames <= 2: | |
| raise ValueError(f"target is {frames} latents ({17 * (frames - 2) // 5 + 5} frames), too short to be a video") | |
| tensors = {"target_latents": target.to(torch.bfloat16).contiguous()} | |
| observations = [] | |
| items = sample.get("observations") or [] | |
| if items: | |
| state = conditioned( | |
| prompt=sample["text"], | |
| references=[ | |
| reference( | |
| o["media"], | |
| shape, | |
| crop=o.get("frames") is not None, | |
| start=start if o.get("frames") is not None else 0.0, | |
| num_frames=num_frames, | |
| box=box if o.get("frames") is not None else None, | |
| ) | |
| for o in items | |
| ], | |
| **size, | |
| ) | |
| for index, (item, latents) in enumerate( | |
| zip(items, state.get("condition_latents"), strict=True) | |
| ): | |
| latents = latents[0] | |
| if item["media"].lower().endswith(IMAGE_EXT) and item.get("frames") is not None: | |
| latents = latents[:, :1] # an aligned still: its first latent (see ``reference``) | |
| aligned = item.get("frames") | |
| if aligned == "all": | |
| aligned = list(range(min(int(latents.shape[1]), frames))) | |
| latents = latents[:, : len(aligned)] | |
| if aligned is not None and len(aligned) != int(latents.shape[1]): | |
| raise ValueError( | |
| f"observation {index} has {int(latents.shape[1])} latent frames " | |
| f"for {len(aligned)} target frames" | |
| ) | |
| tensors[f"observation_{index}_latents"] = latents.to(torch.bfloat16).contiguous() | |
| observations.append( | |
| { | |
| "frames": aligned, | |
| "degree": float(item.get("degree", 0.999)), | |
| "boxes": item.get("boxes"), | |
| } | |
| ) | |
| else: | |
| state = text_only(prompt=sample["text"]) | |
| prompt = state.get("prompt_embeds") | |
| tensors["prompt_embeds"] = prompt.reshape(-1, prompt.shape[-1]).to(torch.bfloat16).contiguous() | |
| tensors["text_token_tags"] = ( | |
| state.get("text_token_tags").reshape(-1).to(torch.int64).contiguous() | |
| ) | |
| if sample.get("camera"): | |
| # The sample hands over absolute camera-to-world poses and normalised intrinsics; the | |
| # reader makes them relative. | |
| loaded = np.load(localize(sample["camera"])) | |
| views_all = torch.as_tensor(np.asarray(loaded["c2w"]), dtype=torch.float32) | |
| k_all = torch.as_tensor(np.asarray(loaded["K"]), dtype=torch.float32) | |
| total = int(views_all.shape[0]) | |
| if k_all.ndim != 2 and int(k_all.shape[0]) < total: | |
| raise ValueError(f"camera holds {total} poses but {int(k_all.shape[0])} intrinsics") | |
| # One camera per latent, taken at the latent's time. A track says its own rate when it was | |
| # made over part of a clip; one without says nothing because it was made one pose per | |
| # frame of the whole clip. Spreading the whole track over the latents instead - as this | |
| # once did - is right only for a clip no longer than the window; for a 20 s clip it laid | |
| # 20 s of camera motion over the 6.6 s the latents hold. | |
| track_fps = float(loaded["fps"]) if "fps" in loaded.files else total * clip_fps / whole_frames | |
| # A track starts where the clip does unless it says otherwise, so a sample that starts | |
| # later reads it from that far in. | |
| offset = start - (float(loaded["start"]) if "start" in loaded.files else 0.0) | |
| if offset < 0: | |
| raise ValueError(f"camera track starts {-offset:.2f} s after the sample does") | |
| if offset + seconds(frames - 1) > (total - 1) / track_fps + 1.0 / FPS: | |
| raise ValueError( | |
| f"camera track covers {(total - 1) / track_fps:.2f} s, the sample needs " | |
| f"{offset + seconds(frames - 1):.2f} s" | |
| ) | |
| picks = [ | |
| min(int(round((offset + seconds(index)) * track_fps)), total - 1) | |
| for index in range(frames) | |
| ] | |
| tensors["camera_c2w"] = views_all[picks].contiguous() | |
| intrinsics = (k_all if k_all.ndim == 2 else k_all[picks]).clone() | |
| # The intrinsics describe the clip as it was shot; the encoder sees a centre-crop of it. | |
| # A normalised focal length is pixels over the frame's own size, so cropping the frame | |
| # raises it by exactly the fraction cropped away - and a field of view that disagrees with | |
| # the picture is worse than no intrinsics at all. | |
| keep_h, keep_w = crop_scale(source_size(sample["target"]), (height, width)) | |
| if keep_h != 1.0 or keep_w != 1.0: | |
| intrinsics[..., 0, 0] /= keep_w | |
| intrinsics[..., 1, 1] /= keep_h | |
| # A long lens on a frame that had to be cropped hard ends up narrower than the released | |
| # contract can express (normalised focal at most 4, about a 14 degree vertical view). The | |
| # sample is dropped by name here rather than written out for the reader to reject later, | |
| # because a rejected member is a silent loss and a skipped sample is a counted one. | |
| focal = torch.stack([intrinsics[..., 0, 0].max(), intrinsics[..., 1, 1].max()]) | |
| if float(focal.max()) > 4.0 or float(focal.min()) <= 0.0: | |
| raise ValueError( | |
| f"cropped focal {float(focal.max()):.2f} is outside the contract's 0..4; " | |
| "this lens is narrower than the canvas can carry" | |
| ) | |
| tensors["camera_K"] = intrinsics.contiguous() | |
| if sample.get("known_mask"): | |
| # Which clock the mask runs on. One that came with the clip runs on the clip's: its frame j | |
| # is the clip at j / mask_fps, so it is put on 24 fps and cut to the window exactly as the | |
| # target was, and frame for frame it covers what the latents cover. One from the generated | |
| # library says what is given over the window, whatever the clip's length, and is stretched | |
| # over it. A dataset's mask marks the thing in the clip, so that is its default. | |
| clock = sample.get("known_mask_clock") or ( | |
| "clip" if sample.get("known_mask_marks_the_region") else "window" | |
| ) | |
| # white = keep (given), black = generate. A mask on the window's clock is read whole. | |
| if sample["known_mask"].endswith(".npz"): | |
| mask, mask_fps = label_mask( | |
| sample["known_mask"], int(sample["known_mask_label"]), clip_fps, start, | |
| num_frames / FPS + 0.5, | |
| ) | |
| skip = 0 # a label map is returned already cut at the start | |
| else: | |
| mask, mask_fps = read_video(sample["known_mask"], span if clock == "clip" else 0.0) | |
| skip = int(round(start * mask_fps)) | |
| if clock == "clip": | |
| # A mask of the clip's own shares its pixels, so the target's crop box cuts it too. | |
| mask = on_grid(region(mask[skip:], box), mask_fps)[:window] | |
| if len(mask) < window: # a mask a frame short of its clip holds its last frame | |
| mask = np.concatenate([mask, np.repeat(mask[-1:], window - len(mask), axis=0)]) | |
| # The target is centre-cropped to the canvas, so a mask that names a real region has to be | |
| # cropped the same way or it points at the wrong part of the picture. A mask already at the | |
| # canvas's aspect - every mask in the generated library is - comes back unchanged. | |
| mask = fit(mask, (height, width)) | |
| keep = torch.from_numpy(mask[..., 0].astype(np.float32) / 255.0)[None, None] | |
| # A mask that came with a dataset marks the thing, not what is kept around it. Saying so | |
| # on the sample is cheaper and more honest than re-encoding the file: the sample records | |
| # which way round its own mask is written, and the convention downstream stays one. | |
| if sample.get("known_mask_marks_the_region"): | |
| keep = 1.0 - keep | |
| grid = torch.nn.functional.interpolate( | |
| keep, | |
| size=(frames, target.shape[2] // 2, target.shape[3] // 2), | |
| mode="trilinear", | |
| align_corners=False, | |
| ) | |
| # A row is given only where the whole patch is: VAE fields leak across a mask edge. | |
| tensors["known_degrees"] = (grid[0, 0] > 0.98).to(torch.float16).reshape(-1).contiguous() | |
| meta = { | |
| "observations": observations, | |
| "latent_frames": frames, | |
| "canvas": [height, width], | |
| # How long the clip was and how much of it the latents hold, so a member that covers only | |
| # the start of a long clip says so. | |
| "clip_seconds": round(whole_frames / clip_fps, 3), | |
| "start_seconds": round(start, 3), | |
| "window_seconds": round(window / FPS, 3), | |
| "source": sample.get("source", ""), | |
| "crop": list(box) if box else None, | |
| # What this sample was made from, kept beside it. The text is the one thing that can be | |
| # rewritten later without re-encoding the pixels - a pass over the members recomputing | |
| # prompt_embeds costs the text encoder only - but that pass can only work if a member | |
| # still says which clip it came from. Without these two lines a rewrite would depend on | |
| # a sample list that has since been regenerated, and the corpus would have to be built | |
| # again to change a word of its text. | |
| "target": sample.get("target", ""), | |
| "media": [item.get("media", "") for item in items], | |
| "known_mask": sample.get("known_mask", ""), | |
| "text": sample["text"][:16000], | |
| } | |
| return tensors, meta | |
| # A row that failed is left alone for this long before anyone tries it again. Most failures are the | |
| # clip itself - a stream that will not decode, a clip too short to be a video - and fail the same way | |
| # every time; retried on every pass by every worker, a set's last thousand such rows held every | |
| # card idle. A day later they are tried once more, for the few that failed on a passing | |
| # hiccup (a read from storage that timed out, a list naming media still being unpacked). | |
| RETRY_FAILED_AFTER = float(os.environ.get("OMNI_RETRY_FAILED_AFTER", 86400)) | |
| def failed_recently(out: str, name: str) -> bool: | |
| """Whether a worker marked this row failed within the retry window: a single request.""" | |
| if "://" not in out: | |
| return False | |
| found = storage.stat(f"{out.rstrip('/')}/failed/{name}.txt") | |
| if found is None: | |
| return False | |
| return found[1] is None or time.time() - found[1] < RETRY_FAILED_AFTER | |
| def not_the_row(error: BaseException) -> bool: | |
| """Whether a failure says nothing about the row: the card ran out of memory or CUDA failed | |
| (another process on the same card is enough for either), or the storage refused the member the | |
| worker had already encoded. A row is left unmarked for these, and tried again next pass.""" | |
| text = str(error) | |
| return ( | |
| isinstance(error, (torch.OutOfMemoryError, WriteRefused)) | |
| or "CUDA" in text | |
| or "out of memory" in text | |
| ) | |
| def mark_failed(out: str, name: str, reason: str) -> None: | |
| """Leave a note in the output location that this row failed and why, so no one tries it again soon.""" | |
| if "://" in out: | |
| with contextlib.suppress(Exception): | |
| storage.write_bytes(f"{out.rstrip('/')}/failed/{name}.txt", reason.encode()) | |
| def written(out: str, name: str) -> bool: | |
| """Whether one member is in the output location now: a single request, asked just before encoding it. | |
| The listing a worker holds is minutes old, and in a set's last minutes that is long enough for | |
| a helper and an owner - or two helpers - to each encode what the other finished a moment ago. | |
| Asking about the one key costs a tenth of a second against half a minute of encoding. | |
| """ | |
| if "://" not in out: | |
| return False | |
| return storage.stat(f"{out.rstrip('/')}/samples/{name}.safetensors") is not None | |
| class Present: | |
| """The names the set's output location holds, kept fresh on a thread of its own. | |
| Two writers can hold the same row in their hands at once - a worker on its own shard and a | |
| sibling that finished early and came to help - and the only thing both can see is the output location. | |
| Listing it takes seconds once a set holds a few hundred thousand members, so the listing runs | |
| beside the encoder instead of in its way, and the encoder reads whatever the last one found. | |
| """ | |
| def __init__(self, out: str, every: float) -> None: | |
| self.out, self.every = out.rstrip("/"), every | |
| self.names: set[str] = set() | |
| self.failed: dict[str, float] = {} # row -> when it last failed | |
| self.ok = False # whether a listing has ever come back; until then it says nothing | |
| if "://" in self.out: | |
| self.refresh() | |
| threading.Thread(target=self.loop, daemon=True).start() | |
| def refresh(self) -> None: | |
| # A listing that failed (None) says nothing about what is there; keep the last one that worked. | |
| samples = storage.listing(f"{self.out}/samples/") | |
| if samples is not None: | |
| self.names = {name[: -len(".safetensors")] for name in samples if name.endswith(".safetensors")} | |
| self.ok = True | |
| marks = storage.listing(f"{self.out}/failed/") | |
| if marks is not None: | |
| now = time.time() | |
| failed = {name[: -len(".txt")]: (stamp or now) for name, stamp in marks.items() if name.endswith(".txt")} | |
| self.failed = {**self.failed, **failed} | |
| def loop(self) -> None: | |
| while True: | |
| time.sleep(self.every) | |
| try: | |
| self.refresh() | |
| except Exception: | |
| pass | |
| def add(self, name: str) -> None: | |
| self.names.add(name) | |
| def fail(self, name: str) -> None: | |
| self.failed[name] = time.time() | |
| def __contains__(self, name: str) -> bool: | |
| return name in self.names or time.time() - self.failed.get(name, 0.0) < RETRY_FAILED_AFTER | |
| def local_weights(root: str) -> str: | |
| """The encoder's weights on this node's own disk, copied there once; the path to load from. | |
| Every pass starts a fresh set of workers, and the weights (74 GB) may live on a shared network | |
| filesystem whose client cache goes cold between passes - so each pass would read them all again, | |
| and workers starting their passes together leave every card idle while they load. A local copy stays in this node's page cache from one pass to the next. The copy is made | |
| by whichever worker gets the lock first; the rest wait for it. Only what the encoder uses is | |
| copied - not the 62 GB transformer - and the index is pointed at the copy. | |
| """ | |
| import fcntl | |
| import re | |
| import shutil | |
| local = os.environ.get("OMNI_LOCAL_WEIGHTS", "/tmp/h3_weights") | |
| if not local: | |
| return root | |
| ready = f"{local}/.ready" | |
| if os.path.exists(ready): | |
| return f"{local}/ref" | |
| os.makedirs(local, exist_ok=True) | |
| with open(f"{local}.lock", "w") as lock: | |
| fcntl.flock(lock, fcntl.LOCK_EX) | |
| if os.path.exists(ready): | |
| return f"{local}/ref" | |
| index = open(f"{root}/modular_model_index.json").read() | |
| pipes = sorted(set(re.findall(r'"pretrained_model_name_or_path":\s*"([^"]+)"', index)) - {root}) | |
| for src in pipes: | |
| dst = f"{local}/pipe{len(pipes) > 1 and '_' + str(pipes.index(src)) or ''}" | |
| names = [n for n in os.listdir(src) if n != "transformer"] | |
| # Parallel copies: one stream from a network filesystem is a fraction of what it can give. | |
| subprocess.run( | |
| ["bash", "-c", f"mkdir -p {dst} && cd {src} && printf '%s\\n' " | |
| + " ".join(f"'{n}'" for n in names) | |
| + f" | xargs -P 16 -I{{}} cp -r {{}} {dst}/"], | |
| check=True, | |
| ) | |
| index = index.replace(f'"{src}"', f'"{dst}"') | |
| shutil.rmtree(f"{local}/ref", ignore_errors=True) | |
| os.makedirs(f"{local}/ref") | |
| for name in os.listdir(root): | |
| if name.startswith("transformer"): | |
| continue # the diffusion transformer is not the encoder's; its entry keeps the original path | |
| if os.path.isdir(f"{root}/{name}"): | |
| shutil.copytree(f"{root}/{name}", f"{local}/ref/{name}") | |
| else: | |
| shutil.copy(f"{root}/{name}", f"{local}/ref/{name}") | |
| open(f"{local}/ref/modular_model_index.json", "w").write(index) | |
| open(ready, "w").write("ok\n") | |
| return f"{local}/ref" | |
| def main() -> None: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--samples", required=True, help="sample JSONL, local or remote path") | |
| parser.add_argument("--out", required=True, help="output prefix (remote object storage or a local directory) for samples/ and index shards") | |
| parser.add_argument("--shard", type=int, default=int(os.environ.get("RANK", 0))) | |
| parser.add_argument("--shards", type=int, default=int(os.environ.get("WORLD_SIZE", 1))) | |
| parser.add_argument("--shard-offset", type=int, default=0, help="global shard = offset + rank") | |
| parser.add_argument( | |
| "--no-help", | |
| action="store_true", | |
| help="stop when this shard is done instead of taking up what is left", | |
| ) | |
| # The length a clip is encoded at is the one thing about a member that cannot be changed | |
| # afterwards without encoding it again, and generation asks for 47 latents (158 frames). | |
| # Training on 37 would hand the rolling schedule a shorter window than it is asked to | |
| # produce. This is a cap, not a demand: a clip with fewer frames yields fewer latents. | |
| parser.add_argument("--num-frames", type=int, default=158) | |
| # A pass that runs until its rows are gone holds the node for as long as the largest set lasts, | |
| # and a set that matters more - one still being derived - waits the whole time. With a limit | |
| # the node stops taking rows, goes back to the top of its queue, and the first set with work | |
| # gets it; whatever it did not reach is still unwritten and is picked up next time. | |
| parser.add_argument("--max-seconds", type=float, default=0, help="stop taking rows after this") | |
| args = parser.parse_args() | |
| from safetensors.torch import save_file | |
| local = int(os.environ.get("LOCAL_RANK", 0)) | |
| device = f"cuda:{local}" | |
| torch.cuda.set_device(local) | |
| torch.zeros(1, device=device) # bind the card at once, before the long load | |
| args.shard += args.shard_offset | |
| with open( | |
| fetch_fresh(args.samples, f"{SCRATCH}/samples.{os.environ.get('RANK', 0)}.jsonl") | |
| ) as lines: | |
| rows = [json.loads(line) for line in lines if line.strip()] | |
| # Shard by the sample's own name, not by where it sits in the list. A list that grows while | |
| # the set is being encoded - a signal still being derived, a pool still being unpacked - puts | |
| # new rows in the middle, and a position-based split would hand rows that were already done to | |
| # a different worker, which would encode them a second time. | |
| def owner(name: str) -> int: | |
| return int(hashlib.sha1(name.encode()).hexdigest()[:8], 16) % args.shards | |
| mine = [row for row in rows if owner(row["id"]) == args.shard] | |
| # Whose rows this worker will pick up once its own are done. A shard is an equal count of | |
| # clips, not an equal amount of work - a 1400-frame clip costs ten times a 130-frame one - so | |
| # the fast workers of a node finish long before the slow ones and their cards sit idle for the | |
| # rest of the pass. Rather than balance the split, which cannot be done without knowing the | |
| # cost in advance, a finished worker goes looking for rows nobody has written yet. | |
| # Strided over the whole list by the worker's place among all the set's workers, not among | |
| # its own node's eight. Striding by the node rank gives worker 3 of every node the same rows, and | |
| # two nodes in their tail at once then walk them side by side, each writing every row the | |
| # other writes. By the global shard the helpers' rows are disjoint across the whole set. | |
| others = [ | |
| row | |
| for index, row in enumerate(rows) | |
| if index % args.shards == args.shard and owner(row["id"]) != args.shard | |
| ] | |
| sink = Sink(args.out, args.shard) | |
| # What this shard wrote, and what the set already holds however it got there. Asking the | |
| # storage as well as the index means a row another worker has already written is never | |
| # written again - whether that worker was a sibling helping out, or this set's split was a | |
| # different number of shards last time it ran. | |
| # Asked again every few minutes, not only at the start: helpers walk this shard's rows from | |
| # the far end while it walks them from the near end, and whatever a helper has already | |
| # written this shard must not write again when it gets there. | |
| present = Present(args.out, float(os.environ.get("OMNI_PRESENT_EVERY", 600))) | |
| # The storage is the account of what is done; this shard's index is only what some worker once | |
| # said it wrote under this shard number. A member moved out of the set to be encoded again is | |
| # still named in an index, and so is one written when the set ran under another split whose | |
| # shard numbers overlap this one's. Only when the storage cannot be listed is the index used. | |
| done = set() if present.ok else sink.done() | |
| global ROOT | |
| try: | |
| ROOT = local_weights(ROOT) | |
| except Exception: # a node whose disk will not take the copy still encodes, from the original path | |
| print(f"local weights unavailable: {traceback.format_exc()[-300:]}", file=sys.stderr, flush=True) | |
| pipes = build_pipelines(device) | |
| start, count = time.time(), 0 | |
| with open(sink.index_path, "a") as index, torch.no_grad(): | |
| for row in mine: | |
| if args.max_seconds and time.time() - start > args.max_seconds: | |
| break | |
| if row["id"] in done or row["id"] in present: | |
| continue | |
| if written(args.out, row["id"]): | |
| present.add(row["id"]) | |
| continue | |
| if failed_recently(args.out, row["id"]): | |
| present.fail(row["id"]) | |
| continue | |
| try: | |
| tensors, meta = encode(row, pipes, args.num_frames) | |
| name = f"samples/{row['id']}.safetensors" | |
| member = sink.member(name) | |
| os.makedirs(os.path.dirname(member), exist_ok=True) | |
| save_file(tensors, member, metadata={"omni": json.dumps(meta)}) | |
| sink.publish(name) | |
| index.write( | |
| json.dumps( | |
| { | |
| "id": row["id"], | |
| "member": name, | |
| **{k: meta[k] for k in ("latent_frames", "canvas", "source")}, | |
| "observations": len(meta["observations"]), | |
| } | |
| ) | |
| + "\n" | |
| ) | |
| index.flush() | |
| count += 1 | |
| if count % 20 == 0: | |
| sink.publish_index() | |
| sweep(f"{SCRATCH}/in") | |
| print( | |
| f"[{args.shard}/{args.shards}] {count} encoded, " | |
| f"{(time.time() - start) / count:.1f} s/sample", | |
| flush=True, | |
| ) | |
| except Exception as error: # one bad sample must not stop the shard | |
| reason = traceback.format_exc()[-600:] | |
| print(f"SKIP {row['id']}: {reason}", file=sys.stderr, flush=True) | |
| if not_the_row(error): | |
| torch.cuda.empty_cache() | |
| else: | |
| mark_failed(args.out, row["id"], reason) | |
| present.fail(row["id"]) | |
| sink.publish_index() | |
| print(f"[{args.shard}/{args.shards}] DONE {count} new samples, helping", flush=True) | |
| # The helper pass. It walks the other shards' rows from the end of the list backwards. Their | |
| # owners are walking the same rows forwards, so a helper takes the rows its owner would reach | |
| # last and the two meet once, in the middle, instead of racing down the same rows side by side | |
| # - which is what a helper walking forwards did, writing each row just before its owner wrote | |
| # it again. Both sides skip whatever the storage already holds; a row encoded twice is never | |
| # wrong, only wasted, and meeting in the middle keeps it to the rows in flight when they meet. | |
| # Then everything else nobody has written yet, so a worker that is done does not sit on an idle | |
| # card while the pass waits for its slowest sibling. Each worker starts at its own place in | |
| # that list, so the ones that arrive together spread out instead of queueing on the same row. | |
| stride = {id(row) for row in others} | |
| rest = [row for row in reversed(rows) if owner(row["id"]) != args.shard and id(row) not in stride] | |
| turn = (args.shard * len(rest)) // max(args.shards, 1) | |
| order = list(reversed(others)) + rest[turn:] + rest[:turn] | |
| helped = 0 | |
| if order and not args.no_help: | |
| with open(sink.index_path, "a") as index, torch.no_grad(): | |
| for row in order: | |
| if args.max_seconds and time.time() - start > args.max_seconds: | |
| break | |
| if row["id"] in present or row["id"] in done: | |
| continue | |
| if written(args.out, row["id"]): | |
| present.add(row["id"]) | |
| continue | |
| if failed_recently(args.out, row["id"]): | |
| present.fail(row["id"]) | |
| continue | |
| try: | |
| tensors, meta = encode(row, pipes, args.num_frames) | |
| member = sink.member(f"samples/{row['id']}.safetensors") | |
| os.makedirs(os.path.dirname(member), exist_ok=True) | |
| save_file(tensors, member, metadata={"omni": json.dumps(meta)}) | |
| sink.publish(f"samples/{row['id']}.safetensors") | |
| index.write( | |
| json.dumps( | |
| { | |
| "id": row["id"], | |
| "member": f"samples/{row['id']}.safetensors", | |
| "latent_frames": meta["latent_frames"], | |
| "canvas": meta["canvas"], | |
| "source": meta["source"], | |
| "observations": len(meta["observations"]), | |
| } | |
| ) | |
| + "\n" | |
| ) | |
| index.flush() | |
| present.add(row["id"]) | |
| helped += 1 | |
| if helped % 20 == 0: | |
| sink.publish_index() | |
| sweep(f"{SCRATCH}/in") | |
| print(f"[{args.shard}/{args.shards}] helped {helped}", flush=True) | |
| except Exception as error: | |
| reason = traceback.format_exc()[-600:] | |
| print(f"SKIP {row['id']}: {reason}", file=sys.stderr, flush=True) | |
| if not_the_row(error): | |
| torch.cuda.empty_cache() | |
| else: | |
| mark_failed(args.out, row["id"], reason) | |
| present.fail(row["id"]) | |
| sink.publish_index() | |
| print(f"[{args.shard}/{args.shards}] DONE {count} of its own, {helped} of others", flush=True) | |
| if __name__ == "__main__": | |
| main() | |