vllm_omni.diffusion.models.dreamid_omni.wan2_2 ¶

eps = eps

weight = Parameter(ones(hidden_size))

cross_attn = WanI2VCrossAttention(
    dim,
    num_heads,
    (-1, -1),
    qk_norm,
    eps,
    additional_emb_length,
)

cross_attn_norm = cross_attn_norm

dim = dim

eps = eps

ffn = Sequential(
    ColumnParallelLinear(
        dim,
        ffn_dim,
        bias=True,
        gather_output=False,
        return_bias=False,
        quant_config=quant_config,
        prefix=f"{prefix}.ffn.0" if prefix else "ffn.0",
    ),
    GELU(approximate="tanh"),
    RowParallelLinear(
        ffn_dim,
        dim,
        bias=True,
        input_is_parallel=True,
        return_bias=False,
        quant_config=quant_config,
        prefix=f"{prefix}.ffn.2" if prefix else "ffn.2",
    ),
)

ffn_dim = ffn_dim

modulation = ModulationAdd(dim, 6)

norm1 = WanLayerNorm(dim, eps)

norm2 = WanLayerNorm(dim, eps)

norm3 = (
    WanLayerNorm(dim, eps, elementwise_affine=True)
    if cross_attn_norm
    else Identity()
)

num_heads = num_heads

qk_norm = qk_norm

self_attn = WanSelfAttention(
    dim,
    num_heads,
    window_size,
    qk_norm,
    eps,
    quant_config=quant_config,
    prefix=f"{prefix}.self_attn" if prefix else "self_attn",
)

window_size = window_size

blocks = ModuleList(
    [
        (
            WanAttentionBlock(
                cross_attn_type,
                dim,
                ffn_dim,
                num_heads,
                window_size,
                qk_norm,
                cross_attn_norm,
                eps,
                additional_emb_length,
                quant_config=quant_config,
                prefix=f"{prefix}.blocks.{layer_idx}"
                if prefix
                else f"blocks.{layer_idx}",
            )
        )
        for layer_idx in (range(num_layers))
    ]
)

cross_attn_norm = cross_attn_norm

dim = dim

eps = eps

ffn_dim = ffn_dim

freq_dim = freq_dim

head = Head(dim, out_dim, patch_size, eps)

img_emb = MLPProj(additional_emb_dim, dim)

in_dim = in_dim

is_audio_type = is_audio_type

is_video_type = is_video_type

model_type = model_type

num_heads = num_heads

num_layers = num_layers

out_dim = out_dim

patch_embedding = Sequential(
    ChannelLastConv1d(
        in_dim, dim, kernel_size=7, padding=3
    ),
    SiLU(),
    ConvMLP(dim, dim * 4, kernel_size=7, padding=3),
)

patch_size = patch_size

qk_norm = qk_norm

temporal_rope_scaling_factor = temporal_rope_scaling_factor

text_dim = text_dim

text_embedding = Sequential(
    Linear(text_dim, dim),
    GELU(approximate="tanh"),
    Linear(dim, dim),
)

text_len = text_len

time_embedding = Sequential(
    Linear(freq_dim, dim), SiLU(), Linear(dim, dim)
)

time_projection = Sequential(SiLU(), Linear(dim, dim * 6))

window_size = window_size

attn = Attention(
    num_heads=num_heads,
    head_size=head_dim,
    num_kv_heads=num_kv_heads,
    softmax_scale=1.0 / head_dim**0.5,
    causal=False,
    prefix=prefix,
)

dim = dim

eps = eps

head_dim = dim // num_heads

norm_k = (
    DistributedRMSNorm(tp_kv_dim, eps=eps)
    if qk_norm
    else Identity()
)

norm_q = (
    DistributedRMSNorm(tp_inner_dim, eps=eps)
    if qk_norm
    else Identity()
)

num_heads = num_heads

num_kv_heads = num_kv_heads

o = RowParallelLinear(
    dim,
    dim,
    bias=True,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.o" if prefix else "o",
)

qk_norm = qk_norm

to_qkv = QKVParallelLinear(
    hidden_size=dim,
    head_size=head_dim,
    total_num_heads=num_heads,
    bias=True,
    quant_config=quant_config,
    prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
)

tp_inner_dim = num_heads * head_dim

tp_kv_dim = num_kv_heads * head_dim

window_size = window_size

attn = Attention(
    num_heads=num_heads,
    head_size=head_dim,
    num_kv_heads=num_kv_heads,
    softmax_scale=1.0 / head_dim**0.5,
    causal=False,
    prefix=prefix,
    disable_kv_quant=True,
)

dim = dim

eps = eps

head_dim = dim // num_heads

k = ColumnParallelLinear(
    dim,
    dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.k" if prefix else "k",
)

norm_k = (
    DistributedRMSNorm(tp_inner_dim, eps=eps)
    if qk_norm
    else Identity()
)

norm_q = (
    DistributedRMSNorm(tp_inner_dim, eps=eps)
    if qk_norm
    else Identity()
)

num_heads = num_heads // tp_size

num_kv_heads = num_heads

o = RowParallelLinear(
    dim,
    dim,
    bias=True,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.o" if prefix else "o",
)

q = ColumnParallelLinear(
    dim,
    dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.q" if prefix else "q",
)

qk_norm = qk_norm

tp_inner_dim = num_heads * head_dim

v = ColumnParallelLinear(
    dim,
    dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.v" if prefix else "v",
)

window_size = window_size