vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep ¶

boi = boi

eoi = eoi

fm_head_batch_mul = fm_head_batch_mul

fm_head_dim = fm_head_dim

fm_head_layers = fm_head_layers

im_loss_weight = im_loss_weight

image_placeholder_id = image_placeholder_id

latent_channels = latent_channels

latent_patch_size = latent_patch_size

latent_size = latent_size

lm_loss_weight = lm_loss_weight

model_type = 'nextstep'

o_attention_bias = (
    self.attention_bias
    if o_attention_bias is None
    else o_attention_bias
)

pad_token_id_added = pad_token_id_added

vae_name_or_path = vae_name_or_path

config = config

embed_tokens = nn.Embedding(
    config.vocab_size, config.hidden_size, self.padding_idx
)

image_head = FlowMatchingHead(
    input_dim=token_dim,
    cond_dim=config.hidden_size,
    dim=config.fm_head_dim,
    layers=config.fm_head_layers,
)

image_in_projector = nn.Linear(
    token_dim, config.hidden_size
)

image_out_projector = nn.Linear(
    config.hidden_size, config.hidden_size
)

layers = nn.ModuleList(
    [
        (LlamaDecoderLayer(config, layer_idx))
        for layer_idx in (range(config.num_hidden_layers))
    ]
)

lm_head = nn.Linear(
    config.hidden_size, config.vocab_size, bias=False
)

norm = LlamaRMSNorm(
    config.hidden_size, eps=config.rms_norm_eps
)

padding_idx = config.pad_token_id

rotary_emb = LlamaRotaryEmbedding(config=config)

vocab_size = config.vocab_size