Update models/gen_pipeline.py

models/gen_pipeline.py

@@ -387,7 +387,7 @@ class NextStepPipeline:
         )
 
         # 7. unpatchify
-        latents = self.model.unpatchify(tokens)
+        latents = self.model.unpatchify(tokens, h=hw[0] // self.down_factor, w=hw[1] // self.down_factor)
         latents = (latents / self.scaling_factor) + self.shift_factor
 
         # 8. decode latents