[taking-lying-flat]

LTX first-frame conditioning replaces the corresponding noisy video latents with clean inputs and sets their token timesteps to zero, but FlowMatchSFTAudioVideoLoss still regresses those outputs against noise - clean. Changing only the discarded noise therefore changes the loss and gradients even when all model inputs are identical. In a two-token example, the current loss changes from 2.5 to 6.5 and the fixed-position output gradient from -1 to -3.

When video latents are replaced through input_latents_video, exclude positions where denoise_mask_video == 0 from the video loss. Average over channels and normalize over each sample's supervised positions before averaging the batch, matching the official LTX masked-loss convention. Nonzero mask values retain full supervision, including soft conditioning; they are not treated as fractional loss weights. A sample with no supervised video positions contributes zero video loss. The unconditioned path and audio loss keep their existing behavior.

The example now returns loss 4.0 for both noise samples and zero direct output gradient at the fixed position. This concerns replaced target positions; appended reference tokens used by end_image are already removed from the returned prediction.

Validation: 64 local CPU cases passed using the actual loss, conditioning unit, model wrapper, scheduler methods and patchifiers, with pre-encoded VAE fixtures and recorder/trainable-linear DiT substitutes:

• First-frame, end-frame, combined, unconditional and soft-conditioning paths with paired noise samples.
• FP32, BF16 and FP16 loss and parameter-gradient comparisons against the official LTX compute_loss, with and without audio.
• Unequal supervised counts across samples, all-fixed inputs, an empty-supervision sample within a batch, and masks without active latent replacement.
• Maximum loss difference versus the reference: 9.54e-7. Maximum parameter-gradient difference: 1.20e-7 in FP32 and zero in BF16/FP16. Audio output gradients are unchanged.
• Syntax compilation and git diff --check pass.

Only diffsynth/diffusion/loss.py is changed; the local validation scripts are not included. No full-model or CUDA training was run.

Official reference: https://github.com/Lightricks/LTX-2/blob/be6d6d0a69d11a0a399c462ed92606692e5f8697/packages/ltx-trainer/src/ltx_trainer/training_strategies/text_to_video.py#L256-L278