Model-Optimizer/examples/diffusers/fastgen/preprocess/preprocessing_multiprocess.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Unified preprocessing tool for images and videos.

Supports:
- Images: FLUX (and other image models)
- Videos: Wan2.1, HunyuanVideo-1.5

Usage:
    # Image preprocessing
    python examples/diffusers/fastgen/preprocess_qwen_image.py image \\
        --image_dir /path/to/images \\
        --output_dir /path/to/cache \\
        --processor flux

    # Video preprocessing
    python examples/diffusers/fastgen/preprocess_qwen_image.py video \\
        --video_dir /path/to/videos \\
        --output_dir /path/to/cache \\
        --processor wan \\
        --resolution_preset 512p

    # List available processors
    python examples/diffusers/fastgen/preprocess_qwen_image.py --list_processors
"""

import argparse
import hashlib
import json
import logging
import os
import traceback
from multiprocessing import Pool
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import torch
from nemo_automodel.components.datasets.diffusion.multi_tier_bucketing import (
    MultiTierBucketCalculator,
)
from PIL import Image
from tqdm import tqdm

from .processors import BaseModelProcessor, ProcessorRegistry, get_caption_loader

logger = logging.getLogger(__name__)

# =============================================================================
# Constants
# =============================================================================
IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "webp", "bmp"}
VIDEO_EXTENSIONS = {"mp4", "avi", "mov", "mkv", "webm"}

# =============================================================================
# Global worker state (initialized once per process)
# =============================================================================
_worker_models: dict[str, Any] | None = None
_worker_processor: BaseModelProcessor | None = None
_worker_calculator: MultiTierBucketCalculator | None = None
_worker_device: str | None = None
_worker_config: dict[str, Any] | None = None


# =============================================================================
# Common Utility Functions
# =============================================================================


def _get_media_files(media_dir: Path, extensions: set) -> list[Path]:
    """Recursively get all media files with given extensions using os.walk()."""
    media_files = []
    for root, dirs, files in os.walk(media_dir):
        root_path = Path(root)
        for file in files:
            if "." in file:
                ext = file.lower().rsplit(".", 1)[-1]
                if ext in extensions:
                    media_files.append(root_path / file)
    return sorted(media_files)


def _save_metadata_shards(
    all_metadata: list[dict],
    output_dir: Path,
    processor_name: str,
    model_name: str,
    model_type: str,
    shard_size: int,
    extra_fields: dict[str, Any],
    shard_rank: int = 0,
    shard_world: int = 1,
) -> None:
    """Save metadata in shards and write config file.

    When shard_world > 1, the index file and shard filenames are namespaced with
    the rank so that multiple jobs sharing an output directory don't overwrite
    each other. Merge the per-rank index files afterwards with a separate
    script to produce a single unified metadata.json.
    """
    sharded = shard_world > 1
    shard_prefix = f"r{shard_rank:02d}_" if sharded else ""
    index_filename = f"metadata_r{shard_rank:02d}.json" if sharded else "metadata.json"

    shard_files = []
    for chunk_start in range(0, len(all_metadata), shard_size):
        chunk_data = all_metadata[chunk_start : chunk_start + shard_size]
        chunk_idx = chunk_start // shard_size
        shard_file = output_dir / f"metadata_shard_{shard_prefix}s{chunk_idx:04d}.json"
        with open(shard_file, "w") as f:
            json.dump(chunk_data, f, indent=2)
        shard_files.append(shard_file.name)

    metadata = {
        "processor": processor_name,
        "model_name": model_name,
        "model_type": model_type,
        "total_items": len(all_metadata),
        "num_shards": len(shard_files),
        "shard_size": shard_size,
        "shards": shard_files,
        **extra_fields,
    }
    if sharded:
        metadata["shard_rank"] = shard_rank
        metadata["shard_world"] = shard_world

    with open(output_dir / index_filename, "w") as f:
        json.dump(metadata, f, indent=2)


def _print_bucket_distribution(all_metadata: list[dict]) -> None:
    """Print bucket resolution distribution."""
    bucket_counts: dict[str, int] = {}
    for item in all_metadata:
        res = f"{item['bucket_resolution'][0]}x{item['bucket_resolution'][1]}"
        bucket_counts[res] = bucket_counts.get(res, 0) + 1

    logger.info("Bucket distribution:")
    for res in sorted(bucket_counts.keys()):
        logger.info("  %s: %d", res, bucket_counts[res])


# =============================================================================
# Image Preprocessing Functions
# =============================================================================


def _init_worker(processor_name: str, model_name: str, gpu_id: int, max_pixels: int):
    """Initialize worker process with models on assigned GPU."""
    global _worker_models, _worker_processor, _worker_calculator, _worker_device

    # Set CUDA_VISIBLE_DEVICES to isolate this GPU for the worker process.
    # After this, the selected GPU becomes cuda:0 (not cuda:{gpu_id}).
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    _worker_device = "cuda:0"

    _worker_processor = ProcessorRegistry.get(processor_name)
    _worker_models = _worker_processor.load_models(model_name, _worker_device)
    _worker_calculator = MultiTierBucketCalculator(quantization=64, max_pixels=max_pixels)

    logger.info("Worker initialized on GPU %d", gpu_id)


def _load_all_captions(
    image_files: list[Path],
    caption_field: str = "internvl",
    caption_format: str = "jsonl",
    verbose: bool = True,
) -> dict[str, str]:
    """Pre-load all captions from caption files. Returns filename->caption dict.

    Args:
        image_files: List of image paths to look up captions for.
        caption_field: Field name inside each caption record (e.g. "internvl", "caption").
        caption_format: One of "sidecar", "meta_json", "jsonl". Selects which CaptionLoader to use.
        verbose: If True, log progress and statistics.
    """
    loader = get_caption_loader(caption_format)
    captions, stats = loader.load_captions_with_stats(image_files, caption_field, verbose=verbose)

    if verbose:
        logger.info(
            "Loaded %d captions from %d caption files (format=%s)",
            stats.loaded_count,
            stats.files_parsed,
            caption_format,
        )
        if stats.files_missing > 0:
            logger.info(
                "  %d caption files not found (will use filename fallback)", stats.files_missing
            )
        if stats.captions_missing > 0:
            logger.info("  %d images will use filename as caption", stats.captions_missing)

    return captions


def _process_image(args: tuple) -> dict | None:
    """Process a single image using pre-initialized worker state."""
    image_path, output_dir, verify, caption = args

    try:
        image = Image.open(image_path).convert("RGB")
        orig_width, orig_height = image.size

        bucket = _worker_calculator.get_bucket_for_image(orig_width, orig_height)
        target_width, target_height = bucket["resolution"]

        resized_image, crop_offset = _worker_calculator.resize_and_crop(
            image, target_width, target_height, crop_mode="center"
        )

        image_tensor = _worker_processor.preprocess_image(resized_image)
        latent = _worker_processor.encode_image(image_tensor, _worker_models, _worker_device)

        if verify and not _worker_processor.verify_latent(latent, _worker_models, _worker_device):
            logger.warning("Verification failed: %s", image_path)
            return None

        # Use pre-loaded caption with fallback to filename
        if not caption:
            caption = Path(image_path).stem.replace("_", " ")

        text_encodings = _worker_processor.encode_text(caption, _worker_models, _worker_device)

        # Save cache file
        resolution = f"{target_width}x{target_height}"
        cache_subdir = Path(output_dir) / resolution
        cache_subdir.mkdir(parents=True, exist_ok=True)

        cache_hash = hashlib.md5(f"{Path(image_path).absolute()}_{resolution}".encode()).hexdigest()
        cache_file = cache_subdir / f"{cache_hash}.pt"

        metadata = {
            "original_resolution": (orig_width, orig_height),
            "bucket_resolution": (target_width, target_height),
            "crop_offset": crop_offset,
            "prompt": caption,
            "image_path": str(Path(image_path).absolute()),
            "bucket_id": bucket["id"],
            "aspect_ratio": bucket["aspect_ratio"],
        }

        cache_data = _worker_processor.get_cache_data(latent, text_encodings, metadata)
        torch.save(cache_data, cache_file)

        return {
            "cache_file": str(cache_file),
            "image_path": str(Path(image_path).absolute()),
            "bucket_resolution": [target_width, target_height],
            "original_resolution": [orig_width, orig_height],
            "prompt": caption,
            "bucket_id": bucket["id"],
            "aspect_ratio": bucket["aspect_ratio"],
            "pixels": target_width * target_height,
            "model_type": _worker_processor.model_type,
        }

    except Exception as e:
        logger.error("Error processing %s: %s", image_path, e)
        logger.debug(traceback.format_exc())
        return None


def _process_shard_on_gpu(
    gpu_id: int,
    image_files: list[Path],
    output_dir: str,
    processor_name: str,
    model_name: str,
    verify: bool,
    caption_cache: dict[str, str],
    max_pixels: int,
) -> list[dict]:
    """Process a shard of images on a specific GPU."""
    _init_worker(processor_name, model_name, gpu_id, max_pixels)

    results = []
    for image_path in tqdm(image_files, desc=f"GPU {gpu_id}", position=gpu_id):
        # Get caption from cache (or None if not found)
        caption = caption_cache.get(image_path.name)
        result = _process_image((str(image_path), output_dir, verify, caption))
        if result:
            results.append(result)

    return results


def preprocess_dataset(
    image_dir: str,
    output_dir: str,
    processor_name: str,
    model_name: str | None = None,
    shard_size: int = 10000,
    verify: bool = False,
    caption_field: str = "internvl",
    caption_format: str = "jsonl",
    max_images: int | None = None,
    max_pixels: int = 256 * 256,
    shard_idx: int = 0,
    shard_count: int = 1,
):
    """
    Preprocess image dataset with one process per GPU.

    Args:
        image_dir: Directory containing images
        output_dir: Output directory for cache
        processor_name: Name of processor to use (e.g., 'flux', 'sdxl')
        model_name: HuggingFace model name (uses processor default if None)
        shard_size: Number of images per metadata shard
        verify: Whether to verify latents can be decoded
        caption_field: Field name inside each caption record (e.g. 'internvl', 'caption')
        caption_format: One of 'sidecar', 'meta_json', 'jsonl' (selects the CaptionLoader)
        max_images: Maximum number of images to process
        max_pixels: Maximum pixels per image
        shard_idx: Rank of this job within a multi-job sweep (0-indexed). Each rank
            processes image_files[shard_idx::shard_count].
        shard_count: Total number of jobs in the sweep. Default 1 (single-job mode).
    """
    image_dir = Path(image_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Get processor and resolve model name
    processor = ProcessorRegistry.get(processor_name)
    if model_name is None:
        model_name = processor.default_model_name

    num_gpus = torch.cuda.device_count()
    if num_gpus == 0:
        raise RuntimeError("No GPUs available")

    logger.info("Processor: %s (%s)", processor_name, processor.model_type)
    logger.info("Model: %s", model_name)
    logger.info("GPUs: %d", num_gpus)
    logger.info("Max pixels: %d", max_pixels)

    # Get all image files
    logger.info("Scanning for images...")
    image_files = _get_media_files(image_dir, IMAGE_EXTENSIONS)

    if max_images is not None:
        image_files = image_files[:max_images]

    if shard_count > 1:
        image_files = image_files[shard_idx::shard_count]
        logger.info("Shard %d/%d: %d images on this rank", shard_idx, shard_count, len(image_files))

    logger.info("Processing %d images", len(image_files))

    if not image_files:
        return

    caption_cache = _load_all_captions(
        image_files, caption_field, caption_format=caption_format, verbose=True
    )

    # Split images across GPUs
    chunks = [image_files[i::num_gpus] for i in range(num_gpus)]

    # Process with one worker per GPU
    all_metadata = []

    with Pool(processes=num_gpus) as pool:
        args = [
            (
                gpu_id,
                chunks[gpu_id],
                str(output_dir),
                processor_name,
                model_name,
                verify,
                caption_cache,
                max_pixels,
            )
            for gpu_id in range(num_gpus)
        ]

        results = pool.starmap(_process_shard_on_gpu, args)

        for gpu_results in results:
            all_metadata.extend(gpu_results)

    # Save metadata
    _save_metadata_shards(
        all_metadata,
        output_dir,
        processor_name,
        model_name,
        processor.model_type,
        shard_size,
        {
            "caption_field": caption_field,
            "caption_format": caption_format,
            "max_pixels": max_pixels,
        },
        shard_rank=shard_idx,
        shard_world=shard_count,
    )

    # Print summary
    logger.info("=" * 50)
    logger.info("COMPLETE: %d/%d images", len(all_metadata), len(image_files))
    logger.info("Output: %s", output_dir)
    _print_bucket_distribution(all_metadata)


# =============================================================================
# Video Preprocessing Functions
# =============================================================================


def _init_video_worker(
    processor_name: str,
    model_name: str,
    gpu_id: int,
    max_pixels: int,
    video_config: dict[str, Any],
):
    """Initialize video worker process with models on assigned GPU."""
    global _worker_models, _worker_processor, _worker_calculator, _worker_device, _worker_config

    # Set CUDA_VISIBLE_DEVICES to isolate this GPU for the worker process.
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    _worker_device = "cuda:0"
    _worker_config = video_config

    _worker_processor = ProcessorRegistry.get(processor_name)
    _worker_models = _worker_processor.load_models(model_name, _worker_device)

    # Create bucket calculator with processor's quantization (8 for video, 64 for image)
    quantization = getattr(_worker_processor, "quantization", 8)
    _worker_calculator = MultiTierBucketCalculator(quantization=quantization, max_pixels=max_pixels)

    logger.info("Video worker initialized on GPU %d (quantization=%d)", gpu_id, quantization)


def _get_video_dimensions(video_path: str) -> tuple[int, int, int]:
    """Get video dimensions and frame count using OpenCV."""
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"Failed to open video: {video_path}")

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()

    return width, height, frame_count


def _extract_evenly_spaced_frames(
    video_path: str,
    num_frames: int,
    target_size: tuple[int, int],
    resize_mode: str = "bilinear",
    center_crop: bool = True,
) -> tuple[list[np.ndarray], list[int]]:
    """Extract evenly-spaced frames. Returns (frames, source_indices)."""
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"Failed to open video: {video_path}")

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    orig_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    orig_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Calculate evenly-spaced frame indices
    if num_frames >= total_frames:
        frame_indices = list(range(total_frames))
    else:
        frame_indices = np.linspace(0, total_frames - 1, num_frames).astype(int).tolist()

    target_height, target_width = target_size

    # Map resize modes to OpenCV interpolation
    interp_map = {
        "bilinear": cv2.INTER_LINEAR,
        "bicubic": cv2.INTER_CUBIC,
        "nearest": cv2.INTER_NEAREST,
        "area": cv2.INTER_AREA,
        "lanczos": cv2.INTER_LANCZOS4,
    }
    interpolation = interp_map.get(resize_mode, cv2.INTER_LINEAR)

    frames = []
    actual_indices = []

    for target_idx in frame_indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, target_idx)
        ret, frame = cap.read()
        if not ret:
            continue

        # Convert BGR to RGB
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        # Resize and optionally center crop
        if center_crop:
            # Calculate scale to cover target area
            scale = max(target_width / orig_width, target_height / orig_height)
            new_width = int(orig_width * scale)
            new_height = int(orig_height * scale)

            frame = cv2.resize(frame, (new_width, new_height), interpolation=interpolation)

            # Center crop
            start_x = (new_width - target_width) // 2
            start_y = (new_height - target_height) // 2
            frame = frame[start_y : start_y + target_height, start_x : start_x + target_width]
        else:
            # Direct resize (may change aspect ratio)
            frame = cv2.resize(frame, (target_width, target_height), interpolation=interpolation)

        frames.append(frame)
        actual_indices.append(target_idx)

    cap.release()
    return frames, actual_indices


def _frame_to_video_tensor(frame: np.ndarray, dtype: torch.dtype = torch.float16) -> torch.Tensor:
    """Convert frame (H,W,C) to video tensor (1,C,1,H,W) normalized to [-1,1]."""
    # (H, W, C) -> (C, H, W)
    tensor = torch.from_numpy(frame).float().permute(2, 0, 1)

    # Normalize to [-1, 1]
    tensor = tensor / 255.0
    tensor = (tensor - 0.5) / 0.5

    # Add batch and temporal dimensions: (C, H, W) -> (1, C, 1, H, W)
    tensor = tensor.unsqueeze(0).unsqueeze(2)

    return tensor.to(dtype)


# =============================================================================
# Video Processing Helper Functions
# =============================================================================


def _resolve_video_resolution(
    orig_width: int,
    orig_height: int,
    config: dict[str, Any],
) -> tuple[int, int, str | None, float]:
    """Resolve target resolution. Returns (width, height, bucket_id, aspect_ratio)."""
    target_height = config.get("target_height")
    target_width = config.get("target_width")

    if target_height is not None and target_width is not None:
        # Explicit size: no bucketing
        return target_width, target_height, None, target_width / target_height
    else:
        # Use bucket calculator to find best resolution
        bucket = _worker_calculator.get_bucket_for_image(orig_width, orig_height)
        return (
            bucket["resolution"][0],
            bucket["resolution"][1],
            bucket["id"],
            bucket["aspect_ratio"],
        )


def _save_cache_file(
    cache_data: dict[str, Any],
    output_dir: str,
    resolution: str,
    cache_hash: str,
    output_format: str,
) -> Path:
    """Save cache data to file. Returns path to saved file."""
    cache_subdir = Path(output_dir) / resolution
    cache_subdir.mkdir(parents=True, exist_ok=True)

    if output_format == "meta":
        cache_file = cache_subdir / f"{cache_hash}.meta"
        torch.save(cache_data, cache_file)
    else:  # pt format
        cache_file = cache_subdir / f"{cache_hash}.pt"
        torch.save(cache_data, cache_file)

    return cache_file


def _build_result_dict(
    cache_file: Path,
    video_path: str,
    target_width: int,
    target_height: int,
    orig_width: int,
    orig_height: int,
    caption: str,
    bucket_id: str | None,
    aspect_ratio: float,
    num_frames: int = 1,
    frame_index: int | None = None,
    total_frames_extracted: int | None = None,
    source_frame_index: int | None = None,
) -> dict[str, Any]:
    """Build a result dictionary for a processed video/frame."""
    result = {
        "cache_file": str(cache_file),
        "video_path": str(Path(video_path).absolute()),
        "bucket_resolution": [target_width, target_height],
        "original_resolution": [orig_width, orig_height],
        "num_frames": num_frames,
        "prompt": caption,
        "bucket_id": bucket_id,
        "aspect_ratio": aspect_ratio,
        "pixels": target_width * target_height,
        "model_type": _worker_processor.model_type,
    }

    # Add frame-specific fields if provided
    if frame_index is not None:
        result["frame_index"] = frame_index
    if total_frames_extracted is not None:
        result["total_frames_extracted"] = total_frames_extracted
    if source_frame_index is not None:
        result["source_frame_index"] = source_frame_index

    return result


def _process_video_frames_mode(args: tuple) -> list[dict]:
    """Process video in frames mode - each frame becomes a separate sample."""
    video_path, output_dir, caption, config = args

    try:
        # Get video dimensions
        orig_width, orig_height, total_frames = _get_video_dimensions(video_path)

        # Resolve target resolution (handles bucketing vs explicit size)
        target_width, target_height, bucket_id, aspect_ratio = _resolve_video_resolution(
            orig_width, orig_height, config
        )

        # Extract evenly-spaced frames
        num_frames = config.get("num_frames", 10)
        frames, source_frame_indices = _extract_evenly_spaced_frames(
            video_path,
            num_frames=num_frames,
            target_size=(target_height, target_width),
            resize_mode=config.get("resize_mode", "bilinear"),
            center_crop=config.get("center_crop", True),
        )

        if not frames:
            logger.warning("No frames extracted from %s", video_path)
            return []

        total_frames_extracted = len(frames)

        # Use caption with fallback to filename
        if not caption:
            caption = Path(video_path).stem.replace("_", " ")

        # Encode text ONCE (reuse for all frames)
        text_encodings = _worker_processor.encode_text(caption, _worker_models, _worker_device)

        # Process each frame individually
        results = []
        deterministic = config.get("deterministic", True)
        output_format = config.get("output_format", "meta")
        resolution = f"{target_width}x{target_height}"

        for frame_idx, (frame, source_idx) in enumerate(zip(frames, source_frame_indices)):
            # Convert single frame to 1-frame video tensor
            video_tensor = _frame_to_video_tensor(frame)

            # Encode with VAE
            latent = _worker_processor.encode_video(
                video_tensor,
                _worker_models,
                _worker_device,
                deterministic=deterministic,
            )

            # Prepare metadata for this frame
            # Note: first_frame and image_embeds are omitted in frames mode
            # (frames mode is intended for t2v training, not i2v conditioning)
            metadata = {
                "original_resolution": (orig_width, orig_height),
                "bucket_resolution": (target_width, target_height),
                "bucket_id": bucket_id,
                "aspect_ratio": aspect_ratio,
                "num_frames": 1,  # Always 1 for frame mode
                "total_original_frames": total_frames,
                "prompt": caption,
                "video_path": str(Path(video_path).absolute()),
                "deterministic": deterministic,
                "mode": "frames",
                "frame_index": frame_idx + 1,  # 1-based index
                "total_frames_extracted": total_frames_extracted,
                "source_frame_index": source_idx,  # 0-based index in source video
            }

            # Get cache data from processor
            cache_data = _worker_processor.get_cache_data(latent, text_encodings, metadata)

            # Include frame index in hash to ensure unique filenames
            cache_hash = hashlib.md5(
                f"{Path(video_path).absolute()}_{resolution}_frame{frame_idx}".encode()
            ).hexdigest()

            # Save cache file using helper
            cache_file = _save_cache_file(
                cache_data, output_dir, resolution, cache_hash, output_format
            )

            # Build result dict using helper
            results.append(
                _build_result_dict(
                    cache_file=cache_file,
                    video_path=video_path,
                    target_width=target_width,
                    target_height=target_height,
                    orig_width=orig_width,
                    orig_height=orig_height,
                    caption=caption,
                    bucket_id=bucket_id,
                    aspect_ratio=aspect_ratio,
                    num_frames=1,
                    frame_index=frame_idx + 1,
                    total_frames_extracted=total_frames_extracted,
                    source_frame_index=source_idx,
                )
            )

        return results

    except Exception as e:
        logger.error("Error processing %s in frames mode: %s", video_path, e)
        logger.debug(traceback.format_exc())
        return []


def _process_video_video_mode(args: tuple) -> dict | None:
    """Process video in video mode - multi-frame encoding as single sample."""
    video_path, output_dir, caption, config = args

    try:
        # Get video dimensions
        orig_width, orig_height, total_frames = _get_video_dimensions(video_path)

        # Resolve target resolution (handles bucketing vs explicit size)
        target_width, target_height, bucket_id, aspect_ratio = _resolve_video_resolution(
            orig_width, orig_height, config
        )

        # Load video with target resolution
        num_frames = config.get("num_frames")
        target_frames = config.get("target_frames")

        video_tensor, first_frame = _worker_processor.load_video(
            video_path,
            target_size=(target_height, target_width),
            num_frames=target_frames or num_frames,
            resize_mode=config.get("resize_mode", "bilinear"),
            center_crop=config.get("center_crop", True),
        )

        actual_frames = video_tensor.shape[2]  # (1, C, T, H, W)

        # Use caption with fallback to filename
        if not caption:
            caption = Path(video_path).stem.replace("_", " ")

        # Encode video
        deterministic = config.get("deterministic", True)
        latent = _worker_processor.encode_video(
            video_tensor,
            _worker_models,
            _worker_device,
            deterministic=deterministic,
        )

        # Encode text
        text_encodings = _worker_processor.encode_text(caption, _worker_models, _worker_device)

        # Encode first frame for i2v (if processor supports it)
        image_embeds = None
        if hasattr(_worker_processor, "encode_first_frame"):
            image_embeds = _worker_processor.encode_first_frame(
                first_frame, _worker_models, _worker_device
            )

        # Prepare metadata
        metadata = {
            "original_resolution": (orig_width, orig_height),
            "bucket_resolution": (target_width, target_height),
            "bucket_id": bucket_id,
            "aspect_ratio": aspect_ratio,
            "num_frames": actual_frames,
            "total_original_frames": total_frames,
            "prompt": caption,
            "video_path": str(Path(video_path).absolute()),
            "first_frame": first_frame,
            "image_embeds": image_embeds,
            "deterministic": deterministic,
            "mode": config.get("mode", "video"),
        }

        # Get cache data from processor
        cache_data = _worker_processor.get_cache_data(latent, text_encodings, metadata)

        # Save cache file using helper
        output_format = config.get("output_format", "meta")
        resolution = f"{target_width}x{target_height}"
        cache_hash = hashlib.md5(
            f"{Path(video_path).absolute()}_{resolution}_{actual_frames}".encode()
        ).hexdigest()
        cache_file = _save_cache_file(cache_data, output_dir, resolution, cache_hash, output_format)

        # Build result dict using helper
        return _build_result_dict(
            cache_file=cache_file,
            video_path=video_path,
            target_width=target_width,
            target_height=target_height,
            orig_width=orig_width,
            orig_height=orig_height,
            caption=caption,
            bucket_id=bucket_id,
            aspect_ratio=aspect_ratio,
            num_frames=actual_frames,
        )

    except Exception as e:
        logger.error("Error processing %s: %s", video_path, e)
        logger.debug(traceback.format_exc())
        return None


def _process_video(args: tuple) -> list[dict]:
    """Process a single video. Dispatches to frames or video mode based on config."""
    video_path, output_dir, caption, config = args
    mode = config.get("mode", "video")

    if mode == "frames":
        return _process_video_frames_mode(args)
    else:
        # Wrap single result in a list for consistent return type
        result = _process_video_video_mode(args)
        return [result] if result is not None else []


def _process_video_shard_on_gpu(
    gpu_id: int,
    video_files: list[Path],
    output_dir: str,
    processor_name: str,
    model_name: str,
    caption_cache: dict[str, str],
    max_pixels: int,
    video_config: dict[str, Any],
) -> list[dict]:
    """Process a shard of videos on a specific GPU."""
    _init_video_worker(processor_name, model_name, gpu_id, max_pixels, video_config)

    results = []
    for video_path in tqdm(video_files, desc=f"GPU {gpu_id}", position=gpu_id):
        caption = caption_cache.get(video_path.name)
        # _process_video now always returns List[Dict] for consistent handling
        results.extend(_process_video((str(video_path), output_dir, caption, video_config)))

    return results


def preprocess_video_dataset(
    video_dir: str,
    output_dir: str,
    processor_name: str,
    model_name: str | None = None,
    mode: str = "video",
    num_frames: int = 10,
    target_frames: int | None = None,
    resolution_preset: str | None = None,
    max_pixels: int | None = None,
    target_height: int | None = None,
    target_width: int | None = None,
    resize_mode: str = "bilinear",
    center_crop: bool = True,
    deterministic: bool = True,
    output_format: str = "meta",
    caption_format: str = "sidecar",
    caption_field: str = "caption",
    shard_size: int = 10000,
    max_videos: int | None = None,
):
    """
    Preprocess video dataset with one process per GPU.

    Args:
        video_dir: Directory containing videos
        output_dir: Output directory for cache
        processor_name: Name of processor ('wan', 'hunyuan')
        model_name: HuggingFace model name (uses processor default if None)
        mode: Processing mode ('video' or 'frames')
        num_frames: Number of frames for 'frames' mode
        target_frames: Target frame count (for HunyuanVideo 4n+1)
        resolution_preset: Resolution preset ('256p', '512p', '768p', '1024p', '1536p')
        max_pixels: Custom pixel budget (mutually exclusive with resolution_preset)
        target_height: Explicit target height (disables bucketing)
        target_width: Explicit target width (disables bucketing)
        resize_mode: Interpolation mode for resizing
        center_crop: Whether to center crop
        deterministic: Use deterministic latent encoding
        output_format: Output format ('meta' or 'pt')
        caption_format: Caption format ('sidecar', 'meta_json', 'jsonl')
        caption_field: Field name for captions
        shard_size: Number of videos per metadata shard
        max_videos: Maximum number of videos to process
    """
    video_dir = Path(video_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Get processor and resolve model name
    processor = ProcessorRegistry.get(processor_name)
    if model_name is None:
        model_name = processor.default_model_name

    # Determine max_pixels
    if resolution_preset:
        if resolution_preset not in MultiTierBucketCalculator.RESOLUTION_PRESETS:
            raise ValueError(
                f"Unknown preset '{resolution_preset}'. "
                f"Available: {list(MultiTierBucketCalculator.RESOLUTION_PRESETS.keys())}"
            )
        max_pixels = MultiTierBucketCalculator.RESOLUTION_PRESETS[resolution_preset]
    elif max_pixels is None and target_height is None:
        # Default to 512p for videos
        max_pixels = 512 * 512

    # If explicit size given, disable bucketing
    use_bucketing = target_height is None or target_width is None
    if not use_bucketing and max_pixels is None:
        max_pixels = target_height * target_width  # Use explicit size as pixel budget

    num_gpus = torch.cuda.device_count()
    if num_gpus == 0:
        raise RuntimeError("No GPUs available")

    logger.info("Processor: %s (%s)", processor_name, processor.model_type)
    logger.info("Model: %s", model_name)
    logger.info("GPUs: %d", num_gpus)
    logger.info("Mode: %s", mode)
    if use_bucketing:
        logger.info("Max pixels: %d (bucketing enabled)", max_pixels)
        logger.info("Quantization: %d", getattr(processor, "quantization", 8))
    else:
        logger.info("Target size: %dx%d (bucketing disabled)", target_width, target_height)

    if hasattr(processor, "frame_constraint") and processor.frame_constraint:
        logger.info("Frame constraint: %s", processor.frame_constraint)

    # Get all video files
    logger.info("Scanning for videos...")
    video_files = _get_media_files(video_dir, VIDEO_EXTENSIONS)

    if max_videos is not None:
        video_files = video_files[:max_videos]

    logger.info("Found %d videos", len(video_files))

    if not video_files:
        return

    # Load captions using appropriate loader
    logger.info("Loading captions (format: %s, field: %s)...", caption_format, caption_field)
    caption_loader = get_caption_loader(caption_format)
    caption_cache = caption_loader.load_captions(video_files, caption_field)
    logger.info("  Loaded %d captions", len(caption_cache))

    # Video config for workers
    video_config = {
        "mode": mode,
        "num_frames": num_frames,
        "target_frames": target_frames,
        "target_height": target_height if not use_bucketing else None,
        "target_width": target_width if not use_bucketing else None,
        "resize_mode": resize_mode,
        "center_crop": center_crop,
        "deterministic": deterministic,
        "output_format": output_format,
    }

    # Split videos across GPUs
    chunks = [video_files[i::num_gpus] for i in range(num_gpus)]

    # Process with one worker per GPU
    all_metadata = []

    with Pool(processes=num_gpus) as pool:
        args = [
            (
                gpu_id,
                chunks[gpu_id],
                str(output_dir),
                processor_name,
                model_name,
                caption_cache,
                max_pixels,
                video_config,
            )
            for gpu_id in range(num_gpus)
        ]

        results = pool.starmap(_process_video_shard_on_gpu, args)

        for gpu_results in results:
            all_metadata.extend(gpu_results)

    # Save metadata
    _save_metadata_shards(
        all_metadata,
        output_dir,
        processor_name,
        model_name,
        processor.model_type,
        shard_size,
        {
            "caption_format": caption_format,
            "caption_field": caption_field,
            "max_pixels": max_pixels,
            "mode": mode,
            "target_frames": target_frames,
        },
    )

    # Print summary
    logger.info("=" * 50)
    logger.info("COMPLETE: %d/%d videos", len(all_metadata), len(video_files))
    logger.info("Output: %s", output_dir)
    _print_bucket_distribution(all_metadata)


# =============================================================================
# CLI Entry Point
# =============================================================================


def main():
    """Run image or video preprocessing from the command line."""

    parser = argparse.ArgumentParser(
        description="Unified preprocessing tool for images and videos",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Image preprocessing with FLUX
  python examples/diffusers/fastgen/preprocess_qwen_image.py image \\
      --image_dir /data/images --output_dir /cache --processor flux

  # Video preprocessing with Wan2.1
  python examples/diffusers/fastgen/preprocess_qwen_image.py video \\
      --video_dir /data/videos --output_dir /cache --processor wan \\
      --resolution_preset 512p --caption_format sidecar

  # Video preprocessing with HunyuanVideo
  python examples/diffusers/fastgen/preprocess_qwen_image.py video \\
      --video_dir /data/videos --output_dir /cache --processor hunyuan \\
      --target_frames 121 --caption_format meta_json
        """,
    )

    parser.add_argument(
        "--list_processors", action="store_true", help="List available processors and exit"
    )

    subparsers = parser.add_subparsers(dest="command", help="Preprocessing type")

    # ===================
    # Image subcommand
    # ===================
    image_parser = subparsers.add_parser("image", help="Preprocess images")
    image_parser.add_argument("--image_dir", type=str, required=True, help="Input image directory")
    image_parser.add_argument(
        "--output_dir", type=str, required=True, help="Output cache directory"
    )
    image_parser.add_argument(
        "--processor", type=str, default="qwen_image", help="Processor name (default: qwen_image)"
    )
    image_parser.add_argument(
        "--model_name", type=str, default=None, help="Model name (uses processor default)"
    )
    image_parser.add_argument("--shard_size", type=int, default=10000, help="Metadata shard size")
    image_parser.add_argument("--verify", action="store_true", help="Verify latents can be decoded")
    image_parser.add_argument(
        "--caption_field",
        type=str,
        default="internvl",
        help="Field name inside each caption record (e.g. 'internvl', 'caption')",
    )
    image_parser.add_argument(
        "--caption_format",
        type=str,
        default="jsonl",
        choices=["sidecar", "meta_json", "jsonl"],
        help="Caption file format (default: jsonl for backward compat)",
    )
    image_parser.add_argument("--max_images", type=int, default=None, help="Max images to process")
    image_parser.add_argument(
        "--shard_idx",
        type=int,
        default=0,
        help="Rank within a multi-job sweep (0-indexed). Default 0.",
    )
    image_parser.add_argument(
        "--shard_count",
        type=int,
        default=1,
        help="Total jobs in the sweep. When >1, this rank processes image_files[shard_idx::shard_count].",
    )

    # Resolution options (mutually exclusive)
    image_res_group = image_parser.add_mutually_exclusive_group()
    image_res_group.add_argument(
        "--resolution_preset",
        type=str,
        choices=["256p", "512p", "768p", "1024p", "1536p"],
        help="Resolution preset for bucketing",
    )
    image_res_group.add_argument("--max_pixels", type=int, help="Custom max pixel budget")

    # ===================
    # Video subcommand
    # ===================
    video_parser = subparsers.add_parser("video", help="Preprocess videos")
    video_parser.add_argument("--video_dir", type=str, required=True, help="Input video directory")
    video_parser.add_argument(
        "--output_dir", type=str, required=True, help="Output cache directory"
    )
    video_parser.add_argument(
        "--processor",
        type=str,
        required=True,
        choices=["wan", "wan2.1", "hunyuan", "hunyuanvideo", "hunyuanvideo-1.5"],
    )
    video_parser.add_argument(
        "--model_name", type=str, default=None, help="Model name (uses processor default)"
    )
    video_parser.add_argument(
        "--mode", type=str, default="video", choices=["video", "frames"], help="Processing mode"
    )
    video_parser.add_argument(
        "--num_frames", type=int, default=10, help="Frames to extract in 'frames' mode"
    )
    video_parser.add_argument(
        "--target_frames",
        type=int,
        default=None,
        help="Target frame count (e.g., 121 for HunyuanVideo)",
    )

    # Resolution options
    video_res_group = video_parser.add_mutually_exclusive_group()
    video_res_group.add_argument(
        "--resolution_preset",
        type=str,
        choices=["256p", "512p", "768p", "1024p", "1536p"],
        help="Resolution preset (videos bucketed by aspect ratio)",
    )
    video_res_group.add_argument("--max_pixels", type=int, help="Custom pixel budget for bucketing")

    # Explicit size options (disables bucketing)
    video_parser.add_argument(
        "--height", type=int, default=None, help="Explicit height (disables bucketing)"
    )
    video_parser.add_argument(
        "--width", type=int, default=None, help="Explicit width (disables bucketing)"
    )

    video_parser.add_argument(
        "--resize_mode",
        type=str,
        default="bilinear",
        choices=["bilinear", "bicubic", "nearest", "area", "lanczos"],
        help="Interpolation mode",
    )
    video_parser.add_argument(
        "--center_crop", action="store_true", default=True, help="Center crop (default: True)"
    )
    video_parser.add_argument(
        "--no_center_crop", dest="center_crop", action="store_false", help="Disable center crop"
    )
    video_parser.add_argument(
        "--deterministic",
        action="store_true",
        default=True,
        help="Use deterministic encoding (default: True)",
    )
    video_parser.add_argument(
        "--stochastic",
        dest="deterministic",
        action="store_false",
        help="Use stochastic (sampled) encoding",
    )
    video_parser.add_argument(
        "--caption_format",
        type=str,
        default="sidecar",
        choices=["sidecar", "meta_json", "jsonl"],
        help="Caption format",
    )
    video_parser.add_argument(
        "--caption_field", type=str, default="caption", help="Caption field name"
    )
    video_parser.add_argument(
        "--output_format",
        type=str,
        default="meta",
        choices=["meta", "pt"],
        help="Output file format",
    )
    video_parser.add_argument("--shard_size", type=int, default=10000, help="Metadata shard size")
    video_parser.add_argument("--max_videos", type=int, default=None, help="Max videos to process")

    args = parser.parse_args()

    # Handle --list_processors
    if args.list_processors:
        logger.info("Available processors:")
        for name in ProcessorRegistry.list_available():
            proc = ProcessorRegistry.get(name)
            quantization = getattr(proc, "quantization", 64)
            logger.info("  %s:", name)
            logger.info("    type: %s", proc.model_type)
            logger.info("    media: image")
            logger.info("    quantization: %d", quantization)
        return

    # Handle subcommands
    if args.command == "image":
        if args.resolution_preset:
            max_pixels = MultiTierBucketCalculator.RESOLUTION_PRESETS[args.resolution_preset]
        elif args.max_pixels:
            max_pixels = args.max_pixels
        else:
            max_pixels = 256 * 256

        preprocess_dataset(
            args.image_dir,
            args.output_dir,
            args.processor,
            args.model_name,
            args.shard_size,
            args.verify,
            args.caption_field,
            args.caption_format,
            args.max_images,
            max_pixels,
            args.shard_idx,
            args.shard_count,
        )

    elif args.command == "video":
        # Validate explicit size args
        if (args.height is None) != (args.width is None):
            parser.error("Both --height and --width must be specified together")

        preprocess_video_dataset(
            video_dir=args.video_dir,
            output_dir=args.output_dir,
            processor_name=args.processor,
            model_name=args.model_name,
            mode=args.mode,
            num_frames=args.num_frames,
            target_frames=args.target_frames,
            resolution_preset=args.resolution_preset,
            max_pixels=args.max_pixels,
            target_height=args.height,
            target_width=args.width,
            resize_mode=args.resize_mode,
            center_crop=args.center_crop,
            deterministic=args.deterministic,
            output_format=args.output_format,
            caption_format=args.caption_format,
            caption_field=args.caption_field,
            shard_size=args.shard_size,
            max_videos=args.max_videos,
        )
    else:
        parser.print_help()


if __name__ == "__main__":
    main()