[diffusion] chore: remove unreachable cosmos3 transfer encoding (#37805)

This commit is contained in:
Mick
2026-09-04 12:14:03 +08:00
committed by GitHub
parent ff1285cc28
commit 97d081ac76
@@ -808,36 +808,6 @@ class Cosmos3LatentPreparationStage(PipelineStage):
self.log_info(f"Prepared latents with shape {shape}")
# Transfer (control-video) conditioning: VAE-encode each control clip
# into clean latents the transformer prepends to the GEN sequence. Stored
# as a list (one block per hint) so multi-hint transfer (edge + depth …)
# threads through the denoiser uniformly with the single-hint case.
preprocessed_control = batch.extra.get("preprocessed_control")
if preprocessed_control is not None:
control_blocks = (
preprocessed_control
if isinstance(preprocessed_control, list)
else [preprocessed_control]
)
vae_dtype = next(self.vae.parameters()).dtype
control_latents_list: list[torch.Tensor] = []
for control_pixels_t in control_blocks:
control_pixels = control_pixels_t.to(device=device, dtype=vae_dtype)
with torch.no_grad():
control_latent = self._vae_encode(control_pixels).to(dtype)
if control_latent.shape[-2:] != latents.shape[-2:]:
raise ValueError(
"control latent spatial dims "
f"{tuple(control_latent.shape[-2:])} must match the target "
f"latents {tuple(latents.shape[-2:])}"
)
control_latents_list.append(control_latent)
batch.extra["control_latents"] = control_latents_list
self.log_info(
f"Prepared {len(control_latents_list)} control latent block(s) "
f"with shape {tuple(control_latents_list[0].shape)}"
)
sound_duration = float(getattr(batch, "sound_duration", 0.0) or 0.0)
if sound_duration > 0.0:
if not getattr(self.transformer, "sound_gen", False):