gemma
fastvideo.models.encoders.gemma.BaseEncoderOutput
BaseEncoderOutput(last_hidden_state: torch.FloatTensor | None = None, pooler_output: torch.FloatTensor | None = None, hidden_states: tuple[torch.FloatTensor, ...] | None = None, attentions: tuple[torch.FloatTensor, ...] | None = None, attention_mask: torch.Tensor | None = None)fastvideo.models.encoders.gemma.TextEncoderConfig
TextEncoderConfig(arch_config: ArchConfig = TextEncoderArchConfig(), prefix: str = '', quant_config: QuantizationConfig | None = None, lora_config: Any | None = None, is_chat_model: bool = False, treat_empty_as_dot: bool = False, *, _resolved_attention_backend: AttentionBackendEnum | None = None, chat_template_enable_thinking: bool = False)fastvideo.models.encoders.gemma.TextEncoder
TextEncoder(config: TextEncoderConfig)Base for native encoders with a model-specific forward output contract.
fastvideo.models.encoders.gemma.TextEncoder.forward
forward(*args: Any, **kwargs: Any) -> TextEncoderOutputTfastvideo.models.encoders.gemma.FeedForward
FeedForward(dim: int, dim_out: int, mult: int = 4, bias: bool = True, quant_config: QuantizationConfig | None = None, prefix: str = '')LTX-2 FFN: GELUApprox -> Identity -> Linear.
fastvideo.models.encoders.gemma.FeedForward.forward
forward(x: torch.Tensor) -> torch.Tensorfastvideo.models.encoders.gemma.LTXRopeType
LTXRopeType()LTX-2 rotary variants (interleaved vs split).
fastvideo.models.encoders.gemma.apply_ltx_rotary_emb
apply_ltx_rotary_emb(input_tensor: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], rope_type: LTXRopeType = LTXRopeType.INTERLEAVED) -> torch.TensorApply LTX-2 rotary embeddings to a tensor.
fastvideo.models.encoders.gemma.generate_ltx_freq_grid_float64
generate_ltx_freq_grid_float64(positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int) -> torch.TensorGenerate an eager-cached, compile-safe float64 LTX-2 rotary grid.
fastvideo.models.encoders.gemma.generate_ltx_freq_grid_pytorch
generate_ltx_freq_grid_pytorch(positional_embedding_theta: float, positional_embedding_max_pos_count: int, inner_dim: int) -> torch.TensorGenerate an eager-cached, compile-safe float32 LTX-2 rotary grid.
fastvideo.models.encoders.gemma.precompute_ltx_freqs_cis
precompute_ltx_freqs_cis(indices_grid: torch.Tensor, dim: int, out_dtype: torch.dtype, theta: float = 10000.0, max_pos: list[int] | None = None, use_middle_indices_grid: bool = False, num_attention_heads: int = 32, rope_type: LTXRopeType = LTXRopeType.INTERLEAVED, freq_grid_generator: Callable[[float, int, int], torch.Tensor] = generate_ltx_freq_grid_pytorch) -> tuple[torch.Tensor, torch.Tensor]Precompute LTX-2 rotary cos/sin grids for (t, x, y) positions.
fastvideo.models.encoders.gemma.default_weight_loader
default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> NoneDefault weight loader.
fastvideo.models.encoders.gemma.effective_attention_backend
effective_attention_backend(config: object) -> AttentionBackendEnum | NoneReturn the attention backend that a component built from config follows.
Some components choose part of their own structure by backend: Wan picks
its VSA transformer block, LTX-2 its distributed-attention path, and the
Gemma and T5-Gemma text encoders switch Hugging Face attention to SDPA.
That choice must match the attention layers inside the component, and
those layers resolve through :func:get_attn_backend. This function
applies the same order:
- The decision a loader recorded on
configwith :func:record_resolved_attention_backend. Every component built by a loader has one when a backend was requested. - While a loader is building a component (a
_component_attention_backend_scopeis active), that scope's rule: its backend, and the environment variable only when the scope consults it. A scope that deliberately ignores the environment variable, such as the dense DMD teacher and critic, keeps ignoring it here. - With no loader involved, the
FASTVIDEO_ATTENTION_BACKENDenvironment variable. This covers components constructed directly, such as MagiHuman's lazily built T5-Gemma encoder or a script that builds a transformer itself.
Reading only the recorded decision is not enough: a directly constructed
component has no recorded decision (None), so it would ignore the
environment variable while its attention layers still follow it, and the
component's structure would not match its layers.
fastvideo.models.encoders.gemma.AttentionBackendEnum
AttentionBackendEnum()fastvideo.models.encoders.gemma.get_local_torch_device
get_local_torch_device() -> torch.deviceReturn the torch device for the current rank.
fastvideo.models.encoders.gemma.GemmaConnectorConfig
GemmaConnectorConfig(num_attention_heads: int, attention_head_dim: int, num_layers: int, positional_embedding_theta: float, positional_embedding_max_pos: list[int], rope_type: LTXRopeType, double_precision_rope: bool, num_learnable_registers: int | None, apply_gated_attention: bool = False, ff_bias: bool = True)Source
@dataclass(frozen=True)
class GemmaConnectorConfig:
num_attention_heads: int
attention_head_dim: int
num_layers: int
positional_embedding_theta: float
positional_embedding_max_pos: list[int]
rope_type: LTXRopeType
double_precision_rope: bool
num_learnable_registers: int | None
# LTX-2.3 connector gated attention (default False == LTX-2.0 behavior).
apply_gated_attention: bool = False
ff_bias: bool = Truefastvideo.models.encoders.gemma.GemmaFeaturesExtractorProjLinear
GemmaFeaturesExtractorProjLinear(in_features: int, out_features: int)Linear projection that aggregates stacked Gemma hidden states.
Source
class GemmaFeaturesExtractorProjLinear(nn.Module):
"""Linear projection that aggregates stacked Gemma hidden states."""
def __init__(self, in_features: int, out_features: int) -> None:
super().__init__()
self.aggregate_embed = nn.Linear(in_features, out_features, bias=False)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.aggregate_embed(x)</code></pre></details>
fastvideo.models.encoders.gemma.GemmaFeaturesExtractorProjLinear.forward
forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.aggregate_embed(x)
fastvideo.models.encoders.gemma.Embeddings1DConnector
Embeddings1DConnector(config: GemmaConnectorConfig)
Transformer connector that refines Gemma embeddings for LTX-2.
Source
class Embeddings1DConnector(nn.Module):
"""Transformer connector that refines Gemma embeddings for LTX-2."""
_supports_gradient_checkpointing = True
def __init__(self, config: GemmaConnectorConfig) -> None:
super().__init__()
self.num_attention_heads = config.num_attention_heads
self.inner_dim = config.num_attention_heads * config.attention_head_dim
self.positional_embedding_theta = config.positional_embedding_theta
self.positional_embedding_max_pos = config.positional_embedding_max_pos
self.rope_type = config.rope_type
self.double_precision_rope = config.double_precision_rope
self.transformer_1d_blocks = nn.ModuleList([
_BasicTransformerBlock1D(
dim=self.inner_dim,
heads=config.num_attention_heads,
dim_head=config.attention_head_dim,
rope_type=config.rope_type,
apply_gated_attention=config.apply_gated_attention,
ff_bias=config.ff_bias,
) for _ in range(config.num_layers)
])
self.num_learnable_registers = config.num_learnable_registers
if self.num_learnable_registers:
self.learnable_registers = nn.Parameter(
torch.rand(
self.num_learnable_registers,
self.inner_dim,
dtype=torch.bfloat16,
) * 2.0 - 1.0)
def _replace_padded_with_learnable_registers(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
assert hidden_states.shape[1] % self.num_learnable_registers == 0, (
f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by "
f"num_learnable_registers {self.num_learnable_registers}.")
num_registers_duplications = (hidden_states.shape[1] // self.num_learnable_registers)
learnable_registers = torch.tile(self.learnable_registers, (num_registers_duplications, 1))
seq_len = hidden_states.shape[1]
valid = attention_mask.squeeze(1).squeeze(1) >= -9000.0 # (batch, seq)
# Left-align each row's valid tokens, then fill the tail with registers.
# A boolean index cannot express this: every row may keep a different
# number of tokens, so the result is not rectangular. Gathering a stable
# descending argsort of the mask moves the valid tokens to the front of
# each row while preserving their original relative order.
order = torch.argsort(valid.to(torch.int8), dim=1, descending=True, stable=True)
adjusted_hidden_states = torch.gather(hidden_states, 1, order.unsqueeze(-1).expand_as(hidden_states))
# After the gather, row b holds its valid tokens in slots
# [0, valid[b].sum()). Deriving the keep-mask from that per-row count is
# what makes this correct when rows have different token counts.
keep = (torch.arange(seq_len, device=hidden_states.device).unsqueeze(0)
< valid.sum(dim=1, keepdim=True)).unsqueeze(-1).to(hidden_states.dtype)
hidden_states = keep * adjusted_hidden_states + (1 - keep) * learnable_registers
attention_mask = torch.full_like(
attention_mask,
0.0,
dtype=attention_mask.dtype,
device=attention_mask.device,
)
return hidden_states, attention_mask
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
if self.num_learnable_registers:
hidden_states, attention_mask = (self._replace_padded_with_learnable_registers(
hidden_states, attention_mask))
indices_grid = torch.arange(
hidden_states.shape[1],
dtype=torch.float32,
device=hidden_states.device,
)
indices_grid = indices_grid[None, None, :]
freq_grid_generator = (generate_ltx_freq_grid_float64
if self.double_precision_rope else generate_ltx_freq_grid_pytorch)
freqs_cis = precompute_ltx_freqs_cis(
indices_grid=indices_grid,
dim=self.inner_dim,
out_dtype=hidden_states.dtype,
theta=self.positional_embedding_theta,
max_pos=self.positional_embedding_max_pos,
num_attention_heads=self.num_attention_heads,
rope_type=self.rope_type,
freq_grid_generator=freq_grid_generator,
)
for block in self.transformer_1d_blocks:
hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis)
hidden_states = torch.nn.functional.rms_norm(hidden_states, (hidden_states.shape[-1], ), eps=1e-6)
return hidden_states, attention_mask</code></pre></details>
fastvideo.models.encoders.gemma.Embeddings1DConnector.forward
forward(hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]
Source
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
if self.num_learnable_registers:
hidden_states, attention_mask = (self._replace_padded_with_learnable_registers(
hidden_states, attention_mask))
indices_grid = torch.arange(
hidden_states.shape[1],
dtype=torch.float32,
device=hidden_states.device,
)
indices_grid = indices_grid[None, None, :]
freq_grid_generator = (generate_ltx_freq_grid_float64
if self.double_precision_rope else generate_ltx_freq_grid_pytorch)
freqs_cis = precompute_ltx_freqs_cis(
indices_grid=indices_grid,
dim=self.inner_dim,
out_dtype=hidden_states.dtype,
theta=self.positional_embedding_theta,
max_pos=self.positional_embedding_max_pos,
num_attention_heads=self.num_attention_heads,
rope_type=self.rope_type,
freq_grid_generator=freq_grid_generator,
)
for block in self.transformer_1d_blocks:
hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis)
hidden_states = torch.nn.functional.rms_norm(hidden_states, (hidden_states.shape[-1], ), eps=1e-6)
return hidden_states, attention_mask</code></pre></details>
fastvideo.models.encoders.gemma.LTX2GemmaTextEncoderModel
LTX2GemmaTextEncoderModel(config: TextEncoderConfig)
Source
class LTX2GemmaTextEncoderModel(TextEncoder):
_supported_attention_backends = (
AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA,
)
def __init__(self, config: TextEncoderConfig) -> None:
super().__init__(config)
arch = config.arch_config
# LTX-2.3 routes the caption projection before the connector and uses
# separate per-modality feature extractor linears. Default (False)
# keeps the single LTX-2.0 GemmaFeaturesExtractorProjLinear.
self.use_v2_feature_extractor = bool(getattr(arch, "caption_proj_before_connector", False))
if self.use_v2_feature_extractor:
video_out_features = (getattr(arch, "video_feature_extractor_out_features", None)
or arch.feature_extractor_out_features)
audio_out_features = (getattr(arch, "audio_feature_extractor_out_features", None) or video_out_features)
self.video_feature_extractor_linear = nn.Linear(
arch.feature_extractor_in_features,
video_out_features,
bias=True,
)
self.audio_feature_extractor_linear = nn.Linear(
arch.feature_extractor_in_features,
audio_out_features,
bias=True,
)
else:
self.feature_extractor_linear = GemmaFeaturesExtractorProjLinear(
in_features=arch.feature_extractor_in_features,
out_features=arch.feature_extractor_out_features,
)
connector_apply_gated_attention = bool(getattr(arch, "connector_apply_gated_attention", False))
connector_ff_bias = bool(getattr(arch, "connector_ff_bias", True))
video_connector_config = GemmaConnectorConfig(
num_attention_heads=arch.connector_num_attention_heads,
attention_head_dim=arch.connector_attention_head_dim,
num_layers=arch.connector_num_layers,
positional_embedding_theta=arch.connector_positional_embedding_theta,
positional_embedding_max_pos=arch.connector_positional_embedding_max_pos,
rope_type=LTXRopeType(arch.connector_rope_type),
double_precision_rope=arch.connector_double_precision_rope,
num_learnable_registers=arch.connector_num_learnable_registers,
apply_gated_attention=connector_apply_gated_attention,
ff_bias=connector_ff_bias,
)
audio_connector_config = GemmaConnectorConfig(
num_attention_heads=getattr(arch, "audio_connector_num_attention_heads", None)
or arch.connector_num_attention_heads,
attention_head_dim=getattr(arch, "audio_connector_attention_head_dim", None)
or arch.connector_attention_head_dim,
num_layers=getattr(arch, "audio_connector_num_layers", None) or arch.connector_num_layers,
positional_embedding_theta=arch.connector_positional_embedding_theta,
positional_embedding_max_pos=arch.connector_positional_embedding_max_pos,
rope_type=LTXRopeType(arch.connector_rope_type),
double_precision_rope=arch.connector_double_precision_rope,
num_learnable_registers=arch.connector_num_learnable_registers,
apply_gated_attention=connector_apply_gated_attention,
ff_bias=connector_ff_bias,
)
self.embeddings_connector = Embeddings1DConnector(video_connector_config)
self.audio_embeddings_connector = Embeddings1DConnector(audio_connector_config)
self.gemma_model_path = arch.gemma_model_path
self.gemma_dtype = arch.gemma_dtype
self.padding_side = arch.padding_side
self._gemma_model: PreTrainedModel | None = None
def named_parameters(self, prefix: str = "", recurse: bool = True):
for name, param in super().named_parameters(prefix=prefix, recurse=recurse):
if name.startswith("_gemma_model."):
continue
yield name, param
def prepare_for_compile(self) -> None:
# Load Gemma outside Dynamo so torch.compile does not trace HF file-system checks.
model = self.gemma_model
# Run one tiny eager forward before torch.compile. FastVideo calls the text encoder with
# Transformers' output_hidden_states path, which wraps each Gemma layer
# forward method and restores it by setting an instance-level forward
# attribute. If that attribute appears after Dynamo captures guards, the
# next text-encoder call recompiles; doing it here makes the first
# compiled call see the stable layer state.
token_id = getattr(model.config, "eos_token_id", None)
if isinstance(token_id, (list, tuple)):
token_id = token_id[0] if token_id else None
if token_id is None:
token_id = getattr(model.config, "pad_token_id", 0)
input_ids = torch.full(
(1, 1),
int(token_id or 0),
dtype=torch.long,
device=model.device,
)
attention_mask = torch.ones_like(input_ids)
with torch.no_grad():
model(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
return_dict=True,
)
@property
def gemma_model(self) -> PreTrainedModel:
if self._gemma_model is None:
gemma_path = self.gemma_model_path
if not gemma_path:
raise ValueError("gemma_model_path must be set (expected text_encoder/gemma).")
dtype = getattr(torch, self.gemma_dtype, torch.bfloat16)
self._gemma_model = AutoModelForImageTextToText.from_pretrained(
gemma_path,
local_files_only=True,
torch_dtype=dtype,
)
# Configure model-level attention implementation when using TORCH_SDPA.
# Note: torch.backends.cuda.enable_*_sdp() settings should be configured
# at application/pipeline initialization level, not here, to avoid
# unexpected side effects across the application.
if effective_attention_backend(self.config) == AttentionBackendEnum.TORCH_SDPA:
if hasattr(self._gemma_model.config, "attn_implementation"):
self._gemma_model.config.attn_implementation = "sdpa"
if hasattr(self._gemma_model.config, "_attn_implementation"):
self._gemma_model.config._attn_implementation = "sdpa"
if self.use_v2_feature_extractor:
device = next(self.video_feature_extractor_linear.parameters()).device
else:
device = next(self.feature_extractor_linear.parameters()).device
self._gemma_model.to(device=device)
self._gemma_model.eval()
return self._gemma_model
def _run_feature_extractor(
self,
hidden_states: tuple[torch.Tensor, ...],
attention_mask: torch.Tensor,
padding_side: str,
) -> tuple[torch.Tensor, torch.Tensor]:
encoded_text_features = torch.stack(hidden_states, dim=-1)
if envs.FASTVIDEO_LTX2_GEMMA_LOG.get():
for idx, layer in enumerate(hidden_states):
_debug_gemma_log_line(f"fastvideo:gemma_hidden_state_{idx}"
f":sum={layer.float().sum().item():.6f}")
_debug_gemma_log_line("fastvideo:gemma_hidden_states_stack"
f":sum={encoded_text_features.float().sum().item():.6f}")
encoded_text_features_dtype = encoded_text_features.dtype
if self.use_v2_feature_extractor:
# LTX-2.3: per-token RMS norm then separate video/audio linears.
normed_text_features = _norm_and_concat_per_token_rms(encoded_text_features,
attention_mask).to(encoded_text_features_dtype)
video_out_dim = self.video_feature_extractor_linear.out_features
audio_out_dim = self.audio_feature_extractor_linear.out_features
video_features = self.video_feature_extractor_linear(
_rescale_norm(
normed_text_features,
target_dim=video_out_dim,
source_dim=self.config.arch_config.hidden_size,
))
audio_features = self.audio_feature_extractor_linear(
_rescale_norm(
normed_text_features,
target_dim=audio_out_dim,
source_dim=self.config.arch_config.hidden_size,
))
return video_features, audio_features
# LTX-2.0: shared min-max norm + single aggregate linear. The two
# returned tensors are the same object (video == audio source).
sequence_lengths = attention_mask.sum(dim=-1)
normed_text_features = _norm_and_concat_padded_batch(encoded_text_features,
sequence_lengths,
padding_side=padding_side)
shared_features = self.feature_extractor_linear(normed_text_features.to(encoded_text_features_dtype))
return shared_features, shared_features
def _convert_to_additive_mask(self, attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
return (attention_mask - 1).to(dtype).reshape(
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])) * torch.finfo(dtype).max
def _run_connectors(
self,
video_features: torch.Tensor,
audio_features: torch.Tensor | None,
attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
connector_attention_mask = self._convert_to_additive_mask(attention_mask, video_features.dtype)
encoded, encoded_connector_attention_mask = self.embeddings_connector(video_features, connector_attention_mask)
attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)
attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
encoded = encoded * attention_mask
# For LTX-2.0 audio_features is the same tensor as video_features, so
# this reproduces the prior behavior (audio connector fed the shared
# features). For LTX-2.3 it is fed the separate audio features.
audio_features = (audio_features if audio_features is not None else video_features)
encoded_for_audio, _ = self.audio_embeddings_connector(audio_features, connector_attention_mask)
return encoded, encoded_for_audio, attention_mask.squeeze(-1)
@torch.no_grad()
def preprocess_text_embeddings(
self,
prompts: str | list[str],
tokenizer: AutoTokenizer,
tokenizer_kwargs: dict[str, Any] | None = None,
padding_side: str | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Compute pre-connector text embeddings for LTX-2 training preprocessing."""
if isinstance(prompts, str):
prompts = [prompts]
model = self.gemma_model
kwargs: dict[str, Any] = {
"padding": "max_length",
"truncation": True,
"return_tensors": "pt",
}
if tokenizer_kwargs is not None:
kwargs.update(tokenizer_kwargs)
if "max_length" not in kwargs:
kwargs["max_length"] = self.config.arch_config.text_len
original_padding_side = tokenizer.padding_side
target_padding_side = padding_side or self.padding_side
tokenizer.padding_side = target_padding_side
try:
text_inputs = tokenizer(prompts, **kwargs)
finally:
tokenizer.padding_side = original_padding_side
input_ids = text_inputs["input_ids"].to(device=model.device)
attention_mask = text_inputs["attention_mask"].to(device=model.device)
input_ids, attention_mask = _ensure_leading_bos(
input_ids,
attention_mask,
bos_token_id=_get_bos_token_id(model.config),
)
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
return_dict=True,
)
video_features, _ = self._run_feature_extractor(
outputs.hidden_states,
attention_mask,
padding_side=target_padding_side,
)
return video_features, attention_mask
def run_connectors(
self,
encoded_input: torch.Tensor,
attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Apply embedding connectors to precomputed Gemma features."""
return self._run_connectors(encoded_input, encoded_input, attention_mask)
def forward(
self,
input_ids: torch.Tensor | None,
position_ids: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
inputs_embeds: torch.Tensor | None = None,
output_hidden_states: bool | None = None,
**kwargs,
) -> BaseEncoderOutput:
if input_ids is None:
raise ValueError("input_ids is required for Gemma text encoding.")
if attention_mask is None:
attention_mask = torch.ones_like(input_ids)
model = self.gemma_model
input_ids, attention_mask = _ensure_leading_bos(
input_ids,
attention_mask,
bos_token_id=_get_bos_token_id(model.config),
)
target_device = get_local_torch_device()
# Do not invoke model.to() inside the compiled forward path.
# _parse_to returns a non-Tensor torch.device, which Dynamo cannot
# trace under fullgraph=True. The model is already moved to device
# when first loaded (see gemma_model property + prepare_for_compile),
# so this guard is a runtime no-op and Dynamo can DCE it.
if model.device != target_device:
model.to(device=target_device)
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
return_dict=True,
)
encoded_video, encoded_audio = self._run_feature_extractor(
outputs.hidden_states,
attention_mask,
padding_side=self.padding_side,
)
video_encoding, audio_encoding, attention_mask = self._run_connectors(encoded_video, encoded_audio,
attention_mask)
hidden_states = (audio_encoding, ) if output_hidden_states else None
return BaseEncoderOutput(
last_hidden_state=video_encoding,
hidden_states=hidden_states,
attention_mask=attention_mask,
)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
for name, loaded_weight in weights:
if name == "aggregate_embed.weight":
name = "feature_extractor_linear.aggregate_embed.weight"
# LTX-2.3 separate video/audio feature extractors. These names
# do not appear in LTX-2.0 checkpoints, so they are no-ops there.
elif name == "video_aggregate_embed.weight":
name = "video_feature_extractor_linear.weight"
elif name == "video_aggregate_embed.bias":
name = "video_feature_extractor_linear.bias"
elif name == "audio_aggregate_embed.weight":
name = "audio_feature_extractor_linear.weight"
elif name == "audio_aggregate_embed.bias":
name = "audio_feature_extractor_linear.bias"
elif name == "feature_extractor.video_aggregate_embed.weight":
name = "video_feature_extractor_linear.weight"
elif name == "feature_extractor.video_aggregate_embed.bias":
name = "video_feature_extractor_linear.bias"
elif name == "feature_extractor.audio_aggregate_embed.weight":
name = "audio_feature_extractor_linear.weight"
elif name == "feature_extractor.audio_aggregate_embed.bias":
name = "audio_feature_extractor_linear.bias"
elif name == "feature_extractor.aggregate_embed.weight":
name = "feature_extractor_linear.aggregate_embed.weight"
elif name.startswith("video_connector."):
name = name.replace("video_connector.", "embeddings_connector.", 1)
elif name.startswith("audio_connector."):
name = name.replace("audio_connector.", "audio_embeddings_connector.", 1)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
return loaded_params</code></pre></details>
fastvideo.models.encoders.gemma.LTX2GemmaTextEncoderModel.named_parameters
named_parameters(prefix: str = '', recurse: bool = True)
Source
def named_parameters(self, prefix: str = "", recurse: bool = True):
for name, param in super().named_parameters(prefix=prefix, recurse=recurse):
if name.startswith("_gemma_model."):
continue
yield name, param
fastvideo.models.encoders.gemma.LTX2GemmaTextEncoderModel.prepare_for_compile
prepare_for_compile() -> None
Source
def prepare_for_compile(self) -> None:
# Load Gemma outside Dynamo so torch.compile does not trace HF file-system checks.
model = self.gemma_model
# Run one tiny eager forward before torch.compile. FastVideo calls the text encoder with
# Transformers' output_hidden_states path, which wraps each Gemma layer
# forward method and restores it by setting an instance-level forward
# attribute. If that attribute appears after Dynamo captures guards, the
# next text-encoder call recompiles; doing it here makes the first
# compiled call see the stable layer state.
token_id = getattr(model.config, "eos_token_id", None)
if isinstance(token_id, (list, tuple)):
token_id = token_id[0] if token_id else None
if token_id is None:
token_id = getattr(model.config, "pad_token_id", 0)
input_ids = torch.full(
(1, 1),
int(token_id or 0),
dtype=torch.long,
device=model.device,
)
attention_mask = torch.ones_like(input_ids)
with torch.no_grad():
model(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
return_dict=True,
)
fastvideo.models.encoders.gemma.LTX2GemmaTextEncoderModel.preprocess_text_embeddings
preprocess_text_embeddings(prompts: str | list[str], tokenizer: AutoTokenizer, tokenizer_kwargs: dict[str, Any] | None = None, padding_side: str | None = None) -> tuple[torch.Tensor, torch.Tensor]
Compute pre-connector text embeddings for LTX-2 training preprocessing.
Source
@torch.no_grad()
def preprocess_text_embeddings(
self,
prompts: str | list[str],
tokenizer: AutoTokenizer,
tokenizer_kwargs: dict[str, Any] | None = None,
padding_side: str | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Compute pre-connector text embeddings for LTX-2 training preprocessing."""
if isinstance(prompts, str):
prompts = [prompts]
model = self.gemma_model
kwargs: dict[str, Any] = {
"padding": "max_length",
"truncation": True,
"return_tensors": "pt",
}
if tokenizer_kwargs is not None:
kwargs.update(tokenizer_kwargs)
if "max_length" not in kwargs:
kwargs["max_length"] = self.config.arch_config.text_len
original_padding_side = tokenizer.padding_side
target_padding_side = padding_side or self.padding_side
tokenizer.padding_side = target_padding_side
try:
text_inputs = tokenizer(prompts, **kwargs)
finally:
tokenizer.padding_side = original_padding_side
input_ids = text_inputs["input_ids"].to(device=model.device)
attention_mask = text_inputs["attention_mask"].to(device=model.device)
input_ids, attention_mask = _ensure_leading_bos(
input_ids,
attention_mask,
bos_token_id=_get_bos_token_id(model.config),
)
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
return_dict=True,
)
video_features, _ = self._run_feature_extractor(
outputs.hidden_states,
attention_mask,
padding_side=target_padding_side,
)
return video_features, attention_mask</code></pre></details>
fastvideo.models.encoders.gemma.LTX2GemmaTextEncoderModel.run_connectors
run_connectors(encoded_input: torch.Tensor, attention_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
Apply embedding connectors to precomputed Gemma features.
Source
def run_connectors(
self,
encoded_input: torch.Tensor,
attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Apply embedding connectors to precomputed Gemma features."""
return self._run_connectors(encoded_input, encoded_input, attention_mask)
fastvideo.models.encoders.gemma.LTX2GemmaTextEncoderModel.forward
forward(input_ids: torch.Tensor | None, position_ids: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None, output_hidden_states: bool | None = None, **kwargs) -> BaseEncoderOutput
Source
def forward(
self,
input_ids: torch.Tensor | None,
position_ids: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
inputs_embeds: torch.Tensor | None = None,
output_hidden_states: bool | None = None,
**kwargs,
) -> BaseEncoderOutput:
if input_ids is None:
raise ValueError("input_ids is required for Gemma text encoding.")
if attention_mask is None:
attention_mask = torch.ones_like(input_ids)
model = self.gemma_model
input_ids, attention_mask = _ensure_leading_bos(
input_ids,
attention_mask,
bos_token_id=_get_bos_token_id(model.config),
)
target_device = get_local_torch_device()
# Do not invoke model.to() inside the compiled forward path.
# _parse_to returns a non-Tensor torch.device, which Dynamo cannot
# trace under fullgraph=True. The model is already moved to device
# when first loaded (see gemma_model property + prepare_for_compile),
# so this guard is a runtime no-op and Dynamo can DCE it.
if model.device != target_device:
model.to(device=target_device)
outputs = model(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
return_dict=True,
)
encoded_video, encoded_audio = self._run_feature_extractor(
outputs.hidden_states,
attention_mask,
padding_side=self.padding_side,
)
video_encoding, audio_encoding, attention_mask = self._run_connectors(encoded_video, encoded_audio,
attention_mask)
hidden_states = (audio_encoding, ) if output_hidden_states else None
return BaseEncoderOutput(
last_hidden_state=video_encoding,
hidden_states=hidden_states,
attention_mask=attention_mask,
)</code></pre></details>
fastvideo.models.encoders.gemma.LTX2GemmaTextEncoderModel.load_weights
load_weights(weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]
Source
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
params_dict = dict(self.named_parameters())
loaded_params: set[str] = set()
for name, loaded_weight in weights:
if name == "aggregate_embed.weight":
name = "feature_extractor_linear.aggregate_embed.weight"
# LTX-2.3 separate video/audio feature extractors. These names
# do not appear in LTX-2.0 checkpoints, so they are no-ops there.
elif name == "video_aggregate_embed.weight":
name = "video_feature_extractor_linear.weight"
elif name == "video_aggregate_embed.bias":
name = "video_feature_extractor_linear.bias"
elif name == "audio_aggregate_embed.weight":
name = "audio_feature_extractor_linear.weight"
elif name == "audio_aggregate_embed.bias":
name = "audio_feature_extractor_linear.bias"
elif name == "feature_extractor.video_aggregate_embed.weight":
name = "video_feature_extractor_linear.weight"
elif name == "feature_extractor.video_aggregate_embed.bias":
name = "video_feature_extractor_linear.bias"
elif name == "feature_extractor.audio_aggregate_embed.weight":
name = "audio_feature_extractor_linear.weight"
elif name == "feature_extractor.audio_aggregate_embed.bias":
name = "audio_feature_extractor_linear.bias"
elif name == "feature_extractor.aggregate_embed.weight":
name = "feature_extractor_linear.aggregate_embed.weight"
elif name.startswith("video_connector."):
name = name.replace("video_connector.", "embeddings_connector.", 1)
elif name.startswith("audio_connector."):
name = name.replace("audio_connector.", "audio_embeddings_connector.", 1)
if name not in params_dict:
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
return loaded_params