Skip to content

feat(ltx2): make run_text_encoder_on_tpu default False and dynamically load torchax#415

Merged
copybara-service[bot] merged 1 commit into
mainfrom
ltx2_tpu_encoder_disable
Jun 2, 2026
Merged

feat(ltx2): make run_text_encoder_on_tpu default False and dynamically load torchax#415
copybara-service[bot] merged 1 commit into
mainfrom
ltx2_tpu_encoder_disable

Conversation

@Perseus14
Copy link
Copy Markdown
Collaborator

@Perseus14 Perseus14 commented May 30, 2026

This PR isolates the torchax dependency in the LTX2 and LTX2.3 pipelines, making TPU-based text encoding opt-in rather than opt-out.
Changes:

  • Lazy Imports: Moved torchax and TorchaxGemma3TextEncoder imports inside the if run_text_encoder_on_tpu: block in ltx2_pipeline.py. This removes the hard dependency on torchax at module import time, allowing users without it to run the pipeline seamlessly.
  • Config Defaults Updated: Changed the default value of run_text_encoder_on_tpu from True to False in both the Python getattr fallbacks and the ltx2_video.yml / ltx2_3_video.yml configs.
  • LTX-2.3 Config Parity: Explicitly added run_text_encoder_on_tpu: False and enable_dynamic_vae_sharding: True to ltx2_3_video.yml to match the base LTX2 config.

@Perseus14 Perseus14 requested a review from entrpn as a code owner May 30, 2026 11:29
@github-actions
Copy link
Copy Markdown

@Perseus14 Perseus14 requested review from mbohlool and prishajain1 May 30, 2026 11:29
@Perseus14 Perseus14 self-assigned this May 30, 2026
@copybara-service copybara-service Bot merged commit 4a3ec4f into main Jun 2, 2026
23 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