Lance transformer pieces.
The Lance LLM is BAGEL's Qwen2-MoT transformer verbatim — the released Lance_3B checkpoint uses the identical *_moe_gen / q_norm / vae2llm / llm2vae / time_embedder / latent_pos_embed layout, so Bagel / Qwen2MoTForCausalLM / Qwen2MoTConfig / NaiveCache are re-exported unchanged. Only the understanding ViT differs (Qwen2.5-VL vision instead of SigLIP) and the video path adds a 3D latent position embedding.
This module also provides :class:LanceBagel, a thin :class:Bagel subclass that overrides the two ViT entry points to consume Qwen2.5-VL's packed pixel_values + image_grid_thw layout directly (BAGEL itself assumes SigLIP-style (C, H, W) tensors that get patchified inside the model).
NOTE ON mRoPE: Lance's backbone is Qwen2.5-VL and its understanding / video paths use multimodal RoPE (mrope_section=[16,24,24]). BAGEL's BagelRotaryEmbedding is plain 1-D RoPE on scalar position ids. For the text2img generation path Lance assigns scalar positions (the gen latent block shares a single rope position, same as BAGEL), so the reused rotary is correct there. Full mRoPE for the x2t / video understanding path is a follow-up — see LancePositionEmbedding3D.
module-attribute ¶
Bases: CFGParallelMixin, Module
instance-attribute ¶
connector = MLPconnector(
self.vit_hidden_size,
self.hidden_size,
config.connector_act,
quant_config=quant_config,
prefix=f"{prefix}.connector",
)
instance-attribute ¶
get_flattened_position_ids = (
get_flattened_position_ids_extrapolate
)
instance-attribute ¶
instance-attribute ¶
latent_pos_embed = PositionEmbedding(
self.max_latent_size, self.hidden_size
)
instance-attribute ¶
instance-attribute ¶
instance-attribute ¶
vit_pos_embed = PositionEmbedding(
self.vit_max_num_patch_per_side, self.hidden_size
)
combine_multi_branch_cfg_noise(
predictions: list[Tensor],
true_cfg_scale: dict[str, float],
cfg_normalize: bool = False,
) -> Tensor
Combine gen/text/img branch velocities via Bagel's renorm CFG.
predictions[0] is the gen branch, [1] the text-CFG branch, and [2] (when present) the image-CFG branch. cfg_normalize is unused because renormalization is folded into _combine_cfg itself.
forward(
x_t: Tensor,
timestep: LongTensor,
packed_vae_token_indexes: LongTensor,
packed_vae_position_ids: LongTensor,
packed_text_ids: LongTensor,
packed_text_indexes: LongTensor,
packed_position_ids: LongTensor,
packed_seqlens: IntTensor,
past_key_values: NaiveCache,
cfg_renorm_min: float = 0.0,
cfg_renorm_type: str = "global",
cfg_text_scale: float = 1.0,
cfg_img_scale: float = 1.0,
cfg_branch_pids: list[Tensor] | None = None,
cfg_branch_caches: list[NaiveCache] | None = None,
)
forward_cache_update_text(
past_key_values: NaiveCache,
packed_text_ids: IntTensor,
packed_text_position_ids: LongTensor,
text_token_lens: LongTensor,
)
forward_cache_update_vae(
vae_model,
past_key_values: NaiveCache,
padded_images: Tensor,
patchified_vae_latent_shapes: list,
packed_vae_position_ids: LongTensor,
packed_timesteps: Tensor,
packed_vae_token_indexes: LongTensor,
packed_text_ids: LongTensor,
packed_text_indexes: LongTensor,
packed_position_ids: LongTensor,
packed_seqlens: IntTensor,
)
forward_cache_update_vit(
past_key_values: NaiveCache,
packed_text_ids: LongTensor,
packed_text_indexes: LongTensor,
packed_vit_tokens: Tensor,
packed_vit_token_indexes: LongTensor,
packed_vit_position_ids: LongTensor,
vit_token_seqlens: IntTensor,
packed_position_ids: LongTensor,
packed_seqlens: IntTensor,
)
forward_single_branch(
x_t: Tensor,
timestep: LongTensor,
packed_vae_token_indexes: LongTensor,
packed_vae_position_ids: LongTensor,
packed_text_ids: LongTensor,
packed_text_indexes: LongTensor,
packed_position_ids: LongTensor,
packed_seqlens: IntTensor,
past_key_values: NaiveCache,
) -> Tensor
Run a single-branch forward pass (no CFG batching).
Used by CFG parallel mode where each rank computes one branch. Returns the velocity v_t for the given branch. Supports Ulysses / Ring SP when parallel_config.sequence_parallel_size > 1.
generate_image(
packed_text_ids: LongTensor,
packed_text_indexes: LongTensor,
packed_init_noises: Tensor,
packed_vae_position_ids: LongTensor,
packed_vae_token_indexes: LongTensor,
packed_seqlens: IntTensor,
packed_position_ids: LongTensor,
past_key_values: NaiveCache,
num_timesteps: int = 24,
timestep_shift: float = 1.0,
cfg_renorm_min: float = 0.0,
cfg_renorm_type: str = "global",
cfg_interval: tuple[float, float] = [0, 1],
cfg_text_scale: float = 1.0,
cfg_text_packed_position_ids: LongTensor | None = None,
cfg_text_past_key_values: NaiveCache | None = None,
cfg_img_scale: float = 1.0,
cfg_img_packed_position_ids: LongTensor | None = None,
cfg_img_past_key_values: NaiveCache | None = None,
return_trajectory_latents: bool = False,
scheduler: object | None = None,
scheduler_kwargs: dict | None = None,
frame_condition_token_indexes: LongTensor | None = None,
)
generate_text(
past_key_values: NaiveCache,
packed_start_tokens: LongTensor,
packed_query_position_ids: LongTensor,
max_length: int,
do_sample: bool = False,
temperature: float = 1.0,
end_token_id: int | None = None,
)
Autoregressive text generation (ported from original BAGEL).
Decodes tokens one at a time, appending to past_key_values until max_length is reached or end_token_id is generated.
Single-branch velocity prediction for CFGParallelMixin.
Each CFG branch differs only by packed_position_ids and past_key_values (carried in kwargs); the heavy lifting is the per-branch forward_single_branch pass.
Prepare start tokens for autoregressive text generation.
Ported from the original BAGEL Bagel.prepare_start_tokens.
dataclass ¶
Bases: ModelOutput
class-attribute instance-attribute ¶
Bases: Bagel
Bagel subclass with Lance-specific ViT handling.
The released Lance checkpoint pairs BAGEL's Qwen2-MoT trunk with the Qwen2.5-VL vision tower (whose merger already projects to LLM hidden_size and which carries its own rotary positional encoding). Two BAGEL assumptions therefore break for Lance:
prepare_vit_imagescallstransforms(image) -> (C, H, W)and thenpatchify(...); Lance's image processor returns(num_patches_flat, patch_features) + image_grid_thwalready.forward_cache_update_vitadds aconnectorprojection plus a 2-Dvit_pos_embedthat has no checkpoint weights and would also double-count the ViT's own positional encoding.
This subclass overrides exactly those two methods. Everything else — LLM trunk, VAE flow, generation loop, forward_cache_update_text / forward_cache_update_vae — is reused unchanged.
forward_cache_update_vae(
vae_model,
past_key_values,
padded_images=None,
patchified_vae_latent_shapes=None,
packed_vae_position_ids=None,
packed_timesteps=None,
packed_vae_token_indexes=None,
packed_text_ids=None,
packed_text_indexes=None,
packed_position_ids=None,
packed_seqlens=None,
packed_indexes=None,
key_values_lens=None,
packed_key_value_indexes=None,
precomputed_latent=None,
)
Lance-native VAE prefill that actually scatters the encoded latents into the LLM query sequence.
:meth:Bagel.forward_cache_update_vae in vllm-omni computes packed_latent = vae2llm(...) + time_embed + pos_embed and then passes only packed_text_ids to the LLM — the VAE embeddings never enter the query sequence (the LLM builds it from embed_tokens(packed_text_ids) which is just the 2 framing tokens). The mismatch between query_lens = [num_vae + 2] and the resulting 2-token sequence is what crashes the gather inside attention.
We scatter both pieces explicitly: text framing tokens at packed_text_indexes and the VAE latent embeddings at packed_vae_token_indexes, producing a full-length (sum(packed_seqlens), hidden) sequence the LLM can attend over. Empty prep data (legacy x2t / x2t_video no-op path) is short-circuited.
VAE prefill router.
- When
imagesis non-empty (image_edit / video_edit path on the image side): delegate to :meth:_lance_native_prepare_vae_imageswhich emits Lance's 3-D mRoPE positions and a real VAE prefill, letting BAGEL's parentimage_editflow handle the rest. - When
imagesis empty (t2i / x2t paths): short-circuit with the no-op output, mirroring BAGEL's "no image to prefill" sentinel.
3-D analogue of :meth:prepare_vae_latent.
video_shapes is a list of (T, H, W) per request (RGB pixel space). We package one packed-init-noise tensor over T_lat × H_lat × W_lat latent tokens per video, plus 1-D indices into the 3-D position embedding table maintained by :class:LancePositionEmbedding3D (bagel.latent_pos_embed). Latent geometry:
- spatial:
H_lat = H // latent_downsample(=16for Lance) - temporal:
T_lat = (T - 1) // downsample_temporal + 1(=4for Wan2.2) - channels:
latent_channel = 48
Position ids are flattened t * max_per_side² + h * max_per_side + w so they index directly into the (max_num_frames * max_per_side², hidden_size) table.
3-D analogue of :meth:prepare_vae_latent_cfg (CFG side).
Mirrors :meth:prepare_video_latent's mRoPE 3-D position layout EXACTLY, including the LANCE_TOKENS_PER_SECOND * LANCE_SECONDS_PER_GRID temporal scaling. Without that scaling, the cfg_text branch attends with different rope coordinates than the cond branch (and than upstream's get_rope_index), which makes cfg_text_v_t diverge from upstream and the CFG combination amplifies the error every denoise step.
Multi-frame ViT prefill for the x2t_video / video_edit paths.
videos is a list of per-request video tensors / numpy arrays of shape (T, H, W, 3). By default the Qwen2-VL video processor is used to convert each video to (pixel_values_videos, video_grid_thw). For video_edit precision matching, the pipeline may pre-compute the upstream-style BucketResize output and pass it via precomputed_vit — a list of (pixel_values, grid_thw) per video, in which case the processor call is skipped.
Bases: Module
No-op connector for Lance.
BAGEL's connector projects the ViT hidden size to the LLM hidden size. Qwen2.5-VL's vision tower (which Lance uses) already projects to the LLM hidden size internally via merger (out_hidden_size = hidden_size), and the released Lance safetensors carry no connector.* weights. We therefore plug in an Identity connector so forward_cache_update_vit keeps its existing call site without a separate code path.
Bases: Module
Frozen 3-D sin-cos latent position embedding for the video path.
BAGEL only ships a 2-D PositionEmbedding (image latents). Lance's Lance_3B_Video checkpoint adds a temporal axis; this mirrors upstream modeling/lance/modeling_utils.py::PositionEmbedding3D. The image path uses t=1 and is numerically equivalent to the 2-D embedding.
instance-attribute ¶
Bases: Module
Packed (NaViT-style) wrapper around the Qwen2.5-VL vision tower.
Bridges BAGEL's vit(packed_pixel_values, ...) -> [num_tokens, vit_hidden] surface to the HF Qwen2_5_VisionTransformerPretrainedModel which consumes (hidden_states, grid_thw). The packed call additionally needs a per-image image_grid_thw so non-square images (and the spatial-merge token count) line up — :class:LanceBagel stashes the grid on the wrapper before invoking the ViT.
Bases: Module
No-op positional embedding for Lance's ViT tokens.
BAGEL adds an extra 2-D sin-cos vit_pos_embed on top of the ViT output. Qwen2.5-VL's vision tower already carries its own (rotary) positional encoding, and the released Lance safetensors carry no vit_pos_embed.* weights. This module returns a broadcast-friendly zero so the addition in forward_cache_update_vit is a no-op without requiring a code-path branch.
Bases: Module
instance-attribute ¶
fc1 = ColumnParallelLinear(
input_dim,
output_dim,
bias=True,
gather_output=False,
quant_config=quant_config,
prefix=f"{prefix}.fc1",
)
instance-attribute ¶
fc2 = RowParallelLinear(
output_dim,
output_dim,
bias=True,
input_is_parallel=True,
quant_config=quant_config,
prefix=f"{prefix}.fc2",
)
classmethod ¶
from_object(obj) -> NaiveCache
Convert a duck-typed cache (e.g., SimpleNamespace from KV transfer) to NaiveCache; in the future, we should find a better way to handle this, e.g., a model agnostic abstraction for key cache transfer instead of having this cache live in bagel.
NOTE: If a NaiveCache is provided, the object is just returned. Otherwise, we enumerate over the key/value cache values and map layer indices to the corresponding tensors.
staticmethod ¶
merge(caches: Sequence[NaiveCache]) -> NaiveCache
Merge per-branch NaiveCaches into one for batched attention; this lets us do the forward passes for CFG in one batched pass, although it's worth noting that this is currently only used for single request. We need this so that gen mode knows the respective kv lengths, and can split things back out as needed.
staticmethod ¶
Split tensor by lengths, which may include 0 entries, e.g., for splitting cfg branches out, since text_cfg may have 0 kv length.
0 lengths will be replaced with None in the returned list.
Bases: Module
Packed attention with Mixture-of-Tokens routing for understanding/generation.
Uses MoTQKVParallelLinear and MoTRowParallelLinear for tensor parallelism. Text and vae weights are held within the same MoT layer (text on self, vae on self.gen_exp). Token routing is driven by text_indices / vae_indices.
instance-attribute ¶
attn_causal = DiffusionAttention(
num_heads=self.total_num_heads,
head_size=self.head_dim,
softmax_scale=1.0 / self.head_dim**0.5,
causal=True,
num_kv_heads=self.total_num_kv_heads,
)
instance-attribute ¶
attn_noncausal = DiffusionAttention(
num_heads=self.total_num_heads,
head_size=self.head_dim,
softmax_scale=1.0 / self.head_dim**0.5,
causal=False,
num_kv_heads=self.total_num_kv_heads,
)
instance-attribute ¶
k_norm = MoTRMSNorm(
self.head_dim, head_norm=True, eps=config.rms_norm_eps
)
instance-attribute ¶
o_proj = MoTRowParallelLinear(
input_size=self.total_num_heads * self.head_dim,
output_size=self.hidden_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.o_proj",
)
instance-attribute ¶
q_norm = MoTRMSNorm(
self.head_dim, head_norm=True, eps=config.rms_norm_eps
)
instance-attribute ¶
qkv_proj = MoTQKVParallelLinear(
self.hidden_size,
self.head_dim,
self.total_num_heads,
self.total_num_kv_heads,
bias=True,
vae_bias=True,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
forward(
packed_query_sequence: Tensor,
query_lens: Tensor,
packed_query_position_embeddings: Tensor,
past_key_values: NaiveCache | None = None,
update_past_key_values=True,
is_causal=True,
mode="und",
packed_vae_token_indexes=None,
packed_text_indexes=None,
)
Bases: Module
instance-attribute ¶
pos_embed = nn.Parameter(
torch.zeros(max_num_patch_per_side**2, hidden_size),
requires_grad=False,
)
Bases: Qwen2Config
Configuration for Qwen2MoT (Mixture of Tokens) model.
This is fundamentally different from Qwen2, hence the distinct name.
Bases: Module
instance-attribute ¶
input_layernorm = MoTRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
instance-attribute ¶
mlp = BagelMLP(
config.hidden_size,
config.intermediate_size,
config.hidden_act,
quant_config=quant_config,
prefix=f"{prefix}.mlp",
)
instance-attribute ¶
mlp_moe_gen = BagelMLP(
config.hidden_size,
config.intermediate_size,
config.hidden_act,
quant_config=quant_config,
prefix=f"{prefix}.mlp_moe_gen",
)
instance-attribute ¶
post_attention_layernorm = MoTRMSNorm(
config.hidden_size, eps=config.rms_norm_eps
)
instance-attribute ¶
self_attn = attn_module(
config,
layer_idx,
parallel_config=parallel_config,
quant_config=quant_config,
prefix=f"{prefix}.self_attn",
)
forward(
hidden_states: Tensor,
encoder_hidden_states: Tensor | None = None,
packed_query_sequence: Tensor | None = None,
query_lens: Tensor = None,
packed_query_position_embeddings: Tensor = None,
past_key_values: NaiveCache | None = None,
update_past_key_values=True,
is_causal=True,
mode="und",
packed_vae_token_indexes=None,
packed_text_indexes=None,
) -> BaseNavitOutputWithPast
Bases: Qwen2PreTrainedModel
instance-attribute ¶
instance-attribute ¶
model = Qwen2MoTModel(
config,
parallel_config=parallel_config,
quant_config=quant_config,
prefix=f"{prefix}.model",
)
forward(
packed_query_sequence: Tensor | None = None,
query_lens: Tensor | None = None,
packed_query_position_ids: Tensor | None = None,
past_key_values: NaiveCache | None = None,
update_past_key_values=True,
is_causal=True,
mode="und",
packed_vae_token_indexes=None,
packed_text_indexes=None,
packed_text_ids: Tensor | None = None,
return_embeddings_only: bool = False,
) -> BaseNavitOutputWithPast
Load weights for MoT parallel layers.
Stacked parameter remapping (checkpoint name → model parameter): - q/k/v_proj → qkv_proj (text, shard q/k/v) - q/k/v_proj_moe_gen → qkv_proj.gen_exp (gen, shard q/k/v)
Direct remapping (no shard dimension): - o_proj_moe_gen → o_proj.gen_exp - {norm}_moe_gen.weight → {norm}.gen_weight (all MoTRMSNorm layers)
Text norm weights (input_layernorm.weight, q_norm.weight, etc.) and other names (embed_tokens, lm_head) pass through unchanged.
Bases: Qwen2PreTrainedModel
instance-attribute ¶
instance-attribute ¶
layers = nn.ModuleList(
[
(
Qwen2MoTDecoderLayer(
config,
layer_idx,
attn_module=PackedAttentionMoT,
parallel_config=parallel_config,
quant_config=quant_config,
prefix=f"{prefix}.layers.{layer_idx}",
)
)
for layer_idx in (range(config.num_hidden_layers))
]
)
forward(
packed_query_sequence: Tensor | None = None,
query_lens: Tensor | None = None,
packed_query_position_ids: Tensor | None = None,
past_key_values: NaiveCache | None = None,
update_past_key_values=True,
is_causal=True,
mode="und",
packed_vae_token_indexes=None,
packed_text_indexes=None,
packed_text_ids: Tensor | None = None,
return_embeddings_only: bool = False,
) -> BaseNavitOutputWithPast
Bases: Module
Embeds scalar timesteps into vector representations.
3-D sin-cos positional embedding (t, h, w), matching the upstream Lance modeling/lance/modeling_utils.py dimension split exactly.
imgs: (N, 3, H, W) or (3, H, W) x: (N, L, patch_size2 *3) or (L, patch_size2 *3)