mirror of
https://github.com/Wan-Video/Wan2.1.git
synced 2025-11-04 06:15:17 +00:00
2046 lines
87 KiB
Python
2046 lines
87 KiB
Python
# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
|
|
import copy
|
|
import inspect
|
|
import math
|
|
import re
|
|
from contextlib import nullcontext
|
|
from dataclasses import dataclass
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from diffusers.image_processor import VaeImageProcessor
|
|
from diffusers.models import AutoencoderKL
|
|
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
|
|
from diffusers.schedulers import DPMSolverMultistepScheduler
|
|
from diffusers.utils import deprecate, logging
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
from einops import rearrange
|
|
from transformers import (
|
|
T5EncoderModel,
|
|
T5Tokenizer,
|
|
AutoModelForCausalLM,
|
|
AutoProcessor,
|
|
AutoTokenizer,
|
|
)
|
|
|
|
from ltx_video.models.autoencoders.causal_video_autoencoder import (
|
|
CausalVideoAutoencoder,
|
|
)
|
|
from ltx_video.models.autoencoders.vae_encode import (
|
|
get_vae_size_scale_factor,
|
|
latent_to_pixel_coords,
|
|
vae_decode,
|
|
vae_encode,
|
|
)
|
|
from ltx_video.models.transformers.symmetric_patchifier import Patchifier
|
|
from ltx_video.models.transformers.transformer3d import Transformer3DModel
|
|
from ltx_video.schedulers.rf import TimestepShifter
|
|
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
|
|
from ltx_video.utils.prompt_enhance_utils import generate_cinematic_prompt
|
|
from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
|
|
from ltx_video.models.autoencoders.vae_encode import (
|
|
un_normalize_latents,
|
|
normalize_latents,
|
|
)
|
|
|
|
|
|
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
|
|
|
|
|
ASPECT_RATIO_1024_BIN = {
|
|
"0.25": [512.0, 2048.0],
|
|
"0.28": [512.0, 1856.0],
|
|
"0.32": [576.0, 1792.0],
|
|
"0.33": [576.0, 1728.0],
|
|
"0.35": [576.0, 1664.0],
|
|
"0.4": [640.0, 1600.0],
|
|
"0.42": [640.0, 1536.0],
|
|
"0.48": [704.0, 1472.0],
|
|
"0.5": [704.0, 1408.0],
|
|
"0.52": [704.0, 1344.0],
|
|
"0.57": [768.0, 1344.0],
|
|
"0.6": [768.0, 1280.0],
|
|
"0.68": [832.0, 1216.0],
|
|
"0.72": [832.0, 1152.0],
|
|
"0.78": [896.0, 1152.0],
|
|
"0.82": [896.0, 1088.0],
|
|
"0.88": [960.0, 1088.0],
|
|
"0.94": [960.0, 1024.0],
|
|
"1.0": [1024.0, 1024.0],
|
|
"1.07": [1024.0, 960.0],
|
|
"1.13": [1088.0, 960.0],
|
|
"1.21": [1088.0, 896.0],
|
|
"1.29": [1152.0, 896.0],
|
|
"1.38": [1152.0, 832.0],
|
|
"1.46": [1216.0, 832.0],
|
|
"1.67": [1280.0, 768.0],
|
|
"1.75": [1344.0, 768.0],
|
|
"2.0": [1408.0, 704.0],
|
|
"2.09": [1472.0, 704.0],
|
|
"2.4": [1536.0, 640.0],
|
|
"2.5": [1600.0, 640.0],
|
|
"3.0": [1728.0, 576.0],
|
|
"4.0": [2048.0, 512.0],
|
|
}
|
|
|
|
ASPECT_RATIO_512_BIN = {
|
|
"0.25": [256.0, 1024.0],
|
|
"0.28": [256.0, 928.0],
|
|
"0.32": [288.0, 896.0],
|
|
"0.33": [288.0, 864.0],
|
|
"0.35": [288.0, 832.0],
|
|
"0.4": [320.0, 800.0],
|
|
"0.42": [320.0, 768.0],
|
|
"0.48": [352.0, 736.0],
|
|
"0.5": [352.0, 704.0],
|
|
"0.52": [352.0, 672.0],
|
|
"0.57": [384.0, 672.0],
|
|
"0.6": [384.0, 640.0],
|
|
"0.68": [416.0, 608.0],
|
|
"0.72": [416.0, 576.0],
|
|
"0.78": [448.0, 576.0],
|
|
"0.82": [448.0, 544.0],
|
|
"0.88": [480.0, 544.0],
|
|
"0.94": [480.0, 512.0],
|
|
"1.0": [512.0, 512.0],
|
|
"1.07": [512.0, 480.0],
|
|
"1.13": [544.0, 480.0],
|
|
"1.21": [544.0, 448.0],
|
|
"1.29": [576.0, 448.0],
|
|
"1.38": [576.0, 416.0],
|
|
"1.46": [608.0, 416.0],
|
|
"1.67": [640.0, 384.0],
|
|
"1.75": [672.0, 384.0],
|
|
"2.0": [704.0, 352.0],
|
|
"2.09": [736.0, 352.0],
|
|
"2.4": [768.0, 320.0],
|
|
"2.5": [800.0, 320.0],
|
|
"3.0": [864.0, 288.0],
|
|
"4.0": [1024.0, 256.0],
|
|
}
|
|
|
|
class MomentumBuffer:
|
|
def __init__(self, momentum: float):
|
|
self.momentum = momentum
|
|
self.running_average = 0
|
|
|
|
def update(self, update_value: torch.Tensor):
|
|
new_average = self.momentum * self.running_average
|
|
self.running_average = update_value + new_average
|
|
|
|
|
|
|
|
def project(
|
|
v0: torch.Tensor, # [B, C, T, H, W]
|
|
v1: torch.Tensor, # [B, C, T, H, W]
|
|
):
|
|
dtype = v0.dtype
|
|
v0, v1 = v0.double(), v1.double()
|
|
v1 = torch.nn.functional.normalize(v1, dim=[-2, -1])
|
|
v0_parallel = (v0 * v1).sum(dim=[-2, -1], keepdim=True) * v1
|
|
v0_orthogonal = v0 - v0_parallel
|
|
return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
|
|
|
|
|
|
def adaptive_projected_guidance(
|
|
diff: torch.Tensor, # [B, C, T, H, W]
|
|
pred_cond: torch.Tensor, # [B, C, T, H, W]
|
|
momentum_buffer: MomentumBuffer = None,
|
|
eta: float = 0.0,
|
|
norm_threshold: float = 55,
|
|
):
|
|
if momentum_buffer is not None:
|
|
momentum_buffer.update(diff)
|
|
diff = momentum_buffer.running_average
|
|
if norm_threshold > 0:
|
|
ones = torch.ones_like(diff)
|
|
diff_norm = diff.norm(p=2, dim=[-2, -1], keepdim=True)
|
|
print(f"diff_norm: {diff_norm}")
|
|
scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
|
|
diff = diff * scale_factor
|
|
diff_parallel, diff_orthogonal = project(diff, pred_cond)
|
|
normalized_update = diff_orthogonal + eta * diff_parallel
|
|
return normalized_update
|
|
|
|
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
|
|
def retrieve_timesteps(
|
|
scheduler,
|
|
num_inference_steps: Optional[int] = None,
|
|
device: Optional[Union[str, torch.device]] = None,
|
|
timesteps: Optional[List[int]] = None,
|
|
max_timestep: Optional[float] = 1.0,
|
|
skip_initial_inference_steps: int = 0,
|
|
skip_final_inference_steps: int = 0,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
|
|
custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
|
|
|
|
Args:
|
|
scheduler (`SchedulerMixin`):
|
|
The scheduler to get timesteps from.
|
|
num_inference_steps (`int`):
|
|
The number of diffusion steps used when generating samples with a pre-trained model. If used,
|
|
`timesteps` must be `None`.
|
|
device (`str` or `torch.device`, *optional*):
|
|
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
|
|
timesteps (`List[int]`, *optional*):
|
|
Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
|
|
timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
|
|
must be `None`.
|
|
max_timestep ('float', *optional*, defaults to 1.0):
|
|
The initial noising level for image-to-image/video-to-video. The list if timestamps will be
|
|
truncated to start with a timestamp greater or equal to this.
|
|
|
|
Returns:
|
|
`Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
|
|
second element is the number of inference steps.
|
|
"""
|
|
if timesteps is not None:
|
|
accepts_timesteps = "timesteps" in set(
|
|
inspect.signature(scheduler.set_timesteps).parameters.keys()
|
|
)
|
|
if not accepts_timesteps:
|
|
raise ValueError(
|
|
f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
|
|
f" timestep schedules. Please check whether you are using the correct scheduler."
|
|
)
|
|
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
|
|
timesteps = scheduler.timesteps
|
|
num_inference_steps = len(timesteps)
|
|
else:
|
|
scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
|
|
timesteps = scheduler.timesteps
|
|
|
|
if (
|
|
skip_initial_inference_steps < 0
|
|
or skip_final_inference_steps < 0
|
|
or skip_initial_inference_steps + skip_final_inference_steps
|
|
>= num_inference_steps
|
|
):
|
|
raise ValueError(
|
|
f"max_timestep {max_timestep} is smaller than the minimum timestep {timesteps.min()}"
|
|
"invalid skip inference step values: must be non-negative and the sum of skip_initial_inference_steps and skip_final_inference_steps must be less than the number of inference steps"
|
|
)
|
|
|
|
timesteps = timesteps[
|
|
skip_initial_inference_steps : len(timesteps) - skip_final_inference_steps
|
|
]
|
|
|
|
if max_timestep < 1.0:
|
|
if max_timestep < timesteps.min():
|
|
raise ValueError(
|
|
f"max_timestep {max_timestep} is smaller than the minimum timestep {timesteps.min()}"
|
|
)
|
|
timesteps = timesteps[timesteps <= max_timestep]
|
|
num_inference_steps = len(timesteps)
|
|
scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
|
|
|
|
return timesteps, num_inference_steps
|
|
|
|
|
|
@dataclass
|
|
class ConditioningItem:
|
|
"""
|
|
Defines a single frame-conditioning item - a single frame or a sequence of frames.
|
|
|
|
Attributes:
|
|
media_item (torch.Tensor): shape=(b, 3, f, h, w). The media item to condition on.
|
|
media_frame_number (int): The start-frame number of the media item in the generated video.
|
|
conditioning_strength (float): The strength of the conditioning (1.0 = full conditioning).
|
|
media_x (Optional[int]): Optional left x coordinate of the media item in the generated frame.
|
|
media_y (Optional[int]): Optional top y coordinate of the media item in the generated frame.
|
|
"""
|
|
|
|
media_item: torch.Tensor
|
|
media_frame_number: int
|
|
conditioning_strength: float
|
|
control_frames: bool = False
|
|
media_x: Optional[int] = None
|
|
media_y: Optional[int] = None
|
|
|
|
|
|
class LTXVideoPipeline(DiffusionPipeline):
|
|
r"""
|
|
Pipeline for text-to-image generation using LTX-Video.
|
|
|
|
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
|
|
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
|
|
|
|
Args:
|
|
vae ([`AutoencoderKL`]):
|
|
Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
|
|
text_encoder ([`T5EncoderModel`]):
|
|
Frozen text-encoder. This uses
|
|
[T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
|
|
[t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant.
|
|
tokenizer (`T5Tokenizer`):
|
|
Tokenizer of class
|
|
[T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
|
|
transformer ([`Transformer2DModel`]):
|
|
A text conditioned `Transformer2DModel` to denoise the encoded image latents.
|
|
scheduler ([`SchedulerMixin`]):
|
|
A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
|
|
"""
|
|
|
|
bad_punct_regex = re.compile(
|
|
r"["
|
|
+ "#®•©™&@·º½¾¿¡§~"
|
|
+ r"\)"
|
|
+ r"\("
|
|
+ r"\]"
|
|
+ r"\["
|
|
+ r"\}"
|
|
+ r"\{"
|
|
+ r"\|"
|
|
+ "\\"
|
|
+ r"\/"
|
|
+ r"\*"
|
|
+ r"]{1,}"
|
|
) # noqa
|
|
|
|
_optional_components = [
|
|
"tokenizer",
|
|
"text_encoder",
|
|
"prompt_enhancer_image_caption_model",
|
|
"prompt_enhancer_image_caption_processor",
|
|
"prompt_enhancer_llm_model",
|
|
"prompt_enhancer_llm_tokenizer",
|
|
]
|
|
model_cpu_offload_seq = "prompt_enhancer_image_caption_model->prompt_enhancer_llm_model->text_encoder->transformer->vae"
|
|
|
|
def __init__(
|
|
self,
|
|
tokenizer: T5Tokenizer,
|
|
text_encoder: T5EncoderModel,
|
|
vae: AutoencoderKL,
|
|
transformer: Transformer3DModel,
|
|
scheduler: DPMSolverMultistepScheduler,
|
|
patchifier: Patchifier,
|
|
prompt_enhancer_image_caption_model: AutoModelForCausalLM,
|
|
prompt_enhancer_image_caption_processor: AutoProcessor,
|
|
prompt_enhancer_llm_model: AutoModelForCausalLM,
|
|
prompt_enhancer_llm_tokenizer: AutoTokenizer,
|
|
allowed_inference_steps: Optional[List[float]] = None,
|
|
):
|
|
super().__init__()
|
|
|
|
self.register_modules(
|
|
tokenizer=tokenizer,
|
|
text_encoder=text_encoder,
|
|
vae=vae,
|
|
transformer=transformer,
|
|
scheduler=scheduler,
|
|
patchifier=patchifier,
|
|
prompt_enhancer_image_caption_model=prompt_enhancer_image_caption_model,
|
|
prompt_enhancer_image_caption_processor=prompt_enhancer_image_caption_processor,
|
|
prompt_enhancer_llm_model=prompt_enhancer_llm_model,
|
|
prompt_enhancer_llm_tokenizer=prompt_enhancer_llm_tokenizer,
|
|
)
|
|
|
|
self.video_scale_factor, self.vae_scale_factor, _ = get_vae_size_scale_factor(
|
|
self.vae
|
|
)
|
|
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
|
|
|
self.allowed_inference_steps = allowed_inference_steps
|
|
|
|
def mask_text_embeddings(self, emb, mask):
|
|
if emb.shape[0] == 1:
|
|
keep_index = mask.sum().item()
|
|
return emb[:, :, :keep_index, :], keep_index
|
|
else:
|
|
masked_feature = emb * mask[:, None, :, None]
|
|
return masked_feature, emb.shape[2]
|
|
|
|
# Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt
|
|
def encode_prompt(
|
|
self,
|
|
prompt: Union[str, List[str]],
|
|
do_classifier_free_guidance: bool = True,
|
|
negative_prompt: str = "",
|
|
num_images_per_prompt: int = 1,
|
|
device: Optional[torch.device] = None,
|
|
prompt_embeds: Optional[torch.FloatTensor] = None,
|
|
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
|
prompt_attention_mask: Optional[torch.FloatTensor] = None,
|
|
negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
|
|
text_encoder_max_tokens: int = 256,
|
|
**kwargs,
|
|
):
|
|
r"""
|
|
Encodes the prompt into text encoder hidden states.
|
|
|
|
Args:
|
|
prompt (`str` or `List[str]`, *optional*):
|
|
prompt to be encoded
|
|
negative_prompt (`str` or `List[str]`, *optional*):
|
|
The prompt not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds`
|
|
instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). For
|
|
This should be "".
|
|
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
|
|
whether to use classifier free guidance or not
|
|
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
|
number of images that should be generated per prompt
|
|
device: (`torch.device`, *optional*):
|
|
torch device to place the resulting embeddings on
|
|
prompt_embeds (`torch.FloatTensor`, *optional*):
|
|
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
|
|
provided, text embeddings will be generated from `prompt` input argument.
|
|
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
|
Pre-generated negative text embeddings.
|
|
"""
|
|
|
|
if "mask_feature" in kwargs:
|
|
deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
|
|
deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
|
|
|
|
if device is None:
|
|
device = self._execution_device
|
|
|
|
if prompt is not None and isinstance(prompt, str):
|
|
batch_size = 1
|
|
elif prompt is not None and isinstance(prompt, list):
|
|
batch_size = len(prompt)
|
|
else:
|
|
batch_size = prompt_embeds.shape[0]
|
|
|
|
# See Section 3.1. of the paper.
|
|
max_length = (
|
|
text_encoder_max_tokens # TPU supports only lengths multiple of 128
|
|
)
|
|
if prompt_embeds is None:
|
|
assert (
|
|
self.text_encoder is not None
|
|
), "You should provide either prompt_embeds or self.text_encoder should not be None,"
|
|
text_enc_device = next(self.text_encoder.parameters()).device
|
|
prompt = self._text_preprocessing(prompt)
|
|
text_inputs = self.tokenizer(
|
|
prompt,
|
|
padding="max_length",
|
|
max_length=max_length,
|
|
truncation=True,
|
|
add_special_tokens=True,
|
|
return_tensors="pt",
|
|
)
|
|
text_input_ids = text_inputs.input_ids
|
|
untruncated_ids = self.tokenizer(
|
|
prompt, padding="longest", return_tensors="pt"
|
|
).input_ids
|
|
|
|
if untruncated_ids.shape[-1] >= text_input_ids.shape[
|
|
-1
|
|
] and not torch.equal(text_input_ids, untruncated_ids):
|
|
removed_text = self.tokenizer.batch_decode(
|
|
untruncated_ids[:, max_length - 1 : -1]
|
|
)
|
|
logger.warning(
|
|
"The following part of your input was truncated because CLIP can only handle sequences up to"
|
|
f" {max_length} tokens: {removed_text}"
|
|
)
|
|
|
|
prompt_attention_mask = text_inputs.attention_mask
|
|
prompt_attention_mask = prompt_attention_mask.to(text_enc_device)
|
|
prompt_attention_mask = prompt_attention_mask.to(device)
|
|
|
|
prompt_embeds = self.text_encoder(
|
|
text_input_ids.to(text_enc_device), attention_mask=prompt_attention_mask
|
|
)
|
|
prompt_embeds = prompt_embeds[0]
|
|
|
|
if self.text_encoder is not None:
|
|
dtype = self.text_encoder.dtype
|
|
elif self.transformer is not None:
|
|
dtype = self.transformer.dtype
|
|
else:
|
|
dtype = None
|
|
|
|
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
|
|
|
|
bs_embed, seq_len, _ = prompt_embeds.shape
|
|
# duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
|
|
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
|
|
prompt_embeds = prompt_embeds.view(
|
|
bs_embed * num_images_per_prompt, seq_len, -1
|
|
)
|
|
prompt_attention_mask = prompt_attention_mask.repeat(1, num_images_per_prompt)
|
|
prompt_attention_mask = prompt_attention_mask.view(
|
|
bs_embed * num_images_per_prompt, -1
|
|
)
|
|
|
|
# get unconditional embeddings for classifier free guidance
|
|
if do_classifier_free_guidance and negative_prompt_embeds is None:
|
|
uncond_tokens = self._text_preprocessing(negative_prompt)
|
|
uncond_tokens = uncond_tokens * batch_size
|
|
max_length = prompt_embeds.shape[1]
|
|
uncond_input = self.tokenizer(
|
|
uncond_tokens,
|
|
padding="max_length",
|
|
max_length=max_length,
|
|
truncation=True,
|
|
return_attention_mask=True,
|
|
add_special_tokens=True,
|
|
return_tensors="pt",
|
|
)
|
|
negative_prompt_attention_mask = uncond_input.attention_mask
|
|
negative_prompt_attention_mask = negative_prompt_attention_mask.to(
|
|
text_enc_device
|
|
)
|
|
|
|
negative_prompt_embeds = self.text_encoder(
|
|
uncond_input.input_ids.to(text_enc_device),
|
|
attention_mask=negative_prompt_attention_mask,
|
|
)
|
|
negative_prompt_embeds = negative_prompt_embeds[0]
|
|
|
|
if do_classifier_free_guidance:
|
|
# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
|
|
seq_len = negative_prompt_embeds.shape[1]
|
|
|
|
negative_prompt_embeds = negative_prompt_embeds.to(
|
|
dtype=dtype, device=device
|
|
)
|
|
|
|
negative_prompt_embeds = negative_prompt_embeds.repeat(
|
|
1, num_images_per_prompt, 1
|
|
)
|
|
negative_prompt_embeds = negative_prompt_embeds.view(
|
|
batch_size * num_images_per_prompt, seq_len, -1
|
|
)
|
|
|
|
negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(
|
|
1, num_images_per_prompt
|
|
)
|
|
negative_prompt_attention_mask = negative_prompt_attention_mask.view(
|
|
bs_embed * num_images_per_prompt, -1
|
|
)
|
|
else:
|
|
negative_prompt_embeds = None
|
|
negative_prompt_attention_mask = None
|
|
|
|
return (
|
|
prompt_embeds,
|
|
prompt_attention_mask,
|
|
negative_prompt_embeds,
|
|
negative_prompt_attention_mask,
|
|
)
|
|
|
|
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
|
|
def prepare_extra_step_kwargs(self, generator, eta):
|
|
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
|
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
|
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
|
# and should be between [0, 1]
|
|
|
|
accepts_eta = "eta" in set(
|
|
inspect.signature(self.scheduler.step).parameters.keys()
|
|
)
|
|
extra_step_kwargs = {}
|
|
if accepts_eta:
|
|
extra_step_kwargs["eta"] = eta
|
|
|
|
# check if the scheduler accepts generator
|
|
accepts_generator = "generator" in set(
|
|
inspect.signature(self.scheduler.step).parameters.keys()
|
|
)
|
|
if accepts_generator:
|
|
extra_step_kwargs["generator"] = generator
|
|
return extra_step_kwargs
|
|
|
|
def check_inputs(
|
|
self,
|
|
prompt,
|
|
height,
|
|
width,
|
|
negative_prompt,
|
|
prompt_embeds=None,
|
|
negative_prompt_embeds=None,
|
|
prompt_attention_mask=None,
|
|
negative_prompt_attention_mask=None,
|
|
enhance_prompt=False,
|
|
):
|
|
if height % 8 != 0 or width % 8 != 0:
|
|
raise ValueError(
|
|
f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
|
|
)
|
|
|
|
if prompt is not None and prompt_embeds is not None:
|
|
raise ValueError(
|
|
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
|
|
" only forward one of the two."
|
|
)
|
|
elif prompt is None and prompt_embeds is None:
|
|
raise ValueError(
|
|
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
|
|
)
|
|
elif prompt is not None and (
|
|
not isinstance(prompt, str) and not isinstance(prompt, list)
|
|
):
|
|
raise ValueError(
|
|
f"`prompt` has to be of type `str` or `list` but is {type(prompt)}"
|
|
)
|
|
|
|
if prompt is not None and negative_prompt_embeds is not None:
|
|
raise ValueError(
|
|
f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
|
|
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
|
|
)
|
|
|
|
if negative_prompt is not None and negative_prompt_embeds is not None:
|
|
raise ValueError(
|
|
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
|
|
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
|
|
)
|
|
|
|
if prompt_embeds is not None and prompt_attention_mask is None:
|
|
raise ValueError(
|
|
"Must provide `prompt_attention_mask` when specifying `prompt_embeds`."
|
|
)
|
|
|
|
if (
|
|
negative_prompt_embeds is not None
|
|
and negative_prompt_attention_mask is None
|
|
):
|
|
raise ValueError(
|
|
"Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`."
|
|
)
|
|
|
|
if prompt_embeds is not None and negative_prompt_embeds is not None:
|
|
if prompt_embeds.shape != negative_prompt_embeds.shape:
|
|
raise ValueError(
|
|
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
|
|
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
|
|
f" {negative_prompt_embeds.shape}."
|
|
)
|
|
if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
|
|
raise ValueError(
|
|
"`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
|
|
f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
|
|
f" {negative_prompt_attention_mask.shape}."
|
|
)
|
|
|
|
if enhance_prompt:
|
|
assert (
|
|
self.prompt_enhancer_image_caption_model is not None
|
|
), "Image caption model must be initialized if enhance_prompt is True"
|
|
assert (
|
|
self.prompt_enhancer_image_caption_processor is not None
|
|
), "Image caption processor must be initialized if enhance_prompt is True"
|
|
assert (
|
|
self.prompt_enhancer_llm_model is not None
|
|
), "Text prompt enhancer model must be initialized if enhance_prompt is True"
|
|
assert (
|
|
self.prompt_enhancer_llm_tokenizer is not None
|
|
), "Text prompt enhancer tokenizer must be initialized if enhance_prompt is True"
|
|
|
|
def _text_preprocessing(self, text):
|
|
if not isinstance(text, (tuple, list)):
|
|
text = [text]
|
|
|
|
def process(text: str):
|
|
text = text.strip()
|
|
return text
|
|
|
|
return [process(t) for t in text]
|
|
|
|
@staticmethod
|
|
def add_noise_to_image_conditioning_latents(
|
|
t: float,
|
|
init_latents: torch.Tensor,
|
|
latents: torch.Tensor,
|
|
noise_scale: float,
|
|
conditioning_mask: torch.Tensor,
|
|
generator,
|
|
eps=1e-6,
|
|
):
|
|
"""
|
|
Add timestep-dependent noise to the hard-conditioning latents.
|
|
This helps with motion continuity, especially when conditioned on a single frame.
|
|
"""
|
|
noise = randn_tensor(
|
|
latents.shape,
|
|
generator=generator,
|
|
device=latents.device,
|
|
dtype=latents.dtype,
|
|
)
|
|
# Add noise only to hard-conditioning latents (conditioning_mask = 1.0)
|
|
need_to_noise = (conditioning_mask > 1.0 - eps).unsqueeze(-1)
|
|
noised_latents = init_latents + noise_scale * noise * (t**2)
|
|
latents = torch.where(need_to_noise, noised_latents, latents)
|
|
return latents
|
|
|
|
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
|
|
def prepare_latents(
|
|
self,
|
|
latents: torch.Tensor | None,
|
|
media_items: torch.Tensor | None,
|
|
timestep: float,
|
|
latent_shape: torch.Size | Tuple[Any, ...],
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
generator: torch.Generator | List[torch.Generator],
|
|
vae_per_channel_normalize: bool = True,
|
|
):
|
|
"""
|
|
Prepare the initial latent tensor to be denoised.
|
|
The latents are either pure noise or a noised version of the encoded media items.
|
|
Args:
|
|
latents (`torch.FloatTensor` or `None`):
|
|
The latents to use (provided by the user) or `None` to create new latents.
|
|
media_items (`torch.FloatTensor` or `None`):
|
|
An image or video to be updated using img2img or vid2vid. The media item is encoded and noised.
|
|
timestep (`float`):
|
|
The timestep to noise the encoded media_items to.
|
|
latent_shape (`torch.Size`):
|
|
The target latent shape.
|
|
dtype (`torch.dtype`):
|
|
The target dtype.
|
|
device (`torch.device`):
|
|
The target device.
|
|
generator (`torch.Generator` or `List[torch.Generator]`):
|
|
Generator(s) to be used for the noising process.
|
|
vae_per_channel_normalize ('bool'):
|
|
When encoding the media_items, whether to normalize the latents per-channel.
|
|
Returns:
|
|
`torch.FloatTensor`: The latents to be used for the denoising process. This is a tensor of shape
|
|
(batch_size, num_channels, height, width).
|
|
"""
|
|
if isinstance(generator, list) and len(generator) != latent_shape[0]:
|
|
raise ValueError(
|
|
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
|
f" size of {latent_shape[0]}. Make sure the batch size matches the length of the generators."
|
|
)
|
|
|
|
# Initialize the latents with the given latents or encoded media item, if provided
|
|
assert (
|
|
latents is None or media_items is None
|
|
), "Cannot provide both latents and media_items. Please provide only one of the two."
|
|
|
|
assert (
|
|
latents is None and media_items is None or timestep < 1.0
|
|
), "Input media_item or latents are provided, but they will be replaced with noise."
|
|
|
|
if media_items is not None:
|
|
latents = vae_encode(
|
|
media_items.to(dtype=self.vae.dtype, device=self.vae.device),
|
|
self.vae,
|
|
vae_per_channel_normalize=vae_per_channel_normalize,
|
|
)
|
|
if latents is not None:
|
|
assert (
|
|
latents.shape == latent_shape
|
|
), f"Latents have to be of shape {latent_shape} but are {latents.shape}."
|
|
latents = latents.to(device=device, dtype=dtype)
|
|
|
|
# For backward compatibility, generate in the "patchified" shape and rearrange
|
|
b, c, f, h, w = latent_shape
|
|
noise = randn_tensor(
|
|
(b, f * h * w, c), generator=generator, device=device, dtype=dtype
|
|
)
|
|
noise = rearrange(noise, "b (f h w) c -> b c f h w", f=f, h=h, w=w)
|
|
|
|
# scale the initial noise by the standard deviation required by the scheduler
|
|
noise = noise * self.scheduler.init_noise_sigma
|
|
|
|
if latents is None:
|
|
latents = noise
|
|
else:
|
|
# Noise the latents to the required (first) timestep
|
|
latents = timestep * noise + (1 - timestep) * latents
|
|
|
|
return latents
|
|
|
|
@staticmethod
|
|
def classify_height_width_bin(
|
|
height: int, width: int, ratios: dict
|
|
) -> Tuple[int, int]:
|
|
"""Returns binned height and width."""
|
|
ar = float(height / width)
|
|
closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))
|
|
default_hw = ratios[closest_ratio]
|
|
return int(default_hw[0]), int(default_hw[1])
|
|
|
|
@staticmethod
|
|
def resize_and_crop_tensor(
|
|
samples: torch.Tensor, new_width: int, new_height: int
|
|
) -> torch.Tensor:
|
|
n_frames, orig_height, orig_width = samples.shape[-3:]
|
|
|
|
# Check if resizing is needed
|
|
if orig_height != new_height or orig_width != new_width:
|
|
ratio = max(new_height / orig_height, new_width / orig_width)
|
|
resized_width = int(orig_width * ratio)
|
|
resized_height = int(orig_height * ratio)
|
|
|
|
# Resize
|
|
samples = LTXVideoPipeline.resize_tensor(
|
|
samples, resized_height, resized_width
|
|
)
|
|
|
|
# Center Crop
|
|
start_x = (resized_width - new_width) // 2
|
|
end_x = start_x + new_width
|
|
start_y = (resized_height - new_height) // 2
|
|
end_y = start_y + new_height
|
|
samples = samples[..., start_y:end_y, start_x:end_x]
|
|
|
|
return samples
|
|
|
|
@staticmethod
|
|
def resize_tensor(media_items, height, width):
|
|
n_frames = media_items.shape[2]
|
|
if media_items.shape[-2:] != (height, width):
|
|
media_items = rearrange(media_items, "b c n h w -> (b n) c h w")
|
|
media_items = F.interpolate(
|
|
media_items,
|
|
size=(height, width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
media_items = rearrange(media_items, "(b n) c h w -> b c n h w", n=n_frames)
|
|
return media_items
|
|
|
|
@torch.no_grad()
|
|
def __call__(
|
|
self,
|
|
height: int,
|
|
width: int,
|
|
num_frames: int,
|
|
frame_rate: float,
|
|
prompt: Union[str, List[str]] = None,
|
|
negative_prompt: str = None,
|
|
num_inference_steps: int = 20,
|
|
timesteps: List[int] = None,
|
|
guidance_scale: Union[float, List[float]] = 4.5,
|
|
skip_layer_strategy: Optional[SkipLayerStrategy] = None,
|
|
skip_block_list: Optional[Union[List[List[int]], List[int]]] = None,
|
|
stg_scale: Union[float, List[float]] = 1.0,
|
|
rescaling_scale: Union[float, List[float]] = 0.7,
|
|
guidance_timesteps: Optional[List[int]] = None,
|
|
num_images_per_prompt: Optional[int] = 1,
|
|
eta: float = 0.0,
|
|
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
|
latents: Optional[torch.FloatTensor] = None,
|
|
prompt_embeds: Optional[torch.FloatTensor] = None,
|
|
prompt_attention_mask: Optional[torch.FloatTensor] = None,
|
|
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
|
|
negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
|
|
output_type: Optional[str] = "pil",
|
|
return_dict: bool = True,
|
|
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
|
|
conditioning_items: Optional[List[ConditioningItem]] = None,
|
|
decode_timestep: Union[List[float], float] = 0.0,
|
|
decode_noise_scale: Optional[List[float]] = None,
|
|
mixed_precision: bool = False,
|
|
offload_to_cpu: bool = False,
|
|
enhance_prompt: bool = False,
|
|
text_encoder_max_tokens: int = 256,
|
|
stochastic_sampling: bool = False,
|
|
media_items: Optional[torch.Tensor] = None,
|
|
tone_map_compression_ratio: float = 0.0,
|
|
strength: Optional[float] = 1.0,
|
|
skip_initial_inference_steps: int = 0,
|
|
skip_final_inference_steps: int = 0,
|
|
joint_pass: bool = False,
|
|
pass_no: int = -1,
|
|
ltxv_model = None,
|
|
callback=None,
|
|
apg_switch = 0,
|
|
**kwargs,
|
|
) -> Union[ImagePipelineOutput, Tuple]:
|
|
"""
|
|
Function invoked when calling the pipeline for generation.
|
|
|
|
Args:
|
|
prompt (`str` or `List[str]`, *optional*):
|
|
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
|
|
instead.
|
|
negative_prompt (`str` or `List[str]`, *optional*):
|
|
The prompt or prompts not to guide the image generation. If not defined, one has to pass
|
|
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
|
|
less than `1`).
|
|
num_inference_steps (`int`, *optional*, defaults to 100):
|
|
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
|
expense of slower inference. If `timesteps` is provided, this parameter is ignored.
|
|
timesteps (`List[int]`, *optional*):
|
|
Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
|
|
timesteps are used. Must be in descending order.
|
|
guidance_scale (`float`, *optional*, defaults to 4.5):
|
|
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
|
|
`guidance_scale` is defined as `w` of equation 2. of [Imagen
|
|
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
|
|
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
|
|
usually at the expense of lower image quality.
|
|
num_images_per_prompt (`int`, *optional*, defaults to 1):
|
|
The number of images to generate per prompt.
|
|
height (`int`, *optional*, defaults to self.unet.config.sample_size):
|
|
The height in pixels of the generated image.
|
|
width (`int`, *optional*, defaults to self.unet.config.sample_size):
|
|
The width in pixels of the generated image.
|
|
eta (`float`, *optional*, defaults to 0.0):
|
|
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
|
|
[`schedulers.DDIMScheduler`], will be ignored for others.
|
|
generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
|
|
One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
|
|
to make generation deterministic.
|
|
latents (`torch.FloatTensor`, *optional*):
|
|
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
|
|
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
|
tensor will ge generated by sampling using the supplied random `generator`.
|
|
prompt_embeds (`torch.FloatTensor`, *optional*):
|
|
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
|
|
provided, text embeddings will be generated from `prompt` input argument.
|
|
prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
|
|
negative_prompt_embeds (`torch.FloatTensor`, *optional*):
|
|
Pre-generated negative text embeddings. This negative prompt should be "". If not
|
|
provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
|
|
negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
|
|
Pre-generated attention mask for negative text embeddings.
|
|
output_type (`str`, *optional*, defaults to `"pil"`):
|
|
The output format of the generate image. Choose between
|
|
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
|
|
return_dict (`bool`, *optional*, defaults to `True`):
|
|
Whether to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
|
|
callback_on_step_end (`Callable`, *optional*):
|
|
A function that calls at the end of each denoising steps during the inference. The function is called
|
|
with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
|
|
callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
|
|
`callback_on_step_end_tensor_inputs`.
|
|
use_resolution_binning (`bool` defaults to `True`):
|
|
If set to `True`, the requested height and width are first mapped to the closest resolutions using
|
|
`ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
|
|
the requested resolution. Useful for generating non-square images.
|
|
enhance_prompt (`bool`, *optional*, defaults to `False`):
|
|
If set to `True`, the prompt is enhanced using a LLM model.
|
|
text_encoder_max_tokens (`int`, *optional*, defaults to `256`):
|
|
The maximum number of tokens to use for the text encoder.
|
|
stochastic_sampling (`bool`, *optional*, defaults to `False`):
|
|
If set to `True`, the sampling is stochastic. If set to `False`, the sampling is deterministic.
|
|
media_items ('torch.Tensor', *optional*):
|
|
The input media item used for image-to-image / video-to-video.
|
|
When provided, they will be noised according to 'strength' and then fully denoised.
|
|
tone_map_compression_ratio: compression ratio for tone mapping, defaults to 0.0.
|
|
If set to 0.0, no tone mapping is applied. If set to 1.0 - full compression is applied.
|
|
strength ('floaty', *optional* defaults to 1.0):
|
|
The editing level in image-to-image / video-to-video. The provided input will be noised
|
|
to this level.
|
|
Examples:
|
|
|
|
Returns:
|
|
[`~pipelines.ImagePipelineOutput`] or `tuple`:
|
|
If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
|
|
returned where the first element is a list with the generated images
|
|
"""
|
|
if "mask_feature" in kwargs:
|
|
deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
|
|
deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
|
|
|
|
is_video = kwargs.get("is_video", False)
|
|
self.check_inputs(
|
|
prompt,
|
|
height,
|
|
width,
|
|
negative_prompt,
|
|
prompt_embeds,
|
|
negative_prompt_embeds,
|
|
prompt_attention_mask,
|
|
negative_prompt_attention_mask,
|
|
)
|
|
|
|
# 2. Default height and width to transformer
|
|
if prompt is not None and isinstance(prompt, str):
|
|
batch_size = 1
|
|
elif prompt is not None and isinstance(prompt, list):
|
|
batch_size = len(prompt)
|
|
else:
|
|
batch_size = prompt_embeds.shape[0]
|
|
|
|
device = self._execution_device
|
|
|
|
self.video_scale_factor = self.video_scale_factor if is_video else 1
|
|
vae_per_channel_normalize = kwargs.get("vae_per_channel_normalize", True)
|
|
image_cond_noise_scale = kwargs.get("image_cond_noise_scale", 0.0)
|
|
|
|
latent_height = height // self.vae_scale_factor
|
|
latent_width = width // self.vae_scale_factor
|
|
latent_num_frames = num_frames // self.video_scale_factor
|
|
if isinstance(self.vae, CausalVideoAutoencoder) and is_video:
|
|
latent_num_frames += 1
|
|
latent_shape = (
|
|
batch_size * num_images_per_prompt,
|
|
self.transformer.config.in_channels,
|
|
latent_num_frames,
|
|
latent_height,
|
|
latent_width,
|
|
)
|
|
|
|
# Prepare the list of denoising time-steps
|
|
|
|
retrieve_timesteps_kwargs = {}
|
|
if isinstance(self.scheduler, TimestepShifter):
|
|
retrieve_timesteps_kwargs["samples_shape"] = latent_shape
|
|
|
|
assert strength == 1.0 or latents is not None or media_items is not None, (
|
|
"strength < 1 is used for image-to-image/video-to-video - "
|
|
"media_item or latents should be provided."
|
|
)
|
|
|
|
timesteps, num_inference_steps = retrieve_timesteps(
|
|
self.scheduler,
|
|
num_inference_steps,
|
|
device,
|
|
timesteps,
|
|
max_timestep=strength,
|
|
skip_initial_inference_steps=skip_initial_inference_steps,
|
|
skip_final_inference_steps=skip_final_inference_steps,
|
|
**retrieve_timesteps_kwargs,
|
|
)
|
|
if self.allowed_inference_steps is not None:
|
|
for timestep in [round(x, 4) for x in timesteps.tolist()]:
|
|
assert (
|
|
timestep in self.allowed_inference_steps
|
|
), f"Invalid inference timestep {timestep}. Allowed timesteps are {self.allowed_inference_steps}."
|
|
|
|
if guidance_timesteps:
|
|
guidance_mapping = []
|
|
for timestep in timesteps:
|
|
indices = [
|
|
i for i, val in enumerate(guidance_timesteps) if val <= timestep
|
|
]
|
|
# assert len(indices) > 0, f"No guidance timestep found for {timestep}"
|
|
guidance_mapping.append(
|
|
indices[0] if len(indices) > 0 else (len(guidance_timesteps) - 1)
|
|
)
|
|
|
|
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
|
|
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
|
|
# corresponds to doing no classifier free guidance.
|
|
if not isinstance(guidance_scale, List):
|
|
guidance_scale = [guidance_scale] * len(timesteps)
|
|
else:
|
|
guidance_scale = [
|
|
guidance_scale[guidance_mapping[i]] for i in range(len(timesteps))
|
|
]
|
|
|
|
# For simplicity, we are using a constant num_conds for all timesteps, so we need to zero
|
|
# out cases where the guidance scale should not be applied.
|
|
guidance_scale = [x if x > 1.0 else 0.0 for x in guidance_scale]
|
|
|
|
if not isinstance(stg_scale, List):
|
|
stg_scale = [stg_scale] * len(timesteps)
|
|
else:
|
|
stg_scale = [stg_scale[guidance_mapping[i]] for i in range(len(timesteps))]
|
|
|
|
if not isinstance(rescaling_scale, List):
|
|
rescaling_scale = [rescaling_scale] * len(timesteps)
|
|
else:
|
|
rescaling_scale = [
|
|
rescaling_scale[guidance_mapping[i]] for i in range(len(timesteps))
|
|
]
|
|
|
|
do_classifier_free_guidance = any(x > 1.0 for x in guidance_scale)
|
|
do_spatio_temporal_guidance = any(x > 0.0 for x in stg_scale)
|
|
do_rescaling = any(x != 1.0 for x in rescaling_scale)
|
|
|
|
num_conds = 1
|
|
if do_classifier_free_guidance:
|
|
num_conds += 1
|
|
if do_spatio_temporal_guidance:
|
|
num_conds += 1
|
|
|
|
# Normalize skip_block_list to always be None or a list of lists matching timesteps
|
|
if skip_block_list is not None:
|
|
# Convert single list to list of lists if needed
|
|
if len(skip_block_list) == 0 or not isinstance(skip_block_list[0], list):
|
|
skip_block_list = [skip_block_list] * len(timesteps)
|
|
else:
|
|
new_skip_block_list = []
|
|
for i, timestep in enumerate(timesteps):
|
|
new_skip_block_list.append(skip_block_list[guidance_mapping[i]])
|
|
skip_block_list = new_skip_block_list
|
|
|
|
# Prepare skip layer masks
|
|
skip_layer_masks: Optional[List[torch.Tensor]] = None
|
|
if do_spatio_temporal_guidance:
|
|
if skip_block_list is not None:
|
|
skip_layer_masks = [
|
|
self.transformer.create_skip_layer_mask(
|
|
batch_size, num_conds, num_conds - 1, skip_blocks
|
|
)
|
|
for skip_blocks in skip_block_list
|
|
]
|
|
|
|
|
|
# if offload_to_cpu and self.text_encoder is not None:
|
|
# self.text_encoder = self.text_encoder.cpu()
|
|
|
|
# self.transformer = self.transformer.to(self._execution_device)
|
|
|
|
prompt_embeds_batch = prompt_embeds
|
|
prompt_attention_mask_batch = prompt_attention_mask
|
|
if do_classifier_free_guidance:
|
|
prompt_embeds_batch = torch.cat(
|
|
[negative_prompt_embeds, prompt_embeds], dim=0
|
|
)
|
|
prompt_attention_mask_batch = torch.cat(
|
|
[negative_prompt_attention_mask.to("cuda"), prompt_attention_mask], dim=0
|
|
)
|
|
if do_spatio_temporal_guidance:
|
|
prompt_embeds_batch = torch.cat([prompt_embeds_batch, prompt_embeds], dim=0)
|
|
prompt_attention_mask_batch = torch.cat(
|
|
[
|
|
prompt_attention_mask_batch,
|
|
prompt_attention_mask,
|
|
],
|
|
dim=0,
|
|
)
|
|
|
|
# 4. Prepare the initial latents using the provided media and conditioning items
|
|
|
|
# Prepare the initial latents tensor, shape = (b, c, f, h, w)
|
|
latents = self.prepare_latents(
|
|
latents=latents,
|
|
media_items=media_items,
|
|
timestep=timesteps[0],
|
|
latent_shape=latent_shape,
|
|
dtype=torch.float32 if mixed_precision else prompt_embeds_batch.dtype,
|
|
device=device,
|
|
generator=generator,
|
|
vae_per_channel_normalize=vae_per_channel_normalize,
|
|
)
|
|
|
|
# Update the latents with the conditioning items and patchify them into (b, n, c)
|
|
latents, pixel_coords, conditioning_mask, num_cond_latents = (
|
|
self.prepare_conditioning(
|
|
conditioning_items=conditioning_items,
|
|
init_latents=latents,
|
|
num_frames=num_frames,
|
|
height=height,
|
|
width=width,
|
|
vae_per_channel_normalize=vae_per_channel_normalize,
|
|
generator=generator,
|
|
)
|
|
)
|
|
init_latents = latents.clone() # Used for image_cond_noise_update
|
|
if conditioning_items is not None and len(conditioning_items) > 0 and not conditioning_items[0].control_frames and conditioning_items[0].media_frame_number == 0:
|
|
prefix_latent_frames = (conditioning_items[0].media_item.shape[2] - 1)// 8 + 1
|
|
else:
|
|
prefix_latent_frames = 0
|
|
# pixel_coords = torch.cat([pixel_coords] * num_conds)
|
|
orig_conditioning_mask = conditioning_mask
|
|
if conditioning_mask is not None and is_video:
|
|
assert num_images_per_prompt == 1
|
|
conditioning_mask = torch.cat([conditioning_mask] * num_conds)
|
|
fractional_coords = pixel_coords.to(torch.float32)
|
|
fractional_coords[:, 0] = fractional_coords[:, 0] * (1.0 / frame_rate)
|
|
freqs_cis = self.transformer.precompute_freqs_cis(fractional_coords)
|
|
|
|
# 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
|
|
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
|
|
|
# 7. Denoising loop
|
|
num_warmup_steps = max(
|
|
len(timesteps) - num_inference_steps * self.scheduler.order, 0
|
|
)
|
|
cfg_star_rescale = True
|
|
|
|
if apg_switch != 0:
|
|
apg_momentum = -0.75
|
|
apg_norm_threshold = 55
|
|
text_momentumbuffer = MomentumBuffer(apg_momentum)
|
|
audio_momentumbuffer = MomentumBuffer(apg_momentum)
|
|
|
|
|
|
if callback != None:
|
|
callback(-1, None, True, override_num_inference_steps = num_inference_steps, pass_no =pass_no)
|
|
|
|
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
|
for i, t in enumerate(timesteps):
|
|
if conditioning_mask is not None and image_cond_noise_scale > 0.0:
|
|
latents = self.add_noise_to_image_conditioning_latents(
|
|
t,
|
|
init_latents,
|
|
latents,
|
|
image_cond_noise_scale,
|
|
orig_conditioning_mask,
|
|
generator,
|
|
)
|
|
|
|
latent_model_input = (
|
|
torch.cat([latents] * num_conds) if num_conds > 1 else latents
|
|
)
|
|
latent_model_input = self.scheduler.scale_model_input(
|
|
latent_model_input, t
|
|
)
|
|
|
|
current_timestep = t
|
|
if not torch.is_tensor(current_timestep):
|
|
# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
|
|
# This would be a good case for the `match` statement (Python 3.10+)
|
|
is_mps = latent_model_input.device.type == "mps"
|
|
if isinstance(current_timestep, float):
|
|
dtype = torch.float32 if is_mps else torch.float64
|
|
else:
|
|
dtype = torch.int32 if is_mps else torch.int64
|
|
current_timestep = torch.tensor(
|
|
[current_timestep],
|
|
dtype=dtype,
|
|
device=latent_model_input.device,
|
|
)
|
|
elif len(current_timestep.shape) == 0:
|
|
current_timestep = current_timestep[None].to(
|
|
latent_model_input.device
|
|
)
|
|
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
|
current_timestep = current_timestep.expand(
|
|
latent_model_input.shape[0]
|
|
).unsqueeze(-1)
|
|
|
|
if conditioning_mask is not None:
|
|
# Conditioning latents have an initial timestep and noising level of (1.0 - conditioning_mask)
|
|
# and will start to be denoised when the current timestep is lower than their conditioning timestep.
|
|
current_timestep = torch.min(
|
|
current_timestep, 1.0 - conditioning_mask
|
|
)
|
|
|
|
# Choose the appropriate context manager based on `mixed_precision`
|
|
if mixed_precision:
|
|
context_manager = torch.autocast(device.type, dtype=self.transformer.dtype)
|
|
else:
|
|
context_manager = nullcontext() # Dummy context manager
|
|
|
|
# predict noise model_output
|
|
with context_manager:
|
|
noise_pred = self.transformer(
|
|
latent_model_input.to(self.transformer.dtype),
|
|
freqs_cis=freqs_cis,
|
|
encoder_hidden_states=prompt_embeds_batch.to(
|
|
self.transformer.dtype
|
|
),
|
|
encoder_attention_mask=prompt_attention_mask_batch,
|
|
timestep=current_timestep,
|
|
skip_layer_mask=(
|
|
skip_layer_masks[i]
|
|
if skip_layer_masks is not None
|
|
else None
|
|
),
|
|
skip_layer_strategy=skip_layer_strategy,
|
|
latent_shape = latent_shape[2:],
|
|
joint_pass = joint_pass,
|
|
ltxv_model = ltxv_model,
|
|
mixed = mixed_precision,
|
|
return_dict=False,
|
|
)[0]
|
|
if noise_pred == None:
|
|
return None
|
|
# perform guidance
|
|
if do_spatio_temporal_guidance:
|
|
noise_pred_text, noise_pred_text_perturb = noise_pred.chunk(
|
|
num_conds
|
|
)[-2:]
|
|
if do_classifier_free_guidance and guidance_scale[i] !=0 and guidance_scale[i] !=1 :
|
|
noise_pred_uncond, noise_pred_text = noise_pred.chunk(num_conds)[:2]
|
|
|
|
if apg_switch != 0:
|
|
noise_pred = noise_pred_text + (guidance_scale[i] - 1) * adaptive_projected_guidance(noise_pred_text - noise_pred_uncond,
|
|
noise_pred_text,
|
|
momentum_buffer=text_momentumbuffer,
|
|
norm_threshold=apg_norm_threshold)
|
|
|
|
else:
|
|
if cfg_star_rescale:
|
|
batch_size = noise_pred_text.shape[0]
|
|
|
|
positive_flat = noise_pred_text.view(batch_size, -1)
|
|
negative_flat = noise_pred_uncond.view(batch_size, -1)
|
|
dot_product = torch.sum(
|
|
positive_flat * negative_flat, dim=1, keepdim=True
|
|
)
|
|
squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
|
|
alpha = dot_product / squared_norm
|
|
noise_pred_uncond = alpha * noise_pred_uncond
|
|
|
|
|
|
noise_pred = noise_pred_uncond + guidance_scale[i] * (
|
|
noise_pred_text - noise_pred_uncond
|
|
)
|
|
elif do_spatio_temporal_guidance:
|
|
noise_pred = noise_pred_text
|
|
if do_spatio_temporal_guidance:
|
|
noise_pred = noise_pred + stg_scale[i] * (
|
|
noise_pred_text - noise_pred_text_perturb
|
|
)
|
|
if do_rescaling and stg_scale[i] > 0.0:
|
|
noise_pred_text_std = noise_pred_text.view(batch_size, -1).std(
|
|
dim=1, keepdim=True
|
|
)
|
|
noise_pred_std = noise_pred.view(batch_size, -1).std(
|
|
dim=1, keepdim=True
|
|
)
|
|
|
|
factor = noise_pred_text_std / noise_pred_std
|
|
factor = rescaling_scale[i] * factor + (1 - rescaling_scale[i])
|
|
|
|
noise_pred = noise_pred * factor.view(batch_size, 1, 1)
|
|
|
|
current_timestep = current_timestep[:1]
|
|
# learned sigma
|
|
if (
|
|
self.transformer.config.out_channels // 2
|
|
== self.transformer.config.in_channels
|
|
):
|
|
noise_pred = noise_pred.chunk(2, dim=1)[0]
|
|
|
|
# compute previous image: x_t -> x_t-1
|
|
latents = self.denoising_step(
|
|
latents,
|
|
noise_pred,
|
|
current_timestep,
|
|
orig_conditioning_mask,
|
|
t,
|
|
extra_step_kwargs,
|
|
stochastic_sampling=stochastic_sampling,
|
|
)
|
|
|
|
if callback is not None:
|
|
# callback(i, None, False, pass_no =pass_no)
|
|
preview_latents= latents[:, num_cond_latents:].squeeze(0).transpose(0, 1)
|
|
preview_latents= preview_latents.reshape(preview_latents.shape[0], latent_num_frames, latent_height, latent_width)
|
|
callback(i, preview_latents, False, pass_no =pass_no)
|
|
preview_latents = None
|
|
# call the callback, if provided
|
|
if i == len(timesteps) - 1 or (
|
|
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
|
|
):
|
|
progress_bar.update()
|
|
|
|
if callback_on_step_end is not None:
|
|
callback_on_step_end(self, i, t, {})
|
|
|
|
|
|
# Remove the added conditioning latents
|
|
latents = latents[:, num_cond_latents:]
|
|
|
|
latents = self.patchifier.unpatchify(
|
|
latents=latents,
|
|
output_height=latent_height,
|
|
output_width=latent_width,
|
|
out_channels=self.transformer.in_channels
|
|
// math.prod(self.patchifier.patch_size),
|
|
)
|
|
if output_type != "latent":
|
|
if self.vae.decoder.timestep_conditioning:
|
|
noise = torch.randn_like(latents)
|
|
if not isinstance(decode_timestep, list):
|
|
decode_timestep = [decode_timestep] * latents.shape[0]
|
|
if decode_noise_scale is None:
|
|
decode_noise_scale = decode_timestep
|
|
elif not isinstance(decode_noise_scale, list):
|
|
decode_noise_scale = [decode_noise_scale] * latents.shape[0]
|
|
|
|
decode_timestep = torch.tensor(decode_timestep).to(latents.device)
|
|
decode_noise_scale = torch.tensor(decode_noise_scale).to(
|
|
latents.device
|
|
)[:, None, None, None, None]
|
|
latents = (
|
|
latents * (1 - decode_noise_scale) + noise * decode_noise_scale
|
|
)
|
|
else:
|
|
decode_timestep = None
|
|
# torch.save(latents, "lala.pt")
|
|
# latents = torch.load("lala.pt")
|
|
latents = self.tone_map_latents(latents, tone_map_compression_ratio, start = prefix_latent_frames)
|
|
image = vae_decode(
|
|
latents,
|
|
self.vae,
|
|
is_video,
|
|
vae_per_channel_normalize=kwargs["vae_per_channel_normalize"],
|
|
timestep=decode_timestep,
|
|
)
|
|
|
|
image = self.image_processor.postprocess(image, output_type=output_type)
|
|
|
|
else:
|
|
image = latents
|
|
|
|
|
|
if not return_dict:
|
|
return (image,)
|
|
|
|
return image
|
|
|
|
@staticmethod
|
|
def tone_map_latents(
|
|
latents: torch.Tensor,
|
|
compression: float,
|
|
start: int = 0
|
|
) -> torch.Tensor:
|
|
"""
|
|
Applies a non-linear tone-mapping function to latent values to reduce their dynamic range
|
|
in a perceptually smooth way using a sigmoid-based compression.
|
|
|
|
This is useful for regularizing high-variance latents or for conditioning outputs
|
|
during generation, especially when controlling dynamic behavior with a `compression` factor.
|
|
|
|
Parameters:
|
|
----------
|
|
latents : torch.Tensor
|
|
Input latent tensor with arbitrary shape. Expected to be roughly in [-1, 1] or [0, 1] range.
|
|
compression : float
|
|
Compression strength in the range [0, 1].
|
|
- 0.0: No tone-mapping (identity transform)
|
|
- 1.0: Full compression effect
|
|
|
|
Returns:
|
|
-------
|
|
torch.Tensor
|
|
The tone-mapped latent tensor of the same shape as input.
|
|
"""
|
|
if compression ==0:
|
|
return latents
|
|
if not (0 <= compression <= 1):
|
|
raise ValueError("Compression must be in the range [0, 1]")
|
|
|
|
# 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))
|
|
# DeepBeepMeep special touch to allow a smooth transition with tone mapping
|
|
if start > 0:
|
|
gradient_tensor = torch.linspace(0, 1, latents.shape[2])
|
|
gradient_tensor = gradient_tensor ** 0.5
|
|
gradient_tensor = gradient_tensor[ None, None, :, None, None ]
|
|
sigmoid_term *= gradient_tensor
|
|
scales = 1.0 - 0.8 * scale_factor * sigmoid_term
|
|
|
|
|
|
filtered = latents * scales
|
|
return filtered
|
|
|
|
def denoising_step(
|
|
self,
|
|
latents: torch.Tensor,
|
|
noise_pred: torch.Tensor,
|
|
current_timestep: torch.Tensor,
|
|
conditioning_mask: torch.Tensor,
|
|
t: float,
|
|
extra_step_kwargs,
|
|
t_eps=1e-6,
|
|
stochastic_sampling=False,
|
|
):
|
|
"""
|
|
Perform the denoising step for the required tokens, based on the current timestep and
|
|
conditioning mask:
|
|
Conditioning latents have an initial timestep and noising level of (1.0 - conditioning_mask)
|
|
and will start to be denoised when the current timestep is equal or lower than their
|
|
conditioning timestep.
|
|
(hard-conditioning latents with conditioning_mask = 1.0 are never denoised)
|
|
"""
|
|
# Denoise the latents using the scheduler
|
|
denoised_latents = self.scheduler.step(
|
|
noise_pred,
|
|
t if current_timestep is None else current_timestep,
|
|
latents,
|
|
**extra_step_kwargs,
|
|
return_dict=False,
|
|
stochastic_sampling=stochastic_sampling,
|
|
)[0]
|
|
|
|
if conditioning_mask is None:
|
|
return denoised_latents
|
|
|
|
tokens_to_denoise_mask = (t - t_eps < (1.0 - conditioning_mask)).unsqueeze(-1)
|
|
return torch.where(tokens_to_denoise_mask, denoised_latents, latents)
|
|
|
|
def prepare_conditioning(
|
|
self,
|
|
conditioning_items: Optional[List[ConditioningItem]],
|
|
init_latents: torch.Tensor,
|
|
num_frames: int,
|
|
height: int,
|
|
width: int,
|
|
vae_per_channel_normalize: bool = False,
|
|
generator=None,
|
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
|
|
"""
|
|
Prepare conditioning tokens based on the provided conditioning items.
|
|
|
|
This method encodes provided conditioning items (video frames or single frames) into latents
|
|
and integrates them with the initial latent tensor. It also calculates corresponding pixel
|
|
coordinates, a mask indicating the influence of conditioning latents, and the total number of
|
|
conditioning latents.
|
|
|
|
Args:
|
|
conditioning_items (Optional[List[ConditioningItem]]): A list of ConditioningItem objects.
|
|
init_latents (torch.Tensor): The initial latent tensor of shape (b, c, f_l, h_l, w_l), where
|
|
`f_l` is the number of latent frames, and `h_l` and `w_l` are latent spatial dimensions.
|
|
num_frames, height, width: The dimensions of the generated video.
|
|
vae_per_channel_normalize (bool, optional): Whether to normalize channels during VAE encoding.
|
|
Defaults to `False`.
|
|
generator: The random generator
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
|
|
- `init_latents` (torch.Tensor): The updated latent tensor including conditioning latents,
|
|
patchified into (b, n, c) shape.
|
|
- `init_pixel_coords` (torch.Tensor): The pixel coordinates corresponding to the updated
|
|
latent tensor.
|
|
- `conditioning_mask` (torch.Tensor): A mask indicating the conditioning-strength of each
|
|
latent token.
|
|
- `num_cond_latents` (int): The total number of latent tokens added from conditioning items.
|
|
|
|
Raises:
|
|
AssertionError: If input shapes, dimensions, or conditions for applying conditioning are invalid.
|
|
"""
|
|
assert isinstance(self.vae, CausalVideoAutoencoder)
|
|
|
|
if conditioning_items:
|
|
batch_size, _, num_latent_frames = init_latents.shape[:3]
|
|
|
|
init_conditioning_mask = torch.zeros(
|
|
init_latents[:, 0, :, :, :].shape,
|
|
dtype=torch.float32,
|
|
device=init_latents.device,
|
|
)
|
|
|
|
extra_conditioning_latents = []
|
|
extra_conditioning_pixel_coords = []
|
|
extra_conditioning_mask = []
|
|
extra_conditioning_num_latents = 0 # Number of extra conditioning latents added (should be removed before decoding)
|
|
|
|
# Process each conditioning item
|
|
for conditioning_item in conditioning_items:
|
|
conditioning_item = self._resize_conditioning_item(
|
|
conditioning_item, height, width
|
|
)
|
|
media_item = conditioning_item.media_item
|
|
media_frame_number = conditioning_item.media_frame_number
|
|
strength = conditioning_item.conditioning_strength
|
|
control_frames = conditioning_item.control_frames
|
|
assert media_item.ndim == 5 # (b, c, f, h, w)
|
|
b, c, n_frames, h, w = media_item.shape
|
|
assert (
|
|
height == h and width == w
|
|
) or media_frame_number == 0, f"Dimensions do not match: {height}x{width} != {h}x{w} - allowed only when media_frame_number == 0"
|
|
# assert n_frames % 8 == 1
|
|
# assert (
|
|
# media_frame_number >= 0
|
|
# and media_frame_number + n_frames <= num_frames
|
|
# )
|
|
|
|
media_item_latents = vae_encode(
|
|
media_item.to(dtype=self.vae.dtype, device=self.vae.device),
|
|
self.vae,
|
|
vae_per_channel_normalize=vae_per_channel_normalize,
|
|
).to(dtype=init_latents.dtype)
|
|
|
|
# Handle the different conditioning cases
|
|
if control_frames:
|
|
#control frames sequence is assumed to start one frame before the actual location so that we can properly insert the prefix latent
|
|
if media_frame_number > 0:
|
|
media_frame_number = media_frame_number -1
|
|
media_item_latents, media_latent_coords = self.patchifier.patchify(
|
|
latents=media_item_latents
|
|
)
|
|
media_pixel_coords = latent_to_pixel_coords(
|
|
media_latent_coords,
|
|
self.vae,
|
|
causal_fix=self.transformer.config.causal_temporal_positioning,
|
|
)
|
|
|
|
media_conditioning_mask = torch.full(
|
|
media_item_latents.shape[:2],
|
|
strength,
|
|
dtype=torch.float32,
|
|
device=init_latents.device,
|
|
)
|
|
|
|
# Update the frame numbers to match the target frame number
|
|
media_pixel_coords[:, 0] += media_frame_number
|
|
extra_conditioning_num_latents += media_item_latents.shape[1]
|
|
extra_conditioning_latents.append(media_item_latents)
|
|
extra_conditioning_pixel_coords.append(media_pixel_coords)
|
|
extra_conditioning_mask.append(media_conditioning_mask)
|
|
elif media_frame_number == 0:
|
|
# Get the target spatial position of the latent conditioning item
|
|
media_item_latents, l_x, l_y = self._get_latent_spatial_position(
|
|
media_item_latents,
|
|
conditioning_item,
|
|
height,
|
|
width,
|
|
strip_latent_border=True,
|
|
)
|
|
b, c_l, f_l, h_l, w_l = media_item_latents.shape
|
|
|
|
# First frame or sequence - just update the initial noise latents and the mask
|
|
init_latents[:, :, :f_l, l_y : l_y + h_l, l_x : l_x + w_l] = (
|
|
torch.lerp(
|
|
init_latents[:, :, :f_l, l_y : l_y + h_l, l_x : l_x + w_l],
|
|
media_item_latents,
|
|
strength,
|
|
)
|
|
)
|
|
init_conditioning_mask[
|
|
:, :f_l, l_y : l_y + h_l, l_x : l_x + w_l
|
|
] = strength
|
|
else:
|
|
# Non-first frame or sequence
|
|
if n_frames > 1:
|
|
# Handle non-first sequence.
|
|
# Encoded latents are either fully consumed, or the prefix is handled separately below.
|
|
(
|
|
init_latents,
|
|
init_conditioning_mask,
|
|
media_item_latents,
|
|
) = self._handle_non_first_conditioning_sequence(
|
|
init_latents,
|
|
init_conditioning_mask,
|
|
media_item_latents,
|
|
media_frame_number,
|
|
strength,
|
|
)
|
|
|
|
# Single frame or sequence-prefix latents
|
|
if media_item_latents is not None:
|
|
noise = randn_tensor(
|
|
media_item_latents.shape,
|
|
generator=generator,
|
|
device=media_item_latents.device,
|
|
dtype=media_item_latents.dtype,
|
|
)
|
|
|
|
media_item_latents = torch.lerp(
|
|
noise, media_item_latents, strength
|
|
)
|
|
|
|
# Patchify the extra conditioning latents and calculate their pixel coordinates
|
|
media_item_latents, latent_coords = self.patchifier.patchify(
|
|
latents=media_item_latents
|
|
)
|
|
pixel_coords = latent_to_pixel_coords(
|
|
latent_coords,
|
|
self.vae,
|
|
causal_fix=self.transformer.config.causal_temporal_positioning,
|
|
)
|
|
|
|
# Update the frame numbers to match the target frame number
|
|
pixel_coords[:, 0] += media_frame_number
|
|
extra_conditioning_num_latents += media_item_latents.shape[1]
|
|
|
|
conditioning_mask = torch.full(
|
|
media_item_latents.shape[:2],
|
|
strength,
|
|
dtype=torch.float32,
|
|
device=init_latents.device,
|
|
)
|
|
|
|
extra_conditioning_latents.append(media_item_latents)
|
|
extra_conditioning_pixel_coords.append(pixel_coords)
|
|
extra_conditioning_mask.append(conditioning_mask)
|
|
|
|
# Patchify the updated latents and calculate their pixel coordinates
|
|
init_latents, init_latent_coords = self.patchifier.patchify(
|
|
latents=init_latents
|
|
)
|
|
init_pixel_coords = latent_to_pixel_coords(
|
|
init_latent_coords,
|
|
self.vae,
|
|
causal_fix=self.transformer.config.causal_temporal_positioning,
|
|
)
|
|
|
|
if not conditioning_items:
|
|
return init_latents, init_pixel_coords, None, 0
|
|
|
|
init_conditioning_mask, _ = self.patchifier.patchify(
|
|
latents=init_conditioning_mask.unsqueeze(1)
|
|
)
|
|
init_conditioning_mask = init_conditioning_mask.squeeze(-1)
|
|
|
|
if extra_conditioning_latents:
|
|
# Stack the extra conditioning latents, pixel coordinates and mask
|
|
init_latents = torch.cat([*extra_conditioning_latents, init_latents], dim=1)
|
|
init_pixel_coords = torch.cat(
|
|
[*extra_conditioning_pixel_coords, init_pixel_coords], dim=2
|
|
)
|
|
init_conditioning_mask = torch.cat(
|
|
[*extra_conditioning_mask, init_conditioning_mask], dim=1
|
|
)
|
|
|
|
if self.transformer.use_tpu_flash_attention:
|
|
# When flash attention is used, keep the original number of tokens by removing
|
|
# tokens from the end.
|
|
init_latents = init_latents[:, :-extra_conditioning_num_latents]
|
|
init_pixel_coords = init_pixel_coords[
|
|
:, :, :-extra_conditioning_num_latents
|
|
]
|
|
init_conditioning_mask = init_conditioning_mask[
|
|
:, :-extra_conditioning_num_latents
|
|
]
|
|
|
|
return (
|
|
init_latents,
|
|
init_pixel_coords,
|
|
init_conditioning_mask,
|
|
extra_conditioning_num_latents,
|
|
)
|
|
|
|
@staticmethod
|
|
def _resize_conditioning_item(
|
|
conditioning_item: ConditioningItem,
|
|
height: int,
|
|
width: int,
|
|
):
|
|
if conditioning_item.media_x or conditioning_item.media_y:
|
|
raise ValueError(
|
|
"Provide media_item in the target size for spatial conditioning."
|
|
)
|
|
new_conditioning_item = copy.copy(conditioning_item)
|
|
new_conditioning_item.media_item = LTXVideoPipeline.resize_tensor(
|
|
conditioning_item.media_item, height, width
|
|
)
|
|
return new_conditioning_item
|
|
|
|
def _get_latent_spatial_position(
|
|
self,
|
|
latents: torch.Tensor,
|
|
conditioning_item: ConditioningItem,
|
|
height: int,
|
|
width: int,
|
|
strip_latent_border,
|
|
):
|
|
"""
|
|
Get the spatial position of the conditioning item in the latent space.
|
|
If requested, strip the conditioning latent borders that do not align with target borders.
|
|
(border latents look different then other latents and might confuse the model)
|
|
"""
|
|
scale = self.vae_scale_factor
|
|
h, w = conditioning_item.media_item.shape[-2:]
|
|
assert (
|
|
h <= height and w <= width
|
|
), f"Conditioning item size {h}x{w} is larger than target size {height}x{width}"
|
|
assert h % scale == 0 and w % scale == 0
|
|
|
|
# Compute the start and end spatial positions of the media item
|
|
x_start, y_start = conditioning_item.media_x, conditioning_item.media_y
|
|
x_start = (width - w) // 2 if x_start is None else x_start
|
|
y_start = (height - h) // 2 if y_start is None else y_start
|
|
x_end, y_end = x_start + w, y_start + h
|
|
assert (
|
|
x_end <= width and y_end <= height
|
|
), f"Conditioning item {x_start}:{x_end}x{y_start}:{y_end} is out of bounds for target size {width}x{height}"
|
|
|
|
if strip_latent_border:
|
|
# Strip one latent from left/right and/or top/bottom, update x, y accordingly
|
|
if x_start > 0:
|
|
x_start += scale
|
|
latents = latents[:, :, :, :, 1:]
|
|
|
|
if y_start > 0:
|
|
y_start += scale
|
|
latents = latents[:, :, :, 1:, :]
|
|
|
|
if x_end < width:
|
|
latents = latents[:, :, :, :, :-1]
|
|
|
|
if y_end < height:
|
|
latents = latents[:, :, :, :-1, :]
|
|
|
|
return latents, x_start // scale, y_start // scale
|
|
|
|
@staticmethod
|
|
def _handle_non_first_conditioning_sequence(
|
|
init_latents: torch.Tensor,
|
|
init_conditioning_mask: torch.Tensor,
|
|
latents: torch.Tensor,
|
|
media_frame_number: int,
|
|
strength: float,
|
|
num_prefix_latent_frames: int = 2,
|
|
prefix_latents_mode: str = "concat",
|
|
prefix_soft_conditioning_strength: float = 0.15,
|
|
):
|
|
"""
|
|
Special handling for a conditioning sequence that does not start on the first frame.
|
|
The special handling is required to allow a short encoded video to be used as middle
|
|
(or last) sequence in a longer video.
|
|
Args:
|
|
init_latents (torch.Tensor): The initial noise latents to be updated.
|
|
init_conditioning_mask (torch.Tensor): The initial conditioning mask to be updated.
|
|
latents (torch.Tensor): The encoded conditioning item.
|
|
media_frame_number (int): The target frame number of the first frame in the conditioning sequence.
|
|
strength (float): The conditioning strength for the conditioning latents.
|
|
num_prefix_latent_frames (int, optional): The length of the sequence prefix, to be handled
|
|
separately. Defaults to 2.
|
|
prefix_latents_mode (str, optional): Special treatment for prefix (boundary) latents.
|
|
- "drop": Drop the prefix latents.
|
|
- "soft": Use the prefix latents, but with soft-conditioning
|
|
- "concat": Add the prefix latents as extra tokens (like single frames)
|
|
prefix_soft_conditioning_strength (float, optional): The strength of the soft-conditioning for
|
|
the prefix latents, relevant if `prefix_latents_mode` is "soft". Defaults to 0.1.
|
|
|
|
"""
|
|
f_l = latents.shape[2]
|
|
f_l_p = num_prefix_latent_frames
|
|
assert f_l >= f_l_p
|
|
assert media_frame_number % 8 == 0
|
|
if f_l > f_l_p:
|
|
# Insert the conditioning latents **excluding the prefix** into the sequence
|
|
f_l_start = media_frame_number // 8 + f_l_p
|
|
f_l_end = f_l_start + f_l - f_l_p
|
|
init_latents[:, :, f_l_start:f_l_end] = torch.lerp(
|
|
init_latents[:, :, f_l_start:f_l_end],
|
|
latents[:, :, f_l_p:],
|
|
strength,
|
|
)
|
|
# Mark these latent frames as conditioning latents
|
|
init_conditioning_mask[:, f_l_start:f_l_end] = strength
|
|
|
|
# Handle the prefix-latents
|
|
if prefix_latents_mode == "soft":
|
|
if f_l_p > 1:
|
|
# Drop the first (single-frame) latent and soft-condition the remaining prefix
|
|
f_l_start = media_frame_number // 8 + 1
|
|
f_l_end = f_l_start + f_l_p - 1
|
|
strength = min(prefix_soft_conditioning_strength, strength)
|
|
init_latents[:, :, f_l_start:f_l_end] = torch.lerp(
|
|
init_latents[:, :, f_l_start:f_l_end],
|
|
latents[:, :, 1:f_l_p],
|
|
strength,
|
|
)
|
|
# Mark these latent frames as conditioning latents
|
|
init_conditioning_mask[:, f_l_start:f_l_end] = strength
|
|
latents = None # No more latents to handle
|
|
elif prefix_latents_mode == "drop":
|
|
# Drop the prefix latents
|
|
latents = None
|
|
elif prefix_latents_mode == "concat":
|
|
# Pass-on the prefix latents to be handled as extra conditioning frames
|
|
latents = latents[:, :, :f_l_p]
|
|
else:
|
|
raise ValueError(f"Invalid prefix_latents_mode: {prefix_latents_mode}")
|
|
return (
|
|
init_latents,
|
|
init_conditioning_mask,
|
|
latents,
|
|
)
|
|
|
|
def trim_conditioning_sequence(
|
|
self, start_frame: int, sequence_num_frames: int, target_num_frames: int
|
|
):
|
|
"""
|
|
Trim a conditioning sequence to the allowed number of frames.
|
|
|
|
Args:
|
|
start_frame (int): The target frame number of the first frame in the sequence.
|
|
sequence_num_frames (int): The number of frames in the sequence.
|
|
target_num_frames (int): The target number of frames in the generated video.
|
|
|
|
Returns:
|
|
int: updated sequence length
|
|
"""
|
|
scale_factor = self.video_scale_factor
|
|
num_frames = min(sequence_num_frames, target_num_frames - start_frame)
|
|
# Trim down to a multiple of temporal_scale_factor frames plus 1
|
|
num_frames = (num_frames - 1) // scale_factor * scale_factor + 1
|
|
return num_frames
|
|
|
|
def adain_filter_latent(
|
|
latents: torch.Tensor, reference_latents: torch.Tensor, factor=1.0
|
|
):
|
|
"""
|
|
Applies Adaptive Instance Normalization (AdaIN) to a latent tensor based on
|
|
statistics from a reference latent tensor.
|
|
|
|
Args:
|
|
latent (torch.Tensor): Input latents to normalize
|
|
reference_latent (torch.Tensor): The reference latents providing style statistics.
|
|
factor (float): Blending factor between original and transformed latent.
|
|
Range: -10.0 to 10.0, Default: 1.0
|
|
|
|
Returns:
|
|
torch.Tensor: The transformed latent tensor
|
|
"""
|
|
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
|
|
|
|
|
|
|
|
class LTXMultiScalePipeline:
|
|
@staticmethod
|
|
def batch_normalize(latents, reference, factor = 0.25):
|
|
latents_copy = latents.clone()
|
|
t = latents_copy # B x C x F x H x W
|
|
|
|
for i in range(t.size(0)): # batch
|
|
for c in range(t.size(1)): # channel
|
|
r_sd, r_mean = torch.std_mean(
|
|
reference[i, c], dim=None
|
|
) # index by original dim order
|
|
i_sd, i_mean = torch.std_mean(t[i, c], dim=None)
|
|
|
|
t[i, c] = ((t[i, c] - i_mean) / i_sd) * r_sd + r_mean
|
|
|
|
latents_copy = torch.lerp(latents, t, factor)
|
|
return latents_copy
|
|
|
|
|
|
def _upsample_latents(
|
|
self, latest_upsampler: LatentUpsampler, latents: torch.Tensor
|
|
):
|
|
# assert latents.device == latest_upsampler.device
|
|
|
|
latents = un_normalize_latents(
|
|
latents, self.vae, vae_per_channel_normalize=True
|
|
)
|
|
upsampled_latents = latest_upsampler(latents)
|
|
upsampled_latents = normalize_latents(
|
|
upsampled_latents, self.vae, vae_per_channel_normalize=True
|
|
)
|
|
return upsampled_latents
|
|
|
|
|
|
def __init__(
|
|
self, video_pipeline: LTXVideoPipeline, latent_upsampler: LatentUpsampler
|
|
):
|
|
self.video_pipeline = video_pipeline
|
|
self.vae = video_pipeline.vae
|
|
self.latent_upsampler = latent_upsampler
|
|
|
|
def __call__(
|
|
self,
|
|
downscale_factor: float,
|
|
first_pass: dict,
|
|
second_pass: dict,
|
|
*args: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
video_pipeline = self.video_pipeline
|
|
|
|
original_kwargs = kwargs.copy()
|
|
original_output_type = kwargs["output_type"]
|
|
original_width = kwargs["width"]
|
|
original_height = kwargs["height"]
|
|
|
|
x_width = int(kwargs["width"] * downscale_factor)
|
|
downscaled_width = x_width - (x_width % self.video_pipeline.vae_scale_factor)
|
|
x_height = int(kwargs["height"] * downscale_factor)
|
|
downscaled_height = x_height - (x_height % self.video_pipeline.vae_scale_factor)
|
|
trans = video_pipeline.transformer
|
|
kwargs["output_type"] = "latent"
|
|
kwargs["width"] = downscaled_width
|
|
kwargs["height"] = downscaled_height
|
|
|
|
|
|
VAE_tile_size = kwargs["VAE_tile_size"]
|
|
|
|
z_tile, hw_tile = VAE_tile_size
|
|
|
|
if z_tile > 0:
|
|
self.vae.enable_z_tiling(z_tile)
|
|
if hw_tile > 0:
|
|
self.vae.enable_hw_tiling()
|
|
self.vae.set_tiling_params(hw_tile)
|
|
|
|
ltxv_model = kwargs["ltxv_model"]
|
|
text_encoder_max_tokens = 256
|
|
prompt = kwargs.pop("prompt")
|
|
negative_prompt = kwargs.pop("negative_prompt")
|
|
if False and kwargs["enhance_prompt"]:
|
|
prompt = generate_cinematic_prompt(
|
|
video_pipeline.prompt_enhancer_image_caption_model,
|
|
video_pipeline.prompt_enhancer_image_caption_processor,
|
|
video_pipeline.prompt_enhancer_llm_model,
|
|
video_pipeline.prompt_enhancer_llm_tokenizer,
|
|
prompt,
|
|
kwargs["conditioning_items"],
|
|
max_new_tokens=text_encoder_max_tokens,
|
|
)
|
|
print("Enhanced prompt: " + prompt[0])
|
|
|
|
# Encode input prompt
|
|
|
|
(
|
|
prompt_embeds,
|
|
prompt_attention_mask,
|
|
negative_prompt_embeds,
|
|
negative_prompt_attention_mask,
|
|
) = video_pipeline.encode_prompt(
|
|
prompt,
|
|
True,
|
|
negative_prompt=negative_prompt,
|
|
device=kwargs["device"],
|
|
text_encoder_max_tokens=text_encoder_max_tokens,
|
|
)
|
|
if ltxv_model._interrupt:
|
|
return None
|
|
|
|
kwargs["prompt_embeds"] = prompt_embeds
|
|
kwargs["prompt_attention_mask"] = prompt_attention_mask
|
|
kwargs["negative_prompt_embeds"] = negative_prompt_embeds
|
|
kwargs["negative_prompt_attention_mask"] = negative_prompt_attention_mask
|
|
|
|
original_kwargs = kwargs.copy()
|
|
|
|
kwargs["joint_pass"] = True
|
|
kwargs["pass_no"] = 1
|
|
|
|
|
|
kwargs.update(**first_pass)
|
|
kwargs["num_inference_steps"] = kwargs["num_inference_steps1"]
|
|
result = video_pipeline(*args, **kwargs)
|
|
if result == None:
|
|
return None
|
|
|
|
latents = result
|
|
|
|
upsampled_latents = self._upsample_latents(self.latent_upsampler, latents)
|
|
|
|
upsampled_latents = adain_filter_latent(
|
|
latents=upsampled_latents, reference_latents=latents
|
|
)
|
|
# upsampled_latents = self.batch_normalize(upsampled_latents, latents)
|
|
|
|
kwargs = original_kwargs
|
|
kwargs["latents"] = upsampled_latents
|
|
kwargs["output_type"] = original_output_type
|
|
kwargs["width"] = downscaled_width * 2
|
|
kwargs["height"] = downscaled_height * 2
|
|
kwargs["joint_pass"] = False
|
|
kwargs["pass_no"] = 2
|
|
|
|
kwargs.update(**second_pass)
|
|
kwargs["num_inference_steps"] = kwargs["num_inference_steps2"]
|
|
|
|
result = video_pipeline(*args, **kwargs)
|
|
if result == None:
|
|
return None
|
|
if original_output_type != "latent":
|
|
num_frames = result.shape[2]
|
|
videos = rearrange(result, "b c f h w -> (b f) c h w")
|
|
|
|
videos = F.interpolate(
|
|
videos,
|
|
size=(original_height, original_width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
videos = rearrange(videos, "(b f) c h w -> b c f h w", f=num_frames)
|
|
result = videos
|
|
|
|
return result
|