hymm_sp/vae/__init__.py

import torch
from pathlib import Path
from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D
from ..constants import VAE_PATH, PRECISION_TO_TYPE

def load_vae(vae_type,
             vae_precision=None,
             sample_size=None,
             vae_path=None,
             logger=None,
             device=None
             ):
    """
    Load and configure a Variational Autoencoder (VAE) model.
    
    This function handles loading 3D causal VAE models, including configuration,
    weight loading, precision setting, and device placement. It ensures the model
    is properly initialized for inference.

    Parameters:
        vae_type (str): Type identifier for the VAE, must follow '???-*' format for 3D VAEs
        vae_precision (str, optional): Desired precision type (e.g., 'fp16', 'fp32'). 
                                     Uses model's default if not specified.
        sample_size (tuple, optional): Input sample dimensions to override config defaults
        vae_path (str, optional): Path to VAE model files. Uses predefined path from
                                VAE_PATH constant if not specified.
        logger (logging.Logger, optional): Logger instance for progress/debug messages
        device (torch.device, optional): Target device to place the model (e.g., 'cuda' or 'cpu')

    Returns:
        tuple: Contains:
            - vae (AutoencoderKLCausal3D): Loaded and configured VAE model
            - vae_path (str): Actual path used to load the VAE
            - spatial_compression_ratio (int): Spatial dimension compression factor
            - time_compression_ratio (int): Temporal dimension compression factor

    Raises:
        ValueError: If vae_type does not follow the required 3D VAE format '???-*'
    """
    if vae_path is None:
        vae_path = VAE_PATH[vae_type]
    vae_compress_spec, _, _ = vae_type.split("-")
    length = len(vae_compress_spec)
    # Process 3D VAE (valid format with 3-character compression spec)
    if 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"]
        vae_ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")}
        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 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)

    # Ensure model is in evaluation mode (disables dropout/batch norm training behavior)
    # Note: Even with dropout rate 0, eval mode is recommended for consistent inference
    vae.eval()

    return vae, vae_path, spatial_compression_ratio, time_compression_ratio