From b88fe1a4b2e0ddf589b8f43aa54b767d340f0543 Mon Sep 17 00:00:00 2001
From: Arison591 <2234819536@qq.com>
Date: Fri, 17 Jul 2026 16:50:15 +0800
Subject: [PATCH] fix Matrix-Game-2 VAE decoding

---
 Matrix-Game-2/utils/wan_wrapper.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Matrix-Game-2/utils/wan_wrapper.py b/Matrix-Game-2/utils/wan_wrapper.py
index 537a1f4..ccd3976 100644
--- a/Matrix-Game-2/utils/wan_wrapper.py
+++ b/Matrix-Game-2/utils/wan_wrapper.py
@@ -58,7 +58,7 @@ def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torc
             decode_function = self.model.decode
 
         output = []
-        for u in zs:
+        for u in latent:
             output.append(decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0))
         output = torch.stack(output, dim=0)
         return output
@@ -206,4 +206,3 @@ def post_init(self):
         We can gradually add more methods here if needed.
         """
         self.get_scheduler()
-