Remove hard-coded batch size configuration in LTXModel

[Linyou]

No description provided.

[a-r-r-o-w]

Shouldn't broadcasting automatically handle batch_size > 1 case?

[Linyou]

The broadcasting isn't the issue here.

latent_mean and latent_std have a shape of [batch_size] when batch_size > 1. When these tensors are reshaped via .view(1, -1, 1, 1, 1), the batch_size dimension is flattened into the second axis (resulting in shape [1, batch_size * (-1), 1, 1, 1]).

[a-r-r-o-w]

Oh, I see. Makes sense, thanks a lot! I'll try to handle this better in the future since we shouldn't need to have the same latents_mean and latents_std to be replicated multiple times for batch_size > 1.

[a-r-r-o-w]

@Linyou Are you able to pass the batch_size > 1 test for LTX Video? I merged the PR earlier since it seemed correct, but I was running tests today and this fails.

torchrun --nnodes=1 --nproc_per_node 2 -m pytest -s tests/trainer/test_sft_trainer.py -k "test___dp_degree_2___batch_size_2 and ___PTD and LoRA and LTXVideo"

I see that we ignore latents_mean and latents_std for collation (see here), so it should not be replicated when batch_size > 1 and work with broadcasting. Am I missing something important in understanding this?