vllm_omni.diffusion.models.omnigen2.omnigen2_transformer ¶

caption_embedder = Sequential(
    RMSNorm(text_feat_dim, eps=norm_eps),
    Linear(text_feat_dim, hidden_size, bias=True),
)

time_proj = Timesteps(
    num_channels=frequency_embedding_size,
    flip_sin_to_cos=True,
    downscale_freq_shift=0.0,
    scale=timestep_scale,
)

timestep_embedder = TimestepEmbedding(
    in_channels=frequency_embedding_size,
    time_embed_dim=min(hidden_size, 1024),
)

act_fn = get_act_and_mul_fn('silu')

down_proj = RowParallelLinear(
    inner_dim,
    dim,
    bias=False,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.down_proj",
)

gate_up_proj = MergedColumnParallelLinear(
    dim,
    [inner_dim, inner_dim],
    bias=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.gate_up_proj",
)

linear_1 = Linear(
    conditioning_embedding_dim, embedding_dim, bias=bias
)

linear_2 = None

norm = LayerNorm(
    embedding_dim, eps, elementwise_affine, bias
)

silu = SiLU()

linear = Linear(
    min(embedding_dim, 1024), 4 * embedding_dim, bias=True
)

norm = RMSNorm(embedding_dim, eps=norm_eps)

silu = SiLU()

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

dim = dim

eps = eps

head_dim = dim // num_heads

norm_k = RMSNorm(head_dim, eps=eps)

norm_q = RMSNorm(head_dim, eps=eps)

num_heads = num_heads

num_kv_heads = num_kv_heads

to_out = ModuleList(
    [
        RowParallelLinear(
            dim,
            dim,
            bias=False,
            input_is_parallel=False,
            quant_config=quant_config,
            return_bias=False,
            prefix=f"{prefix}.to_out.0",
        )
    ]
)

to_qkv = QKVParallelLinear(
    hidden_size=dim,
    head_size=head_dim,
    total_num_heads=num_heads,
    total_num_kv_heads=num_kv_heads,
    disable_tp=True,
    bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.to_qkv",
)

axes_dim = axes_dim

axes_lens = axes_lens

patch_size = patch_size

theta = theta

config = SimpleNamespace(
    patch_size=patch_size,
    in_channels=in_channels,
    out_channels=out_channels or in_channels,
    hidden_size=hidden_size,
    axes_dim_rope=axes_dim_rope,
    axes_lens=axes_lens,
)

context_refiner = ModuleList(
    [
        (
            OmniGen2TransformerBlock(
                hidden_size,
                num_attention_heads,
                num_kv_heads,
                multiple_of,
                ffn_dim_multiplier,
                norm_eps,
                modulation=False,
                quant_config=quant_config,
                prefix=f"context_refiner.{i}",
            )
        )
        for i in (range(num_refiner_layers))
    ]
)

image_index_embedding = Parameter(randn(5, hidden_size))

layers = ModuleList(
    [
        (
            OmniGen2TransformerBlock(
                hidden_size,
                num_attention_heads,
                num_kv_heads,
                multiple_of,
                ffn_dim_multiplier,
                norm_eps,
                modulation=True,
                quant_config=quant_config,
                prefix=f"layers.{i}",
            )
        )
        for i in (range(num_layers))
    ]
)

noise_refiner = ModuleList(
    [
        (
            OmniGen2TransformerBlock(
                hidden_size,
                num_attention_heads,
                num_kv_heads,
                multiple_of,
                ffn_dim_multiplier,
                norm_eps,
                modulation=True,
                quant_config=quant_config,
                prefix=f"noise_refiner.{i}",
            )
        )
        for i in (range(num_refiner_layers))
    ]
)

norm_out = LuminaLayerNormContinuous(
    embedding_dim=hidden_size,
    conditioning_embedding_dim=min(hidden_size, 1024),
    elementwise_affine=False,
    eps=1e-06,
    bias=True,
    out_dim=patch_size * patch_size * out_channels,
)

out_channels = out_channels

ref_image_patch_embedder = Linear(
    in_features=patch_size * patch_size * in_channels,
    out_features=hidden_size,
)

ref_image_refiner = ModuleList(
    [
        (
            OmniGen2TransformerBlock(
                hidden_size,
                num_attention_heads,
                num_kv_heads,
                multiple_of,
                ffn_dim_multiplier,
                norm_eps,
                modulation=True,
                quant_config=quant_config,
                prefix=f"ref_image_refiner.{i}",
            )
        )
        for i in (range(num_refiner_layers))
    ]
)

rope_embedder = OmniGen2RotaryPosEmbed(
    theta=10000,
    axes_dim=axes_dim_rope,
    axes_lens=axes_lens,
    patch_size=patch_size,
)

time_caption_embed = (
    Lumina2CombinedTimestepCaptionEmbedding(
        hidden_size=hidden_size,
        text_feat_dim=text_feat_dim,
        norm_eps=norm_eps,
        timestep_scale=timestep_scale,
    )
)

x_embedder = Linear(
    in_features=patch_size * patch_size * in_channels,
    out_features=hidden_size,
)

attn = OmniGen2Attention(
    dim=dim,
    num_heads=num_attention_heads,
    num_kv_heads=num_kv_heads,
    eps=1e-05,
    quant_config=quant_config,
    prefix=f"{prefix}.attn",
)

feed_forward = LuminaFeedForward(
    dim=dim,
    inner_dim=4 * dim,
    multiple_of=multiple_of,
    ffn_dim_multiplier=ffn_dim_multiplier,
    quant_config=quant_config,
    prefix=f"{prefix}.feed_forward",
)

ffn_norm1 = RMSNorm(dim, eps=norm_eps)

ffn_norm2 = RMSNorm(dim, eps=norm_eps)

head_dim = dim // num_attention_heads

modulation = modulation

norm1 = LuminaRMSNormZero(
    embedding_dim=dim,
    norm_eps=norm_eps,
    norm_elementwise_affine=True,
    quant_config=quant_config,
    prefix=f"{prefix}.norm1",
)

norm2 = RMSNorm(dim, eps=norm_eps)

act = get_activation(act_fn)

cond_proj = Linear(cond_proj_dim, in_channels, bias=False)

linear_1 = Linear(
    in_channels, time_embed_dim, sample_proj_bias
)

linear_2 = Linear(
    time_embed_dim, time_embed_dim_out, sample_proj_bias
)

post_act = None