URSA/scripts/eval_distill_dimo.py

#!/usr/bin/env python3
# Copyright (c) 2024-present, BAAI. All Rights Reserved.
#
# 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.
# -----------------------------------------------------------------------
"""Evaluation script for distill_dimo checkpoints.

Generates videos from both the student (1-step) and teacher (multi-step)
using checkpoints saved by train_distill_dimo.py.

Verified native inference regime (from A/B testing — ground truth):
  height=320, width=512, num_frames=49, guidance_scale=7, teacher_steps=50.
  no_cfg (guidance_scale=1) does NOT produce valid output for this URSA
  checkpoint — outputs are blank or blurry.

Student generation modes
------------------------
  cfg      : 1-step, guidance_scale=7   (2× forward, inference-time CFG)

Teacher generation modes
------------------------
  cfg      : 50-step, guidance_scale=7  (official working regime)

Usage:
  python scripts/eval_distill_dimo.py \
      --teacher_ckpt /gfs/space/private/fengzl/World_Model/URSA-1.7B \
      --student_ckpt ./experiments/distill_dimo_v3/checkpoints/checkpoint-200/student.pt \
      --out_dir ./outputs/eval_distill_v3_200steps_49frames
"""

import argparse
import os
import sys

import numpy as np
import torch

_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if _REPO_ROOT not in sys.path:
    sys.path.insert(0, _REPO_ROOT)

from diffnext.pipelines import URSAPipeline
from diffnext.pipelines.ursa.pipeline_ursa_distill_dimo import (
    VERIFIED_NATIVE_DEFAULTS,
    check_verified_regime,
)
from diffnext.utils import export_to_video


# ---------------------------------------------------------------------------
# Default prompts and seeds
# ---------------------------------------------------------------------------

DEFAULT_PROMPTS = [
    "a lone grizzly bear walks through a misty forest at dawn, sunlight catching its fur.",
    "beautiful fireworks in the sky with red, white and blue.",
    "a wave crashes on a rocky shoreline at sunset, slow motion.",
    "a hummingbird hovers in front of a red flower, wings a blur.",
    "timelapse of clouds rolling over mountain peaks.",
    "a neon-lit city street at night with rain-soaked reflections.",
    "a kitten playing with a ball of yarn on a wooden floor.",
    "astronaut floating weightlessly inside a space station.",
]

# Official URSA negative prompt (from README / app scripts)
DEFAULT_NEGATIVE_PROMPT = (
    "worst quality, low quality, inconsistent motion, static, still, "
    "blurry, jittery, distorted, ugly"
)

DEFAULT_SEEDS = [0, 1, 2, 3]


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser(
        description="Evaluate distill_dimo student (1-step) vs teacher (multi-step)"
    )

    p.add_argument("--teacher_ckpt", required=True,
                   help="URSA diffusers pipeline directory (teacher weights)")
    p.add_argument("--student_ckpt", required=True,
                   help="student.pt from train_distill_dimo.py checkpoint")
    p.add_argument("--out_dir", default="./outputs/eval_distill")

    # Geometry — verified native: 320×512×49 (from A/B testing)
    p.add_argument("--num_frames", type=int, default=49)
    p.add_argument("--height", type=int, default=320)
    p.add_argument("--width", type=int, default=512)
    p.add_argument("--fps", type=int, default=12)

    # Student generation — default: cfg only (no_cfg is known to fail)
    p.add_argument("--student_modes", nargs="+", default=["cfg"],
                   choices=["no_cfg", "cfg", "baked"],
                   help="Student generation modes to evaluate. "
                        "Default: ['cfg']. no_cfg is known to produce blank/blurry "
                        "output for this checkpoint.")
    p.add_argument("--eval_cfg_scale", type=float, default=7.0,
                   help="Guidance scale for 'cfg' mode (verified working value=7)")

    # Teacher generation — default: cfg only (no_cfg is known to fail)
    p.add_argument("--teacher_modes", nargs="+", default=["cfg"],
                   choices=["no_cfg", "cfg"],
                   help="Teacher generation modes. Default: ['cfg']. "
                        "no_cfg is NOT a valid baseline for this URSA checkpoint.")
    p.add_argument("--teacher_steps", type=int, default=50,
                   help="Number of inference steps for teacher (verified default=50)")

    # Shared generation params (match verified official defaults)
    p.add_argument("--guidance_trunc", type=float, default=0.9,
                   help="Truncation threshold for inference CFG")
    p.add_argument("--negative_prompt", type=str, default=DEFAULT_NEGATIVE_PROMPT,
                   help="Negative prompt for CFG (official URSA uses one)")
    p.add_argument("--max_prompt_length", type=int, default=320)
    p.add_argument("--vae_batch_size", type=int, default=1)

    # Safety override for no_cfg
    p.add_argument("--allow_bad_nocfg", action="store_true", default=False,
                   help="Suppress the no_cfg warning/block. Use at your own risk.")

    # Data
    p.add_argument("--prompt_file", default=None,
                   help="Text file with one prompt per line (overrides defaults)")
    p.add_argument("--seeds", nargs="*", type=int, default=DEFAULT_SEEDS)

    # Device
    p.add_argument("--device", type=int, default=0)
    p.add_argument("--mixed_precision", default="bf16",
                   choices=["fp16", "bf16", "fp32"])

    return p.parse_args()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def slug(text: str, max_len: int = 40) -> str:
    s = text.lower()
    s = "".join(c if c.isalnum() or c == " " else "" for c in s)
    s = "_".join(s.split())[:max_len]
    return s or "prompt"


