runtime/orbitquant_wan_a2/official_runtime.py

from __future__ import annotations

import gc
import json
import os
from pathlib import Path
import socket
import subprocess

import torch
import torch.distributed as dist
import yaml
from easydict import EasyDict

from .attention_integration import install_dense_attention_acceleration
from .diffusers_components import CachedT5Adapter, DiffusersCLIPAdapter, DiffusersWanVAEAdapter
from .loader_official import build_packed_official_transformer, place_transformer
from .kv_cache_streamer import Animate2KVCacheCPUOffloader
from .para_bridge import install_para_context_parallel


def _init_dist():
    if dist.is_available() and dist.is_initialized():
        return int(os.environ.get('LOCAL_RANK', '0')), dist.get_world_size()
    rank = int(os.environ.get('RANK', '0'))
    world = int(os.environ.get('WORLD_SIZE', '1'))
    local = int(os.environ.get('LOCAL_RANK', rank))
    os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
    os.environ.setdefault('MASTER_PORT', '29671')
    torch.cuda.set_device(local)
    dist.init_process_group('nccl', rank=rank, world_size=world)
    return local, world




def _install_distilled_euler_scheduler(pipeline_module):
    """Honor Wan-Animate-2 Distillation's documented Euler solver.

    The low-level upstream `inference_core` constructs its historical
    FlowDPMSolver symbol directly even though the distilled release documents
    `flow_solver=euler`.  Replace that module-global constructor with a small
    compatibility subclass of Diffusers' FlowMatchEulerDiscreteScheduler.  The
    upstream code already supplies the shift-5 sigma schedule, so this wrapper
    uses shift=1 to avoid shifting it a second time.
    """
    from diffusers import FlowMatchEulerDiscreteScheduler

    class _Animate2DistilledEuler(FlowMatchEulerDiscreteScheduler):
        def __init__(
            self,
            num_train_timesteps=1000,
            shift=1.0,
            use_dynamic_shifting=False,
            **_ignored,
        ):
            super().__init__(
                num_train_timesteps=int(num_train_timesteps),
                shift=float(shift),
                use_dynamic_shifting=bool(use_dynamic_shifting),
            )

    pipeline_module.FlowDPMSolverMultistepScheduler = _Animate2DistilledEuler
    return _Animate2DistilledEuler.__name__


def _load_cfg(official_repo: Path) -> EasyDict:
    cfg_path = official_repo / 'infer' / 'wan_animate_2_distillation.yaml'
    cfg = EasyDict(yaml.safe_load(cfg_path.read_text()))
    cfg.test_cfg.sp_size = 1
    cfg.test_cfg.sharding_size = 1
    cfg.test_cfg.world_size = 1
    return cfg


