Skip to content

[Text Pipeline] Implement Text Encoders Wrappers with mesh support#339

Merged
copybara-service[bot] merged 1 commit intoAI-Hypercomputer:mainfrom
syhuang22:feat/ltx2-text-encoders-wrappers
Mar 9, 2026
Merged

[Text Pipeline] Implement Text Encoders Wrappers with mesh support#339
copybara-service[bot] merged 1 commit intoAI-Hypercomputer:mainfrom
syhuang22:feat/ltx2-text-encoders-wrappers

Conversation

@syhuang22
Copy link
Collaborator

This PR introduces the high-level nnx.Module wrappers (LTX2VideoGemmaTextEncoder and LTX2AudioVideoGemmaTextEncoder) for the LTX-2 text pipeline. These wrappers serve as the orchestrators that connect the output of the raw Gemma 3 text encoder (hidden states) to the GemmaFeaturesExtractorProjLinear and subsequently to the Embeddings1DConnector.

@syhuang22 syhuang22 requested a review from entrpn as a code owner February 26, 2026 00:31
# Feature Extractor Config
gemma_dim: int = 3840, # Gemma-3-12b
gemma_layers: int = 49, # Gemma-3 has 48 layers + 1 embedding layer output = 49 hidden states
projection_dim: int = 4096, # LTX-2 conditioning dim
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Change needed 4096 -> 3840

# Feature Extractor Config (Shared)
gemma_dim: int = 3840, # Gemma-3-12b
gemma_layers: int = 49, # Gemma-3 has 48 layers + 1 embedding layer output = 49 hidden states
projection_dim: int = 4096,
Copy link
Collaborator

@prishajain1 prishajain1 Feb 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

4096 -> 3840

gemma_layers: int = 49, # Gemma-3 has 48 layers + 1 embedding layer output = 49 hidden states
projection_dim: int = 4096,
# Connector Config
connector_heads: int = 32,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@syhuang22 syhuang22 force-pushed the feat/ltx2-text-encoders-wrappers branch 2 times, most recently from 4c2baff to 2d4337c Compare March 9, 2026 18:23
Signed-off-by: James Huang <syhuang1201@gmail.com>
@syhuang22 syhuang22 force-pushed the feat/ltx2-text-encoders-wrappers branch from 2d4337c to 02dbc99 Compare March 9, 2026 18:24
@copybara-service copybara-service bot merged commit 4afed9f into AI-Hypercomputer:main Mar 9, 2026
20 of 21 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants