mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-15 00:04:26 +00:00
CrossProduct
81 lines
3.2 KiB
Python
81 lines
3.2 KiB
Python
from typing import Optional
|
|
import torch
|
|
from diffusers import AutoencoderKLLTXVideo
|
|
|
|
def _normalize_latents(
|
|
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
|
|
) -> torch.Tensor:
|
|
# Normalize latents across the channel dimension [B, C, F, H, W]
|
|
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
|
|
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
|
|
latents = (latents - latents_mean) * scaling_factor / latents_std
|
|
return latents
|
|
|
|
def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor:
|
|
# Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p].
|
|
# The patch dimensions are then permuted and collapsed into the channel dimension of shape:
|
|
# [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor).
|
|
# dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features
|
|
batch_size, num_channels, num_frames, height, width = latents.shape
|
|
post_patch_num_frames = num_frames // patch_size_t
|
|
post_patch_height = height // patch_size
|
|
post_patch_width = width // patch_size
|
|
|
|
# to check math lol
|
|
# dim1 = num_frames // patch_size_t * height // patch_size * width // patch_size
|
|
# dim2 = num_channels * patch_size_t * patch_size * patch_size
|
|
|
|
latents = latents.reshape(
|
|
batch_size,
|
|
-1,
|
|
post_patch_num_frames,
|
|
patch_size_t,
|
|
post_patch_height,
|
|
patch_size,
|
|
post_patch_width,
|
|
patch_size,
|
|
)
|
|
latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
|
|
return latents
|
|
|
|
def post_conditioned_latent_patchify(
|
|
latents: torch.Tensor,
|
|
num_frames: int,
|
|
height: int,
|
|
width: int,
|
|
patch_size: int = 1,
|
|
patch_size_t: int = 1,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
latents = _pack_latents(latents, patch_size, patch_size_t)
|
|
return {"latents": latents, "num_frames": num_frames, "height": height, "width": width}
|
|
|
|
def prepare_latents_for_conditioning(
|
|
vae: AutoencoderKLLTXVideo,
|
|
image_or_video: torch.Tensor,
|
|
patch_size: int = 1,
|
|
patch_size_t: int = 1,
|
|
device: Optional[torch.device] = None,
|
|
dtype: Optional[torch.dtype] = None,
|
|
generator: Optional[torch.Generator] = None,
|
|
|
|
) -> torch.Tensor:
|
|
device = device or vae.device
|
|
|
|
if image_or_video.ndim == 4:
|
|
image_or_video = image_or_video.unsqueeze(2)
|
|
assert image_or_video.ndim == 5, f"Expected 5D tensor, got {image_or_video.ndim}D tensor"
|
|
|
|
image_or_video = image_or_video.to(device=device, dtype=vae.dtype)
|
|
image_or_video = image_or_video.permute(0, 2, 1, 3, 4).contiguous() # [B, C, F, H, W] -> [B, F, C, H, W]
|
|
|
|
latents = vae.encode(image_or_video).latent_dist.sample(generator=generator)
|
|
|
|
latents = latents.to(dtype=dtype)
|
|
_, _, num_frames, height, width = latents.shape
|
|
latents = _normalize_latents(latents, vae.latents_mean, vae.latents_std)
|
|
# latents = _pack_latents(latents, patch_size, patch_size_t)
|
|
return {"latents": latents,
|
|
"num_frames": num_frames,
|
|
"height": height,
|
|
"width": width } |