`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) -> TextEncoderOutputT

`fastvideo.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.Tensor

`fastvideo.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.Tensor

Apply 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.Tensor

Generate 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.Tensor

Generate 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) -> None

Default weight loader.

`fastvideo.models.encoders.gemma.effective_attention_backend`

effective_attention_backend(config: object) -> AttentionBackendEnum | None

Return 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 config with :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_scope is 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_BACKEND environment 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.device

Return 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 = True

`fastvideo.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)

`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

`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

`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

`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

`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,
    )

`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