263 lines
12 KiB
Python
263 lines
12 KiB
Python
import os
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from diffusers import AutoencoderKLLTX2Video
|
|
from diffusers.image_processor import PipelineImageInput
|
|
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
|
|
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
from diffusers.video_processor import VideoProcessor
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
|
from vllm_omni.diffusion.distributed.utils import get_local_device
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class LTX2LatentUpsamplePipeline(nn.Module):
|
|
def __init__(
|
|
self,
|
|
od_config: OmniDiffusionConfig,
|
|
vae: AutoencoderKLLTX2Video,
|
|
latent_upsampler: LTX2LatentUpsamplerModel = None,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
if vae is None:
|
|
raise ValueError("vae must be provided")
|
|
self.vae = vae
|
|
|
|
self.device = get_local_device()
|
|
model = od_config.model
|
|
local_files_only = os.path.exists(model)
|
|
|
|
if latent_upsampler is None:
|
|
# Use cpu context to create latent upsampler. The code k[:, None] @ k[None, :] in
|
|
# diffuser's BlurDownsample is not supported on GPU as k is type of torch.Int64
|
|
with torch.device("cpu"):
|
|
latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
|
|
model,
|
|
subfolder="latent_upsampler",
|
|
torch_dtype=torch.bfloat16,
|
|
local_files_only=local_files_only,
|
|
).to(self.device)
|
|
self.latent_upsampler = latent_upsampler
|
|
|
|
self.vae_spatial_compression_ratio = (
|
|
self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32
|
|
)
|
|
self.vae_temporal_compression_ratio = (
|
|
self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8
|
|
)
|
|
self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio)
|
|
|
|
def prepare_latents(
|
|
self,
|
|
video: torch.Tensor | None = None,
|
|
batch_size: int = 1,
|
|
num_frames: int = 121,
|
|
height: int = 512,
|
|
width: int = 768,
|
|
spatial_patch_size: int = 1,
|
|
temporal_patch_size: int = 1,
|
|
dtype: torch.dtype | None = None,
|
|
device: torch.device | None = None,
|
|
generator: torch.Generator | None = None,
|
|
latents: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
if latents is not None:
|
|
if latents.ndim == 3:
|
|
# Convert token seq [B, S, D] to latent video [B, C, F, H, W]
|
|
latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
|
|
latent_height = height // self.vae_spatial_compression_ratio
|
|
latent_width = width // self.vae_spatial_compression_ratio
|
|
latents = self._unpack_latents(
|
|
latents, latent_num_frames, latent_height, latent_width, spatial_patch_size, temporal_patch_size
|
|
)
|
|
return latents.to(device=device, dtype=dtype)
|
|
|
|
video = video.to(device=device, dtype=self.vae.dtype)
|
|
if isinstance(generator, list):
|
|
if len(generator) != batch_size:
|
|
raise ValueError(
|
|
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
|
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
|
)
|
|
|
|
init_latents = [
|
|
retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator[i]) for i in range(batch_size)
|
|
]
|
|
else:
|
|
init_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video]
|
|
|
|
init_latents = torch.cat(init_latents, dim=0).to(dtype)
|
|
# NOTE: latent upsampler operates on the unnormalized latents, so don't normalize here
|
|
# init_latents = self._normalize_latents(init_latents, self.vae.latents_mean, self.vae.latents_std)
|
|
return init_latents
|
|
|
|
def adain_filter_latent(self, latents: torch.Tensor, reference_latents: torch.Tensor, factor: float = 1.0):
|
|
result = latents.clone()
|
|
|
|
for i in range(latents.size(0)):
|
|
for c in range(latents.size(1)):
|
|
r_sd, r_mean = torch.std_mean(reference_latents[i, c], dim=None) # index by original dim order
|
|
i_sd, i_mean = torch.std_mean(result[i, c], dim=None)
|
|
|
|
result[i, c] = ((result[i, c] - i_mean) / i_sd) * r_sd + r_mean
|
|
|
|
result = torch.lerp(latents, result, factor)
|
|
return result
|
|
|
|
def tone_map_latents(self, latents: torch.Tensor, compression: float) -> torch.Tensor:
|
|
# Remap [0-1] to [0-0.75] and apply sigmoid compression in one shot
|
|
scale_factor = compression * 0.75
|
|
abs_latents = torch.abs(latents)
|
|
|
|
# Sigmoid compression: sigmoid shifts large values toward 0.2, small values stay ~1.0
|
|
# When scale_factor=0, sigmoid term vanishes, when scale_factor=0.75, full effect
|
|
sigmoid_term = torch.sigmoid(4.0 * scale_factor * (abs_latents - 1.0))
|
|
scales = 1.0 - 0.8 * scale_factor * sigmoid_term
|
|
|
|
filtered = latents * scales
|
|
return filtered
|
|
|
|
@staticmethod
|
|
# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_latents
|
|
def _denormalize_latents(
|
|
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
|
|
) -> torch.Tensor:
|
|
# Denormalize 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_std / scaling_factor + latents_mean
|
|
return latents
|
|
|
|
@staticmethod
|
|
# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_latents
|
|
def _unpack_latents(
|
|
latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1
|
|
) -> torch.Tensor:
|
|
# Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) # noqa
|
|
# are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of
|
|
# what happens in the `_pack_latents` method.
|
|
batch_size = latents.size(0)
|
|
latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size)
|
|
latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
|
|
return latents
|
|
|
|
def check_inputs(self, video, height, width, latents, tone_map_compression_ratio):
|
|
if height % self.vae_spatial_compression_ratio != 0 or width % self.vae_spatial_compression_ratio != 0:
|
|
raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
|
|
|
|
if video is not None and latents is not None:
|
|
raise ValueError("Only one of `video` or `latents` can be provided.")
|
|
if video is None and latents is None:
|
|
raise ValueError("One of `video` or `latents` has to be provided.")
|
|
|
|
if not (0 <= tone_map_compression_ratio <= 1):
|
|
raise ValueError("`tone_map_compression_ratio` must be in the range [0, 1]")
|
|
|
|
def forward(
|
|
self,
|
|
video: list[PipelineImageInput] | None = None,
|
|
height: int = 512,
|
|
width: int = 768,
|
|
num_frames: int = 121,
|
|
spatial_patch_size: int = 1,
|
|
temporal_patch_size: int = 1,
|
|
latents: torch.Tensor | None = None,
|
|
latents_normalized: bool = False,
|
|
decode_timestep: float | list[float] = 0.0,
|
|
decode_noise_scale: float | list[float] | None = None,
|
|
adain_factor: float = 0.0,
|
|
tone_map_compression_ratio: float = 0.0,
|
|
generator: torch.Generator | list[torch.Generator] | None = None,
|
|
output_type: str | None = "pil",
|
|
return_dict: bool = True,
|
|
):
|
|
self.check_inputs(
|
|
video=video,
|
|
height=height,
|
|
width=width,
|
|
latents=latents,
|
|
tone_map_compression_ratio=tone_map_compression_ratio,
|
|
)
|
|
|
|
if video is not None:
|
|
# Batched video input is not yet tested/supported. TODO: take a look later
|
|
batch_size = 1
|
|
else:
|
|
batch_size = latents.shape[0]
|
|
device = self.device
|
|
|
|
if video is not None:
|
|
num_frames = len(video)
|
|
if num_frames % self.vae_temporal_compression_ratio != 1:
|
|
num_frames = num_frames // self.vae_temporal_compression_ratio * self.vae_temporal_compression_ratio + 1
|
|
video = video[:num_frames]
|
|
logger.warning(
|
|
f"Video length expected to be of the form `k * {self.vae_temporal_compression_ratio} + 1` but is {len(video)}. Truncating to {num_frames} frames." # noqa
|
|
)
|
|
video = self.video_processor.preprocess_video(video, height=height, width=width)
|
|
video = video.to(device=device, dtype=torch.float32)
|
|
|
|
latents_supplied = latents is not None
|
|
latents = self.prepare_latents(
|
|
video=video,
|
|
batch_size=batch_size,
|
|
num_frames=num_frames,
|
|
height=height,
|
|
width=width,
|
|
spatial_patch_size=spatial_patch_size,
|
|
temporal_patch_size=temporal_patch_size,
|
|
dtype=torch.float32,
|
|
device=device,
|
|
generator=generator,
|
|
latents=latents,
|
|
)
|
|
|
|
if latents_supplied and latents_normalized:
|
|
latents = self._denormalize_latents(
|
|
latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
|
|
)
|
|
latents = latents.to(self.latent_upsampler.dtype)
|
|
latents_upsampled = self.latent_upsampler(latents)
|
|
|
|
if adain_factor > 0.0:
|
|
latents = self.adain_filter_latent(latents_upsampled, latents, adain_factor)
|
|
else:
|
|
latents = latents_upsampled
|
|
|
|
if tone_map_compression_ratio > 0.0:
|
|
latents = self.tone_map_latents(latents, tone_map_compression_ratio)
|
|
|
|
if output_type == "latent":
|
|
video = latents
|
|
else:
|
|
if not self.vae.config.timestep_conditioning:
|
|
timestep = None
|
|
else:
|
|
noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
|
|
if not isinstance(decode_timestep, list):
|
|
decode_timestep = [decode_timestep] * batch_size
|
|
if decode_noise_scale is None:
|
|
decode_noise_scale = decode_timestep
|
|
elif not isinstance(decode_noise_scale, list):
|
|
decode_noise_scale = [decode_noise_scale] * batch_size
|
|
|
|
timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
|
|
decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
|
|
:, None, None, None, None
|
|
]
|
|
latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
|
|
|
|
video = self.vae.decode(latents, timestep, return_dict=False)[0]
|
|
video = self.video_processor.postprocess_video(video, output_type=output_type)
|
|
|
|
if not return_dict:
|
|
return (video,)
|
|
|
|
return DiffusionOutput(output=(video,))
|