"""
LoomVideo inference script.

Supports:
- t2v: Text-to-Video / Text-to-Image (num_frames=1)
- edit: Image/Video editing with source image or video + text instruction
- ref_edit: Image/Video editing with source + reference image(s) + text instruction
- mi2v: Multi-Image to Video generation from reference image(s) + text instruction

All output dimensions are aligned to the combined Wan VAE + DiT patch grid.
"""
import argparse
import os
import sys
import datetime
import logging
import gc
import types
from typing import List, Optional

# Reduce CUDA allocator fragmentation before torch initializes CUDA.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import numpy as np
import torch
from torchvision import transforms
from omegaconf import OmegaConf
from PIL import Image
from decord import VideoReader, cpu

from accelerate import Accelerator
from accelerate.utils import set_seed
from diffusers.utils.export_utils import export_to_video

from transformers import AutoProcessor


# Add project root to path
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
sys.path.insert(0, PROJECT_ROOT)

from src.models.utils import load_model, load_checkpoint
from src.models.transformers.loomvideo import LoomVideo

logger = logging.getLogger(__name__)


VAE_TEMPORAL_SCALE = 4
DEFAULT_MAX_EDIT_FRAMES = 121
DEFAULT_SPATIAL_ALIGNMENT = 32


def _module_device(module: torch.nn.Module) -> torch.device:
    """Return the device of the first parameter or buffer in a module."""
    try:
        return next(module.parameters()).device
    except StopIteration:
        try:
            return next(module.buffers()).device
        except StopIteration:
            return torch.device("cpu")


def get_und_device(model: LoomVideo) -> torch.device:
    """Device used by the Qwen3-VL understanding backbone."""
    return _module_device(model.und_model)


def get_gen_device(model: LoomVideo) -> torch.device:
    """Device used by the Wan diffusion transformer."""
    return _module_device(model.gen_model)


def get_vae_device(model: LoomVideo) -> torch.device:
    """Device used by the Wan VAE."""
    return _module_device(model.gen_vae)


def _cuda_memory_summary(device: torch.device) -> str:
    if device.type != "cuda":
        return str(device)
    index = device.index if device.index is not None else torch.cuda.current_device()
    free_bytes, total_bytes = torch.cuda.mem_get_info(index)
    return (
        f"cuda:{index}, free={free_bytes / 1024**3:.2f} GiB, "
        f"total={total_bytes / 1024**3:.2f} GiB"
    )


def configure_two_gpu_inference(
    model: LoomVideo,
    und_device: torch.device,
    gen_device: torch.device,
) -> None:
    """Place Qwen on one GPU and Wan DiT/VAE on the other in one process."""
    und_device = torch.device(und_device)
    gen_device = torch.device(gen_device)

    if und_device.type != "cuda" or gen_device.type != "cuda":
        raise ValueError("Two-GPU mode requires CUDA devices for both backbones.")
    if und_device == gen_device:
        raise ValueError("und_device and gen_device must be different CUDA devices.")

    logger.info("Placing understanding backbone on %s", und_device)
    model.und_model.to(device=und_device, dtype=torch.bfloat16)

    logger.info("Placing generation backbone and VAE on %s", gen_device)
    model.gen_model.to(device=gen_device, dtype=torch.bfloat16)
    model.gen_vae.to(device=gen_device, dtype=torch.bfloat16)

    # LoomVideo.generate() consumes the precomputed fixed T5 embeddings during
    # inference, so the full text encoder can stay on CPU.
    if hasattr(model, "text_encoder") and model.text_encoder is not None:
        model.text_encoder.to(device="cpu", dtype=torch.bfloat16)

    if hasattr(model, "fixed_t5_embeds") and model.fixed_t5_embeds is not None:
        model.fixed_t5_embeds = model.fixed_t5_embeds.to(
            device=gen_device,
            dtype=torch.bfloat16,
        )

    hidden_state_cache = {}
    original_forward_gen = model.forward_gen
    original_generate = model.generate

    def forward_gen_two_gpu(self, all_und_hidden_states, *args, **kwargs):
        # The same CFG branches are reused at every denoising step. Move each
        # branch to the generation GPU once per generate() call.
        cache_key = id(all_und_hidden_states)
        moved_hidden_states = hidden_state_cache.get(cache_key)
        if moved_hidden_states is None:
            moved_values = [
                value.to(
                    device=gen_device,
                    dtype=torch.bfloat16,
                    non_blocking=True,
                )
                if torch.is_tensor(value)
                else value
                for value in all_und_hidden_states
            ]
            moved_hidden_states = (
                tuple(moved_values)
                if isinstance(all_und_hidden_states, tuple)
                else moved_values
            )
            hidden_state_cache[cache_key] = moved_hidden_states

        return original_forward_gen(moved_hidden_states, *args, **kwargs)

    def generate_two_gpu(self, *args, **kwargs):
        hidden_state_cache.clear()
        try:
            return original_generate(*args, **kwargs)
        finally:
            hidden_state_cache.clear()
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    model.forward_gen = types.MethodType(forward_gen_two_gpu, model)
    model.generate = types.MethodType(generate_two_gpu, model)
    model.eval()

    logger.info("Understanding GPU: %s", _cuda_memory_summary(und_device))
    logger.info("Generation GPU: %s", _cuda_memory_summary(gen_device))


