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