def run_official_packed_w4a4(
    *,
    official_repo: str | Path,
    model_root: str | Path,
    packed_dir: str | Path,
    reference_image: str | Path,
    driving_video: str | Path,
    output_dir: str | Path,
    prompt: str,
    prompt_ref: str = '人物动作的参考视频',
    negative_prompt: str | None = None,
    width: int = 256,
    height: int = 320,
    fps: int = 24,
    clip_len: int = 17,
    steps: int = 10,
    guidance_scale: float = 1.0,
    seed: int = 42,
    transformer_placement: str = 'auto',
    attention: str = 'hybrid',
    sol_tau: float = 1.0,
    kv_cache_placement: str = 'cpu',
    max_input_frames: int | None = None,
) -> dict:
    official_repo = Path(official_repo).resolve()
    model_root = Path(model_root).resolve()
    pdir = Path(packed_dir).resolve()
    outdir = Path(output_dir).resolve()
    outdir.mkdir(parents=True, exist_ok=True)

    # Hard identity gate: this runtime is only for the exact Animate-2
    # distilled Diffusers release.  Refuse older Wan-Animate/Wan I2V layouts.
    model_index_path = model_root / 'model_index.json'
    if not model_index_path.is_file():
        raise RuntimeError(f'model_index.json missing from target model root: {model_root}')
    model_index = json.loads(model_index_path.read_text())
    if model_index.get('_class_name') != 'WanAnimate2Pipeline':
        raise RuntimeError(
            f"wrong pipeline class: {model_index.get('_class_name')!r}; expected 'WanAnimate2Pipeline'"
        )
    transformer_entry = model_index.get('transformer')
    if not (isinstance(transformer_entry, list) and 'WanAnimate2Transformer3DModel' in transformer_entry):
        raise RuntimeError(f'wrong Animate-2 transformer entry in model_index.json: {transformer_entry!r}')
    pmanifest = json.loads((pdir / 'packed_manifest.json').read_text())
    if pmanifest.get('model') != 'Wan-AI/Wan2.2-Animate-2-14B-Distilled-Diffusers':
        raise RuntimeError(f"wrong packed-model identity: {pmanifest.get('model')!r}")

    local_rank, world = _init_dist()
    device = torch.device(f'cuda:{local_rank}')

    # Prompt encode before transformer residency: UMT5-XXL is ~11 GB BF16.
    neg = negative_prompt or ''
    t5 = CachedT5Adapter(model_root)
    prompts = [prompt, prompt_ref] + ([neg] if guidance_scale > 1 else [])
    t5.precompute(prompts, device=device)
    t5.release_model()

    model, load_report = build_packed_official_transformer(official_repo, pdir)
    attn_stats = install_dense_attention_acceleration(
        attention,
        sol_tau=sol_tau,
        sol_thresh_type='exact',
        strict=attention in {'sol', 'hybrid'},
    )

    para_report = {'enabled': False, 'world_size': world}
    if world > 1:
        para_report = install_para_context_parallel(model)
        # Official model will split sequence using ParaAttention collectives.
        cfg_sp = world
    else:
        cfg_sp = 1

    placement = place_transformer(model, mode=transformer_placement, device=device)

    kv_cache_placement = str(kv_cache_placement).lower()
    if kv_cache_placement not in {'cpu', 'cpu-pinned', 'gpu'}:
        raise ValueError("kv_cache_placement must be 'cpu', 'cpu-pinned', or 'gpu'")
    kv_offloader = None
    if kv_cache_placement in {'cpu', 'cpu-pinned'}:
        kv_offloader = Animate2KVCacheCPUOffloader(
            model, device=device, pin_memory=(kv_cache_placement == 'cpu-pinned')
        ).install()

    # VAE + CLIP from the exact Diffusers model repo; these are converted versions
    # of the same Wan components and preserve the official transformer's interfaces.
    clip = DiffusersCLIPAdapter(model_root).to(device)
    vae = DiffusersWanVAEAdapter(model_root).to(device)

    cfg = _load_cfg(official_repo)
    cfg.model.use_t5 = True
    cfg.test_cfg.sp_size = cfg_sp
    cfg.test_cfg.world_size = world
    cfg.test_cfg.sample_guide_scale = float(guidance_scale)

    import pipelines.wan_animate_2_pipeline as official_pipeline
    scheduler_name = _install_distilled_euler_scheduler(official_pipeline)
    inference_core = official_pipeline.inference_core

    runtime_video = Path(driving_video).resolve()
    if max_input_frames is not None and int(max_input_frames) > 0:
        # The official pipeline otherwise processes the entire driving video.
        # For the mandatory smoke gate, make a short standalone clip so memory,
        # KV-cache size and wall time are bounded before attempting a full run.
        import imageio_ffmpeg
        smoke_video = outdir / 'driving_smoke.mp4'
        duration = float(max_input_frames) / float(fps)
        cmd = [
            imageio_ffmpeg.get_ffmpeg_exe(), '-y', '-i', str(runtime_video),
            '-t', f'{duration:.6f}', '-c:v', 'libx264', '-preset', 'veryfast',
            '-crf', '18', '-c:a', 'aac', str(smoke_video),
        ]
        subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        runtime_video = smoke_video

    item = EasyDict({
        'height': int(height), 'width': int(width), 'clip_len': int(clip_len), 'first_num': 1,
        'fps': int(fps), 'seed': int(seed), 'output_path': str(outdir),
        'prompt': prompt, 'prompt_ref': prompt_ref,
        'neg_prompt': neg,
        'sample_guide_scale': float(guidance_scale), 'step': int(steps),
        'tpl_video_path': str(runtime_video),
        'refer_img_path': str(Path(reference_image).resolve()),
    })

    torch.cuda.reset_peak_memory_stats(device)
    result = inference_core(item, local_rank, cfg, clip, vae, model, t5, True, None, None, device)
    dist.barrier()

    packed_modules = [m for m in model.modules() if getattr(m, '_orbitquant_w4a4', False)]
    active_modules = [m for m in packed_modules if int(getattr(m, '_orbitquant_call_count', 0)) > 0]
    if len(packed_modules) != 480:
        raise RuntimeError(f'packed runtime contains {len(packed_modules)} OrbitQuant linears, expected 480')
    if len(active_modules) != 480:
        raise RuntimeError(
            f'only {len(active_modules)}/480 packed W4A4 linears executed; refusing to report a partial W4A4 run'
        )
    engine = model._orbitquant_engine
    if attention in {'sol', 'hybrid'}:
        if attn_stats.sol_calls < 40 or (attn_stats.sol_calls % 40) != 0:
            raise RuntimeError(
                f'attention mode {attention!r} requires one successful Sol-Attn reference self-attention '
                f'per layer/segment; got sol_calls={attn_stats.sol_calls}, expected a positive multiple of 40'
            )
        if attn_stats.fallbacks != 0:
            raise RuntimeError(
                f'Sol-Attn had {attn_stats.fallbacks} eligible-call failures; refusing to report a partial Sol run'
            )
    if world > 1 and attention in {'para', 'hybrid'}:
        if not para_report.get('enabled', False):
            raise RuntimeError('ParaAttention context parallelism was requested for multi-GPU but is not enabled')
        if int(para_report.get('all_to_all_calls', 0)) < 80:
            raise RuntimeError(
                'ParaAttention was enabled but did not execute the expected Animate-2 Ulysses collectives '
                f"(all_to_all_calls={para_report.get('all_to_all_calls', 0)})"
            )

    quant_stats = {
        'packed_modules': len(packed_modules),
        'active_modules': len(active_modules),
        'w4a4_linear_calls': int(engine.linear_calls),
        'a4_pack_cache': engine.activation_cache.stats(),
        'all_targets_executed': len(active_modules) == 480,
    }

    report = {
        'result': result if local_rank == 0 else None,
        'rank': dist.get_rank(), 'world_size': world,
        'packed_load': load_report,
        'placement': {k:v for k,v in placement.items() if k != 'streamer'},
        'quant_runtime': quant_stats,
        'attention': attn_stats.dict(),
        'para_attention': para_report,
        'kv_cache': kv_offloader.report() if kv_offloader is not None else {
            'mode': 'gpu', 'full_cache_gpu_residency': True
        },
        'peak_vram_gib': torch.cuda.max_memory_allocated(device) / 2**30,
        'w4_target_count': 480,
        'a4_target_count': 480,
        'steps': int(steps), 'guidance_scale': float(guidance_scale),
        'scheduler': scheduler_name,
        'sigma_shift': float(cfg.test_cfg.sample_shift),
        'input_frame_cap': None if max_input_frames is None else int(max_input_frames),
        'runtime_driving_video': str(runtime_video),
    }
    if local_rank == 0:
        (outdir / 'orbitquant_w4a4_report.json').write_text(json.dumps(report, indent=2, default=str))
    return report