def retrieve_latents(
    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
):
    """Retrieve latents from an encoder output.

    Args:
        encoder_output (torch.Tensor): Output from an encoder (e.g., VAE).
        generator (torch.Generator, optional): Random generator for sampling. Defaults to None.
        sample_mode (str): Sampling mode ("sample" or "argmax"). Defaults to "sample".

    Returns:
        torch.Tensor: Retrieved latent tensor.
    """
    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
        return encoder_output.latent_dist.sample(generator)
    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
        return encoder_output.latent_dist.mode()
    elif hasattr(encoder_output, "latents"):
        return encoder_output.latents
    else:
        raise AttributeError("Could not access latents of provided encoder_output")
Retrieve latents from an encoder output.

    Args:
        encoder_output (torch.Tensor): Output from an encoder (e.g., VAE).
        generator (torch.Generator, optional): Random generator for sampling. Defaults to None.
        sample_mode (str): Sampling mode ("sample" or "argmax"). Defaults to "sample".

    Returns:
        torch.Tensor: Retrieved latent tensor.
    
examples/community/pipeline_faithdiff_stable_diffusion_xl.py