vllm_omni.diffusion.models.magi_human.magi_human_dit ¶

audio_embedder = Linear(
    audio_in_channels, hidden_size, bias=True, dtype=float32
)

config: AdapterConfig = config

rope = ElementWiseFourierEmbed(
    hidden_size // num_attention_heads,
    in_pixels=False,
    learnable=False,
)

text_embedder = Linear(
    text_in_channels, hidden_size, bias=True, dtype=float32
)

video_embedder = Linear(
    video_in_channels, hidden_size, bias=True, dtype=float32
)

audio_in_channels: int

hidden_size: int

num_attention_heads: int

params_dtype: dtype

text_in_channels: int

video_in_channels: int

config: AttentionConfig = config

gating_size = num_heads_q if enable_attn_gating else 0

k_norm = MultiModalityRMSNorm(
    head_dim, num_modality=num_modality
)

kv_size = num_kv_heads * head_dim

linear_gating = ColumnParallelLinear(
    input_size=hidden_size,
    output_size=num_heads_q,
    bias=False,
    gather_output=False,
    return_bias=False,
)

linear_proj = RowParallelLinear(
    input_size=num_heads_q * head_dim,
    output_size=hidden_size,
    bias=False,
    input_is_parallel=True,
    return_bias=False,
)

linear_qkv = QKVParallelLinear(
    hidden_size=hidden_size,
    head_size=head_dim,
    total_num_heads=num_heads_q,
    total_num_kv_heads=num_heads_kv,
    bias=False,
    return_bias=False,
)

pre_norm = MultiModalityRMSNorm(
    hidden_size, eps=1e-06, num_modality=num_modality
)

q_norm = MultiModalityRMSNorm(
    head_dim, num_modality=num_modality
)

q_size = num_heads * head_dim

checkpoint_qk_layernorm_rope: bool

enable_attn_gating: bool = False

head_dim: int

hidden_size: int

num_heads_kv: int

num_heads_q: int

num_layers: int

num_modality: int

params_dtype: dtype

use_local_attn: bool = False

bias = Parameter(
    empty(out_features * num_experts, **factory_kwargs)
)

in_features = in_features

num_experts = num_experts

num_layers_for_initialization = (
    num_layers_for_initialization
)

out_features = out_features

use_bias = bias

weight = Parameter(
    empty(
        (out_features * num_experts, in_features),
        **factory_kwargs,
    )
)

adapter: Adapter = Adapter(adapter_config)

block: TransformerBlock = TransformerBlock(
    model_config=model_config
)

blocks: ModuleList

config: TransformerConfig = TransformerConfig(
    hidden_size=hidden_size,
    video_in_channels=video_in_channels,
    audio_in_channels=audio_in_channels,
    text_in_channels=text_in_channels,
    params_dtype=params_dtype,
    post_process_dtype=float32,
)

final_linear_audio = Linear(
    hidden_size,
    audio_in_channels,
    bias=False,
    dtype=float32,
)

final_linear_video = Linear(
    hidden_size,
    video_in_channels,
    bias=False,
    dtype=float32,
)

final_norm_audio = MultiModalityRMSNorm(hidden_size)

final_norm_video = MultiModalityRMSNorm(hidden_size)

bands = Parameter(bands, requires_grad=learnable)

device = device

dim = dim

dtype = dtype

in_pixels = in_pixels

learnable = learnable

linear_bands = linear_bands

max_res = max_res

temperature = temperature

attn_type_map: Tensor

k_ranges: Tensor

max_seqlen_k: int

max_seqlen_q: int

q_ranges: Tensor

softmax_scale: float

activation_func = create_activation_func(activation_type)

config: MLPConfig

down_proj = RowParallelLinear(
    input_size=intermediate_size,
    output_size=hidden_size,
    bias=False,
    input_is_parallel=True,
    return_bias=False,
)

pre_norm = MultiModalityRMSNorm(
    hidden_size, num_modality=num_modality
)

up_gate_proj = ColumnParallelLinear(
    input_size=hidden_size,
    output_size=intermediate_size_up,
    bias=False,
    gather_output=False,
    return_bias=False,
)

