101 lines
4.4 KiB
Python
101 lines
4.4 KiB
Python
"""
|
|
Reversible bridge between ComfyUI's raw Wan VAE latent space and the
|
|
normalized + temporally-scaled latent space that the STARFlow-V transformer
|
|
was trained on.
|
|
|
|
Verified facts this bridge depends on (checked directly against the
|
|
installed ComfyUI source and against Apple ml-starflow's
|
|
misc/wan_vae2.py::video_vae2 at commit 20984834c1a68ee679faa8ed9cda2842f02b042d):
|
|
|
|
* comfy's core VAE.encode()/VAE.decode() (comfy.sd.VAE / comfy.ldm.wan.vae)
|
|
return/consume the *raw* Wan VAE latent. They do NOT apply
|
|
comfy.latent_formats.Wan22's mean/std normalization themselves -- that
|
|
normalization is only ever applied by consumers of the latent (e.g. the
|
|
Wan diffusion samplers), via LatentFormat.process_in/process_out.
|
|
|
|
* comfy.latent_formats.Wan22.latents_mean / latents_std are byte-for-byte
|
|
identical to the `self.mean` / `self.std` tensors baked into upstream's
|
|
misc/wan_vae2.py `video_vae2` (WanVAE_) class -- both are the official
|
|
Wan2.2-TI2V-5B 48-channel latent statistics.
|
|
|
|
* Upstream's VAE wrapper (utils/model_setup.py::VAE, Wan-AI/Wan2.2 branch)
|
|
calls `self.vae.sample(x, self.vae.scale)`, where `scale = [mean, 1/std]`
|
|
is applied *inside* `encode()`/`decode()` (i.e. before the function
|
|
returns), and only *after that* does the outer wrapper apply the
|
|
`config.vae_temporal_scale` (0.6) multiply/divide to frames after the
|
|
first:
|
|
|
|
encode: mu = (raw - mean) / std # inside video_vae2.encode
|
|
z[:, 1:] = z[:, 1:] * 0.6 # outer wrapper, applied after
|
|
decode: z[:, 1:] = z[:, 1:] / 0.6 # outer wrapper, applied before
|
|
raw = z * std + mean # inside video_vae2.decode
|
|
|
|
So the correct order is NORMALIZE THEN SCALE on encode, and UNSCALE THEN
|
|
DENORMALIZE on decode. (This is the opposite order than "scale then
|
|
normalize" -- scaling before normalizing would not reproduce the
|
|
checkpoint's trained latent distribution, since the affine mean-subtract
|
|
does not commute with a bare multiplicative scale.)
|
|
|
|
Layout: ComfyUI's standard video LATENT tensor is (B, C, T, H, W), which
|
|
matches comfy.latent_formats.Wan22.process_in/out's expected shape directly
|
|
(no permute needed for the normalize/scale step itself). The STARFlow-V
|
|
transformer works in (B, T, H, W, C) layout; nodes are responsible for the
|
|
rearrange at the point they hand tensors to/from the model.
|
|
"""
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
from ..config import STARFlowVConfig
|
|
|
|
_wan22_format = None
|
|
|
|
|
|
def _get_wan22_format():
|
|
global _wan22_format
|
|
if _wan22_format is None:
|
|
from comfy.latent_formats import Wan22
|
|
_wan22_format = Wan22()
|
|
return _wan22_format
|
|
|
|
|
|
def validate_video_latent_shape(latent: torch.Tensor, cfg: Optional[STARFlowVConfig] = None) -> None:
|
|
cfg = cfg or STARFlowVConfig()
|
|
if latent.dim() != 5:
|
|
raise ValueError(f"expected a 5D (B,C,T,H,W) video latent, got shape {tuple(latent.shape)}")
|
|
b, c, t, h, w = latent.shape
|
|
if c != cfg.channel_size:
|
|
raise ValueError(f"expected {cfg.channel_size} latent channels (Wan2.2 TI2V), got {c}")
|
|
spatial = cfg.vae_downsample_factor // cfg.patch_size
|
|
if h % max(spatial // cfg.vae_downsample_factor, 1) != 0:
|
|
pass # spatial patch-size divisibility is checked at the pixel level upstream of the VAE
|
|
if (t - 1) % 1 != 0 or t < 1:
|
|
raise ValueError(f"invalid temporal latent length {t}")
|
|
|
|
|
|
def comfy_latent_to_starflow(latent: torch.Tensor, cfg: Optional[STARFlowVConfig] = None) -> torch.Tensor:
|
|
"""Raw comfy Wan2.2 VAE latent (B,C,T,H,W) -> normalized + temporally
|
|
scaled STARFlow-V internal latent (B,C,T,H,W)."""
|
|
cfg = cfg or STARFlowVConfig()
|
|
validate_video_latent_shape(latent, cfg)
|
|
wan22 = _get_wan22_format()
|
|
z = wan22.process_in(latent)
|
|
if cfg.vae_temporal_scale != 1.0 and z.size(2) > 1:
|
|
z = z.clone()
|
|
z[:, :, 1:] = z[:, :, 1:] * cfg.vae_temporal_scale
|
|
return z
|
|
|
|
|
|
def starflow_to_comfy_latent(latent: torch.Tensor, cfg: Optional[STARFlowVConfig] = None) -> torch.Tensor:
|
|
"""STARFlow-V internal latent (B,C,T,H,W) -> raw comfy Wan2.2 VAE latent
|
|
(B,C,T,H,W), the exact inverse of comfy_latent_to_starflow."""
|
|
cfg = cfg or STARFlowVConfig()
|
|
validate_video_latent_shape(latent, cfg)
|
|
wan22 = _get_wan22_format()
|
|
z = latent
|
|
if cfg.vae_temporal_scale != 1.0 and z.size(2) > 1:
|
|
z = z.clone()
|
|
z[:, :, 1:] = z[:, :, 1:] / cfg.vae_temporal_scale
|
|
raw = wan22.process_out(z)
|
|
return raw
|