def frames_to_mp4(frames, path: str, fps: int = 12):
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    if isinstance(frames, np.ndarray) and frames.ndim == 4:
        frames = list(frames)
    export_to_video(frames, output_video_path=path, fps=fps)


def _extract_frames(frames_output):
    """Normalise pipeline output → list of uint8 numpy arrays [H, W, 3]."""
    if isinstance(frames_output, np.ndarray):
        frames_output = frames_output[0] if frames_output.ndim == 5 else frames_output
        frames = list(frames_output)
    elif isinstance(frames_output, list):
        frames = [np.array(f) if not isinstance(f, np.ndarray) else f
                  for f in frames_output]
    else:
        raise TypeError(f"Unexpected frames type: {type(frames_output)}")
    result = []
    for f in frames:
        if f.dtype != np.uint8:
            f = ((f * 255).clip(0, 255).astype(np.uint8)
                 if f.max() <= 1.0 else f.astype(np.uint8))
        result.append(f)
    return result


def _gen(pipe, prompt, negative_prompt, seed, num_frames, height, width,
         guidance_scale, num_inference_steps, guidance_trunc,
         max_prompt_length, vae_batch_size, device):
    """Single generation call, returns list of uint8 frames."""
    gen = torch.Generator(device=device).manual_seed(seed)
    with torch.inference_mode():
        out = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_frames=num_frames,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            guidance_trunc=guidance_trunc,
            max_prompt_length=max_prompt_length,
            vae_batch_size=vae_batch_size,
            output_type="np",
            generator=gen,
        )
    return _extract_frames(out.frames)


def _debug_pipeline(pipe, label=""):
    """Print diagnostic info about the pipeline state."""
    print(f"\n{'='*60}")
    print(f"  Pipeline diagnostics {label}")
    print(f"{'='*60}")
    print(f"  scheduler class     : {type(pipe.scheduler).__name__}")
    print(f"  scheduler type      : {type(pipe.scheduler)}")
    if hasattr(pipe.scheduler, 'config'):
        print(f"  scheduler.config    : {dict(pipe.scheduler.config)}")
    if hasattr(pipe.scheduler, 'path'):
        print(f"  scheduler.path      : {type(pipe.scheduler.path).__name__}")
        if hasattr(pipe.scheduler.path, 'emb'):
            emb = pipe.scheduler.path.emb
            print(f"  path.emb shape      : {emb.shape}")
            print(f"  path.emb device     : {emb.device}")
            print(f"  path.emb dtype      : {emb.dtype}")
    else:
        print(f"  scheduler.path      : MISSING (scheduler not fully loaded!)")
    print(f"  codebook_size       : {getattr(pipe.scheduler, 'codebook_size', 'N/A')}")
    print(f"  transformer class   : {type(pipe.transformer).__name__}")
    print(f"  transformer device  : {next(pipe.transformer.parameters()).device}")
    print(f"  vae class           : {type(pipe.vae).__name__}")
    if hasattr(pipe, 'image_processor'):
        print(f"  image_processor     : {type(pipe.image_processor).__name__}")
    print(f"{'='*60}\n")


def _debug_frames(frames, label=""):
    """Print diagnostic info about generated frames."""
    if not frames:
        print(f"  [{label}] No frames generated!")
        return
    f0 = frames[0]
    print(f"  [{label}] n_frames={len(frames)}  shape={f0.shape}  "
          f"dtype={f0.dtype}  min={f0.min()}  max={f0.max()}")


