feat: add LTX2 smoke test and update setup.sh environment

[Perseus14]

Description

This PR adds an end-to-end smoke test for the LTX2 model in MaxDiffusion, modify the litx2 pipeline script and updates the environment setup script to handle virtual environments more robustly.

Key Changes

1. LTX2 Smoke Test

• Added src/maxdiffusion/tests/ltx2_smoke_test.py to provide a basic verification that LTX2 generation runs successfully on the configured hardware. This follows similar patterns used for the Wan model tests.

2. Environment Setup (setup.sh)

• Package Index URLs: Explicitly set PIP_INDEX_URL and UV_INDEX_URL to https://pypi.org/simple to ensure standard resolution from PyPI.
• Conditional UV_SYSTEM_PYTHON: Instead of forcing uv to use the system Python by default, added a check to detect if the script is running within a virtual environment.
  • If in a virtual environment, it respects that environment and avoids setting UV_SYSTEM_PYTHON=1.
  • If using system Python, it sets UV_SYSTEM_PYTHON=1 to maintain compatibility with typical bare-metal TPU VM workflows where we want to use the system package space.

3. Pipeline Updates

• Updated ltx2_pipeline.py to use get_value() and set_value() methods for logical state sharding, ensuring correct state handling.

Verification Results

• Verified the virtual environment detection logic in setup.sh works as expected when run inside and outside a venv.
• The LTX2 smoke test has been added and is ready to be run in CI or manually.

[github-actions]

e2e testgrid: https://8bcf50593faf4ea38060e236169827e5-dot-us-central1.composer.googleusercontent.com/dags/maxdiffusion_tpu_e2e/grid