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, Any],
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: Tensor,
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,
cfg_vae_lengths: list[int] | None = None,
cfg_text_scales: list[float] | None = None,
cfg_img_scales: list[float] | 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: Tensor,
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_denoise_schedule(
x_t: Tensor, num_timesteps: int, timestep_shift: float
) -> tuple[Tensor, Tensor]
Build BAGEL's flow-matching timesteps and per-step deltas.
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 NaiveCaches into one for packed request/CFG attention.
Already-merged inputs keep their per-sequence lengths so nested request-by-CFG packing can split cache rows back out correctly.
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)