def _verify_state_dict_swap(pipe, state_dict, label=""):
    """Verify transformer weights actually changed after load_state_dict."""
    sample_key = next(iter(state_dict.keys()))
    loaded_val = state_dict[sample_key].flatten()[:8]
    current_val = pipe.transformer.state_dict()[sample_key].flatten()[:8]
    match = torch.allclose(loaded_val.cpu().float(), current_val.cpu().float(), atol=1e-6)
    print(f"  [{label}] state_dict match for '{sample_key}': {match}")
    if not match:
        print(f"    loaded  : {loaded_val[:4]}")
        print(f"    current : {current_val[:4]}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    args = parse_args()

    dtype_map = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}
    dtype = dtype_map[args.mixed_precision]
    device = (torch.device("cuda", args.device)
              if torch.cuda.is_available() else torch.device("cpu"))
    os.makedirs(args.out_dir, exist_ok=True)

    # -- Verified regime check --------------------------------------------
    is_native = check_verified_regime(
        height=args.height,
        width=args.width,
        num_frames=args.num_frames,
        guidance_scale=args.eval_cfg_scale,
        teacher_steps=args.teacher_steps,
        label="eval",
    )
    print(f"[eval] verified_native_regime={is_native}")
    print(f"[eval] geometry=({args.num_frames},{args.height},{args.width}), "
          f"guidance_scale={args.eval_cfg_scale}, teacher_steps={args.teacher_steps}")

    # -- no_cfg safety gate -----------------------------------------------
    all_modes = list(args.student_modes) + list(args.teacher_modes)
    if "no_cfg" in all_modes:
        if args.allow_bad_nocfg:
            print("[WARN] no_cfg is known to fail for this URSA checkpoint. "
                  "Outputs may be blank or blurry. Proceeding because --allow_bad_nocfg is set.")
        else:
            print("[WARN] no_cfg is known to fail for this URSA checkpoint. "
                  "Outputs may be blank or blurry. "
                  "Pass --allow_bad_nocfg to override this warning.")

    # -- Load prompts -----------------------------------------------------
    if args.prompt_file:
        with open(args.prompt_file, encoding="utf-8") as f:
            prompts = [l.strip() for l in f if l.strip() and not l.startswith("#")]
    else:
        prompts = DEFAULT_PROMPTS

    print(f"[eval] {len(prompts)} prompts × {len(args.seeds)} seeds  "
          f"| student modes={args.student_modes}  "
          f"| teacher modes={args.teacher_modes}")
    print(f"[eval] guidance_scale={args.eval_cfg_scale}  "
          f"guidance_trunc={args.guidance_trunc}  "
          f"teacher_steps={args.teacher_steps}")
    print(f"[eval] negative_prompt='{args.negative_prompt[:60]}...'")

    # -- Load pipeline (teacher) ------------------------------------------
    print(f"[eval] Loading pipeline from {args.teacher_ckpt} …")
    # 【修改点 2】尝试启用 Flash Attention 2
    try:
        pipe = URSAPipeline.from_pretrained(
            args.teacher_ckpt, 
            torch_dtype=dtype, 
            trust_remote_code=True,
            attn_implementation="flash_attention_2" 
        ).to(device)
    except Exception:
        # 如果环境不支持 FA2，降级到默认
        pipe = URSAPipeline.from_pretrained(
            args.teacher_ckpt, torch_dtype=dtype, trust_remote_code=True
        ).to(device)
    
    if hasattr(pipe.vae, "disable_slicing"):
        pipe.vae.disable_slicing()
    if hasattr(pipe.vae, "disable_tiling"):
        pipe.vae.disable_tiling()
        
    # print("[eval] Compiling transformer (this takes ~2 mins for the first time)...")
    # pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead")

    # Diagnostic: verify scheduler loaded correctly
    _debug_pipeline(pipe, label="after from_pretrained + .to(device)")

    # CRITICAL CHECK: scheduler must have .path with embeddings
    if not hasattr(pipe.scheduler, 'path') or pipe.scheduler.path is None:
        print("[ERROR] Scheduler path not loaded! This will cause blurry output.")
        print("[ERROR] The scheduler needs scheduler_model.pth with codebook embeddings.")
        return

    if pipe.scheduler.codebook_size == 0:
        print("[ERROR] codebook_size=0 — scheduler not properly initialized!")
        return

    # Save teacher state for switching back after student inference
    teacher_state = {k: v.clone() for k, v in pipe.transformer.state_dict().items()}

    # -- Load student checkpoint ------------------------------------------
    print(f"[eval] Loading student weights from {args.student_ckpt} …")
    student_state = torch.load(
        args.student_ckpt, map_location=device, weights_only=True
    )
    print(f"[eval] student state_dict keys: {len(student_state)}  "
          f"sample key: {next(iter(student_state.keys()))}")

    # Common kwargs for every pipeline call
    gen_kwargs = dict(
        num_frames=args.num_frames,
        height=args.height,
        width=args.width,
        guidance_trunc=args.guidance_trunc,
        max_prompt_length=args.max_prompt_length,
        vae_batch_size=args.vae_batch_size,
    )

    # Mode → guidance_scale mapping
    student_guidance = {
        "no_cfg": 1.0,
        "cfg":    args.eval_cfg_scale,
        "baked":  1.0,
    }
    teacher_guidance = {
        "no_cfg": 1.0,
        "cfg":    args.eval_cfg_scale,
    }

    # -- Evaluation loop --------------------------------------------------
    for idx, prompt in enumerate(prompts):
        p_slug = slug(prompt)
        print(f"\n[{idx+1}/{len(prompts)}] {prompt[:70]}")

        for seed in args.seeds:
            # ---- Student: 1-step generation -----------------------------
            for mode in args.student_modes:
                g_scale = student_guidance[mode]
                neg = args.negative_prompt if g_scale > 1 else None
                pipe.transformer.load_state_dict(student_state, strict=True)
                pipe.transformer.eval()

                if idx == 0 and seed == args.seeds[0]:
                    _verify_state_dict_swap(pipe, student_state, f"student/{mode}")

                with torch.no_grad():
                    frames = _gen(pipe, prompt, neg, seed,
                                  guidance_scale=g_scale,
                                  num_inference_steps=1,
                                  device=device, **gen_kwargs)

                if idx == 0 and seed == args.seeds[0]:
                    _debug_frames(frames, f"student/{mode}")

                path = os.path.join(
                    args.out_dir,
                    f"{idx:02d}_s{seed}_{p_slug}_student_1step_{mode}.mp4",
                )
                frames_to_mp4(frames, path, fps=args.fps)
                print(f"  [student/{mode:6s}] seed={seed}  scale={g_scale}  → {path}")

            # ---- Teacher: multi-step reference --------------------------
            for t_mode in args.teacher_modes:
                g_scale = teacher_guidance[t_mode]
                neg = args.negative_prompt if g_scale > 1 else None
                pipe.transformer.load_state_dict(teacher_state, strict=True)
                pipe.transformer.eval()

                if idx == 0 and seed == args.seeds[0]:
                    _verify_state_dict_swap(pipe, teacher_state, f"teacher/{t_mode}")

                with torch.no_grad():
                    frames = _gen(pipe, prompt, neg, seed,
                                  guidance_scale=g_scale,
                                  num_inference_steps=args.teacher_steps,
                                  device=device, **gen_kwargs)

                if idx == 0 and seed == args.seeds[0]:
                    _debug_frames(frames, f"teacher/{t_mode}")

                path = os.path.join(
                    args.out_dir,
                    f"{idx:02d}_s{seed}_{p_slug}_teacher_{args.teacher_steps}step_{t_mode}.mp4",
                )
                frames_to_mp4(frames, path, fps=args.fps)
                print(f"  [teacher/{t_mode:6s}] seed={seed}  scale={g_scale}  "
                      f"steps={args.teacher_steps}  → {path}")

    print(f"\n[eval] Done. Results in {args.out_dir}")
    _print_guide(args)


def _print_guide(args):
    print(f"""
╔══════════════════════════════════════════════════════════════╗
║  Interpretation guide                                        ║
╠══════════════════════════════════════════════════════════════╣
║  student_1step_cfg      : 1-step, guidance_scale={args.eval_cfg_scale:<4}        ║
║                           (verified working student mode)    ║
║  student_1step_baked    : 1-step, guidance_scale=1           ║
║                           (for students trained with CFG KD) ║
║  teacher_{args.teacher_steps}step_cfg     : {args.teacher_steps}-step, guidance_scale={args.eval_cfg_scale:<4}  ║
║                           (verified working teacher mode)    ║
╠══════════════════════════════════════════════════════════════╣
║  NOTE: no_cfg (guidance_scale=1) is NOT a valid baseline     ║
║  for this URSA checkpoint. Use --allow_bad_nocfg to override.║
╚══════════════════════════════════════════════════════════════╝""")


if __name__ == "__main__":
    main()