def get_required_spatial_alignment(model: LoomVideo) -> tuple[int, int]:
    """Return required pixel-space (height, width) alignment for Wan DiT.

    LoomVideo first downsamples pixels with the VAE and then patchifies the
    latent with the Wan transformer. Therefore each pixel dimension must be
    divisible by ``vae_scale_factor_spatial * transformer_patch_size``.
    For the released LoomVideo checkpoint this is normally 16 * 2 = 32.
    """
    vae_scale = int(getattr(model, "vae_scale_factor_spatial", 16))
    patch_size = getattr(model.gen_model.config, "patch_size", (1, 2, 2))

    if isinstance(patch_size, int):
        patch_h = patch_w = int(patch_size)
    else:
        patch_values = tuple(int(v) for v in patch_size)
        if len(patch_values) >= 3:
            patch_h, patch_w = patch_values[-2], patch_values[-1]
        elif len(patch_values) == 2:
            patch_h, patch_w = patch_values
        elif len(patch_values) == 1:
            patch_h = patch_w = patch_values[0]
        else:
            patch_h = patch_w = 2

    align_h = max(1, vae_scale * patch_h)
    align_w = max(1, vae_scale * patch_w)
    return align_h, align_w


def normalize_spatial_dimensions(
    model: LoomVideo,
    height: int,
    width: int,
    *,
    round_down: bool = True,
) -> tuple[int, int]:
    """Normalize output dimensions to VAE + DiT patch-compatible values.

    By default values are rounded down so a request never silently consumes
    more VRAM. With the released checkpoint, 240x240 becomes 224x224 while
    480x832 remains unchanged.
    """
    height = int(height)
    width = int(width)
    if height < 1 or width < 1:
        raise ValueError(f"height and width must be positive, got {width}x{height}.")

    align_h, align_w = get_required_spatial_alignment(model)
    if round_down:
        aligned_h = (height // align_h) * align_h
        aligned_w = (width // align_w) * align_w
    else:
        aligned_h = round(height / align_h) * align_h
        aligned_w = round(width / align_w) * align_w

    aligned_h = max(align_h, int(aligned_h))
    aligned_w = max(align_w, int(aligned_w))
    return aligned_h, aligned_w


def normalize_num_frames(
    num_frames: int,
    *,
    allow_single_frame: bool = True,
    temporal_scale: int = VAE_TEMPORAL_SCALE,
) -> int:
    """Normalize a frame count to the VAE-compatible form T = temporal_scale*k + 1.

    Values are rounded down so a request cannot silently use more frames/VRAM
    than requested. For example, 60 becomes 57 and 97 remains 97.
    """
    value = int(num_frames)
    if value < 1:
        raise ValueError(f"num_frames must be >= 1, got {value}.")
    if value == 1:
        if allow_single_frame:
            return 1
        raise ValueError("Video editing requires more than one source frame.")

    aligned = ((value - 1) // temporal_scale) * temporal_scale + 1
    if aligned < temporal_scale + 1:
        aligned = temporal_scale + 1
    return aligned

# Preprocessing: normalize to [-1, 1] for VAE/DiT conditioning
IMAGE_TRANSFORM = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
])

def resize_and_crop(image: Image.Image, target_height: int, target_width: int) -> Image.Image:
    """
    Resize and center-crop image to exact target dimensions.

    Args:
        image: PIL Image.
        target_height: Desired output height.
        target_width: Desired output width.

    Returns:
        Resized and cropped PIL Image.
    """
    width, height = image.size
    scale = max(target_width / width, target_height / height)
    image = transforms.functional.resize(
        image,
        (round(height * scale), round(width * scale)),
        interpolation=transforms.InterpolationMode.BILINEAR,
    )
    image = transforms.functional.center_crop(image, (target_height, target_width))
    return image


def load_image(image_path: str, target_height: int = None, target_width: int = None) -> tuple:
    """
    Load and preprocess a single image, returning both PIL and tensor.

    Args:
        image_path: Path to the image file.
        target_height: If set, resize image to this height.
        target_width: If set, resize image to this width.

    Returns:
        Tuple of (pil_image, tensor):
            pil_image: Resized PIL Image (for VLM input).
            tensor: Tensor of shape [C, H, W] normalized to [-1, 1] (for DiT conditioning).
    """
    image = Image.open(image_path).convert("RGB")
    if target_height is not None and target_width is not None:
        image = resize_and_crop(image, target_height, target_width)
    return image, IMAGE_TRANSFORM(image)

def load_images(
    image_paths: List[str], target_height: int = None, target_width: int = None
) -> tuple:
    """
    Load multiple images, returning both PIL images and tensors.

    Args:
        image_paths: List of image file paths.
        target_height: If set, resize each image to this height.
        target_width: If set, resize each image to this width.

    Returns:
        Tuple of (pil_images, tensors):
            pil_images: List of resized PIL Images.
            tensors: List of tensors, each [C, H, W] normalized to [-1, 1].
    """
    pil_images = []
    tensors = []
    for path in image_paths:
        pil_img, tensor = load_image(path, target_height, target_width)
        pil_images.append(pil_img)
        tensors.append(tensor)
    return pil_images, tensors



def load_video(
    video_path: str,
    target_fps: int = 24,
    target_num_frames: int = DEFAULT_MAX_EDIT_FRAMES,
    scale_factor_temporal: int = VAE_TEMPORAL_SCALE,
    target_height: int = None,
    target_width: int = None,
) -> tuple:
    """Load a source video and return VAE-aligned frames and [C,T,H,W] tensor.

    The returned temporal length is always ``4k+1`` (or generally
    ``scale_factor_temporal*k+1``). The caller must use the returned tensor's
    actual T as the generated ``num_frames``.
    """
    target_num_frames = normalize_num_frames(
        target_num_frames,
        allow_single_frame=True,
        temporal_scale=scale_factor_temporal,
    )
    if target_fps <= 0:
        raise ValueError(f"target_fps must be positive, got {target_fps}.")

    video_reader = VideoReader(video_path, ctx=cpu(0))
    num_video_frames = len(video_reader)
    if num_video_frames < 1:
        raise ValueError(f"Source video has no readable frames: {video_path}")

    video_fps = float(video_reader.get_avg_fps())
    if not np.isfinite(video_fps) or video_fps <= 0:
        video_fps = float(target_fps)

    # Preserve the official target-FPS sampling behaviour, but cap it using the
    # requested, already-aligned output length.
    target_duration = (target_num_frames - 1) / float(target_fps)
    num_video_in_range_frames = min(
        int(video_fps * target_duration) + 1,
        num_video_frames,
    )

    frame_float = 0.0
    interval = video_fps / float(target_fps)
    frame_indices = []
    max_loops = max(target_num_frames * 3, 8)

    for _ in range(max_loops):
        idx = min(round(frame_float), num_video_frames - 1)
        frame_indices.append(idx)
        frame_float += interval

        if (round(frame_float)) >= num_video_in_range_frames:
            break
        if len(frame_indices) >= target_num_frames:
            break

    if not frame_indices:
        frame_indices = [0]

    # Never let a video-conditioning request accidentally become a non-aligned
    # arbitrary length such as 60. 60 -> 57; 97 -> 97.
    needed_len = normalize_num_frames(
        min(len(frame_indices), target_num_frames),
        allow_single_frame=True,
        temporal_scale=scale_factor_temporal,
    )
    needed_len = min(needed_len, len(frame_indices))
    if needed_len > 1:
        needed_len = ((needed_len - 1) // scale_factor_temporal) * scale_factor_temporal + 1
    frame_indices = frame_indices[:needed_len]

    batch = video_reader.get_batch(frame_indices)
    video_data = batch.asnumpy()  # [T, H, W, C]

    pil_frames = []
    frame_tensors = []
    for frame in video_data:
        frame_pil = Image.fromarray(frame)
        if target_height is not None and target_width is not None:
            frame_pil = resize_and_crop(frame_pil, target_height, target_width)
        pil_frames.append(frame_pil)
        frame_tensors.append(IMAGE_TRANSFORM(frame_pil))

    if not frame_tensors:
        raise ValueError(f"Failed to decode source video: {video_path}")

    video_tensor = torch.stack(frame_tensors, dim=0).permute(1, 0, 2, 3)
    actual_frames = int(video_tensor.shape[1])
    if actual_frames > 1 and (actual_frames - 1) % scale_factor_temporal != 0:
        raise RuntimeError(
            "Internal frame-alignment error: source video produced "
            f"{actual_frames} frames, expected {scale_factor_temporal}k+1."
        )
    return pil_frames, video_tensor

def temporal_downsample_for_vlm(pil_frames: List[Image.Image], config) -> List[Image.Image]:
    """
    Downsample video frames for VLM input, matching evaluation pipeline.

    Args:
        pil_frames: Full list of PIL frames after resize.
        config: OmegaConf config object.

    Returns:
        Downsampled list of PIL frames for VLM input.
    """
    downsample_factor = OmegaConf.select(
        config, "data.und_temporal_downsample_factor", default=15
    )
    total_frames = len(pil_frames)
    if total_frames > downsample_factor:
        indices = np.linspace(
            0, total_frames - 1,
            max(total_frames // downsample_factor, 1),
            dtype=int,
        ).tolist()
    else:
        indices = list(range(total_frames))
    return [pil_frames[i] for i in indices]

def _get_negative_prompt(args, config) -> str:
    """Resolve negative prompt from args or config."""
    if args.negative_prompt is not None:
        return args.negative_prompt
    return OmegaConf.select(config, "generation.negative_prompt", default="")


def _get_guidance_scale(args, config, is_edit: bool) -> float:
    """
    Resolve guidance scale: use edit-specific scale for edit tasks.
    """
    if args.guidance_scale is not None:
        return args.guidance_scale
    if is_edit:
        return OmegaConf.select(config, "generation.guidance_scale_edit", default=2.5)
    return OmegaConf.select(config, "generation.guidance_scale", default=5.0)


def _get_guidance_scale_visual(args, config) -> float:
    """Resolve visual guidance scale from args or config."""
    if args.guidance_scale_visual is not None:
        return args.guidance_scale_visual
    return OmegaConf.select(config, "generation.guidance_scale_visual", default=1.5)


def _get_generator(args, device) -> Optional[torch.Generator]:
    """Create a random generator if seed is set."""
    if args.seed >= 0:
        return torch.Generator(device=device).manual_seed(args.seed)
    return None



def prepare_t2v_inputs(model: LoomVideo, config, args) -> dict:
    """Prepare inputs for text-to-video or text-to-image generation."""
    processor = AutoProcessor.from_pretrained(
        config.model.und.pretrained_model_path, trust_remote_code=True
    )

    requested_height = args.height if args.height is not None else 480
    requested_width = args.width if args.width is not None else 832
    height, width = normalize_spatial_dimensions(
        model, requested_height, requested_width
    )
    requested_frames = args.num_frames if args.num_frames is not None else 81
    num_frames = normalize_num_frames(requested_frames, allow_single_frame=True)

    is_image = num_frames == 1
    prefix = "Generate an image: " if is_image else "Generate a video: "
    messages = [
        {
            "role": "user",
            "content": [{"type": "text", "text": f"{prefix}{args.prompt}"}],
        }
    ]

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(get_und_device(model))

    return {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"],
        "negative_prompt": _get_negative_prompt(args, config),
        "height": int(height),
        "width": int(width),
        "num_frames": int(num_frames),
        "num_inference_steps": args.num_inference_steps,
        "guidance_scale": _get_guidance_scale(args, config, is_edit=False),
        "guidance_scale_visual": _get_guidance_scale_visual(args, config),
        "generator": _get_generator(args, get_gen_device(model)),
    }


def prepare_edit_inputs(model: LoomVideo, config, args) -> dict:
    """Prepare image/video editing inputs with strict temporal alignment."""
    processor = AutoProcessor.from_pretrained(
        config.model.und.pretrained_model_path, trust_remote_code=True
    )

    source_path = args.source_video_path
    image_extensions = (".jpg", ".jpeg", ".png", ".bmp", ".webp")
    is_source_image = source_path.lower().endswith(image_extensions)

    if is_source_image:
        source_image = Image.open(source_path).convert("RGB")
        native_width, native_height = source_image.size
        requested_height = int(
            args.height if args.height is not None else native_height
        )
        requested_width = int(
            args.width if args.width is not None else native_width
        )
        height, width = normalize_spatial_dimensions(
            model, requested_height, requested_width
        )

        source_image = resize_and_crop(source_image, height, width)
        # IMAGE_TRANSFORM yields [C, H, W]; the rest of the pipeline (VAE
        # encoding + validation) requires [C, T, H, W] with T=1 for images.
        source_tensor = (
            IMAGE_TRANSFORM(source_image)
            .unsqueeze(1)
            .to(get_vae_device(model))
        )
        num_frames = 1
        content = [
            {"type": "image", "image": source_image},
            {"type": "text", "text": f"Edit this image: {args.prompt}"},
        ]
    else:
        probe_reader = VideoReader(source_path, ctx=cpu(0))
        if len(probe_reader) < 1:
            raise ValueError(f"Source video has no readable frames: {source_path}")
        first_frame = probe_reader[0].asnumpy()
        native_height, native_width = first_frame.shape[:2]
        del probe_reader

        requested_height = int(
            args.height if args.height is not None else native_height
        )
        requested_width = int(
            args.width if args.width is not None else native_width
        )
        height, width = normalize_spatial_dimensions(
            model, requested_height, requested_width
        )

        requested_frames = (
            normalize_num_frames(args.num_frames, allow_single_frame=True)
            if args.num_frames is not None
            else DEFAULT_MAX_EDIT_FRAMES
        )
        pil_frames, source_tensor = load_video(
            source_path,
            target_num_frames=requested_frames,
            scale_factor_temporal=VAE_TEMPORAL_SCALE,
            target_height=height,
            target_width=width,
        )
        source_tensor = source_tensor.to(get_vae_device(model))

        # Critical fix: generated target and clean source condition use the exact
        # same temporal length. Never retain the unaligned user value (e.g. 60).
        num_frames = int(source_tensor.shape[1])
        if num_frames > 1 and (num_frames - 1) % VAE_TEMPORAL_SCALE != 0:
            raise RuntimeError(f"Source tensor has invalid temporal length: {num_frames}")

        vlm_frames = temporal_downsample_for_vlm(pil_frames, config)
        content = [
            {"type": "video", "video": vlm_frames},
            {"type": "text", "text": f"Edit this video: {args.prompt}"},
        ]

    messages = [{"role": "user", "content": content}]
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(get_und_device(model))

    result = {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"],
        "pixel_values": inputs.get("pixel_values"),
        "pixel_values_videos": inputs.get("pixel_values_videos"),
        "image_grid_thw": inputs.get("image_grid_thw"),
        "video_grid_thw": inputs.get("video_grid_thw"),
        "negative_prompt": _get_negative_prompt(args, config),
        "height": height,
        "width": width,
        "num_frames": num_frames,
        "num_inference_steps": args.num_inference_steps,
        "guidance_scale": _get_guidance_scale(args, config, is_edit=True),
        "guidance_scale_visual": _get_guidance_scale_visual(args, config),
        "source_pixel_values": [source_tensor],
        "generator": _get_generator(args, get_gen_device(model)),
    }
    _validate_generation_alignment(result, task="edit", model=model)
    return result


def prepare_ref_edit_inputs(model: LoomVideo, config, args) -> dict:
    """Prepare reference-guided video editing with strict temporal alignment."""
    processor = AutoProcessor.from_pretrained(
        config.model.und.pretrained_model_path, trust_remote_code=True
    )

    probe_reader = VideoReader(args.source_video_path, ctx=cpu(0))
    if len(probe_reader) < 1:
        raise ValueError(f"Source video has no readable frames: {args.source_video_path}")
    first_frame = probe_reader[0].asnumpy()
    native_height, native_width = first_frame.shape[:2]
    del probe_reader

    requested_height = int(
        args.height if args.height is not None else native_height
    )
    requested_width = int(
        args.width if args.width is not None else native_width
    )
    height, width = normalize_spatial_dimensions(
        model, requested_height, requested_width
    )
    requested_frames = (
        normalize_num_frames(args.num_frames, allow_single_frame=True)
        if args.num_frames is not None
        else DEFAULT_MAX_EDIT_FRAMES
    )

    source_pil_frames, source_video_tensor = load_video(
        args.source_video_path,
        target_num_frames=requested_frames,
        scale_factor_temporal=VAE_TEMPORAL_SCALE,
        target_height=height,
        target_width=width,
    )
    source_video_tensor = source_video_tensor.to(get_vae_device(model))
    num_frames = int(source_video_tensor.shape[1])

    ref_pil_images, ref_tensors = load_images(
        args.ref_image_paths, target_height=height, target_width=width
    )
    ref_image_tensors = [t.to(get_vae_device(model)) for t in ref_tensors]

    vlm_frames = temporal_downsample_for_vlm(source_pil_frames, config)
    content = []
    for ref_pil in ref_pil_images:
        content.append({"type": "image", "image": ref_pil})
    content.append({"type": "video", "video": vlm_frames})
    content.append(
        {
            "type": "text",
            "text": f"Edit this video with reference images: {args.prompt}",
        }
    )

    messages = [{"role": "user", "content": content}]
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(get_und_device(model))

    result = {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"],
        "pixel_values": inputs.get("pixel_values"),
        "pixel_values_videos": inputs.get("pixel_values_videos"),
        "image_grid_thw": inputs.get("image_grid_thw"),
        "video_grid_thw": inputs.get("video_grid_thw"),
        "negative_prompt": _get_negative_prompt(args, config),
        "height": height,
        "width": width,
        "num_frames": num_frames,
        "num_inference_steps": args.num_inference_steps,
        "guidance_scale": _get_guidance_scale(args, config, is_edit=True),
        "guidance_scale_visual": _get_guidance_scale_visual(args, config),
        "source_pixel_values": [source_video_tensor],
        "ref_pixel_values": ref_image_tensors,
        "generator": _get_generator(args, get_gen_device(model)),
    }
    _validate_generation_alignment(result, task="ref_edit", model=model)
    return result


def prepare_mi2v_inputs(model: LoomVideo, config, args) -> dict:
    """Prepare inputs for multi-image-to-video generation."""
    processor = AutoProcessor.from_pretrained(
        config.model.und.pretrained_model_path, trust_remote_code=True
    )

    ref_image = Image.open(args.ref_image_paths[0]).convert("RGB")
    native_width, native_height = ref_image.size
    requested_height = int(
        args.height if args.height is not None else native_height
    )
    requested_width = int(
        args.width if args.width is not None else native_width
    )
    height, width = normalize_spatial_dimensions(
        model, requested_height, requested_width
    )
    requested_frames = args.num_frames if args.num_frames is not None else 81
    num_frames = normalize_num_frames(requested_frames, allow_single_frame=True)

    ref_pil_images, ref_tensors = load_images(
        args.ref_image_paths, target_height=height, target_width=width
    )
    ref_image_tensors = [t.to(get_vae_device(model)) for t in ref_tensors]

    content = [
        {"type": "image", "image": ref_pil}
        for ref_pil in ref_pil_images
    ]
    content.append(
        {"type": "text", "text": f"Generate a video with reference images: {args.prompt}"}
    )
    messages = [{"role": "user", "content": content}]

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(get_und_device(model))

    return {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"],
        "pixel_values": inputs.get("pixel_values"),
        "image_grid_thw": inputs.get("image_grid_thw"),
        "negative_prompt": _get_negative_prompt(args, config),
        "height": height,
        "width": width,
        "num_frames": int(num_frames),
        "num_inference_steps": args.num_inference_steps,
        "guidance_scale": _get_guidance_scale(args, config, is_edit=False),
        "guidance_scale_visual": _get_guidance_scale_visual(args, config),
        "ref_pixel_values": ref_image_tensors,
        "generator": _get_generator(args, get_gen_device(model)),
    }


def _validate_generation_alignment(
    generate_kwargs: dict,
    task: str,
    model: Optional[LoomVideo] = None,
) -> None:
    """Fail early on temporal or spatial incompatibilities."""
    num_frames = int(generate_kwargs["num_frames"])
    if num_frames > 1 and (num_frames - 1) % VAE_TEMPORAL_SCALE != 0:
        raise ValueError(
            f"{task}: num_frames={num_frames} is invalid; expected 4k+1 "
            "such as 49, 57, 81, 97 or 121."
        )

    height = int(generate_kwargs["height"])
    width = int(generate_kwargs["width"])
    if model is not None:
        align_h, align_w = get_required_spatial_alignment(model)
    else:
        align_h = align_w = DEFAULT_SPATIAL_ALIGNMENT

    if height % align_h != 0 or width % align_w != 0:
        raise ValueError(
            f"{task}: resolution {width}x{height} is incompatible with the "
            f"VAE/DiT. Width must be divisible by {align_w} and height by "
            f"{align_h}. For example, use 224x224 or 256x256 instead of "
            "240x240."
        )

    sources = generate_kwargs.get("source_pixel_values") or []
    for index, source in enumerate(sources):
        if not torch.is_tensor(source) or source.ndim != 4:
            raise ValueError(
                f"{task}: source_pixel_values[{index}] must be [C,T,H,W], "
                f"got {type(source)!r}."
            )
        source_frames = int(source.shape[1])
        if source_frames != num_frames:
            raise ValueError(
                f"{task}: source conditioning has {source_frames} frames but "
                f"generation requests {num_frames}. These values must be identical."
            )
        source_h, source_w = int(source.shape[-2]), int(source.shape[-1])
        if source_h != height or source_w != width:
            raise ValueError(
                f"{task}: source conditioning is {source_w}x{source_h}, but "
                f"generation requests {width}x{height}. These values must match."
            )

    refs = generate_kwargs.get("ref_pixel_values") or []
    for index, ref in enumerate(refs):
        if not torch.is_tensor(ref) or ref.ndim not in (3, 4):
            raise ValueError(
                f"{task}: ref_pixel_values[{index}] must be [C,H,W] or "
                f"[C,T,H,W], got {type(ref)!r}."
            )
        ref_h, ref_w = int(ref.shape[-2]), int(ref.shape[-1])
        if ref_h != height or ref_w != width:
            raise ValueError(
                f"{task}: reference conditioning is {ref_w}x{ref_h}, but "
                f"generation requests {width}x{height}. These values must match."
            )


def save_output(output_np: np.ndarray, output_path: str, fps: int = 24):
    """
    Save generation output as image (T=1) or video (T>1).

    Args:
        output_np: Output array of shape [T, H, W, C] with values in [0, 1].
        output_path: Output file path. Extension is auto-corrected based on content.
        fps: Frames per second (only used for video).
    """
    num_frames = output_np.shape[0]

    # Single frame → save as image (matching evaluation: .jpg format)
    if num_frames == 1:
        if output_path.endswith(".mp4"):
            output_path = output_path.replace(".mp4", ".jpg")
        os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
        frame = (output_np[0] * 255).clip(0, 255).astype(np.uint8)
        Image.fromarray(frame).save(output_path)
        logger.info(f"Saved image to: {output_path}")
        return

    # Multiple frames → save as video using export_to_video (matching evaluation)
    if not output_path.endswith(".mp4"):
        output_path = output_path.rsplit(".", 1)[0] + ".mp4"
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)

    export_to_video(output_np, output_path, fps=fps)

    logger.info(f"Saved video to: {output_path} ({num_frames} frames, {fps} fps)")

def generate_task(model: LoomVideo, config, args, accelerator: Accelerator):
    """
    Unified generation entry point for all tasks.

    Args:
        model: Loaded LoomVideo model.
        config: OmegaConf config object.
        args: Parsed arguments.
        accelerator: Accelerator instance.
    """
    task = args.task

    # Prepare task-specific inputs
    if task == "t2v":
        generate_kwargs = prepare_t2v_inputs(model, config, args)
    elif task == "edit":
        generate_kwargs = prepare_edit_inputs(model, config, args)
    elif task == "ref_edit":
        generate_kwargs = prepare_ref_edit_inputs(model, config, args)
    elif task == "mi2v":
        generate_kwargs = prepare_mi2v_inputs(model, config, args)
    else:
        raise NotImplementedError(f"Task '{task}' is not yet implemented.")

    _validate_generation_alignment(generate_kwargs, task=task, model=model)

    # Log resolved parameters after preparation
    resolved_height = generate_kwargs["height"]
    resolved_width = generate_kwargs["width"]
    resolved_num_frames = generate_kwargs["num_frames"]

    if accelerator.is_main_process:
        task_display = {
            "t2v": "Text-to-Video",
            "edit": "Video Editing",
            "ref_edit": "Reference-guided Editing",
            "mi2v": "Multi-Image to Video",
        }
        logger.info(f"Task: {task_display.get(task, task)}")
        logger.info(f"Prompt: {args.prompt}")
        if args.source_video_path:
            logger.info(f"Source: {args.source_video_path}")
        if args.ref_image_paths:
            logger.info(f"Reference images: {args.ref_image_paths}")
        logger.info(
            f"Resolution: {resolved_width}x{resolved_height}, "
            f"Frames: {resolved_num_frames}"
        )
        logger.info(
            f"Steps: {args.num_inference_steps}, "
            f"CFG scale: {generate_kwargs['guidance_scale']}, "
            f"Visual CFG: {generate_kwargs['guidance_scale_visual']}"
        )
        logger.info("Starting generation...")

    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        output = model.generate(**generate_kwargs)

    # Save output (only on main process)
    if accelerator.is_main_process:
        # fps: user-specified > default 24 (matching evaluation config.data.train.fps=24)
        fps = args.fps if args.fps is not None else 24

        output_path = args.output_path
        if output_path is None:
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            ext = "png" if resolved_num_frames == 1 else "mp4"
            output_path = f"outputs/loomvideo_{task}_{timestamp}.{ext}"

        save_output(output, output_path, fps=fps)

    accelerator.wait_for_everyone()


def main():
    parser = argparse.ArgumentParser(
        description="LoomVideo Inference",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument("--config_path", type=str, required=True, help="Path to YAML config file.")
    parser.add_argument(
        "--ckpt_path", type=str, default=None,
        help="Path to trained checkpoint (directory or file). "
             "Directory: loads ema/ema.pth or latest/<component>.pth. "
             "File: loads .safetensors or .pt directly.",
    )

    parser.add_argument(
        "--task", type=str, default="t2v",
        choices=["t2v", "edit", "ref_edit", "mi2v"],
        help="Generation task type.",
    )

    parser.add_argument("--prompt", type=str, required=True, help="Text prompt / editing instruction.")
    parser.add_argument("--negative_prompt", type=str, default=None,
                        help="Negative prompt for CFG. If not set, uses value from config.")
    parser.add_argument("--height", type=int, default=None,
                        help="Output height. If not set: t2v uses 480; edit/ref_edit infers from source.")
    parser.add_argument("--width", type=int, default=None,
                        help="Output width. If not set: t2v uses 832; edit/ref_edit infers from source.")
    parser.add_argument("--num_frames", type=int, default=None,
                        help="Number of output frames. If not set: t2v uses 81; edit infers from source.")
    parser.add_argument("--num_inference_steps", type=int, default=50, help="Number of denoising steps.")
    parser.add_argument("--guidance_scale", type=float, default=None,
                        help="Text CFG scale. If not set: t2v uses 5.0; edit/ref_edit uses 2.5 from config.")
    parser.add_argument("--guidance_scale_visual", type=float, default=None,
                        help="Visual CFG scale. If not set, uses 1.5 from config.")
    parser.add_argument("--seed", type=int, default=0, help="Random seed. Set to -1 for random.")
    parser.add_argument("--fps", type=int, default=None, help="Output video FPS. Auto-detected if not set.")

    parser.add_argument(
        "--source_video_path", type=str, default=None,
        help="Path to source video for edit/ref_edit tasks.",
    )
    parser.add_argument(
        "--ref_image_paths", type=str, nargs="+", default=None,
        help="Path(s) to reference image(s) for ref_edit task. Supports multiple images.",
    )

    parser.add_argument(
        "--und_device", type=str, default="cuda:0",
        help="CUDA device for the Qwen understanding backbone.",
    )
    parser.add_argument(
        "--gen_device", type=str, default="cuda:1",
        help="CUDA device for the Wan DiT and VAE generation backbone.",
    )

    parser.add_argument("--output_path", type=str, default=None, help="Output file path.")

    args = parser.parse_args()

    # Validate task-specific arguments
    if args.task == "edit" and not args.source_video_path:
        parser.error("--source_video_path is required for 'edit' task.")
    if args.task == "ref_edit":
        if not args.source_video_path:
            parser.error("--source_video_path is required for 'ref_edit' task.")
        if not args.ref_image_paths:
            parser.error("--ref_image_paths is required for 'ref_edit' task.")
    if args.task == "mi2v":
        if not args.ref_image_paths:
            parser.error("--ref_image_paths is required for 'mi2v' task.")

    # Initialize Accelerator
    accelerator = Accelerator(mixed_precision="bf16")

    # Setup logging: only main process outputs logs to avoid duplicate prints
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%Y/%m/%d %H:%M:%S",
        level=logging.INFO if accelerator.is_main_process else logging.WARNING,
    )

    # Set seed
    if args.seed >= 0:
        set_seed(args.seed)

    # Load config
    if not os.path.exists(args.config_path):
        raise FileNotFoundError(f"Config file not found: {args.config_path}")
    config = OmegaConf.load(args.config_path)
    if accelerator.is_main_process:
        logger.info(f"Loaded config from: {args.config_path}")

    # Build model
    if accelerator.is_main_process:
        logger.info("Building LoomVideo model...")
        logger.info(f"  Understanding backbone: {config.model.und.pretrained_model_path}")
        logger.info(f"  Generation backbone: {config.model.gen.pretrained_model_path}")

    model = load_model(config)

    # Load checkpoint (EMA priority)
    if args.ckpt_path:
        load_checkpoint(model, args.ckpt_path)
    elif accelerator.is_main_process:
        logger.info("No checkpoint provided. Using pretrained backbone weights only.")

    # One Python process controls both GPUs. Do not launch two processes and do
    # not move the whole model to accelerator.device first.
    if accelerator.num_processes != 1:
        raise RuntimeError(
            "Two-GPU LoomVideo inference requires --num_processes=1. "
            "One process controls both CUDA devices."
        )
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for LoomVideo inference.")

    und_device = torch.device(args.und_device)
    gen_device = torch.device(args.gen_device)
    visible_count = torch.cuda.device_count()
    for device in (und_device, gen_device):
        index = 0 if device.index is None else device.index
        if device.type != "cuda" or index >= visible_count:
            raise RuntimeError(
                f"Requested device {device} is unavailable; visible CUDA count is "
                f"{visible_count}."
            )

    configure_two_gpu_inference(model, und_device, gen_device)

    accelerator.wait_for_everyone()

    # Run generation
    generate_task(model, config, args, accelerator)

    if accelerator.is_main_process:
        logger.info("Done.")


if __name__ == "__main__":
    main()

使用以上generate.py代码，可以把模型放在双卡上运行。

Uh oh!

请问下必须确保 checkpoints/Qwen3-VL-8B-Instruct 和 checkpoints/Wan2.2-TI2V-5B-Diffusers 这两个目录存在且包含完整的模型权重文件？ #4

Description

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions