"""Runtime monkey-patches. Import this before launching test_infer_seedance.""" from __future__ import annotations import os import sys import torch import torch.nn.functional as F def install_elastic_schedule(model, meta: dict) -> None: elastic = (meta or {}).get("elastic") or {} if not elastic: return from inference.model.dit.dit_module import SparseMoEMLP quality_steps = set(int(x) for x in elastic.get("quality_steps") or []) skip_layers = set(int(x) for x in elastic.get("skip_layers") or []) n_steps = int(elastic.get("num_steps") or 32) per_layer = elastic.get("layers") or {} fast_steps = [i for i in range(n_steps) if i not in quality_steps] for m in model.modules(): if not isinstance(m, SparseMoEMLP): continue spec = per_layer.get(str(int(m.layer_idx))) or {} quality_width = int(spec.get("quality_width") or m.experts[0].down_proj.in_features) fast_width = int(spec.get("fast_width") or quality_width) skip_steps = set(fast_steps) if int(m.layer_idx) in skip_layers else set() m._elastic = { "quality_width": quality_width, "fast_width": min(fast_width, quality_width), "quality_steps": quality_steps, "skip_steps": skip_steps, } m._skip_mlp = False m._active_mid = quality_width for expert in m.experts: expert._active_mid = quality_width print( f"prune_pro elastic schedule: quality_steps={sorted(quality_steps)} " f"skip_layers={sorted(skip_layers)}" ) def _apply_elastic_step(module, step_index) -> None: from inference.model.dit.dit_module import SparseMoEMLP force = os.getenv("MAGI_ELASTIC_FORCE", "").strip().lower() skip_off = os.getenv("MAGI_ELASTIC_SKIP", "1").strip().lower() in ("0", "off", "false") step = 0 if step_index is None else int(step_index) for m in module.modules(): if not isinstance(m, SparseMoEMLP): continue elastic = getattr(m, "_elastic", None) if not elastic: m._skip_mlp = False m._active_mid = None continue quality_steps = elastic["quality_steps"] skip_steps = set() if skip_off else elastic.get("skip_steps") or set() if force == "off": m._skip_mlp = False m._active_mid = None elif force == "quality" or (force != "fast" and step in quality_steps): m._skip_mlp = False m._active_mid = int(elastic["quality_width"]) elif step in skip_steps: m._skip_mlp = True m._active_mid = int(elastic["fast_width"]) else: m._skip_mlp = False m._active_mid = int(elastic["fast_width"]) for expert in m.experts: expert._active_mid = None if m._skip_mlp else m._active_mid def _patch_expert_forward(ExpertMLP) -> None: orig = ExpertMLP.forward_with_intermediate def forward_with_intermediate(self, x: torch.Tensor): n_mid = getattr(self, "_active_mid", None) full = int(self.down_proj.in_features) if n_mid is None or int(n_mid) >= full: return orig(self, x) n = min(full, max(1, int(n_mid))) x_bf = x.to(torch.bfloat16) projected = F.linear(x_bf, self.up_gate_proj.weight[: 2 * n]).to(torch.bfloat16) z = self.activation_func(projected).to(torch.bfloat16) y = F.linear(z, self.down_proj.weight[:, :n]).to(torch.bfloat16) return y, z ExpertMLP.forward_with_intermediate = forward_with_intermediate def _patch_moe_forward(SparseMoEMLP) -> None: orig = SparseMoEMLP.forward def forward(self, x: torch.Tensor, modality_dispatcher=None): if getattr(self, "_skip_mlp", False): return torch.zeros_like(x) n_mid = getattr(self, "_active_mid", None) for expert in self.experts: expert._active_mid = n_mid return orig(self, x, modality_dispatcher) SparseMoEMLP.forward = forward def _patch_grouped(SparseMoEMLP) -> None: orig = SparseMoEMLP._experts_grouped def _experts_grouped(self, x_norm, topk_idx, topk_vals, mode): n_mid = getattr(self, "_active_mid", None) e0 = self.experts[0] full = int(e0.down_proj.in_features) if n_mid is None or int(n_mid) >= full or mode != "bmm": return orig(self, x_norm, topk_idx, topk_vals, mode) n = min(full, max(1, int(n_mid))) ids = topk_idx[:, 0] g = topk_vals[:, 0] n_tok = x_norm.shape[0] ecount = self.num_experts dtype = x_norm.dtype order = ids.argsort(stable=True) ids_s = ids[order] x_s = x_norm[order] g_s = g.to(x_norm.dtype)[order] counts = torch.bincount(ids_s, minlength=ecount) offsets = torch.cumsum(counts, 0) - counts col = torch.arange(n_tok, device=x_norm.device) - offsets[ids_s] m_max = int(counts.max().item()) if m_max == 0: return torch.zeros_like(x_norm) x_pad = x_norm.new_zeros(ecount, m_max, x_norm.shape[1]) g_pad = x_norm.new_zeros(ecount, m_max) x_pad[ids_s, col] = x_s g_pad[ids_s, col] = g_s w_up = torch.stack([e.up_gate_proj.weight.to(dtype)[: 2 * n] for e in self.experts], dim=0) z = torch.bmm(x_pad, w_up.transpose(-2, -1)).to(dtype) z = e0.activation_func(z) if z.dtype != dtype: z = z.to(dtype) w_dn = torch.stack([e.down_proj.weight.to(dtype)[:, :n] for e in self.experts], dim=0) y = torch.bmm(z, w_dn.transpose(-2, -1)).to(dtype) y = y * g_pad.unsqueeze(-1) y_s = y[ids_s, col] y_orig = y_s.new_empty(n_tok, y_s.shape[1]) y_orig[order] = y_s return y_orig SparseMoEMLP._experts_grouped = _experts_grouped def install_runtime_patches() -> None: """Patch inference classes in-process. Call before importing test_infer_seedance.""" from inference.model.dit import dit_module as dm from inference.pruning import checkpoint as ckpt if getattr(dm, "_prune_pro_elastic_patched", False): return _patch_expert_forward(dm.ExpertMLP) _patch_moe_forward(dm.SparseMoEMLP) _patch_grouped(dm.SparseMoEMLP) orig_set = dm.set_sparse_moe_diffusion_timestep def set_sparse_moe_diffusion_timestep(module, timestep, step_index=None, num_steps=None): orig_set(module, timestep, step_index=step_index, num_steps=num_steps) _apply_elastic_step(module, step_index) dm.set_sparse_moe_diffusion_timestep = set_sparse_moe_diffusion_timestep dit_pkg = sys.modules.get("inference.model.dit") if dit_pkg is not None: dit_pkg.set_sparse_moe_diffusion_timestep = set_sparse_moe_diffusion_timestep vg = sys.modules.get("inference.pipeline.video_generate") if vg is not None: vg.set_sparse_moe_diffusion_timestep = set_sparse_moe_diffusion_timestep orig_rebuild = ckpt.rebuild_mid_shapes def rebuild_mid_shapes(model, meta): orig_rebuild(model, meta) if meta.get("elastic"): install_elastic_schedule(model, meta) ckpt.rebuild_mid_shapes = rebuild_mid_shapes dm._prune_pro_elastic_patched = True print("prune_pro: installed elastic runtime patches") def _dit_bench_enabled() -> bool: return os.getenv("MAGI_DIT_BENCH", "").strip().lower() in ("1", "true", "yes", "on") def install_dit_bench_patches() -> None: """Skip VAE load/decode/save; time DiT evaluate_with_latent only. Call after importing test_infer_seedance so MagiEvaluator exists. """ if not _dit_bench_enabled(): return import time import inference.pipeline.pipeline as pl import inference.pipeline.video_generate as vg if getattr(vg, "_prune_pro_dit_bench", False): return class _SkipVAE(torch.nn.Module): def __init__(self, *args, **kwargs): super().__init__() # Defaults match the published MoE denoise bench (sr=44100, C_a=128). # Native MagiHuman uses C_a=64 and hop=2048; set the env overrides for that run. self.sample_rate = int(os.getenv("MAGI_DIT_BENCH_SAMPLE_RATE", "44100")) self.latent_dim = int(os.getenv("MAGI_DIT_BENCH_LATENT_DIM", "128")) hop = os.getenv("MAGI_DIT_BENCH_HOP", "").strip() if hop: self.hop = int(hop) self.downsampling_ratio = int(hop) orig_init = vg.MagiEvaluator.__init__ def _init(self, *args, **kwargs): saved = ( vg.get_vae2_2, vg.get_turbo_vaed, vg.OmniAudioFeatureExtractor, vg.SAAudioFeatureExtractor, ) vg.get_vae2_2 = lambda *a, **k: _SkipVAE() vg.get_turbo_vaed = lambda *a, **k: _SkipVAE() vg.OmniAudioFeatureExtractor = _SkipVAE vg.SAAudioFeatureExtractor = _SkipVAE try: orig_init(self, *args, **kwargs) finally: ( vg.get_vae2_2, vg.get_turbo_vaed, vg.OmniAudioFeatureExtractor, vg.SAAudioFeatureExtractor, ) = saved print("prune_pro dit_bench: skipped video/audio VAE load", flush=True) vg.MagiEvaluator.__init__ = _init orig_ewl = vg.MagiEvaluator.evaluate_with_latent def _ewl(self, *args, **kwargs): torch.cuda.synchronize() t0 = time.perf_counter() out = orig_ewl(self, *args, **kwargs) torch.cuda.synchronize() print(f"dit_denoise_s={time.perf_counter() - t0:.6f}", flush=True) return out vg.MagiEvaluator.evaluate_with_latent = _ewl orig_eval = vg.MagiEvaluator.evaluate def _eval(self, *args, **kwargs): orig_pp = self.post_process def _skip_pp(*a, **k): from inference.utils import event_path_timer, print_rank_last event_path_timer().synced_record( "Step7: Decode Audio skipped", print_fn=print_rank_last ) event_path_timer().synced_record( "Step8: Post Process skipped", print_fn=print_rank_last ) return None, None self.post_process = _skip_pp try: return orig_eval(self, *args, **kwargs) finally: self.post_process = orig_pp vg.MagiEvaluator.evaluate = _eval def _run( self, prompt, image, audio, save_path_prefix, seed=42, seconds=4.0, br_width=480, br_height=272, sr_width=None, sr_height=None, output_width=None, output_height=None, upsample_mode=None, ): self._validate_offline_request(prompt=prompt, save_path_prefix=save_path_prefix) save_path = f"{save_path_prefix}_{seconds}s_{br_width}x{br_height}.dit_bench" with torch.random.fork_rng(devices=[torch.cuda.current_device()]): torch.random.manual_seed(seed) self.evaluator.evaluate( prompt, image, audio, seconds=seconds, br_width=br_width, br_height=br_height, sr_width=sr_width, sr_height=sr_height, br_num_inference_steps=self.evaluation_config.num_inference_steps, sr_num_inference_steps=self.evaluation_config.sr_num_inference_steps, ) self.evaluator.dump_moe_clip(save_path_prefix) print(f"dit_bench skip_video {save_path}", flush=True) if torch.distributed.is_initialized(): torch.distributed.barrier() return save_path pl.MagiPipeline.run_offline = _run vg._prune_pro_dit_bench = True print("prune_pro: installed DiT-denoise bench (skip VAE decode/save)", flush=True)