refactor(train): unify data classes and add LTX2 support

[gushiqiao]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces significant structural changes to the training and inference pipeline, including the addition of new model components, improved data loading mechanisms, and memory-efficient decoding strategies. The review identified several critical and high-severity issues, including potential crashes during compilation due to type mismatches in attention masks, runtime errors when handling missing video paths, platform-specific overflow errors in CSV processing, and security risks related to insecure deserialization of .pt files. I have filtered out comments that provided validation without actionable feedback or focused on non-code changes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If args.self_attention_mask is a BlockMask (which is returned by the new FlexAttention mask preparation functions), accessing .shape or calling torch._dynamo.mark_dynamic on it will raise an AttributeError because BlockMask is not a Tensor. Adding an isinstance(..., torch.Tensor) check prevents compilation crashes.

        if args.self_attention_mask is not None and isinstance(args.self_attention_mask, torch.Tensor):
            # Dense form is (B, 1, T, T); key-padding form (from the SP wrapper)
            # is (B, 1, 1, T) -- leave the size-1 query dim static so Dynamo
            # keeps the broadcast.
            if args.self_attention_mask.shape[2] > 1:
                torch._dynamo.mark_dynamic(args.self_attention_mask, 2)
            torch._dynamo.mark_dynamic(args.self_attention_mask, 3)