Fix crash in Cosmos2 latent generation

Your Name · Your Name · commit 889e54b58e6c · 2025-08-07T13:49:38.000Z
diff --git a/helpers/caching/vae.py b/helpers/caching/vae.py
@@ -555,7 +555,7 @@ def process_video_latents(self, latents_uncached):
             logger.debug(f"Video latent processing results: {output_cache_entry}")
             # we'll now overwrite the latents after logging.
             output_cache_entry["latents"] = latents_uncached
-        elif StateTracker.get_model_family() in ["wan", "cosmos2image"]:
+        elif StateTracker.get_model_family() in ["wan"]:
             logger.debug(
                 f"Shape for Wan VAE encode: {latents_uncached.shape} with latents_mean: {self.vae.latents_mean} and latents_std: {self.vae.latents_std}"
             )

Original file line number	Diff line number	Diff line change
`@@ -555,7 +555,7 @@ def process_video_latents(self, latents_uncached):`
`555`	`555`	`logger.debug(f"Video latent processing results: {output_cache_entry}")`
`556`	`556`	`# we'll now overwrite the latents after logging.`
`557`	`557`	`output_cache_entry["latents"] = latents_uncached`
`558`		`- elif StateTracker.get_model_family() in ["wan", "cosmos2image"]:`
	`558`	`+ elif StateTracker.get_model_family() in ["wan"]:`
`559`	`559`	`logger.debug(`
`560`	`560`	`f"Shape for Wan VAE encode: {latents_uncached.shape} with latents_mean: {self.vae.latents_mean} and latents_std: {self.vae.latents_std}"`
`561`	`561`	`)`