inference_code/crown/runtime.py

"""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)