mirror of
https://github.com/storytold/FineTrainers-Conditioning.git
synced 2026-10-09 00:09:45 +00:00
CrossProduct
373 lines
16 KiB
Python
373 lines
16 KiB
Python
from typing import Dict, List, Optional, Union
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from accelerate.logging import get_logger
|
|
from diffusers import AutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel
|
|
from finetrainers.conditioning.LTXVideoConditionedTransformer3DModel import LTXVideoConditionedTransformer3DModel
|
|
from PIL import Image
|
|
from transformers import T5EncoderModel,T5TokenizerFast
|
|
|
|
from diffusers.pipelines.ltx.pipeline_ltx import PipelineCallback, MultiPipelineCallbacks,retrieve_timesteps, calculate_shift
|
|
from typing import Callable, Any
|
|
import numpy as np
|
|
from diffusers.utils import is_torch_xla_available
|
|
from diffusers.pipelines.ltx.pipeline_output import LTXPipelineOutput
|
|
from diffusers.utils import is_torch_xla_available
|
|
from torch import torch
|
|
from transformers import T5EncoderModel, T5TokenizerFast
|
|
from typing import Any, Callable, Dict, List, Optional, Union
|
|
|
|
from finetrainers.conditioning import condition_latents_prepare,post_conditioned_latent_patchify
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
|
|
if is_torch_xla_available():
|
|
import torch_xla.core.xla_model as xm
|
|
|
|
XLA_AVAILABLE = True
|
|
else:
|
|
XLA_AVAILABLE = False
|
|
|
|
class LTXConditionedPipeline(LTXPipeline):
|
|
def __init__(self,
|
|
scheduler: FlowMatchEulerDiscreteScheduler,
|
|
vae: AutoencoderKLLTXVideo,
|
|
text_encoder: T5EncoderModel,
|
|
tokenizer: T5TokenizerFast,
|
|
transformer: LTXVideoTransformer3DModel,
|
|
):
|
|
super().__init__(scheduler, vae, text_encoder, tokenizer, transformer)
|
|
|
|
def noise_condition_latent_prepare(self,
|
|
batch_size: int = 1,
|
|
num_channels_latents: int = 128,
|
|
height: int = 512,
|
|
width: int = 704,
|
|
num_frames: int = 161,
|
|
dtype: Optional[torch.dtype] = None,
|
|
device: Optional[torch.device] = None,
|
|
generator: Optional[torch.Generator] = None,
|
|
latents: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
if latents is not None:
|
|
return latents.to(device=device, dtype=dtype)
|
|
|
|
height = height // self.vae_spatial_compression_ratio
|
|
width = width // self.vae_spatial_compression_ratio
|
|
num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
|
|
|
|
shape = (batch_size, num_channels_latents, num_frames, height, width)
|
|
|
|
if isinstance(generator, list) and 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."
|
|
)
|
|
|
|
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
|
|
|
|
return latents
|
|
|
|
def __call__(
|
|
self,
|
|
prompt: Union[str, List[str]] = None,
|
|
negative_prompt: Optional[Union[str, List[str]]] = None,
|
|
height: int = 512,
|
|
width: int = 704,
|
|
num_frames: int = 161,
|
|
frame_rate: int = 25,
|
|
num_inference_steps: int = 50,
|
|
timesteps: List[int] = None,
|
|
guidance_scale: float = 3,
|
|
num_videos_per_prompt: Optional[int] = 1,
|
|
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
|
latents: Optional[torch.Tensor] = None,
|
|
prompt_embeds: Optional[torch.Tensor] = None,
|
|
prompt_attention_mask: Optional[torch.Tensor] = None,
|
|
negative_prompt_embeds: Optional[torch.Tensor] = None,
|
|
negative_prompt_attention_mask: Optional[torch.Tensor] = None,
|
|
decode_timestep: Union[float, List[float]] = 0.0,
|
|
decode_noise_scale: Optional[Union[float, List[float]]] = None,
|
|
output_type: Optional[str] = "pil",
|
|
return_dict: bool = True,
|
|
attention_kwargs: Optional[Dict[str, Any]] = None,
|
|
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
|
|
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
|
|
max_sequence_length: int = 128,
|
|
pose_video=None,
|
|
img_ref_video=None,
|
|
):
|
|
# These are the raw videos
|
|
if pose_video is None:
|
|
raise ValueError("pose_video cannot be None.")
|
|
if img_ref_video is None:
|
|
raise ValueError("image_ref_video cannot be None.")
|
|
|
|
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
|
|
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
|
|
|
|
# 1. Check inputs. Raise error if not correct
|
|
self.check_inputs(
|
|
prompt=prompt,
|
|
height=height,
|
|
width=width,
|
|
callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
|
|
prompt_embeds=prompt_embeds,
|
|
negative_prompt_embeds=negative_prompt_embeds,
|
|
prompt_attention_mask=prompt_attention_mask,
|
|
negative_prompt_attention_mask=negative_prompt_attention_mask,
|
|
)
|
|
|
|
self._guidance_scale = guidance_scale
|
|
self._attention_kwargs = attention_kwargs
|
|
self._interrupt = False
|
|
|
|
# 2. Define call parameters
|
|
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
|
|
|
|
# 3. Prepare text embeddings
|
|
(
|
|
prompt_embeds,
|
|
prompt_attention_mask,
|
|
negative_prompt_embeds,
|
|
negative_prompt_attention_mask,
|
|
) = self.encode_prompt(
|
|
prompt=prompt,
|
|
negative_prompt=negative_prompt,
|
|
do_classifier_free_guidance=self.do_classifier_free_guidance,
|
|
num_videos_per_prompt=num_videos_per_prompt,
|
|
prompt_embeds=prompt_embeds,
|
|
negative_prompt_embeds=negative_prompt_embeds,
|
|
prompt_attention_mask=prompt_attention_mask,
|
|
negative_prompt_attention_mask=negative_prompt_attention_mask,
|
|
max_sequence_length=max_sequence_length,
|
|
device=device,
|
|
)
|
|
if self.do_classifier_free_guidance:
|
|
prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
|
|
prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0)
|
|
|
|
# 4. Prepare latent variables
|
|
# it needs to be the size of the image
|
|
|
|
num_channels_latents = self.transformer.config.in_channels
|
|
|
|
# # TODO this is the noise latent patchified.
|
|
# # we have to use the image size by default.
|
|
# latents = self.prepare_latents(
|
|
# batch_size * num_videos_per_prompt,
|
|
# num_channels_latents,
|
|
# height,
|
|
# width,
|
|
# num_frames,
|
|
# torch.float32,
|
|
# device,
|
|
# generator,
|
|
# latents,
|
|
# dtype=self.text_encoder.dtype
|
|
# ) # creates noise tensor the size suggested by the user
|
|
|
|
noise_latent = self.noise_condition_latent_prepare(
|
|
batch_size * num_videos_per_prompt,
|
|
num_channels_latents,
|
|
height,
|
|
width,
|
|
num_frames,
|
|
torch.bfloat16,
|
|
device,
|
|
generator,
|
|
)
|
|
|
|
# 4. create conditioning latents from the video.
|
|
pose_latent = condition_latents_prepare.prepare_latents_for_conditioning(
|
|
vae=self.vae,
|
|
image_or_video=pose_video,
|
|
patch_size=self.transformer.config.patch_size,
|
|
patch_size_t=self.transformer.config.patch_size_t,
|
|
device=device,
|
|
dtype=torch.bfloat16,
|
|
generator=generator,
|
|
)["latents"]
|
|
|
|
img_ref_latent = condition_latents_prepare.prepare_latents_for_conditioning(
|
|
vae=self.vae,
|
|
image_or_video=img_ref_video,
|
|
patch_size=self.transformer.config.patch_size,
|
|
patch_size_t=self.transformer.config.patch_size_t,
|
|
device=device,
|
|
dtype=torch.bfloat16,
|
|
generator=generator,
|
|
)["latents"]
|
|
|
|
# pose template noisey input [cat] img_ref + pose video
|
|
# pose template ref video latent + patchify video latent.
|
|
# img ref patchify video latent +
|
|
# add them together as input
|
|
# residual x latent
|
|
|
|
# add pose information to both channels
|
|
noisy_latents = noise_latent # + pose_latent
|
|
pose_img_ref_latents = img_ref_latent + pose_latent
|
|
|
|
# expand channel information # B x 2C latent will be projected to adapter to scale it back to 128d using adapter.
|
|
condition_latent = torch.cat([pose_img_ref_latents,noisy_latents], dim=1)
|
|
|
|
noisy_latent_tokens = post_conditioned_latent_patchify(latents=noisy_latents, num_frames=num_frames,
|
|
height=height,
|
|
width=width,
|
|
patch_size = 1,
|
|
patch_size_t = 1)["latents"]
|
|
|
|
condition_tokens = post_conditioned_latent_patchify(latents=condition_latent,
|
|
num_frames=num_frames,
|
|
height=height,
|
|
width=width,
|
|
patch_size = 1,
|
|
patch_size_t = 1)["latents"]
|
|
# target video as a input residual
|
|
noisy_residual_tokens = post_conditioned_latent_patchify(latents=noisy_latents,
|
|
num_frames=num_frames,
|
|
height=height,
|
|
width=width,
|
|
patch_size = 1,
|
|
patch_size_t = 1)["latents"]
|
|
|
|
# Need to change the latents to
|
|
# 5. Prepare timesteps
|
|
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
|
|
video_sequence_length = latent_num_frames * latent_height * latent_width
|
|
sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
|
|
mu = calculate_shift(
|
|
video_sequence_length,
|
|
self.scheduler.config.base_image_seq_len,
|
|
self.scheduler.config.max_image_seq_len,
|
|
self.scheduler.config.base_shift,
|
|
self.scheduler.config.max_shift,
|
|
)
|
|
timesteps, num_inference_steps = retrieve_timesteps(
|
|
self.scheduler,
|
|
num_inference_steps,
|
|
device,
|
|
timesteps,
|
|
sigmas=sigmas,
|
|
mu=mu,
|
|
)
|
|
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
|
|
self._num_timesteps = len(timesteps)
|
|
|
|
# 6. Prepare micro-conditions
|
|
latent_frame_rate = frame_rate / self.vae_temporal_compression_ratio
|
|
rope_interpolation_scale = (
|
|
1 / latent_frame_rate,
|
|
self.vae_spatial_compression_ratio,
|
|
self.vae_spatial_compression_ratio,
|
|
)
|
|
|
|
# 7. Denoising loop
|
|
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
|
for i, t in enumerate(timesteps):
|
|
if self.interrupt:
|
|
continue
|
|
# change this
|
|
# latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
|
|
# latent_model_input = latent_model_input.to(prompt_embeds.dtype)
|
|
|
|
# might be the bug ..
|
|
latent_model_input = noisy_latent_tokens
|
|
latent_model_input = latent_model_input.to(prompt_embeds.dtype)
|
|
|
|
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
|
timestep = t.expand(latent_model_input.shape[0])
|
|
|
|
|
|
noise_pred = self.transformer(
|
|
hidden_states=condition_tokens,
|
|
encoder_hidden_states=prompt_embeds,
|
|
timestep=timestep,
|
|
encoder_attention_mask=prompt_attention_mask,
|
|
num_frames=latent_num_frames,
|
|
height=latent_height,
|
|
width=latent_width,
|
|
rope_interpolation_scale=rope_interpolation_scale,
|
|
attention_kwargs=attention_kwargs,
|
|
return_dict=False,
|
|
residual_x=noisy_residual_tokens,
|
|
)[0]
|
|
|
|
noise_pred = noise_pred.float()
|
|
|
|
if self.do_classifier_free_guidance:
|
|
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
|
noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
|
|
|
|
# compute the previous noisy sample x_t -> x_t-1
|
|
noisy_latents = self.scheduler.step(noise_pred, t, noisy_latents, return_dict=False)[0]
|
|
|
|
if callback_on_step_end is not None:
|
|
callback_kwargs = {}
|
|
for k in callback_on_step_end_tensor_inputs:
|
|
callback_kwargs[k] = locals()[k]
|
|
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
|
|
|
|
noisy_latents = callback_outputs.pop("latents", noisy_latents)
|
|
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
|
|
|
|
# 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 XLA_AVAILABLE:
|
|
xm.mark_step()
|
|
|
|
if output_type == "latent":
|
|
video = latents
|
|
else:
|
|
noisy_latents = self._unpack_latents(
|
|
noisy_latents,
|
|
latent_num_frames,
|
|
latent_height,
|
|
latent_width,
|
|
self.transformer_spatial_patch_size,
|
|
self.transformer_temporal_patch_size,
|
|
)
|
|
noisy_latents = self._denormalize_latents(
|
|
noisy_latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
|
|
)
|
|
noisy_latents = noisy_latents.to(prompt_embeds.dtype)
|
|
|
|
if not self.vae.config.timestep_conditioning:
|
|
timestep = None
|
|
else:
|
|
noise = torch.randn(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)
|
|
|
|
# Offload all models
|
|
self.maybe_free_model_hooks()
|
|
|
|
if not return_dict:
|
|
return (video,)
|
|
|
|
return LTXPipelineOutput(frames=video) |