vllm_omni.model_executor.models.common.ming.dit ¶

attn_mask_enabled = attn_mask_enabled

dim = dim

dropout = dropout

heads = heads

inner_dim = dim_head * heads

k_norm = None

pe_attn_head = pe_attn_head

q_norm = None

to_k = nn.Linear(dim, self.inner_dim)

to_out = nn.ModuleList([])

to_q = nn.Linear(dim, self.inner_dim)

to_v = nn.Linear(dim, self.inner_dim)

cond_embedder = nn.Linear(input_feature_size, hidden_size)

attn = Attention(
    dim=hidden_size,
    heads=num_heads,
    dim_head=hidden_size // num_heads,
    dropout=dropout,
    qk_norm=qk_norm,
    pe_attn_head=pe_attn_head,
    attn_mask_enabled=attn_mask_enabled,
)

mlp = FeedForward(
    dim=hidden_size,
    mult=mlp_ratio,
    dropout=dropout,
    approximate="tanh",
)

norm1 = RMSNorm(hidden_size)

norm2 = RMSNorm(hidden_size)

ff = nn.Sequential(
    project_in,
    nn.Dropout(dropout),
    nn.Linear(inner_dim, dim_out),
)

linear = nn.Linear(hidden_size, out_channels, bias=True)

norm_final = RMSNorm(hidden_size)

eps = eps

weight = nn.Parameter(torch.ones(dim))