aigc-apps-VideoX-Fun/scripts/wan2.1_self_forcing/generate_ode_pairs.py

# Based on https://github.com/guandeh17/Self-Forcing
import argparse
import gc
import json
import math
import os
import sys

import torch
from accelerate import Accelerator
from diffusers import FlowMatchEulerDiscreteScheduler
from einops import rearrange
from omegaconf import OmegaConf
from safetensors.torch import save_file
from tqdm import tqdm
from transformers import AutoTokenizer

current_file_path = os.path.abspath(__file__)
project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
for project_root in project_roots:
    sys.path.insert(0, project_root) if project_root not in sys.path else None

from videox_fun.models import (AutoencoderKLWan, WanT5EncoderModel,
                               WanTransformer3DModel)
from videox_fun.utils.utils import save_videos_grid


def filter_kwargs(cls, kwargs):
    import inspect
    sig = inspect.signature(cls.__init__)
    valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
    return {k: v for k, v in kwargs.items() if k in valid_params}


def load_prompts(caption_path):
    with open(caption_path, encoding="utf-8") as f:
        return [line.rstrip() for line in f if line.strip()]


def main():
    parser = argparse.ArgumentParser(description="Generate ODE trajectory pairs for ODE regression training.")
    parser.add_argument(
        "--pretrained_model_name_or_path",
        type=str,
        required=True,
        help="Path to pretrained model or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--config_path",
        type=str,
        required=True,
        help="Path to the model config YAML file (e.g. config/wan2.1/wan_civitai.yaml).",
    )
    parser.add_argument(
        "--caption_path",
        type=str,
        required=True,
        help="Path to a text file containing prompts, one per line. Download at https://huggingface.co/gdhe17/Self-Forcing/blob/main/vidprom_filtered_extended.txt",
    )
    parser.add_argument(
        "--output_folder",
        type=str,
        required=True,
        help="The output directory where per-prompt .safetensors ODE trajectory files will be saved.",
    )
    parser.add_argument(
        "--guidance_scale",
        type=float,
        default=6.0,
        help="Classifier-free guidance scale for ODE denoising. Default: 6.0.",
    )
    parser.add_argument(
        "--num_inference_steps",
        type=int,
        default=48,
        help="Number of ODE denoising steps for the teacher model. Default: 48.",
    )
    parser.add_argument(
        "--shift",
        type=float,
        default=8.0,
        help="Shift value for FlowMatchEulerDiscreteScheduler. Default: 8.0.",
    )
    parser.add_argument(
        "--video_sample_n_frames",
        type=int,
        default=81,
        help="Number of pixel frames for the generated video. Default: 81.",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=480,
        help="Video height in pixels. Will be divided by VAE spatial ratio (8) for latent size. Default: 480.",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=832,
        help="Video width in pixels. Will be divided by VAE spatial ratio (8) for latent size. Default: 832.",
    )
    parser.add_argument(
        "--negative_prompt",
        type=str,
        default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
        help="The negative prompt for classifier-free guidance.",
    )
    parser.add_argument(
        "--vae_mini_batch",
        type=int,
        default=1,
        help="Mini batch size for VAE decode. Default: 1.",
    )
    parser.add_argument(
        "--sample_every_n_prompts",
        type=int,
        default=0,
        help="Decode and save sample video every N prompts for visualization. 0 to disable. Default: 0.",
    )
    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank.")
    parser.add_argument(
        "--mixed_precision",
        type=str,
        default="bf16",
        choices=["no", "fp16", "bf16"],
        help="Whether to use mixed precision. Default: bf16.",
    )
    args = parser.parse_args()

    # Initialize accelerator for distributed generation
    accelerator = Accelerator(mixed_precision=args.mixed_precision)
    device = accelerator.device
    world_size = accelerator.num_processes
    rank = accelerator.process_index

    # Disable gradients globally since this is inference-only
    torch.set_grad_enabled(False)
    # Enable TF32 for faster computation on Ampere GPUs
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    config = OmegaConf.load(args.config_path)

    # For mixed precision we cast all weights to half-precision
    # as these weights are only used for inference, keeping weights in full precision is not required.
    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16

    # Load tokenizer and text encoder
    tokenizer = AutoTokenizer.from_pretrained(
        os.path.join(args.pretrained_model_name_or_path,
                     config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer'))
    )
    text_encoder = WanT5EncoderModel.from_pretrained(
        os.path.join(args.pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
        additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
        low_cpu_mem_usage=True,
        torch_dtype=weight_dtype,
    ).to(device).eval()
    text_encoder.requires_grad_(False)

    # Load VAE
    vae = AutoencoderKLWan.from_pretrained(
        os.path.join(args.pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')),
        additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
    ).to(device, dtype=weight_dtype).eval()
    vae.requires_grad_(False)

    # Load bidirectional transformer (teacher)
    transformer = WanTransformer3DModel.from_pretrained(
        os.path.join(args.pretrained_model_name_or_path,
                     config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
        low_cpu_mem_usage=True,
    ).to(device, dtype=weight_dtype).eval()
    transformer.requires_grad_(False)

    # Load scheduler and configure shift
    scheduler_kwargs = OmegaConf.to_container(config['scheduler_kwargs'])
    scheduler_kwargs['shift'] = args.shift
    noise_scheduler = FlowMatchEulerDiscreteScheduler(
        **filter_kwargs(FlowMatchEulerDiscreteScheduler, scheduler_kwargs)
    )
    noise_scheduler.set_timesteps(args.num_inference_steps, device=device)
    timesteps = noise_scheduler.timesteps

    # Compute latent shapes from VAE config
    # latent_h/w = pixel_h/w / spatial_compression_ratio
    # num_frames  = (pixel_frames - 1) / temporal_compression_ratio + 1
    latent_h = args.height // vae.spatial_compression_ratio
    latent_w = args.width // vae.spatial_compression_ratio
    latent_channels = vae.latent_channels
    num_frames = (args.video_sample_n_frames - 1) // vae.config.temporal_compression_ratio + 1

    # Compute seq_len for transformer
    patch_size = transformer.config.patch_size
    seq_len = math.ceil((latent_h * latent_w) / (patch_size[1] * patch_size[2]) * num_frames)

    # Load prompts and distribute across ranks
    prompts = load_prompts(args.caption_path)
    os.makedirs(args.output_folder, exist_ok=True)
    total_per_rank = int(math.ceil(len(prompts) / world_size))

    # Negative prompt embedding (unconditional)
    with torch.no_grad():
        neg_inputs = tokenizer(
            [args.negative_prompt], padding="max_length", max_length=512,
            truncation=True, add_special_tokens=True, return_tensors="pt"
        )
        neg_seq_lens = neg_inputs.attention_mask.gt(0).sum(dim=1).long()
        neg_embeds = text_encoder(neg_inputs.input_ids.to(device), attention_mask=neg_inputs.attention_mask.to(device))[0]
        neg_prompt_embeds = [neg_embeds[i, :neg_seq_lens[i]] for i in range(neg_embeds.shape[0])]

    # Main generation loop: each rank processes interleaved prompts
    for index in tqdm(range(total_per_rank), disable=rank != 0, desc="Generating ODE pairs"):
        prompt_index = index * world_size + rank
        if prompt_index >= len(prompts):
            continue

        prompt = prompts[prompt_index]
        output_path = os.path.join(args.output_folder, f"{prompt_index:05d}.safetensors")
        print(rank, output_path)
        if os.path.exists(output_path):
            continue

        # Encode prompt (keep padded [512, D] for saving to safetensors)
        text_inputs = tokenizer(
            [prompt], 
            padding="max_length", 
            max_length=512,
            truncation=True, 
            add_special_tokens=True, 
            return_tensors="pt"
        )
        prompt_attention_mask = text_inputs.attention_mask  # [1, 512]
        text_seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long()
        text_embeds = text_encoder(text_inputs.input_ids.to(device), attention_mask=prompt_attention_mask.to(device))[0]  # [1, 512, D]
        prompt_embeds = [text_embeds[i, :text_seq_lens[i]] for i in range(text_embeds.shape[0])]

        # Sample initial noise: [B, C, F, H, W]
        latents = torch.randn(
            [1, latent_channels, num_frames, latent_h, latent_w],
            dtype=weight_dtype, device=device
        )

        # Run full ODE denoising with CFG, collecting intermediate latents
        noisy_inputs = []

        # Reset scheduler state for each prompt to avoid stale `_step_index`
        # leaking across iterations and causing IndexError on `self.sigmas[sigma_idx + 1]`.
        noise_scheduler._step_index = None
        if hasattr(noise_scheduler, 'model_outputs'):
            noise_scheduler.model_outputs = []

        for progress_id, t in enumerate(timesteps):
            timestep = t.expand(latents.shape[0])  # [B]

            noisy_inputs.append(latents.clone())

            # Conditional prediction
            with torch.cuda.amp.autocast(dtype=weight_dtype):
                flow_pred_cond = transformer(
                    x=latents,
                    context=prompt_embeds,
                    t=timestep,
                    seq_len=seq_len,
                )

                # Unconditional prediction
                flow_pred_uncond = transformer(
                    x=latents,
                    context=neg_prompt_embeds,
                    t=timestep,
                    seq_len=seq_len,
                )

            # CFG
            flow_pred = flow_pred_uncond + args.guidance_scale * (flow_pred_cond - flow_pred_uncond)

            # Scheduler step
            latents = noise_scheduler.step(flow_pred, t, latents, return_dict=False)[0]

        # Append final clean latent
        noisy_inputs.append(latents.clone())

        # Stack all intermediate + final latents: [1, num_steps+1, C, F, H, W]
        noisy_inputs_tensor = torch.stack(noisy_inputs, dim=1)

        # Sparse sample 5 points along the ODE trajectory: [0, 12, 24, 36, -1]
        # This reduces storage while preserving the trajectory shape
        noisy_inputs_tensor = noisy_inputs_tensor[:, [0, 12, 24, 36, -1]]

        # Save as safetensors with latents, prompt_embeds, prompt_attention_mask
        save_file(
            {
                "latents": noisy_inputs_tensor.squeeze(0).cpu(),
                "prompt_embeds": text_embeds.squeeze(0).cpu(),
                "prompt_attention_mask": prompt_attention_mask.squeeze(0).cpu(),
            },
            output_path,
            metadata={"prompt": prompt},
        )

        # Decode and save sample video for visualization
        if args.sample_every_n_prompts > 0 and prompt_index % args.sample_every_n_prompts == 0:
            sample_dir = os.path.join(args.output_folder, "sample")
            os.makedirs(sample_dir, exist_ok=True)
            with torch.no_grad():
                # Decode the final clean latent (last sparse point)
                clean_latent = noisy_inputs_tensor[:, -1]  # [1, C, F, H, W]
                video = vae.decode(clean_latent.to(vae.dtype)).sample
                video = (video / 2 + 0.5).clamp(0, 1)
                save_videos_grid(
                    video.cpu().float(),
                    os.path.join(sample_dir, f"{prompt_index:05d}_clean.mp4"),
                )
            gc.collect()
            torch.cuda.empty_cache()

    accelerator.wait_for_everyone()

    # Write outputs.json annotation file for ImageVideoSafetensorsDataset
    # This JSON lists all generated safetensors files so they can be loaded by train_ode.py
    if accelerator.is_main_process:
        safe_json_writer = []
        for i in range(len(prompts)):
            safetensor_path = os.path.join(args.output_folder, f"{i:05d}.safetensors")
            if os.path.exists(safetensor_path):
                safe_json_writer.append({"file_path": safetensor_path})
        json_path = os.path.join(args.output_folder, "outputs.json")
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(safe_json_writer, f, ensure_ascii=False, indent=4)
        print(f"Done. Generated {len(safe_json_writer)} ODE pairs, saved to {args.output_folder}")
        print(f"Annotation JSON: {json_path}")


if __name__ == "__main__":
    main()