From 3fcd8f90bec3162fe89fc35cf40cdb1814e6223d Mon Sep 17 00:00:00 2001
From: Prince Canuma <prince.gdt@gmail.com>
Date: Thu, 15 Jan 2026 03:47:52 +0100
Subject: [PATCH] Refactor LTXModel: Change transformer_blocks from list to
 dictionary

---
 mlx_video/models/ltx/ltx.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/mlx_video/models/ltx/ltx.py b/mlx_video/models/ltx/ltx.py
index 2b4bec2..dea4089 100644
--- a/mlx_video/models/ltx/ltx.py
+++ b/mlx_video/models/ltx/ltx.py
@@ -384,8 +384,9 @@ class LTXModel(nn.Module):
         video_config = config.get_video_config()
         audio_config = config.get_audio_config()
 
-        self.transformer_blocks = [
-            BasicAVTransformerBlock(
+
+        self.transformer_blocks = {
+            idx: BasicAVTransformerBlock(
                 idx=idx,
                 video=video_config,
                 audio=audio_config,
@@ -393,7 +394,7 @@ class LTXModel(nn.Module):
                 norm_eps=config.norm_eps,
             )
             for idx in range(config.num_layers)
-        ]
+        }
 
     def _process_transformer_blocks(
         self,
@@ -401,7 +402,7 @@ class LTXModel(nn.Module):
         audio: Optional[TransformerArgs],
     ) -> Tuple[Optional[TransformerArgs], Optional[TransformerArgs]]:
         """Process through all transformer blocks."""
-        for block in self.transformer_blocks:
+        for block in self.transformer_blocks.values():
             video, audio = block(video=video, audio=audio)
         return video, audio