#!/usr/bin/env python3 """ Precompute VAE latents for Context-as-Memory dataset with separate ctx and target storage: - ctx: 1 latent per frame (for context/memory frames) - target: 1 latent per 4 frames (time_division_factor=4) 8-GPU distributed: each rank processes a subset of segments. Two modes: 1. With metadata: --metadata_path metadata_full.csv (uses VideoDataset) 2. No metadata: --no_metadata - auto-discovers segments from frames/ + captions.txt (or captions.jsonl) Usage (8 GPUs, no metadata): accelerate launch --num_processes 8 scripts/precompute_ctx_target_latents.py \\ --dataset_base_path /path/to/Context-as-Memory-Dataset \\ --output_dir /path/to/latents \\ --model_paths '["dit.safetensors","t5.pth","VAE.pth"]' \\ --no_metadata """ import argparse import json import os import sys import warnings from datetime import datetime import torch from PIL import Image from tqdm import tqdm # Add project root for imports SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) PROJECT_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, "..")) if PROJECT_ROOT not in sys.path: sys.path.insert(0, PROJECT_ROOT) from accelerate import Accelerator from accelerate.utils import set_seed from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig from diffsynth.trainers.utils import VideoDataset def load_captions_txt(captions_path): """Load captions.txt: video_name/start_end.mp4\\tcaption -> video_name -> caption.""" captions = {} if not os.path.isfile(captions_path): return captions with open(captions_path, "r", encoding="utf-8") as f: for line in f: parts = line.strip().split("\t", 1) if len(parts) != 2: continue video_path, caption = parts[0], parts[1] video_name = video_path.split("/")[0] if video_name not in captions: captions[video_name] = caption return captions def load_captions_jsonl(captions_path): """Load captions.jsonl: each line {"video_name": "...", "prompt": "..."} or similar.""" captions = {} if not os.path.isfile(captions_path): return captions with open(captions_path, "r", encoding="utf-8") as f: for line in f: try: obj = json.loads(line) vn = obj.get("video_name") or obj.get("video") or obj.get("id", "") prompt = obj.get("prompt") or obj.get("caption") or obj.get("text", "") if vn: captions[vn] = prompt except json.JSONDecodeError: continue return captions def build_segments_from_frames( frames_dir, captions_path=None, captions_jsonl_path=None, num_frames=81, segment_stride=None, overlap_labels_dir=None, overlap_labels_dense=False, ): """ Build segment list without metadata CSV. Modes: 1. overlap_labels_dir + dense=False: 1 segment per overlap_labels JSON (~240k) 2. overlap_labels_dir + dense=True: stride=1 within each video (~10x-40x more) 3. else: stride-based from frames/ - stride 40 ~19k, stride 1 ~760k Returns: [(video_name, start_frame, end_frame, frame_paths, prompt), ...] """ captions = {} if captions_path: captions = load_captions_txt(captions_path) if captions_jsonl_path: captions.update(load_captions_jsonl(captions_jsonl_path)) default_prompt = "A video scene." segments = [] stride = segment_stride if segment_stride is not None else max(1, num_frames // 2) if overlap_labels_dir and os.path.isdir(overlap_labels_dir): video_dirs = sorted( [d for d in os.listdir(overlap_labels_dir) if os.path.isdir(os.path.join(overlap_labels_dir, d))] ) for video_name in video_dirs: video_frames_dir = os.path.join(frames_dir, video_name) if not os.path.isdir(video_frames_dir): continue frame_files = sorted([f for f in os.listdir(video_frames_dir) if f.endswith(".png")]) if len(frame_files) < num_frames: continue prompt = captions.get(video_name, default_prompt) if overlap_labels_dense: # Dense: stride=1 for max segments (~10x-40x more). Overridable via segment_stride. seg_stride = segment_stride if segment_stride is not None else 1 else: # 1 segment per overlap_labels JSON video_overlap_dir = os.path.join(overlap_labels_dir, video_name) json_files = sorted([f for f in os.listdir(video_overlap_dir) if f.endswith(".json")]) for jf in json_files: try: start_frame = int(jf.replace(".json", "")) except ValueError: continue end_frame = start_frame + num_frames - 1 frame_paths = [ os.path.join(video_name, f"{start_frame + i:04d}.png") for i in range(num_frames) ] first_path = os.path.join(frames_dir, frame_paths[0]) last_path = os.path.join(frames_dir, frame_paths[-1]) if os.path.isfile(first_path) and os.path.isfile(last_path): segments.append((video_name, start_frame, end_frame, frame_paths, prompt)) continue for start in range(0, len(frame_files) - num_frames + 1, seg_stride): end = start + num_frames - 1 frame_paths = [os.path.join(video_name, frame_files[i]) for i in range(start, end + 1)] segments.append((video_name, start, end, frame_paths, prompt)) else: # Stride-based from frames/ video_dirs = sorted([d for d in os.listdir(frames_dir) if os.path.isdir(os.path.join(frames_dir, d))]) for video_name in video_dirs: video_dir = os.path.join(frames_dir, video_name) frame_files = sorted([f for f in os.listdir(video_dir) if f.endswith(".png")]) if len(frame_files) < num_frames: continue prompt = captions.get(video_name, default_prompt) for start in range(0, len(frame_files) - num_frames + 1, stride): end = start + num_frames - 1 frame_paths = [os.path.join(video_name, frame_files[i]) for i in range(start, end + 1)] segments.append((video_name, start, end, frame_paths, prompt)) return segments class FrameSegmentDataset(torch.utils.data.Dataset): """Dataset that loads frames from segment list (no metadata CSV).""" def __init__(self, base_path, segments, height, width): self.base_path = base_path self.frames_dir = os.path.join(base_path, "frames") self.segments = segments self.height = height self.width = width def __len__(self): return len(self.segments) def _load_image(self, rel_path): path = os.path.join(self.frames_dir, rel_path) img = Image.open(path).convert("RGB") import torchvision.transforms.functional as TF w, h = img.size scale = max(self.width / w, self.height / h) img = TF.resize(img, (round(h * scale), round(w * scale)), interpolation=TF.InterpolationMode.BILINEAR) img = TF.center_crop(img, (self.height, self.width)) return img def __getitem__(self, idx): video_name, start_frame, end_frame, frame_paths, prompt = self.segments[idx] frames = [] for fp in frame_paths: try: frames.append(self._load_image(fp)) except Exception: return None if len(frames) != len(frame_paths): return None return { "video": frames, "prompt": prompt, "video_name": video_name, "start_frame": start_frame, "end_frame": end_frame, } def parse_args(): parser = argparse.ArgumentParser( description="Precompute ctx (1 latent/frame) and target (1 latent/4 frames) latents." ) parser.add_argument( "--dataset_base_path", type=str, required=True, help="Dataset root (contains frames/, metadata).", ) parser.add_argument( "--metadata_path", type=str, default=None, help="Metadata CSV path. Omit when using --no_metadata.", ) parser.add_argument( "--no_metadata", action="store_true", help="Skip metadata CSV; auto-discover segments from frames/ + captions.", ) parser.add_argument( "--captions_path", type=str, default=None, help="captions.txt path (for --no_metadata). Default: {dataset_base_path}/captions.txt", ) parser.add_argument( "--captions_jsonl_path", type=str, default=None, help="Optional captions.jsonl path (for --no_metadata).", ) parser.add_argument( "--segment_stride", type=int, default=None, help="Stride between segments. Default: 40 (stride mode), 1 (--overlap_labels_dense). Use 1 for max.", ) parser.add_argument( "--use_overlap_labels", action="store_true", help="When --no_metadata: use overlap_labels/ to discover segments (matches metadata_full ~240k).", ) parser.add_argument( "--overlap_labels_dense", action="store_true", help="With --use_overlap_labels: stride=1 per video (~10x-40x more segments).", ) parser.add_argument( "--output_dir", type=str, required=True, help="Output directory. Will create ctx_latents/ and target_latents/ subdirs.", ) parser.add_argument( "--model_paths", type=str, required=True, help='JSON array of model paths, e.g. \'["dit.safetensors","t5.pth","Wan2.1_VAE.pth"]\'.', ) parser.add_argument("--tokenizer_path", type=str, default=None) parser.add_argument("--height", type=int, default=352) parser.add_argument("--width", type=int, default=640) parser.add_argument("--num_frames", type=int, default=81) parser.add_argument( "--context_frames", type=int, default=5, help="Number of context frames (each gets 1 latent).", ) parser.add_argument( "--target_frames_per_latent", type=int, default=4, help="Target: 1 latent per N frames (default 4).", ) parser.add_argument("--action_base_path", type=str, default=None) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--skip_existing", action="store_true") parser.add_argument( "--encode_batch_size", type=int, default=8, help="Batch size for VAE encode (frames per call). Higher = better GPU util, more VRAM.", ) parser.add_argument( "--segment_batch_size", type=int, default=8, help="Process N segments per iteration (DataLoader batch). Speeds up I/O + encode.", ) return parser.parse_args() def make_dataset_args(args): return argparse.Namespace( dataset_base_path=args.dataset_base_path, dataset_metadata_path=args.metadata_path, height=args.height, width=args.width, max_pixels=1920 * 1080, num_frames=args.num_frames, dataset_repeat=1, data_file_keys="video,video_name,start_frame,end_frame", action_base_path=args.action_base_path or args.dataset_base_path, ) def crop_and_resize(image, target_height, target_width): import torchvision.transforms.functional as TF width, height = image.size scale = max(target_width / width, target_height / height) image = TF.resize( image, (round(height * scale), round(width * scale)), interpolation=TF.InterpolationMode.BILINEAR, ) image = TF.center_crop(image, (target_height, target_width)) return image def main(): args = parse_args() set_seed(args.seed) if not args.no_metadata and not args.metadata_path: raise ValueError("Either --metadata_path or --no_metadata is required.") accelerator = Accelerator() if accelerator.num_processes != 8 and accelerator.is_main_process: print(f"Info: using {accelerator.num_processes} processes (expected 8).") ctx_dir = os.path.join(args.output_dir, "ctx_latents") target_dir = os.path.join(args.output_dir, "target_latents") os.makedirs(ctx_dir, exist_ok=True) os.makedirs(target_dir, exist_ok=True) # Redirect rank 0 stdout to log file in output_dir _rank0_log_file = None if accelerator.is_main_process: log_path = os.path.join(args.output_dir, "precompute_log.txt") _rank0_log_file = open(log_path, "w", encoding="utf-8") _rank0_log_file.write(f"[{datetime.now().isoformat()}] precompute_ctx_target_latents started\n") _rank0_log_file.write(f"output_dir={args.output_dir}\n") _rank0_log_file.flush() class _Tee: def __init__(self, *files): self.files = files def write(self, obj): for f in self.files: f.write(obj) f.flush() def flush(self): for f in self.files: f.flush() sys.stdout = _Tee(sys.__stdout__, _rank0_log_file) if args.no_metadata: frames_dir = os.path.join(args.dataset_base_path, "frames") captions_path = args.captions_path or os.path.join(args.dataset_base_path, "captions.txt") overlap_labels_dir = None if args.use_overlap_labels: overlap_labels_dir = os.path.join(args.dataset_base_path, "overlap_labels") segments = build_segments_from_frames( frames_dir, captions_path=captions_path, captions_jsonl_path=args.captions_jsonl_path, num_frames=args.num_frames, segment_stride=args.segment_stride, overlap_labels_dir=overlap_labels_dir, overlap_labels_dense=args.overlap_labels_dense, ) dataset = FrameSegmentDataset( args.dataset_base_path, segments, args.height, args.width ) if accelerator.is_main_process: if overlap_labels_dir: src = "overlap_labels (dense)" if args.overlap_labels_dense else "overlap_labels" else: src = "frames (stride)" print(f"No-metadata mode: discovered {len(segments)} segments from {src}") else: dataset_args = make_dataset_args(args) dataset = VideoDataset(args=dataset_args) total = len(dataset) if total == 0: if accelerator.is_main_process: print("Dataset is empty. Exit.") return sampler = torch.utils.data.DistributedSampler( dataset, num_replicas=accelerator.num_processes, rank=accelerator.process_index, shuffle=False, drop_last=False, ) indices = list(sampler) n_local = len(indices) if accelerator.is_main_process: print(f"Dataset size: {total}. Rank 0 processing {n_local} indices.") meta = { "dataset_base_path": args.dataset_base_path, "metadata_path": args.metadata_path, "no_metadata": args.no_metadata, "total_samples": total, "height": args.height, "width": args.width, "num_frames": args.num_frames, "context_frames": args.context_frames, "target_frames_per_latent": args.target_frames_per_latent, } with open(os.path.join(args.output_dir, "metadata_precompute.json"), "w") as f: json.dump(meta, f, indent=2) # Load pipeline (VAE only) model_paths = json.loads(args.model_paths) model_configs = [ModelConfig(path=p) for p in model_paths] from_pretrained_kw = { "torch_dtype": torch.bfloat16, "device": "cpu", "model_configs": model_configs, } if args.tokenizer_path: from_pretrained_kw["tokenizer_config"] = ModelConfig(path=args.tokenizer_path) if accelerator.is_main_process: print("Loading pipeline (VAE)...") pipe = WanVideoPipeline.from_pretrained(**from_pretrained_kw) pipe.vae.to(accelerator.device) pipe.vae.eval() K = args.context_frames step = args.target_frames_per_latent def preprocess_frames(frames): return pipe.preprocess_video(frames) @torch.no_grad() def encode_frame(frame_pil): """Encode single frame -> (C, 1, H//8, W//8)""" vid = preprocess_frames([frame_pil]) if vid.dim() == 5: vid = vid.squeeze(0) lat = pipe.vae.encode([vid], device=accelerator.device, tiled=False, tile_size=None, tile_stride=None) return lat[0].cpu() failed = 0 skipped = 0 for idx in tqdm( indices, desc=f"Rank {accelerator.process_index}", disable=not accelerator.is_local_main_process, ): sample = dataset[idx] ctx_path = os.path.join(ctx_dir, f"{idx:08d}.pt") target_path = os.path.join(target_dir, f"{idx:08d}.pt") if args.skip_existing and os.path.isfile(ctx_path) and os.path.isfile(target_path): skipped += 1 continue try: if sample is None: failed += 1 continue video_frames = sample.get("video") if not video_frames or len(video_frames) == 0: failed += 1 continue if len(video_frames) != args.num_frames: if len(video_frames) > args.num_frames: video_frames = video_frames[: args.num_frames] else: last = video_frames[-1] if video_frames else None while len(video_frames) < args.num_frames and last is not None: video_frames = video_frames + [last] if len(video_frames) < args.num_frames: failed += 1 continue # Context: 1 latent per frame (frames 0..K-1) ctx_latents_list = [] for i in range(min(K, len(video_frames))): lat = encode_frame(video_frames[i]) if isinstance(lat, (list, tuple)): lat = lat[0] ctx_latents_list.append(lat) ctx_latent = torch.cat(ctx_latents_list, dim=1) if ctx_latent.dim() == 5: ctx_latent = ctx_latent.squeeze(0) # Target: 1 latent per `step` frames (frames K, K+step, ...) target_indices = list(range(K, len(video_frames), step)) target_latents_list = [] for i in target_indices: lat = encode_frame(video_frames[i]) if isinstance(lat, (list, tuple)): lat = lat[0] target_latents_list.append(lat) if not target_latents_list: failed += 1 continue target_latent = torch.cat(target_latents_list, dim=1) if target_latent.dim() == 5: target_latent = target_latent.squeeze(0) save_meta = { "prompt": sample.get("prompt", ""), "video_name": sample.get("video_name"), "start_frame": sample.get("start_frame"), "end_frame": sample.get("end_frame"), } if "actions" in sample and sample["actions"] is not None: a = sample["actions"] save_meta["actions"] = torch.tensor(a) if not isinstance(a, torch.Tensor) else a.cpu() torch.save({"latent": ctx_latent, **save_meta}, ctx_path) torch.save({"latent": target_latent, **save_meta}, target_path) except Exception as e: if accelerator.is_local_main_process: tqdm.write(f"Rank {accelerator.process_index} idx {idx}: {e}") failed += 1 accelerator.wait_for_everyone() if accelerator.is_main_process: print( f"Precompute done. ctx_latents/ and target_latents/ under {args.output_dir}. " f"Failed: {failed}, Skipped: {skipped}." ) if _rank0_log_file is not None: sys.stdout = sys.__stdout__ _rank0_log_file.close() if __name__ == "__main__": main()