GELU7 = 'gelu7'

SWIGLU7 = 'swiglu7'

activation_type: MLPActivationType

gated_act: bool = False

hidden_size: int

intermediate_size: int

num_layers: int = 1

num_modality: int = 1

params_dtype: dtype

coords_style = coords_style

frame_receptive_field = frame_receptive_field

patch_size = patch_size

ref_audio_offset = ref_audio_offset

spatial_rope_interpolation = spatial_rope_interpolation

t_patch_size = t_patch_size

text_offset = text_offset

audio_in_channels: int = 64

checkpoint_qk_layernorm_rope: bool = False

enable_attn_gating: bool = True

gelu7_layers: list = field(
    default_factory=lambda: [0, 1, 2, 3]
)

head_dim: int = 128

hidden_size: int = 5120

local_attn_layers: list = field(default_factory=list)

mm_layers: list = field(
    default_factory=lambda: [0, 1, 2, 3, 36, 37, 38, 39]
)

num_layers: int = 40

num_query_groups: int = 8

params_dtype: dtype = float32

post_norm_layers: list = field(default_factory=list)

text_in_channels: int = 3584

video_in_channels: int = 48 * 4

experts = ModuleList(
    [
        (
            ColumnParallelLinear(
                input_size=input_size,
                output_size=output_size,
                bias=bias,
                gather_output=False,
                return_bias=False,
            )
        )
        for _ in (range(num_experts))
    ]
)

num_experts = num_experts

experts = ModuleList(
    [
        (
            QKVParallelLinear(
                hidden_size=hidden_size,
                head_size=head_size,
                total_num_heads=total_num_heads,
                total_num_kv_heads=total_num_kv_heads,
                bias=bias,
                return_bias=False,
            )
        )
        for _ in (range(num_experts))
    ]
)

head_size = head_size

num_experts = num_experts

num_heads = num_heads

num_kv_heads = num_kv_heads

experts = ModuleList(
    [
        (
            RowParallelLinear(
                input_size=input_size,
                output_size=output_size,
                bias=bias,
                input_is_parallel=True,
                return_bias=False,
            )
        )
        for _ in (range(num_experts))
    ]
)

num_experts = num_experts

AUDIO = 1

TEXT = 2

VIDEO = 0

group_size: Tensor = to(int32)

group_size_cpu: list[int] = [(int(x)) for x in (tolist())]

modality_mapping = modality_mapping

num_modalities: int = num_modalities

permuted_modality_mapping: Tensor = (
    _precompute_permute_mapping(modality_mapping)
)

dim = dim

eps = eps

forward = forward_multi_experts

num_modality = num_modality

weight = Parameter(
    zeros(dim * num_modality, device=device, dtype=float32)
)

coords_mapping

cu_seqlen

items: list[SingleData]

max_seqlen

modality_mapping

token_sequence

total_token_num

audio_feat_len: int

audio_x_t: Tensor

coords_mapping

coords_style: Literal['v1', 'v2'] = 'v1'

default_dtype

device

h: int

modality_mapping

patch_size: int

ref_audio_offset: int

spatial_rope_interpolation: Literal['inter', 'extra']

t: int

t_patch_size: int

text_offset: int

token_sequence

total_token_num

txt_feat: Tensor

txt_feat_len: int

video_x_t: Tensor

w: int

attention: Attention = Attention(attention_config)

attn_post_norm = MultiModalityRMSNorm(
    hidden_size, num_modality=num_modality
)

mlp: MLP = MLP(mlp_config)

mlp_post_norm = MultiModalityRMSNorm(
    hidden_size, num_modality=num_modality
)

post_norm = layer_idx in post_norm_layers

layers: list[TransFormerLayer] = ModuleList()

audio_in_channels: int

hidden_size: int

params_dtype: dtype

post_process_dtype: dtype

text_in_channels: int

video_in_channels: int

cu_seqlens_k: Tensor

cu_seqlens_q: Tensor

max_seqlen_k: int

max_seqlen_q: int