1
0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00

Get LTX 2 transformer compile tests passing

This commit is contained in:
Daniel Gu
2025-12-15 03:38:34 +01:00
parent 780fb61d32
commit 5765759cd3

View File

@@ -246,7 +246,6 @@ class LTX2Attention(torch.nn.Module, AttentionModuleMixin):
return hidden_states
@maybe_allow_in_graph
class LTX2VideoTransformerBlock(nn.Module):
r"""
Transformer block used in [LTX-2.0](https://huggingface.co/Lightricks/LTX-Video).
@@ -802,7 +801,6 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
return cos_freqs, sin_freqs
@maybe_allow_in_graph
class LTX2VideoTransformer3DModel(
ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin, PeftAdapterMixin, CacheMixin
):
@@ -834,7 +832,7 @@ class LTX2VideoTransformer3DModel(
_supports_gradient_checkpointing = True
_skip_layerwise_casting_patterns = ["norm"]
_repeated_blocks = ["LTXVideoTransformerBlock"]
_repeated_blocks = ["LTX2VideoTransformerBlock"]
_cp_plan = {
"": {
"hidden_states": ContextParallelInput(split_dim=1, expected_dims=3, split_output=False),