Spaces:
Running on Zero
Running on Zero
| from pathlib import Path | |
| import torch | |
| from diffusers.models import AutoencoderKL | |
| from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D | |
| from .flux_vae import FluxAutoencoderKL | |
| from ..constants import VAE_PATH | |
| from ..utils.torch_utils import PRECISION_TO_TYPE | |
| def transform_pytorch_ckpt_to_safetensors(vae_path): | |
| from safetensors.torch import save_file | |
| vae = AutoencoderKL.from_config(AutoencoderKL.load_config(vae_path)) | |
| ckpt = torch.load(Path(vae_path) / "pytorch_model.pt", map_location=vae.device) | |
| if "state_dict" in ckpt: | |
| ckpt = ckpt["state_dict"] | |
| vae.load_state_dict(ckpt) | |
| save_file(vae.state_dict(), Path(vae_path) / "diffusion_pytorch_model.safetensors") | |
| return vae | |
| def load_vae(vae_type, | |
| vae_precision=None, | |
| sample_size=None, | |
| vae_path=None, | |
| logger=None, | |
| device=None | |
| ): | |
| if vae_path is None: | |
| vae_path = VAE_PATH[vae_type] | |
| vae_compress_spec, vae_latent_channel, vae_name = vae_type.split("-") | |
| length = len(vae_compress_spec) | |
| if length == 2: | |
| if logger is not None: | |
| logger.info(f"Loading 2D VAE model ({vae_type}) from: {vae_path}") | |
| if vae_name == "flux": | |
| vae = FluxAutoencoderKL.from_pretrained(vae_path) | |
| else: | |
| vae = AutoencoderKL.from_pretrained(vae_path) | |
| # vae = transform_pytorch_ckpt_to_safetensors(vae_path) | |
| spatial_compression_ratio = 8 | |
| time_compression_ratio = 1 | |
| elif length == 3: | |
| if logger is not None: | |
| logger.info(f"Loading 3D VAE model ({vae_type}) from: {vae_path}") | |
| config = AutoencoderKLCausal3D.load_config(vae_path) | |
| if sample_size: | |
| vae = AutoencoderKLCausal3D.from_config(config, sample_size=sample_size) | |
| else: | |
| vae = AutoencoderKLCausal3D.from_config(config) | |
| ckpt = torch.load(Path(vae_path) / "pytorch_model.pt", map_location=vae.device) | |
| if "state_dict" in ckpt: | |
| ckpt = ckpt["state_dict"] | |
| # Internal checkpoints stored the VAE as a "vae." submodule; the public | |
| # HunyuanVideo VAE ships the keys unprefixed. Support both. | |
| if any(k.startswith("vae.") for k in ckpt): | |
| vae_ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")} | |
| else: | |
| vae_ckpt = ckpt | |
| vae.load_state_dict(vae_ckpt) | |
| spatial_compression_ratio = vae.config.spatial_compression_ratio | |
| time_compression_ratio = vae.config.time_compression_ratio | |
| else: | |
| raise ValueError(f"Invalid VAE model: {vae_type}. Must be either 2D VAE in the format of '??-*' or " | |
| f"3D VAE in the format of '???-*'.") | |
| if vae_precision is not None: | |
| vae = vae.to(dtype=PRECISION_TO_TYPE[vae_precision]) | |
| vae.requires_grad_(False) | |
| if logger is not None: | |
| logger.info(f"VAE to dtype: {vae.dtype}") | |
| if device is not None: | |
| vae = vae.to(device) | |
| # Set vae to eval mode, even though it's dropout rate is 0. | |
| vae.eval() | |
| return vae, vae_path, spatial_compression_ratio, time_compression_ratio | |