Meissonic/train/extract_features.py

#!/usr/bin/env python3
"""
Extract video codes and text embeddings from video-text pairs for efficient training.

This script pre-extracts:
1. Video codes: Discrete tokens from CosmosVideoTokenizer
2. Text embeddings: Encoder hidden states from T5/UMT5

The extracted features are saved to disk and can be loaded directly during training,
avoiding repeated encoding operations.
"""

import argparse
import os
import sys
import logging
from pathlib import Path
from tqdm import tqdm
import torch
import numpy as np
from torch.utils.data import DataLoader, DistributedSampler
import json

sys.path.append(os.getcwd())

from train.dataset_utils import OpenVid1MDataset, tokenize_prompt, encode_prompt
from src.pipeline_video import CosmosVideoTokenizer
from transformers import T5Tokenizer, T5EncoderModel
from accelerate import Accelerator

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%m/%d/%Y %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)


def get_hierarchical_path(base_dir, index):
    """
    Get hierarchical path for storing features in 3-level directory structure.
    
    Structure: base_dir/level1/level2/level3/filename.npy
    - level1: index // 1000000 (0-999)
    - level2: (index // 1000) % 1000 (0-999)
    - level3: index % 1000 (0-999)
    
    Args:
        base_dir: Base directory for features
        index: Sample index
        
    Returns:
        Full path to the file
    """
    level1 = index // 1000000
    level2 = (index // 1000) % 1000
    level3 = index % 1000
    
    dir_path = os.path.join(
        base_dir,
        f"{level1:03d}",
        f"{level2:03d}",
        f"{level3:03d}"
    )
    file_path = os.path.join(dir_path, f"{index:08d}.npy")
    
    return dir_path, file_path


def parse_args():
    parser = argparse.ArgumentParser(description="Extract video codes and text embeddings")
    
    parser.add_argument(
        "--csv_path",
        type=str,
        required=True,
        help="Path to OpenVid1M CSV file",
    )
    parser.add_argument(
        "--video_root_dir",
        type=str,
        default=None,
        help="Root directory containing video files. If None, will auto-detect.",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        required=True,
        help="Output directory to save extracted features",
    )
    parser.add_argument(
        "--text_encoder_architecture",
        type=str,
        default="umt5-base",
        choices=["umt5-base", "umt5-xxl", "t5"],
        help="Text encoder architecture",
    )
    parser.add_argument(
        "--video_tokenizer_model_id",
        type=str,
        default="Cosmos-1.0-Tokenizer-DV8x16x16",
        help="HuggingFace model ID for Cosmos video tokenizer",
    )
    parser.add_argument(
        "--num_frames",
        type=int,
        default=16,
        help="Number of frames per video",
    )
    parser.add_argument(
        "--video_height",
        type=int,
        default=480,
        help="Video height",
    )
    parser.add_argument(
        "--video_width",
        type=int,
        default=848,
        help="Video width",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=4,
        help="Batch size for feature extraction",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=4,
        help="Number of dataloader workers",
    )
    parser.add_argument(
        "--max_samples",
        type=int,
        default=None,
        help="Maximum number of samples to process (for testing). If None, process all.",
    )
    parser.add_argument(
        "--resume_from_index",
        type=int,
        default=0,
        help="Resume extraction from this index (for resuming interrupted extraction)",
    )
    parser.add_argument(
        "--prompt_prefix",
        type=str,
        default=None,
        help="Prefix to add to prompts",
    )
    parser.add_argument(
        "--extract_video",
        action="store_true",
        default=False,
        help="Extract video codes. Enable this flag to dump video codes.",
    )
    parser.add_argument(
        "--extract_text",
        action="store_true",
        default=False,
        help="Extract text embeddings. Enable this flag to dump text embeddings.",
    )
    
    return parser.parse_args()


def main():
    args = parse_args()
    
    # Initialize accelerator
    accelerator = Accelerator()
    
    # Log GPU info
    logger.info(f"Process {accelerator.process_index}/{accelerator.num_processes} on device {accelerator.device}")
    
    if accelerator.is_main_process:
        os.makedirs(args.output_dir, exist_ok=True)
        logger.info(f"Output directory: {args.output_dir}")
        logger.info(f"Using {accelerator.num_processes} GPUs for parallel extraction")
        logger.info(f"Extract video codes: {args.extract_video}")
        logger.info(f"Extract text embeddings: {args.extract_text}")
    
    # Validate that at least one feature type is selected
    if not args.extract_video and not args.extract_text:
        raise ValueError("At least one feature type must be enabled. Use --extract_video and/or --extract_text.")
    
    device = accelerator.device
    dtype = torch.float32
    
    # Initialize text encoder (only if needed)
    text_encoder = None
    tokenizer = None
    if args.extract_text:
        logger.info(f"Loading text encoder: {args.text_encoder_architecture}")
        if args.text_encoder_architecture == "umt5-base":
            model_id = "google/umt5-base"
        elif args.text_encoder_architecture == "umt5-xxl":
            model_id = "google/umt5-xxl"
        elif args.text_encoder_architecture == "t5":
            model_id = "t5-base"
        else:
            raise ValueError(f"Unknown text encoder: {args.text_encoder_architecture}")
        
        text_encoder = T5EncoderModel.from_pretrained(model_id)
        tokenizer = T5Tokenizer.from_pretrained(model_id)
        text_encoder.to(device=device, dtype=dtype)
        text_encoder.eval()
        text_encoder.requires_grad_(False)
        
        # Extract empty_embeds for conditional dropout (only on main process)
        if accelerator.is_main_process:
            logger.info("Extracting empty_embeds for conditional dropout...")
            with torch.no_grad():
                empty_input_ids = tokenize_prompt(tokenizer, "", args.text_encoder_architecture)
                empty_input_ids = empty_input_ids.to(device)
                empty_embeds, _ = encode_prompt(
                    text_encoder,
                    empty_input_ids,
                    args.text_encoder_architecture
                )
            
            # Convert to CPU and save as .npy (more space-efficient than .pt)
            empty_embeds_np = empty_embeds.cpu().numpy().astype(np.float16)  # Save as float16 to save space
            empty_embeds_path = os.path.join(args.output_dir, "empty_embeds.npy")
            np.save(empty_embeds_path, empty_embeds_np)
            logger.info(f"Saved empty_embeds to: {empty_embeds_path}")
            logger.info(f"  Shape: {empty_embeds_np.shape}, dtype: {empty_embeds_np.dtype}")
    else:
        logger.info("Skipping text encoder loading (--no_extract_text)")
        # Still need tokenizer for dataset if video extraction needs prompts
        # But if only extracting video, we might not need text processing
        # For now, we'll still create a minimal tokenizer for dataset compatibility
        if args.extract_video:
            # Use a simple tokenizer just for dataset compatibility
            tokenizer = T5Tokenizer.from_pretrained("google/umt5-base")
    
    # Initialize video tokenizer (only if needed)
    video_tokenizer = None
    if args.extract_video:
        logger.info(f"Loading video tokenizer: {args.video_tokenizer_model_id}")
        video_tokenizer = CosmosVideoTokenizer(
            model_id=args.video_tokenizer_model_id,
            device=device,
            dtype=dtype
        )
        video_tokenizer.requires_grad_(False)
        video_tokenizer.eval()
    else:
        logger.info("Skipping video tokenizer loading (--no_extract_video)")
    
    # Auto-detect video_root_dir if not provided
    if args.video_root_dir is None:
        csv_dir = os.path.dirname(args.csv_path)
        if os.path.exists(os.path.join(csv_dir, 'video_reorg')):
            video_root_dir = os.path.join(csv_dir, 'video_reorg')
        elif os.path.exists(os.path.join(os.path.dirname(csv_dir), 'video_reorg')):
            video_root_dir = os.path.join(os.path.dirname(csv_dir), 'video_reorg')
        else:
            video_root_dir = csv_dir
            logger.warning(f"Video directory not found, using CSV directory: {video_root_dir}")
    else:
        video_root_dir = args.video_root_dir
    
    # Create dataset
    # For feature extraction, use fixed temporal crop (from beginning) and center crop for consistency
    # Training will use random temporal crop and random crop for data augmentation
    dataset = OpenVid1MDataset(
        csv_path=args.csv_path,
        video_root_dir=video_root_dir,
        tokenizer=tokenizer,
        num_frames=args.num_frames,
        height=args.video_height,
        width=args.video_width,
        text_encoder_architecture=args.text_encoder_architecture,
        prompt_prefix=args.prompt_prefix,
        use_random_temporal_crop=False,  # Fixed sampling for feature extraction consistency
        use_random_crop=False,  # Center crop for feature extraction consistency
    )
    
    # Limit dataset size if specified
    if args.max_samples is not None:
        dataset.data = dataset.data[:args.max_samples]
        logger.info(f"Limited dataset to {len(dataset)} samples")
    
    # Calculate dataset size and per-process info
    num_processes = accelerator.num_processes
    process_index = accelerator.process_index
    
    # Resume from index
    if args.resume_from_index > 0:
        dataset.data = dataset.data[args.resume_from_index:]
        logger.info(f"Resuming from index {args.resume_from_index}, remaining samples: {len(dataset)}")
    
    # Create DistributedSampler for proper data sharding across GPUs
    sampler = DistributedSampler(
        dataset,
        num_replicas=num_processes,
        rank=process_index,
        shuffle=False,  # Don't shuffle for feature extraction
        drop_last=False,
    )
    
    # Get sampler indices before prepare (they won't be accessible after)
    sampler_indices = list(sampler)  # This gives us the global dataset indices for this process
    
    # Create dataloader
    dataloader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        sampler=sampler,  # Use DistributedSampler
        num_workers=args.num_workers,
        pin_memory=True,
    )
    
    dataloader = accelerator.prepare(dataloader)
    
    # Create output directories (only for features that will be extracted)
    video_codes_dir = None
    text_embeddings_dir = None
    if args.extract_video:
        video_codes_dir = os.path.join(args.output_dir, "video_codes")
        os.makedirs(video_codes_dir, exist_ok=True)
    if args.extract_text:
        text_embeddings_dir = os.path.join(args.output_dir, "text_embeddings")
        os.makedirs(text_embeddings_dir, exist_ok=True)
    
    metadata_file = os.path.join(args.output_dir, "metadata.json")
    
    # Calculate dataset size and per-process info
    total_samples = len(dataset)
    
    # Extract features
    logger.info(f"[GPU {process_index}] Starting feature extraction for {total_samples} samples (process {process_index+1}/{num_processes})...")
    logger.info(f"[GPU {process_index}] This process will handle ~{len(dataloader) * args.batch_size} samples")
    
    # Get codebook_size and mask_token_id from video_tokenizer (if extracting video)
    codebook_size = None
    mask_token_id = None
    if args.extract_video and video_tokenizer is not None:
        codebook_size = video_tokenizer.codebook_size
        mask_token_id = video_tokenizer.mask_token_id
        logger.info(f"[GPU {process_index}] Video tokenizer info: codebook_size={codebook_size}, mask_token_id={mask_token_id}")
    
    # Get empty_embeds info (only on main process, will be added to metadata later)
    empty_embeds_shape = None
    empty_embeds_path = os.path.join(args.output_dir, "empty_embeds.npy")
    if args.extract_text and accelerator.is_main_process and os.path.exists(empty_embeds_path):
        empty_embeds_np = np.load(empty_embeds_path)
        empty_embeds_shape = list(empty_embeds_np.shape)
        logger.info(f"Empty embeds shape: {empty_embeds_shape}")
    
    # Per-process metadata (will be merged at the end)
    process_metadata = {
        "process_index": process_index,
        "num_samples": total_samples,
        "extract_video": args.extract_video,
        "extract_text": args.extract_text,
        "text_encoder_architecture": args.text_encoder_architecture if args.extract_text else None,
        "video_tokenizer_model_id": args.video_tokenizer_model_id if args.extract_video else None,
        "codebook_size": codebook_size,
        "mask_token_id": mask_token_id,
        "num_frames": args.num_frames,
        "video_height": args.video_height,
        "video_width": args.video_width,
        "prompt_prefix": args.prompt_prefix,
        "empty_embeds_shape": empty_embeds_shape if process_index == 0 else None,  # Only main process has this
        "empty_embeds_path": "empty_embeds.npy" if args.extract_text else None,
        "samples": []
    }
    
    # Track processed samples for this process
    process_failed_samples = []
    process_samples_processed = 0
    
    with torch.no_grad():
        for batch_idx, batch in enumerate(tqdm(dataloader, desc=f"[GPU {process_index}] Extracting", disable=not accelerator.is_main_process)):
            batch_size = batch["video"].shape[0] if args.extract_video else batch["prompt_input_ids"].shape[0]
            
            # Extract video codes (if needed)
            video_codes = None
            if args.extract_video:
                videos = batch["video"].to(device, non_blocking=True)  # [B, C, F, H, W]
                try:
                    video_codes = video_tokenizer.encode(videos)  # [B, F', H', W']
                    video_codes = video_codes.cpu().numpy()  # Convert to numpy for saving
                except Exception as e:
                    logger.error(f"[GPU {process_index}] Failed to encode video batch {batch_idx}: {e}")
                    continue
            
            # Extract text embeddings (if needed)
            encoder_hidden_states = None
            if args.extract_text:
                prompt_input_ids = batch["prompt_input_ids"].to(device, non_blocking=True)
                try:
                    encoder_hidden_states, _ = encode_prompt(
                        text_encoder,
                        prompt_input_ids,
                        args.text_encoder_architecture
                    )  # [B, L, D]
                    encoder_hidden_states = encoder_hidden_states.cpu().numpy()  # Convert to numpy for saving
                except Exception as e:
                    logger.error(f"[GPU {process_index}] Failed to encode text batch {batch_idx}: {e}")
                    continue
            
            # Get the actual dataset indices for this batch
            # sampler_indices contains the global dataset indices assigned to this process
            local_start_idx = batch_idx * args.batch_size
            
            # Save features for each sample in the batch
            for i in range(batch_size):
                local_idx = local_start_idx + i
                
                if local_idx < len(sampler_indices):
                    # Get the global dataset index from the sampler
                    global_dataset_idx = sampler_indices[local_idx]
                    
                    # Calculate the final global sample index (accounting for resume_from_index)
                    sample_idx = args.resume_from_index + global_dataset_idx
                    
                    # Get original video path and prompt from dataset for metadata
                    row = None
                    if global_dataset_idx < len(dataset.data):
                        row = dataset.data[global_dataset_idx]
                    
                    # Save video codes (if extracted)
                    video_code = None
                    if args.extract_video and video_codes is not None:
                        video_code_dir, video_code_path = get_hierarchical_path(video_codes_dir, sample_idx)
                        os.makedirs(video_code_dir, exist_ok=True)
                        video_code = video_codes[i]  # [F', H', W']
                        # Ensure saved as CPU numpy array with int32 (smaller than int64)
                        if isinstance(video_code, torch.Tensor):
                            video_code = video_code.cpu().numpy()
                        # Convert to int32 to save space (codebook_size is typically < 2^31)
                        video_code = video_code.astype(np.int32)
                        np.save(video_code_path, video_code)
                    
                    # Save text embedding (if extracted)
                    text_embedding = None
                    if args.extract_text and encoder_hidden_states is not None:
                        text_embedding_dir, text_embedding_path = get_hierarchical_path(text_embeddings_dir, sample_idx)
                        os.makedirs(text_embedding_dir, exist_ok=True)
                        text_embedding = encoder_hidden_states[i]  # [L, D]
                        # Ensure saved as CPU numpy array with float16 to save space
                        if isinstance(text_embedding, torch.Tensor):
                            text_embedding = text_embedding.cpu().numpy()
                        # Convert to float16 to save space (half the size of float32)
                        text_embedding = text_embedding.astype(np.float16)
                        np.save(text_embedding_path, text_embedding)
                    
                    # Add to metadata
                    if row is not None:
                        sample_meta = {
                            "index": sample_idx,
                            "video_path": row.get("video", ""),
                            "caption": row.get("caption", ""),
                        }
                        if args.extract_video and video_code is not None:
                            sample_meta["video_code_shape"] = list(video_code.shape)
                        if args.extract_text and text_embedding is not None:
                            sample_meta["text_embedding_shape"] = list(text_embedding.shape)
                        process_metadata["samples"].append(sample_meta)
                    
                    process_samples_processed += 1
            
            # Save per-process metadata periodically (every 1000 samples per process)
            if process_samples_processed % 1000 == 0:
                process_metadata_file = os.path.join(args.output_dir, f"metadata_process_{process_index}.json")
                process_metadata["num_extracted"] = process_samples_processed
                process_metadata["failed_samples"] = process_failed_samples
                with open(process_metadata_file, 'w') as f:
                    json.dump(process_metadata, f, indent=2)
                logger.info(f"[GPU {process_index}] Progress: {process_samples_processed} samples extracted")
    
    # Wait for all processes to finish
    accelerator.wait_for_everyone()
    
    # Save per-process metadata
    process_metadata_file = os.path.join(args.output_dir, f"metadata_process_{process_index}.json")
    process_metadata["num_extracted"] = process_samples_processed
    process_metadata["failed_samples"] = process_failed_samples
    with open(process_metadata_file, 'w') as f:
        json.dump(process_metadata, f, indent=2)
    
    logger.info(f"[GPU {process_index}] Process complete: {process_samples_processed} samples extracted")
    
    # Merge metadata from all processes (only on main process)
    accelerator.wait_for_everyone()
    
    if accelerator.is_main_process:
        logger.info("Merging metadata from all processes...")
        
        # Load all process metadata files
        all_samples = []
        total_extracted = 0
        all_failed = []
        
        for proc_idx in range(num_processes):
            proc_metadata_file = os.path.join(args.output_dir, f"metadata_process_{proc_idx}.json")
            if os.path.exists(proc_metadata_file):
                with open(proc_metadata_file, 'r') as f:
                    proc_meta = json.load(f)
                all_samples.extend(proc_meta.get("samples", []))
                total_extracted += proc_meta.get("num_extracted", 0)
                all_failed.extend(proc_meta.get("failed_samples", []))
        
        # Sort samples by index
        all_samples.sort(key=lambda x: x["index"])
        
        # Get codebook_size, mask_token_id, and empty_embeds info from first process metadata
        codebook_size = None
        mask_token_id = None
        empty_embeds_shape = None
        empty_embeds_path = None
        for proc_idx in range(num_processes):
            proc_metadata_file = os.path.join(args.output_dir, f"metadata_process_{proc_idx}.json")
            if os.path.exists(proc_metadata_file):
                with open(proc_metadata_file, 'r') as f:
                    proc_meta = json.load(f)
                if proc_meta.get("codebook_size") is not None:
                    codebook_size = proc_meta.get("codebook_size")
                if proc_meta.get("mask_token_id") is not None:
                    mask_token_id = proc_meta.get("mask_token_id")
                if proc_meta.get("empty_embeds_shape") is not None:
                    empty_embeds_shape = proc_meta.get("empty_embeds_shape")
                if proc_meta.get("empty_embeds_path") is not None:
                    empty_embeds_path = proc_meta.get("empty_embeds_path")
                if codebook_size is not None and mask_token_id is not None:
                    if not args.extract_text or (empty_embeds_shape is not None and empty_embeds_path is not None):
                        break
        
        # Create merged metadata
        merged_metadata = {
            "num_samples": total_samples,
            "num_extracted": total_extracted,
            "num_processes": num_processes,
            "extract_video": args.extract_video,
            "extract_text": args.extract_text,
            "text_encoder_architecture": args.text_encoder_architecture if args.extract_text else None,
            "video_tokenizer_model_id": args.video_tokenizer_model_id if args.extract_video else None,
            "codebook_size": codebook_size,
            "mask_token_id": mask_token_id,
            "num_frames": args.num_frames,
            "video_height": args.video_height,
            "video_width": args.video_width,
            "prompt_prefix": args.prompt_prefix,
            "empty_embeds_shape": empty_embeds_shape,
            "empty_embeds_path": empty_embeds_path,
            "samples": all_samples,
            "failed_samples": sorted(set(all_failed)),
        }
        
        # Save merged metadata
        with open(metadata_file, 'w') as f:
            json.dump(merged_metadata, f, indent=2)
        
        logger.info(f"Feature extraction complete!")
        logger.info(f"  Total samples: {total_samples}")
        logger.info(f"  Extracted: {total_extracted}")
        logger.info(f"  Failed: {len(merged_metadata['failed_samples'])}")
        if args.extract_video:
            logger.info(f"  Video codes saved to: {video_codes_dir}")
        if args.extract_text:
            logger.info(f"  Text embeddings saved to: {text_embeddings_dir}")
            if empty_embeds_path:
                logger.info(f"  Empty embeds saved to: {os.path.join(args.output_dir, empty_embeds_path)}")
        logger.info(f"  Metadata saved to: {metadata_file}")
        
        # Optionally clean up per-process metadata files
        # Uncomment if you want to remove them after merging
        # for proc_idx in range(num_processes):
        #     proc_metadata_file = os.path.join(args.output_dir, f"metadata_process_{proc_idx}.json")
        #     if os.path.exists(proc_metadata_file):
        #         os.remove(proc_metadata_file)


if __name__ == "__main__":
    main()