vllm_omni.diffusion.models.sdxl.sdxl_unet ¶

act = nn.SiLU()

linear_1 = nn.Linear(
    text_embed_dim + addition_time_embed_dim * 6,
    time_embed_dim,
)

linear_2 = nn.Linear(time_embed_dim, time_embed_dim)

attention_head_dim = attention_head_dim

attn1 = SDXLSelfAttention(
    dim=dim,
    num_attention_heads=num_attention_heads,
    attention_head_dim=attention_head_dim,
    prefix=f"{prefix}.attn1",
)

attn2 = SDXLCrossAttention(
    dim=dim,
    cross_attention_dim=cross_attention_dim,
    num_attention_heads=num_attention_heads,
    attention_head_dim=attention_head_dim,
    prefix=f"{prefix}.attn2",
)

ff = SDXLFeedForward(dim=dim, inner_dim=dim * 4)

norm1 = nn.LayerNorm(dim, elementwise_affine=True)

norm2 = nn.LayerNorm(dim, elementwise_affine=True)

norm3 = nn.LayerNorm(dim, elementwise_affine=True)

num_attention_heads = num_attention_heads

attention = Attention(
    num_heads=num_attention_heads,
    head_size=attention_head_dim,
    causal=False,
    softmax_scale=1.0 / math.sqrt(attention_head_dim),
    role="cross",
    prefix=prefix,
    skip_sequence_parallel=True,
)

to_k = ColumnParallelLinear(
    cross_attention_dim, inner_dim, bias=False
)

to_out = nn.ModuleList(
    [RowParallelLinear(inner_dim, dim, bias=True)]
)

to_q = ColumnParallelLinear(dim, inner_dim, bias=False)

to_v = ColumnParallelLinear(
    cross_attention_dim, inner_dim, bias=False
)

attentions = nn.ModuleList()

downsamplers = None

resnets = nn.ModuleList()

attentions = nn.ModuleList()

resnets = nn.ModuleList()

upsamplers = None

downsamplers = None

resnets = nn.ModuleList()

conv = nn.Conv2d(
    channels, channels, kernel_size=3, stride=2, padding=1
)

geglu = SDXLGEGLU(dim, inner_dim)

out_proj = RowParallelLinear(inner_dim, dim, bias=True)

proj = MergedColumnParallelLinear(
    dim_in, [dim_out, dim_out], bias=True
)

conv1 = nn.Conv2d(
    in_channels, out_channels, kernel_size=3, padding=1
)

conv2 = nn.Conv2d(
    out_channels, out_channels, kernel_size=3, padding=1
)

conv_shortcut = None

nonlinearity = nn.SiLU()

norm1 = nn.GroupNorm(groups, in_channels, eps=1e-05)

norm2 = nn.GroupNorm(groups, out_channels, eps=1e-05)

time_emb_proj = nn.Linear(time_embed_dim, out_channels)

attention = Attention(
    num_heads=num_attention_heads,
    head_size=attention_head_dim,
    causal=False,
    softmax_scale=1.0 / math.sqrt(attention_head_dim),
    role="self",
    prefix=prefix,
)

to_out = nn.ModuleList(
    [RowParallelLinear(inner_dim, dim, bias=True)]
)

to_qkv = QKVParallelLinear(
    hidden_size=dim,
    head_size=attention_head_dim,
    total_num_heads=num_attention_heads,
    bias=False,
)

act = nn.SiLU()

linear_1 = nn.Linear(in_channels, time_embed_dim)

linear_2 = nn.Linear(time_embed_dim, time_embed_dim)

norm = nn.GroupNorm(32, in_channels, eps=1e-06, affine=True)

proj_in = nn.Linear(in_channels, inner_dim)

proj_out = nn.Linear(inner_dim, in_channels)

transformer_blocks = nn.ModuleList(
    [
        (
            SDXLBasicTransformerBlock(
                dim=inner_dim,
                num_attention_heads=num_attention_heads,
                attention_head_dim=attention_head_dim,
                cross_attention_dim=cross_attention_dim,
                prefix=f"{prefix}.transformer_blocks.{i}",
            )
        )
        for i in (range(num_layers))
    ]
)

add_embedding = SDXLAddTimestepEmbedding(
    addition_time_embed_dim=addition_time_embed_dim,
    text_embed_dim=1280,
    time_embed_dim=time_embed_dim,
)

conv_act = nn.SiLU()

conv_in = nn.Conv2d(
    4, model_channels, kernel_size=3, padding=1
)

conv_norm_out = nn.GroupNorm(
    32, block_out_channels[0], eps=1e-05
)

conv_out = nn.Conv2d(
    block_out_channels[0], 4, kernel_size=3, padding=1
)

down_blocks = nn.ModuleList()

in_channels = 4

mid_block = SDXLUNetMidBlock2DCrossAttn(
    in_channels=block_out_channels[2],
    time_embed_dim=time_embed_dim,
    num_attention_heads=block_out_channels[2]
    // num_head_channels,
    cross_attention_dim=cross_attention_dim,
    transformer_layers_per_block=transformer_layers_per_block[
        2
    ],
    prefix="mid_block",
)

od_config = od_config

out_channels = 4

time_embedding = SDXLTimestepEmbedding(
    model_channels, time_embed_dim
)

time_proj_dim = model_channels

up_blocks = nn.ModuleList()

attentions = nn.ModuleList(
    [
        SDXLTransformer2DModel(
            num_attention_heads=num_attention_heads,
            attention_head_dim=attention_head_dim,
            in_channels=in_channels,
            num_layers=transformer_layers_per_block,
            cross_attention_dim=cross_attention_dim,
            prefix=f"{prefix}.attentions.0",
        )
    ]
)

resnets = nn.ModuleList(
    [
        SDXLResnetBlock2D(
            in_channels, in_channels, time_embed_dim
        ),
        SDXLResnetBlock2D(
            in_channels, in_channels, time_embed_dim
        ),
    ]
)

resnets = nn.ModuleList()

upsamplers = None

conv = nn.Conv2d(
    channels, channels, kernel_size=3, padding=1
)