vllm_omni.diffusion.models.helios.helios_transformer ¶

approximate = approximate

proj = ColumnParallelLinear(
    dim_in,
    dim_out,
    bias=bias,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
)

eps = eps

weight = Parameter(ones(hidden_size))

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

dim = dim

dropout = Dropout(dropout)

head_dim = head_dim

inner_dim = num_heads * head_dim

norm_k = DistributedRMSNorm(tp_inner_dim, eps=eps)

norm_q = DistributedRMSNorm(tp_inner_dim, eps=eps)

num_heads = num_heads // tp_size

to_k = ColumnParallelLinear(
    dim,
    inner_dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
)

to_out = RowParallelLinear(
    inner_dim,
    dim,
    bias=True,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
)

to_q = ColumnParallelLinear(
    dim,
    inner_dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
)

to_v = ColumnParallelLinear(
    dim,
    inner_dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
)

tp_inner_dim = num_heads * head_dim

net_0 = ColumnParallelGELU(
    dim,
    inner_dim,
    approximate="tanh",
    bias=bias,
    quant_config=quant_config,
)

net_1 = Identity()

net_2 = RowParallelLinear(
    inner_dim,
    dim_out,
    bias=bias,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
)

norm = FP32LayerNorm(dim, eps, elementwise_affine=False)

scale_shift_table = Parameter(randn(1, 2, dim) / dim**0.5)

theta = theta

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

dim = dim

dropout = Dropout(dropout)

head_dim = head_dim

history_key_scale = Parameter(ones(1))

history_scale_mode = history_scale_mode

inner_dim = num_heads * head_dim

is_amplify_history = is_amplify_history

max_scale = 10.0

norm_k = DistributedRMSNorm(tp_inner_dim, eps=eps)

norm_q = DistributedRMSNorm(tp_inner_dim, eps=eps)

num_heads = num_heads

num_kv_heads = num_kv_heads

to_out = RowParallelLinear(
    inner_dim,
    dim,
    bias=True,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
)

to_qkv = QKVParallelLinear(
    hidden_size=dim,
    head_size=head_dim,
    total_num_heads=num_heads,
    bias=True,
    quant_config=quant_config,
)

tp_inner_dim = num_heads * head_dim

act_fn = SiLU()

text_embedder = PixArtAlphaTextProjection(
    text_embed_dim, dim, act_fn="gelu_tanh"
)

time_embedder = TimestepEmbedding(
    in_channels=time_freq_dim, time_embed_dim=dim
)

time_proj = Linear(dim, time_proj_dim)

timesteps_proj = Timesteps(
    num_channels=time_freq_dim,
    flip_sin_to_cos=True,
    downscale_freq_shift=0,
)

blocks = ModuleList(
    [
        (
            HeliosTransformerBlock(
                inner_dim,
                ffn_dim,
                num_attention_heads,
                eps,
                cross_attn_norm,
                guidance_cross_attn=guidance_cross_attn,
                is_amplify_history=is_amplify_history,
                history_scale_mode=history_scale_mode,
                quant_config=quant_config,
            )
        )
        for _ in (range(num_layers))
    ]
)

condition_embedder = HeliosTimeTextEmbedding(
    dim=inner_dim,
    time_freq_dim=freq_dim,
    time_proj_dim=inner_dim * 6,
    text_embed_dim=text_dim,
)

config = type(
    "Config",
    (),
    {
        "patch_size": patch_size,
        "num_attention_heads": num_attention_heads,
        "attention_head_dim": attention_head_dim,
        "in_channels": in_channels,
        "out_channels": out_channels,
        "text_dim": text_dim,
        "freq_dim": freq_dim,
        "ffn_dim": ffn_dim,
        "num_layers": num_layers,
        "cross_attn_norm": cross_attn_norm,
        "qk_norm": qk_norm,
        "eps": eps,
        "added_kv_proj_dim": added_kv_proj_dim,
        "rope_dim": rope_dim,
        "rope_theta": rope_theta,
        "guidance_cross_attn": guidance_cross_attn,
        "zero_history_timestep": zero_history_timestep,
        "has_multi_term_memory_patch": has_multi_term_memory_patch,
        "is_amplify_history": is_amplify_history,
        "history_scale_mode": history_scale_mode,
    },
)()

dtype: dtype

has_multi_term_memory_patch = has_multi_term_memory_patch

inner_dim = inner_dim

norm_out = HeliosOutputNorm(inner_dim, eps)

packed_modules_mapping = {
    "to_qkv": ["to_q", "to_k", "to_v"]
}

patch_embedding = Conv3dLayer(
    in_channels=in_channels,
    out_channels=inner_dim,
    kernel_size=patch_size,
    stride=patch_size,
)

patch_long = Conv3dLayer(
    in_channels=in_channels,
    out_channels=inner_dim,
    kernel_size=(4, 8, 8),
    stride=(4, 8, 8),
)

patch_mid = Conv3dLayer(
    in_channels=in_channels,
    out_channels=inner_dim,
    kernel_size=(2, 4, 4),
    stride=(2, 4, 4),
)

patch_short = Conv3dLayer(
    in_channels=in_channels,
    out_channels=inner_dim,
    kernel_size=(1, 2, 2),
    stride=(1, 2, 2),
)

proj_out = Linear(
    inner_dim, out_channels * prod(patch_size)
)

rope = HeliosRotaryPosEmbed(
    rope_dim=rope_dim, theta=rope_theta
)

zero_history_timestep = zero_history_timestep

attn1 = HeliosSelfAttention(
    dim=dim,
    num_heads=num_heads,
    head_dim=head_dim,
    eps=eps,
    is_amplify_history=is_amplify_history,
    history_scale_mode=history_scale_mode,
    quant_config=quant_config,
)

attn2 = HeliosCrossAttention(
    dim=dim,
    num_heads=num_heads,
    head_dim=head_dim,
    eps=eps,
    quant_config=quant_config,
)

ffn = HeliosFeedForward(
    dim=dim,
    inner_dim=ffn_dim,
    dim_out=dim,
    quant_config=quant_config,
)

guidance_cross_attn = guidance_cross_attn

norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)

norm2 = (
    FP32LayerNorm(dim, eps, elementwise_affine=True)
    if cross_attn_norm
    else Identity()
)

norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)

scale_shift_table = Parameter(randn(1, 6, dim) / dim**0.5)