476 lines
16 KiB
Python
476 lines
16 KiB
Python
# Adapted from https://github.com/MooreThreads/Moore-AnimateAnyone/blob/master/src/pipelines/pipeline_pose2vid_long.py
|
|
import inspect
|
|
import math
|
|
from enum import Enum
|
|
from typing import Callable, List, Optional, Union
|
|
|
|
import comfy.utils
|
|
|
|
import torch
|
|
from diffusers import DiffusionPipeline, StableDiffusionPipeline
|
|
from diffusers.loaders import LoraLoaderMixin
|
|
from diffusers.schedulers import (
|
|
DDIMScheduler,
|
|
LCMScheduler,
|
|
DPMSolverMultistepScheduler,
|
|
EulerAncestralDiscreteScheduler,
|
|
EulerDiscreteScheduler,
|
|
LMSDiscreteScheduler,
|
|
PNDMScheduler,
|
|
)
|
|
from diffusers.utils.torch_utils import randn_tensor
|
|
from sfast.compilers.diffusion_pipeline_compiler import compile_unet, CompilationConfig
|
|
|
|
from .mutual_self_attention import ReferenceAttentionControl
|
|
from ..utils.util import get_tensor_interpolation_method, get_context_scheduler
|
|
|
|
class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
|
|
def __init__(
|
|
self,
|
|
reference_unet,
|
|
unet,
|
|
scheduler: Union[
|
|
DDIMScheduler,
|
|
LCMScheduler,
|
|
PNDMScheduler,
|
|
LMSDiscreteScheduler,
|
|
EulerDiscreteScheduler,
|
|
EulerAncestralDiscreteScheduler,
|
|
DPMSolverMultistepScheduler,
|
|
],
|
|
):
|
|
super().__init__()
|
|
self.register_modules(
|
|
reference_unet=reference_unet,
|
|
unet=unet,
|
|
scheduler=scheduler,
|
|
)
|
|
self.reference_unet = reference_unet
|
|
self.unet = unet
|
|
self.scheduler = scheduler
|
|
|
|
self.noise_pred_uncond_prev = None
|
|
|
|
def load_lora(self, pretrained_model_name_or_path_or_dict, adapter_name=None, **kwargs):
|
|
# First, ensure that the checkpoint is a compatible one and can be successfully loaded.
|
|
state_dict, network_alphas = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs)
|
|
|
|
is_correct_format = all("lora" in key for key in state_dict.keys())
|
|
if not is_correct_format:
|
|
raise ValueError("Invalid LoRA checkpoint.")
|
|
|
|
low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", True)
|
|
|
|
self.load_lora_into_unet(
|
|
state_dict,
|
|
network_alphas=network_alphas,
|
|
unet=getattr(self, self.denoising_unet_name) if not hasattr(self, "unet") else self.denoising_unet,
|
|
low_cpu_mem_usage=low_cpu_mem_usage,
|
|
adapter_name=adapter_name,
|
|
_pipeline=self,
|
|
)
|
|
|
|
def prepare_latents(
|
|
self,
|
|
batch_size,
|
|
num_channels_latents,
|
|
video_length,
|
|
height,
|
|
width,
|
|
dtype,
|
|
device,
|
|
generator,
|
|
latents=None,
|
|
):
|
|
shape = (
|
|
batch_size,
|
|
num_channels_latents,
|
|
video_length,
|
|
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."
|
|
)
|
|
|
|
if latents is None:
|
|
latents = randn_tensor(
|
|
shape, generator=generator, device=device, dtype=dtype
|
|
)
|
|
else:
|
|
latents = latents.to(device)
|
|
|
|
# scale the initial noise by the standard deviation required by the scheduler
|
|
latents = latents * self.scheduler.init_noise_sigma
|
|
return latents
|
|
|
|
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 interpolate_latents(self, latents: torch.Tensor, interpolation_factor: int, device):
|
|
if interpolation_factor < 2:
|
|
return latents
|
|
|
|
new_latents = torch.zeros(
|
|
(
|
|
latents.shape[0],
|
|
latents.shape[1],
|
|
((latents.shape[2] - 1) * interpolation_factor) + 1,
|
|
latents.shape[3],
|
|
latents.shape[4],
|
|
),
|
|
device=latents.device,
|
|
dtype=latents.dtype,
|
|
)
|
|
|
|
org_video_length = latents.shape[2]
|
|
rate = [i / interpolation_factor for i in range(interpolation_factor)][1:]
|
|
|
|
new_index = 0
|
|
|
|
v0 = None
|
|
v1 = None
|
|
|
|
for i0, i1 in zip(range(org_video_length), range(org_video_length)[1:]):
|
|
v0 = latents[:, :, i0, :, :]
|
|
v1 = latents[:, :, i1, :, :]
|
|
|
|
new_latents[:, :, new_index, :, :] = v0
|
|
new_index += 1
|
|
|
|
for f in rate:
|
|
v = get_tensor_interpolation_method()(
|
|
v0.to(device=device), v1.to(device=device), f
|
|
)
|
|
new_latents[:, :, new_index, :, :] = v.to(latents.device)
|
|
new_index += 1
|
|
|
|
new_latents[:, :, new_index, :, :] = v1
|
|
new_index += 1
|
|
|
|
return new_latents
|
|
|
|
def perform_guidance(
|
|
self,
|
|
noise_pred,
|
|
counter,
|
|
cfg,
|
|
delta,
|
|
do_double_pass,
|
|
):
|
|
|
|
# perform CFG
|
|
if do_double_pass:
|
|
noise_pred_uncond, noise_pred_cond = (noise_pred / counter).chunk(2)
|
|
else:
|
|
noise_pred_cond = noise_pred_uncond = noise_pred / counter
|
|
cfg = 0.0
|
|
noise_pred = delta * noise_pred_uncond + cfg * (noise_pred_cond - delta * noise_pred_uncond)
|
|
|
|
return noise_pred
|
|
|
|
def denoise_all_one_step(
|
|
self,
|
|
t,
|
|
latents,
|
|
encoder_hidden_states,
|
|
pose_latent,
|
|
global_context,
|
|
do_double_pass,
|
|
device,
|
|
):
|
|
noise_pred = torch.zeros(
|
|
(
|
|
latents.shape[0] * (2 if do_double_pass else 1),
|
|
*latents.shape[1:],
|
|
),
|
|
device=device,
|
|
dtype=latents.dtype,
|
|
)
|
|
counter = torch.zeros(
|
|
(1, 1, latents.shape[2], 1, 1),
|
|
device=device,
|
|
dtype=latents.dtype,
|
|
)
|
|
|
|
for context in global_context:
|
|
# 3.1 expand the latents if we are doing classifier free guidance
|
|
latent_model_input = (
|
|
torch.cat([latents[:, :, c] for c in context]) # Concatenates the context in the same batch
|
|
.to(device)
|
|
.repeat(2 if do_double_pass else 1, 1, 1, 1, 1)
|
|
)
|
|
latent_model_input = self.scheduler.scale_model_input(
|
|
latent_model_input, t
|
|
)
|
|
b, c, f, h, w = latent_model_input.shape
|
|
latent_pose_input = torch.cat(
|
|
[pose_latent[:, :, c] for c in context]
|
|
).repeat(2 if do_double_pass else 1, 1, 1, 1, 1)
|
|
pred = self.unet(
|
|
latent_model_input, # torch.Size([2, 4, 24, 96, 64])
|
|
t,
|
|
encoder_hidden_states=encoder_hidden_states[:b],
|
|
pose_cond_fea=latent_pose_input, # torch.Size([2, 320, 24, 96, 64])
|
|
return_dict=False,
|
|
)[0]
|
|
|
|
for j, c in enumerate(context):
|
|
noise_pred[:, :, c] = noise_pred[:, :, c] + pred
|
|
counter[:, :, c] = counter[:, :, c] + 1
|
|
|
|
return noise_pred, counter
|
|
|
|
def solve_one_step(
|
|
self,
|
|
t,
|
|
latents,
|
|
encoder_hidden_states,
|
|
pose_latent,
|
|
global_context,
|
|
cfg,
|
|
delta,
|
|
do_double_pass,
|
|
device,
|
|
extra_step_kwargs,
|
|
):
|
|
|
|
noise_pred, counter = self.denoise_all_one_step(
|
|
t=t,
|
|
latents=latents,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
pose_latent=pose_latent,
|
|
global_context=global_context,
|
|
do_double_pass=do_double_pass,
|
|
device=device,
|
|
)
|
|
|
|
noise_pred = self.perform_guidance(
|
|
noise_pred=noise_pred,
|
|
counter=counter,
|
|
cfg=cfg,
|
|
delta=delta,
|
|
do_double_pass=do_double_pass,
|
|
)
|
|
|
|
scheduler_output = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs)
|
|
|
|
return scheduler_output.prev_sample
|
|
|
|
def denoise_loop(
|
|
self,
|
|
ref_image_latent,
|
|
loop_steps,
|
|
latents,
|
|
encoder_hidden_states,
|
|
pose_latent,
|
|
global_context,
|
|
cfg,
|
|
delta,
|
|
do_double_pass,
|
|
interpolation_factor,
|
|
device,
|
|
extra_step_kwargs,
|
|
):
|
|
# Prepare confitional embeds
|
|
if do_double_pass:
|
|
uncond_encoder_hidden_states = torch.zeros_like(encoder_hidden_states)
|
|
encoder_hidden_states=torch.cat(
|
|
[uncond_encoder_hidden_states, encoder_hidden_states], dim=0
|
|
)
|
|
|
|
# Prepare reference attention control
|
|
reference_control_writer = ReferenceAttentionControl(
|
|
self.reference_unet,
|
|
do_classifier_free_guidance=do_double_pass,
|
|
mode="write",
|
|
fusion_blocks="full",
|
|
)
|
|
reference_control_reader = ReferenceAttentionControl(
|
|
self.unet,
|
|
do_classifier_free_guidance=do_double_pass,
|
|
mode="read",
|
|
fusion_blocks="full",
|
|
)
|
|
|
|
self.reference_unet(
|
|
ref_image_latent.repeat(2 if do_double_pass else 1, 1, 1, 1),
|
|
torch.zeros(size=(1,), dtype=int, device=device),
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
return_dict=False,
|
|
)
|
|
reference_control_reader.update(reference_control_writer)
|
|
|
|
config = CompilationConfig.Default()
|
|
# xformers and Triton are suggested for achieving best performance.
|
|
try:
|
|
import xformers
|
|
config.enable_xformers = True
|
|
except ImportError:
|
|
print('xformers not installed, skip')
|
|
try:
|
|
import triton
|
|
config.enable_triton = True
|
|
except ImportError:
|
|
print('Triton not installed, skip')
|
|
|
|
# CUDA Graph is suggested for small batch sizes and small resolutions to reduce CPU overhead.
|
|
# But it can increase the amount of GPU memory used.
|
|
# For StableVideoDiffusionPipeline it is not needed.
|
|
config.enable_cuda_graph = True
|
|
|
|
self.unet = compile_unet(self.unet, config)
|
|
|
|
# Prepare timesteps
|
|
self.scheduler.set_timesteps(loop_steps, device=device)
|
|
timesteps = self.scheduler.timesteps
|
|
|
|
comfy_pbar = comfy.utils.ProgressBar(loop_steps)
|
|
|
|
with self.progress_bar(total=loop_steps) as progress_bar:
|
|
for i, t in enumerate(timesteps):
|
|
|
|
latents = self.solve_one_step(
|
|
t=t,
|
|
latents=latents,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
pose_latent=pose_latent,
|
|
global_context=global_context,
|
|
cfg=cfg,
|
|
delta=delta,
|
|
do_double_pass=do_double_pass,
|
|
device=device,
|
|
extra_step_kwargs=extra_step_kwargs,
|
|
)
|
|
|
|
progress_bar.update()
|
|
comfy_pbar.update_absolute(i + 1)
|
|
|
|
if interpolation_factor > 0:
|
|
latents = self.interpolate_latents(latents, interpolation_factor, device)
|
|
|
|
reference_control_reader.clear()
|
|
reference_control_writer.clear()
|
|
|
|
return latents
|
|
|
|
@torch.no_grad()
|
|
def __call__(
|
|
self,
|
|
steps,
|
|
cfg,
|
|
delta,
|
|
ref_image_latent,
|
|
pose_latent,
|
|
encoder_hidden_states,
|
|
seed = 999999999,
|
|
eta: float = 0.0,
|
|
context_schedule="uniform",
|
|
context_frames=24,
|
|
context_stride=1,
|
|
context_overlap=4,
|
|
context_batch_size=1,
|
|
interpolation_factor=1,
|
|
):
|
|
|
|
do_double_pass = cfg > 0
|
|
|
|
device = pose_latent.device
|
|
generator=torch.manual_seed(seed)
|
|
video_length = pose_latent.shape[2] # number of frames
|
|
|
|
self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
|
|
if hasattr(self.scheduler, "final_alpha_cumprod"):
|
|
self.scheduler.final_alpha_cumprod = self.scheduler.final_alpha_cumprod.to(device)
|
|
|
|
# Prepare extra step kwargs.
|
|
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
|
|
|
# Prepare Input Batches
|
|
context_scheduler = get_context_scheduler(context_schedule)
|
|
|
|
context_queue = list(
|
|
context_scheduler(
|
|
0,
|
|
steps,
|
|
video_length,
|
|
context_frames,
|
|
context_stride,
|
|
context_overlap,
|
|
)
|
|
)
|
|
|
|
num_context_batches = math.ceil(len(context_queue) / context_batch_size)
|
|
global_context = []
|
|
for i in range(num_context_batches):
|
|
global_context.append(
|
|
context_queue[
|
|
i * context_batch_size : (i + 1) * context_batch_size
|
|
]
|
|
)
|
|
|
|
"""
|
|
Example: num_frames == 72, context_frames == 24, context_stride == 1, context_overlap == 4, context_batch_size == 1
|
|
context_queue: [
|
|
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23],
|
|
[20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43],
|
|
[40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63],
|
|
[60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
|
|
]
|
|
global_context: [ # shape: (num_context_batches, context_batch_size, context_frames)
|
|
[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]],
|
|
[[20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43]],
|
|
[[40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63]],
|
|
[[60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]]
|
|
]
|
|
"""
|
|
|
|
# Prepare initial latents
|
|
latents = self.prepare_latents(
|
|
batch_size=pose_latent.shape[0],
|
|
num_channels_latents=self.unet.config.in_channels,
|
|
video_length=video_length,
|
|
height=pose_latent.shape[-2],
|
|
width=pose_latent.shape[-1],
|
|
dtype=encoder_hidden_states.dtype,
|
|
device=device,
|
|
generator=generator,
|
|
)
|
|
|
|
# sequential double pass
|
|
latents = self.denoise_loop(
|
|
ref_image_latent=ref_image_latent,
|
|
loop_steps=steps,
|
|
latents=latents,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
pose_latent=pose_latent,
|
|
global_context=global_context,
|
|
cfg=cfg,
|
|
delta=delta,
|
|
do_double_pass=do_double_pass,
|
|
interpolation_factor=interpolation_factor,
|
|
device=device,
|
|
extra_step_kwargs=extra_step_kwargs,
|
|
)
|
|
|
|
return latents |