vllm_omni.diffusion.models.flux.flux_transformer ¶

approximate = approximate

proj = ColumnParallelLinear(
    dim_in,
    dim_out,
    bias=bias,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.proj",
)

net = ModuleList(layers)

add_kv_proj = QKVParallelLinear(
    hidden_size=added_kv_proj_dim,
    head_size=head_dim,
    total_num_heads=heads,
    bias=added_proj_bias,
    quant_config=quant_config,
    prefix=f"{prefix}.add_kv_proj",
)

added_kv_proj_dim = added_kv_proj_dim

added_proj_bias = added_proj_bias

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

context_pre_only = context_pre_only

dropout = dropout

head_dim = dim_head

heads = (
    out_dim // dim_head if out_dim is not None else heads
)

inner_dim = (
    out_dim if out_dim is not None else dim_head * heads
)

norm_added_k = RMSNorm(dim_head, eps=eps)

norm_added_q = RMSNorm(dim_head, eps=eps)

norm_k = RMSNorm(dim_head, eps=eps)

norm_q = RMSNorm(dim_head, eps=eps)

out_dim = out_dim if out_dim is not None else query_dim

pre_only = pre_only

query_dim = query_dim

rope = RotaryEmbedding(is_neox_style=False)

to_add_out = RowParallelLinear(
    inner_dim,
    query_dim,
    bias=out_bias,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.to_add_out",
)

to_out = ModuleList(
    [
        RowParallelLinear(
            inner_dim,
            out_dim,
            bias=out_bias,
            input_is_parallel=True,
            return_bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.to_out.0",
        ),
        Dropout(dropout),
    ]
)

to_qkv = QKVParallelLinear(
    hidden_size=query_dim,
    head_size=head_dim,
    total_num_heads=heads,
    bias=bias,
    quant_config=quant_config,
    prefix=f"{prefix}.to_qkv",
)

use_bias = bias

axes_dim = axes_dim

theta = theta

act_mlp = GELU(approximate='tanh')

attn = FluxAttention(
    query_dim=dim,
    dim_head=attention_head_dim,
    heads=num_attention_heads,
    out_dim=dim,
    bias=True,
    eps=1e-06,
    pre_only=True,
    quant_config=quant_config,
    prefix=f"{prefix}.attn",
)

mlp_hidden_dim = int(dim * mlp_ratio)

norm = AdaLayerNormZeroSingle(
    dim,
    quant_config=_safe_quant_config(quant_config),
    prefix=f"{prefix}.norm",
)

proj_mlp = ReplicatedLinear(
    dim,
    mlp_hidden_dim,
    bias=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.proj_mlp",
)

proj_out = ReplicatedLinear(
    dim + mlp_hidden_dim,
    dim,
    bias=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.proj_out",
)

context_embedder = Linear(joint_attention_dim, inner_dim)

guidance_embeds = guidance_embeds

in_channels = in_channels

inner_dim = num_attention_heads * attention_head_dim

norm_out = AdaLayerNormContinuous(
    inner_dim,
    inner_dim,
    elementwise_affine=False,
    eps=1e-06,
    quant_config=_safe_quant_config(quant_config),
    prefix="norm_out",
)

out_channels = out_channels or in_channels

packed_modules_mapping = {
    "to_qkv": ["to_q", "to_k", "to_v"],
    "add_kv_proj": [
        "add_q_proj",
        "add_k_proj",
        "add_v_proj",
    ],
}

parallel_config = parallel_config

pos_embed = FluxPosEmbed(
    theta=theta, axes_dim=axes_dims_rope
)

proj_out = Linear(
    inner_dim,
    patch_size * patch_size * out_channels,
    bias=True,
)

single_transformer_blocks = ModuleList(
    [
        (
            FluxSingleTransformerBlock(
                dim=inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
                quant_config=quant_config,
                prefix=f"single_transformer_blocks.{i}",
            )
        )
        for i in (range(num_single_layers))
    ]
)

time_text_embed = text_time_guidance_cls(
    embedding_dim=inner_dim,
    pooled_projection_dim=pooled_projection_dim,
)

transformer_blocks = ModuleList(
    [
        (
            FluxTransformerBlock(
                dim=inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
                quant_config=_safe_quant_config(
                    quant_config
                ),
                prefix=f"transformer_blocks.{i}",
            )
        )
        for i in (range(num_layers))
    ]
)

x_embedder = Linear(in_channels, inner_dim)

attn = FluxAttention(
    query_dim=dim,
    added_kv_proj_dim=dim,
    dim_head=attention_head_dim,
    heads=num_attention_heads,
    out_dim=dim,
    context_pre_only=False,
    bias=True,
    eps=eps,
    quant_config=quant_config,
    prefix=f"{prefix}.attn",
)

ff = FeedForward(
    dim=dim,
    dim_out=dim,
    quant_config=quant_config,
    prefix=f"{prefix}.ff",
)

ff_context = FeedForward(
    dim=dim,
    dim_out=dim,
    quant_config=quant_config,
    prefix=f"{prefix}.ff_context",
)

norm1 = AdaLayerNormZero(
    dim, quant_config=quant_config, prefix=f"{prefix}.norm1"
)

norm1_context = AdaLayerNormZero(
    dim,
    quant_config=quant_config,
    prefix=f"{prefix}.norm1_context",
)

norm2 = LayerNorm(dim, elementwise_affine=False, eps=1e-06)

norm2_context = LayerNorm(
    dim, elementwise_affine=False, eps=1e-06
)