|
|
from typing import Tuple |
|
|
import torch |
|
|
from diffusers import AutoencoderKL |
|
|
from einops import rearrange |
|
|
from torch import Tensor |
|
|
|
|
|
|
|
|
from ltx_video.models.autoencoders.causal_video_autoencoder import ( |
|
|
CausalVideoAutoencoder, |
|
|
) |
|
|
from ltx_video.models.autoencoders.video_autoencoder import ( |
|
|
Downsample3D, |
|
|
VideoAutoencoder, |
|
|
) |
|
|
|
|
|
try: |
|
|
import torch_xla.core.xla_model as xm |
|
|
except ImportError: |
|
|
xm = None |
|
|
|
|
|
|
|
|
def vae_encode( |
|
|
media_items: Tensor, |
|
|
vae: AutoencoderKL, |
|
|
split_size: int = 1, |
|
|
vae_per_channel_normalize=False, |
|
|
) -> Tensor: |
|
|
""" |
|
|
Encodes media items (images or videos) into latent representations using a specified VAE model. |
|
|
The function supports processing batches of images or video frames and can handle the processing |
|
|
in smaller sub-batches if needed. |
|
|
|
|
|
Args: |
|
|
media_items (Tensor): A torch Tensor containing the media items to encode. The expected |
|
|
shape is (batch_size, channels, height, width) for images or (batch_size, channels, |
|
|
frames, height, width) for videos. |
|
|
vae (AutoencoderKL): An instance of the `AutoencoderKL` class from the `diffusers` library, |
|
|
pre-configured and loaded with the appropriate model weights. |
|
|
split_size (int, optional): The number of sub-batches to split the input batch into for encoding. |
|
|
If set to more than 1, the input media items are processed in smaller batches according to |
|
|
this value. Defaults to 1, which processes all items in a single batch. |
|
|
|
|
|
Returns: |
|
|
Tensor: A torch Tensor of the encoded latent representations. The shape of the tensor is adjusted |
|
|
to match the input shape, scaled by the model's configuration. |
|
|
|
|
|
Examples: |
|
|
>>> import torch |
|
|
>>> from diffusers import AutoencoderKL |
|
|
>>> vae = AutoencoderKL.from_pretrained('your-model-name') |
|
|
>>> images = torch.rand(10, 3, 8 256, 256) # Example tensor with 10 videos of 8 frames. |
|
|
>>> latents = vae_encode(images, vae) |
|
|
>>> print(latents.shape) # Output shape will depend on the model's latent configuration. |
|
|
|
|
|
Note: |
|
|
In case of a video, the function encodes the media item frame-by frame. |
|
|
""" |
|
|
is_video_shaped = media_items.dim() == 5 |
|
|
batch_size, channels = media_items.shape[0:2] |
|
|
|
|
|
if channels != 3: |
|
|
raise ValueError(f"Expects tensors with 3 channels, got {channels}.") |
|
|
|
|
|
if is_video_shaped and not isinstance( |
|
|
vae, (VideoAutoencoder, CausalVideoAutoencoder) |
|
|
): |
|
|
media_items = rearrange(media_items, "b c n h w -> (b n) c h w") |
|
|
if split_size > 1: |
|
|
if len(media_items) % split_size != 0: |
|
|
raise ValueError( |
|
|
"Error: The batch size must be divisible by 'train.vae_bs_split" |
|
|
) |
|
|
encode_bs = len(media_items) // split_size |
|
|
|
|
|
latents = [] |
|
|
if media_items.device.type == "xla": |
|
|
xm.mark_step() |
|
|
for image_batch in media_items.split(encode_bs): |
|
|
latents.append(vae.encode(image_batch).latent_dist.sample()) |
|
|
if media_items.device.type == "xla": |
|
|
xm.mark_step() |
|
|
latents = torch.cat(latents, dim=0) |
|
|
else: |
|
|
latents = vae.encode(media_items).latent_dist.sample() |
|
|
|
|
|
latents = normalize_latents(latents, vae, vae_per_channel_normalize) |
|
|
if is_video_shaped and not isinstance( |
|
|
vae, (VideoAutoencoder, CausalVideoAutoencoder) |
|
|
): |
|
|
latents = rearrange(latents, "(b n) c h w -> b c n h w", b=batch_size) |
|
|
return latents |
|
|
|
|
|
|
|
|
def vae_decode( |
|
|
latents: Tensor, |
|
|
vae: AutoencoderKL, |
|
|
is_video: bool = True, |
|
|
split_size: int = 1, |
|
|
vae_per_channel_normalize=False, |
|
|
timestep=None, |
|
|
) -> Tensor: |
|
|
is_video_shaped = latents.dim() == 5 |
|
|
batch_size = latents.shape[0] |
|
|
|
|
|
if is_video_shaped and not isinstance( |
|
|
vae, (VideoAutoencoder, CausalVideoAutoencoder) |
|
|
): |
|
|
latents = rearrange(latents, "b c n h w -> (b n) c h w") |
|
|
if split_size > 1: |
|
|
if len(latents) % split_size != 0: |
|
|
raise ValueError( |
|
|
"Error: The batch size must be divisible by 'train.vae_bs_split" |
|
|
) |
|
|
encode_bs = len(latents) // split_size |
|
|
image_batch = [ |
|
|
_run_decoder( |
|
|
latent_batch, vae, is_video, vae_per_channel_normalize, timestep |
|
|
) |
|
|
for latent_batch in latents.split(encode_bs) |
|
|
] |
|
|
images = torch.cat(image_batch, dim=0) |
|
|
else: |
|
|
images = _run_decoder( |
|
|
latents, vae, is_video, vae_per_channel_normalize, timestep |
|
|
) |
|
|
|
|
|
if is_video_shaped and not isinstance( |
|
|
vae, (VideoAutoencoder, CausalVideoAutoencoder) |
|
|
): |
|
|
images = rearrange(images, "(b n) c h w -> b c n h w", b=batch_size) |
|
|
return images |
|
|
|
|
|
|
|
|
def _run_decoder( |
|
|
latents: Tensor, |
|
|
vae: AutoencoderKL, |
|
|
is_video: bool, |
|
|
vae_per_channel_normalize=False, |
|
|
timestep=None, |
|
|
) -> Tensor: |
|
|
if isinstance(vae, (VideoAutoencoder, CausalVideoAutoencoder)): |
|
|
*_, fl, hl, wl = latents.shape |
|
|
temporal_scale, spatial_scale, _ = get_vae_size_scale_factor(vae) |
|
|
latents = latents.to(vae.dtype) |
|
|
vae_decode_kwargs = {} |
|
|
if timestep is not None: |
|
|
vae_decode_kwargs["timestep"] = timestep |
|
|
image = vae.decode( |
|
|
un_normalize_latents(latents, vae, vae_per_channel_normalize), |
|
|
return_dict=False, |
|
|
target_shape=( |
|
|
1, |
|
|
3, |
|
|
fl * temporal_scale if is_video else 1, |
|
|
hl * spatial_scale, |
|
|
wl * spatial_scale, |
|
|
), |
|
|
**vae_decode_kwargs, |
|
|
)[0] |
|
|
else: |
|
|
image = vae.decode( |
|
|
un_normalize_latents(latents, vae, vae_per_channel_normalize), |
|
|
return_dict=False, |
|
|
)[0] |
|
|
return image |
|
|
|
|
|
|
|
|
def get_vae_size_scale_factor(vae: AutoencoderKL) -> float: |
|
|
if isinstance(vae, CausalVideoAutoencoder): |
|
|
spatial = vae.spatial_downscale_factor |
|
|
temporal = vae.temporal_downscale_factor |
|
|
else: |
|
|
down_blocks = len( |
|
|
[ |
|
|
block |
|
|
for block in vae.encoder.down_blocks |
|
|
if isinstance(block.downsample, Downsample3D) |
|
|
] |
|
|
) |
|
|
spatial = vae.config.patch_size * 2**down_blocks |
|
|
temporal = ( |
|
|
vae.config.patch_size_t * 2**down_blocks |
|
|
if isinstance(vae, VideoAutoencoder) |
|
|
else 1 |
|
|
) |
|
|
|
|
|
return (temporal, spatial, spatial) |
|
|
|
|
|
|
|
|
def latent_to_pixel_coords( |
|
|
latent_coords: Tensor, vae: AutoencoderKL, causal_fix: bool = False |
|
|
) -> Tensor: |
|
|
""" |
|
|
Converts latent coordinates to pixel coordinates by scaling them according to the VAE's |
|
|
configuration. |
|
|
|
|
|
Args: |
|
|
latent_coords (Tensor): A tensor of shape [batch_size, 3, num_latents] |
|
|
containing the latent corner coordinates of each token. |
|
|
vae (AutoencoderKL): The VAE model |
|
|
causal_fix (bool): Whether to take into account the different temporal scale |
|
|
of the first frame. Default = False for backwards compatibility. |
|
|
Returns: |
|
|
Tensor: A tensor of pixel coordinates corresponding to the input latent coordinates. |
|
|
""" |
|
|
|
|
|
scale_factors = get_vae_size_scale_factor(vae) |
|
|
causal_fix = isinstance(vae, CausalVideoAutoencoder) and causal_fix |
|
|
pixel_coords = latent_to_pixel_coords_from_factors( |
|
|
latent_coords, scale_factors, causal_fix |
|
|
) |
|
|
return pixel_coords |
|
|
|
|
|
|
|
|
def latent_to_pixel_coords_from_factors( |
|
|
latent_coords: Tensor, scale_factors: Tuple, causal_fix: bool = False |
|
|
) -> Tensor: |
|
|
pixel_coords = ( |
|
|
latent_coords |
|
|
* torch.tensor(scale_factors, device=latent_coords.device)[None, :, None] |
|
|
) |
|
|
if causal_fix: |
|
|
|
|
|
pixel_coords[:, 0] = (pixel_coords[:, 0] + 1 - scale_factors[0]).clamp(min=0) |
|
|
return pixel_coords |
|
|
|
|
|
|
|
|
def normalize_latents( |
|
|
latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False |
|
|
) -> Tensor: |
|
|
return ( |
|
|
(latents - vae.mean_of_means.to(latents.dtype).view(1, -1, 1, 1, 1)) |
|
|
/ vae.std_of_means.to(latents.dtype).view(1, -1, 1, 1, 1) |
|
|
if vae_per_channel_normalize |
|
|
else latents * vae.config.scaling_factor |
|
|
) |
|
|
|
|
|
|
|
|
def un_normalize_latents( |
|
|
latents: Tensor, vae: AutoencoderKL, vae_per_channel_normalize: bool = False |
|
|
) -> Tensor: |
|
|
return ( |
|
|
latents * vae.std_of_means.to(latents.dtype).view(1, -1, 1, 1, 1) |
|
|
+ vae.mean_of_means.to(latents.dtype).view(1, -1, 1, 1, 1) |
|
|
if vae_per_channel_normalize |
|
|
else latents / vae.config.scaling_factor |
|
|
) |
|
|
|
