[Bug] [Bagel] Fix kv transfer bug (#1437)

Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>
Co-authored-by: Wang Zhipeng: princepride <wangzhipeng628@gmail.com>

Author: nussejzz

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -1286,8 +1286,14 @@ def prepare_vae_latent_cfg(self, curr_kvlens, curr_rope, image_sizes):
     @staticmethod
     def _merge_naive_caches(caches: list) -> NaiveCache:
         """Merge multiple NaiveCache objects by concatenating KV tensors per layer."""
-        merged = NaiveCache(caches[0].num_layers)
-        for layer_idx in range(merged.num_layers):
+        if not caches:
+            # Handle empty list case gracefully if desired,
+            # though original code also crashed on this.
+            return NaiveCache(0)
+
+        num_layers = len(caches[0].key_cache)
+        merged = NaiveCache(num_layers)
+        for layer_idx in range(num_layers):
             merged.key_cache[layer_idx] = torch.cat([c.key_cache[layer_idx] for c in caches], dim=0)
             merged.value_cache[layer_idx] = torch.cat([c.value_cache[layer_idx] for c in caches], dim=0)
         return merged