magi_human ¶

def __init__(
    self,
    dim: int,
    temperature: float = 10000.0,
    dtype: torch.dtype = torch.float32,
):
    super().__init__()
    self.dim = dim
    bands = _freq_bands(dim // 8, temperature=temperature).to(dtype)
    # `register_buffer` so state_dict keeps it, matching upstream naming.
    self.register_buffer("bands", bands)

def __init__(self, cfg: AttentionSubConfig):
    super().__init__()
    self.cfg = cfg
    self.gating_size = cfg.num_heads_q if cfg.enable_attn_gating else 0
    qkv_out = (
        cfg.num_heads_q * cfg.head_dim
        + 2 * cfg.num_heads_kv * cfg.head_dim
        + self.gating_size
    )
    self.pre_norm = MultiModalityRMSNorm(cfg.hidden_size, num_modality=cfg.num_modality)
    self.linear_qkv = PackedExpertLinear(
        cfg.hidden_size, qkv_out, num_experts=cfg.num_modality, bias=False,
    )
    self.linear_proj = PackedExpertLinear(
        cfg.num_heads_q * cfg.head_dim, cfg.hidden_size,
        num_experts=cfg.num_modality, bias=False,
    )
    self.q_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality)
    self.k_norm = MultiModalityRMSNorm(cfg.head_dim, num_modality=cfg.num_modality)

    self.q_size = cfg.num_heads_q * cfg.head_dim
    self.kv_size = cfg.num_heads_kv * cfg.head_dim

    self.attn = LocalAttention(
        num_heads=cfg.num_heads_q,
        head_size=cfg.head_dim,
        num_kv_heads=cfg.num_heads_kv,
        causal=False,
        supported_attention_backends=(
            AttentionBackendEnum.FLASH_ATTN,
            AttentionBackendEnum.TORCH_SDPA,
        ),
    )

def __init__(self, config: MagiHumanVideoConfig, hf_config: dict | None = None, **kwargs):
    super().__init__(config=config, hf_config=hf_config or {})
    arch: MagiHumanArchConfig = getattr(config, "arch_config", config)
    self.arch = arch

    # BaseDiT contract instance vars.
    self.hidden_size = arch.hidden_size
    self.num_attention_heads = arch.num_attention_heads
    self.num_channels_latents = arch.num_channels_latents

    self.adapter = MagiAdapter(arch)
    self.block = _TransformerBlock(arch)
    self.final_norm_video = MultiModalityRMSNorm(arch.hidden_size)
    self.final_norm_audio = MultiModalityRMSNorm(arch.hidden_size)
    self.final_linear_video = nn.Linear(
        arch.hidden_size, arch.video_in_channels, bias=False, dtype=torch.float32,
    )
    self.final_linear_audio = nn.Linear(
        arch.hidden_size, arch.audio_in_channels, bias=False, dtype=torch.float32,
    )
    # Dispatcher + mask cache: modality_mapping is the same tensor object
    # across all timesteps in the denoising loop; data_ptr()+shape+dtype+device
    # is a fast, collision-free key that avoids rebuilding ModalityDispatcher
    # (which calls argsort + bincount) on every forward call.
    self._cached_dispatcher: Optional[ModalityDispatcher] = None
    self._cached_video_mask: Optional[torch.Tensor] = None
    self._cached_audio_mask: Optional[torch.Tensor] = None
    self._cached_text_mask: Optional[torch.Tensor] = None
    self._cached_modality_key: Optional[tuple] = None

def __init__(self, modality_mapping: torch.Tensor, num_modalities: int):
    self.modality_mapping = modality_mapping
    self.num_modalities = num_modalities
    self.permute_mapping = torch.argsort(modality_mapping)
    self.inv_permute_mapping = torch.argsort(self.permute_mapping)
    permuted = modality_mapping[self.permute_mapping]
    self.group_size = torch.bincount(permuted, minlength=num_modalities).to(torch.int32)
    self.group_size_cpu: list[int] = [int(x) for x in self.group_size.cpu().tolist()]

def __init__(self, dim: int, eps: float = 1e-6, num_modality: int = 1):
    super().__init__()
    self.dim = dim
    self.eps = eps
    self.num_modality = num_modality
    # Always stored in fp32; matches the reference initialization.
    self.weight = nn.Parameter(torch.zeros(dim * num_modality, dtype=torch.float32))

def __init__(
    self,
    in_features: int,
    out_features: int,
    num_experts: int = 1,
    bias: bool = False,
    dtype: torch.dtype = torch.bfloat16,
):
    super().__init__()
    self.in_features = in_features
    self.out_features = out_features
    self.num_experts = num_experts
    self.use_bias = bias
    self.weight = nn.Parameter(
        torch.empty(out_features * num_experts, in_features, dtype=dtype)
    )
    if bias:
        self.bias = nn.Parameter(
            torch.empty(out_features * num_experts, dtype=dtype)
        )
    else:
        self.register_parameter("bias", None)