nodes/MrForExample-ComfyUI-AnimateAnyone-Evolved
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Tested Multi Uncondition Steps for RCFG, No improvements visible

Browse Source
This commit is contained in:
MrForExample
2024-01-21 11:51:43 +01:00
parent c2dd4665a7
commit 3c7bd59d33
4 changed files with 791 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
_test.py
+8 -12
View File
@@ -1,13 +1,9 @@
import inspect
from diffusers.schedulers import (
DDIMScheduler,
DPMSolverMultistepScheduler,
LCMScheduler,
EulerDiscreteScheduler,
EulerAncestralDiscreteScheduler,
LMSDiscreteScheduler,
PNDMScheduler,
)
class MyClass:
def my_function(self):
self.x = 123
pass
print("clip_samples" in inspect.signature(DDIMScheduler.__init__).parameters)
mc = MyClass()
mc.my_function()
if hasattr(mc, 'z'):
print('MyClass has x')
nodes.py
+17 -38
View File
@@ -23,7 +23,7 @@ from .src.models.pose_guider import PoseGuider
from .src.models.unet_2d_condition import UNet2DConditionModel
from .src.models.unet_3d import UNet3DConditionModel
from .src.models.main_model_patcher import AnimateAnyoneModelPatcher
from .src.models.main_diffuser import AADiffusion
from .src.models.main_diffuser import AADiffusion, CFGType
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "./ComfyUI-AnimateAnyone-Evolved")
DEFAULT_CONFIG_PATH = os.path.join(ROOT_PATH, "./configs/default.yaml")
@@ -43,6 +43,7 @@ SCHEDULER_DICT = OrderedDict([
])
class Animate_Anyone_Sampler:
@classmethod
def INPUT_TYPES(s):
return {
@@ -55,7 +56,10 @@ class Animate_Anyone_Sampler:
"seed": ("INT", {"default": 999999999, "min": 0, "max": 0xffffffffffffffff}),
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
"cfg": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
"context_frames": ("INT", {"default": 12, "min": 1}),
"delta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
"cfg_type": (["rcfg_self_negative", "rcfg_one_time_negative", "full_double_pass"],),
"steps_rcfg_uncondition": ("INT", {"default": 3, "min": 1, "max": 10000}),
"context_frames": ("INT", {"default": 24, "min": 1}),
"context_stride": ("INT", {"default": 1, "min": 1}),
"context_overlap": ("INT", {"default": 4, "min": 0}),
"context_batch_size": ("INT", {"default": 1, "min": 1}),
@@ -95,7 +99,10 @@ class Animate_Anyone_Sampler:
pose_latent,
seed,
steps,
cfg,
cfg,
steps_rcfg_uncondition,
delta,
cfg_type,
context_frames,
context_stride,
context_overlap,
@@ -115,41 +122,12 @@ class Animate_Anyone_Sampler:
):
latent_format = latent_formats.SD15()
do_classifier_free_guidance = cfg > 0
# encoder_hidden_states.shape: torch.Size([1, 1, 768]) clip_image_embeds.shape: torch.Size([1, 768])
encoder_hidden_states = clip_image_embeds["image_embeds"].unsqueeze(1).to(DEVICE, dtype=WEIGHT_DETYPE)
if do_classifier_free_guidance:
uncond_encoder_hidden_states = torch.zeros_like(encoder_hidden_states)
encoder_hidden_states = torch.cat(
[uncond_encoder_hidden_states, encoder_hidden_states], dim=0
)
reference_control_writer = ReferenceAttentionControl(
reference_unet,
do_classifier_free_guidance=do_classifier_free_guidance,
mode="write",
fusion_blocks="full",
)
reference_control_reader = ReferenceAttentionControl(
denoising_unet,
do_classifier_free_guidance=do_classifier_free_guidance,
mode="read",
fusion_blocks="full",
)
# forward reference image latent with shape (1, 4, 96, 64) to reference net
ref_image_latent = latent_format.process_in(ref_image_latent["samples"]).to(DEVICE, dtype=WEIGHT_DETYPE)
reference_unet(
ref_image_latent.repeat(
(2 if do_classifier_free_guidance 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)
# setup scheduler from user inputs
scheduler_class = SCHEDULER_DICT[sampler_scheduler_pairs]
@@ -171,7 +149,7 @@ class Animate_Anyone_Sampler:
scheduler = scheduler_class(**sched_kwargs)
# setup diffuser and then denoise
diffuser = AADiffusion(denoising_unet, scheduler)
diffuser = AADiffusion(reference_unet, denoising_unet, scheduler)
if use_lora:
lora_path = comfy_paths.get_full_path("loras", lora_name)
@@ -179,9 +157,13 @@ class Animate_Anyone_Sampler:
samples = diffuser(
steps,
cfg,
cfg,
delta,
CFGType[cfg_type],
steps_rcfg_uncondition,
ref_image_latent,
pose_latent,
encoder_hidden_states,
encoder_hidden_states,
seed,
context_frames=context_frames,
context_stride=context_stride,
@@ -194,9 +176,6 @@ class Animate_Anyone_Sampler:
# (1, 4, f, h, w) -> (f, 4, h, w)
samples = torch.squeeze(samples, 0).permute(1, 0, 2, 3)
reference_control_reader.clear()
reference_control_writer.clear()
return ({"samples":samples}, )
class Load_UNet2D_ConditionModel:
src/models/main_diffuser.py
+269 -73
View File
@@ -1,6 +1,7 @@
# 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 torch
@@ -17,13 +18,19 @@ from diffusers.schedulers import (
)
from diffusers.utils.torch_utils import randn_tensor
from .mutual_self_attention import ReferenceAttentionControl
from ..utils.util import get_tensor_interpolation_method, get_context_scheduler
class CFGType(Enum):
rcfg_self_negative = 1
rcfg_one_time_negative = 2
full_double_pass = 3
class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
def __init__(
self,
unet,
reference_unet,
denoising_unet,
scheduler: Union[
DDIMScheduler,
LCMScheduler,
@@ -36,12 +43,16 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
):
super().__init__()
self.register_modules(
unet=unet,
reference_unet=reference_unet,
denoising_unet=denoising_unet,
scheduler=scheduler,
)
self.unet = unet
self.reference_unet = reference_unet
self.denoising_unet = denoising_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)
@@ -55,7 +66,7 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
self.load_lora_into_unet(
state_dict,
network_alphas=network_alphas,
unet=getattr(self, self.unet_name) if not hasattr(self, "unet") else self.unet,
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,
@@ -160,13 +171,152 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
new_index += 1
return new_latents
def perform_guidance(
self,
noise_pred,
counter,
cfg,
delta,
do_classifier_free_guidance,
do_double_pass,
do_rcfg,
noise_pred_uncond_prev=None,
):
# perform CFG/RCFG guidance
if do_classifier_free_guidance:
if do_double_pass:
noise_pred_uncond, noise_pred_cond = (noise_pred / counter).chunk(2)
elif noise_pred_uncond_prev is not None:
noise_pred_cond = noise_pred / counter
noise_pred_uncond = noise_pred_uncond_prev
else:
noise_pred_cond = noise_pred_uncond = noise_pred / counter
cfg = 0.0
# RCFG of formula (7) from https://arxiv.org/pdf/2312.12491.pdf
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.denoising_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_classifier_free_guidance,
do_double_pass,
do_rcfg,
noise_pred_uncond_prev,
device,
extra_step_kwargs,
):
noise_pred, counter = self.denoise_all_one_step(
t,
latents,
encoder_hidden_states,
pose_latent,
global_context,
do_double_pass,
device,
)
noise_pred = self.perform_guidance(
noise_pred,
counter,
cfg,
delta,
do_classifier_free_guidance,
do_double_pass,
noise_pred_uncond_prev,
)
scheduler_output = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs)
pred_original_sample = scheduler_output.pred_original_sample
# compute the residual empty/negative condition noise sample
if do_rcfg:
alpha_prod_t = self.scheduler.alphas_cumprod[t].to(device)
beta_prod_t = 1 - alpha_prod_t
# Seems following two methods both worked for prediction_type == "v_prediction"
if self.scheduler.config.prediction_type == "epsilon":
# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
# rearranged from formular: pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5)
noise_pred_uncond_prev = (latents - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
elif self.scheduler.config.prediction_type == "v_prediction":
# v_prediction: (see section 2.4 of [Imagen Video](https://imagen.research.google/video/paper.pdf)
# rearranged from formular: pred_original_sample = (alpha_prod_t ** (0.5) * latents - (beta_prod_t ** (0.5) * noise_pred
noise_pred_uncond_prev = (alpha_prod_t ** (0.5) * latents - pred_original_sample) / beta_prod_t ** (0.5)
return scheduler_output.prev_sample, noise_pred_uncond_prev
@torch.no_grad()
def __call__(
self,
steps,
cfg,
pose_fea,
delta,
cfg_type,
steps_uncond,
ref_image_latent,
pose_latent,
encoder_hidden_states,
seed = 999999999,
eta: float = 0.0,
@@ -180,27 +330,21 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
callback_steps: Optional[int] = 1,
):
device = pose_fea.device
do_classifier_free_guidance = cfg > 0
do_double_pass = do_classifier_free_guidance and cfg_type == CFGType.full_double_pass
do_rcfg = do_classifier_free_guidance and (cfg_type == CFGType.rcfg_self_negative or cfg_type == CFGType.rcfg_one_time_negative)
device = pose_latent.device
generator=torch.manual_seed(seed)
latents = self.prepare_latents(
batch_size=pose_fea.shape[0],
num_channels_latents=self.unet.config.in_channels,
video_length=pose_fea.shape[2],
height=pose_fea.shape[3],
width=pose_fea.shape[4],
dtype=encoder_hidden_states.dtype,
device=device,
generator=generator,
)
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 timesteps
self.scheduler.set_timesteps(steps, device=device)
timesteps = self.scheduler.timesteps
# Prepare Input Batches
context_scheduler = get_context_scheduler(context_schedule)
@@ -208,7 +352,7 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
context_scheduler(
0,
steps,
latents.shape[2], # number of frames
video_length,
context_frames,
context_stride,
context_overlap,
@@ -239,64 +383,113 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
[[60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]]
]
"""
do_classifier_free_guidance = cfg > 0
# Prepare initial latents
latents = self.prepare_latents(
batch_size=pose_latent.shape[0],
num_channels_latents=self.denoising_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,
)
# Prepare confitional embeds
if do_classifier_free_guidance:
uncond_encoder_hidden_states = torch.zeros_like(encoder_hidden_states)
if cfg_type == CFGType.full_double_pass:
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.denoising_unet,
do_classifier_free_guidance=do_double_pass,
mode="read",
fusion_blocks="full",
)
if do_rcfg:
# predict initial uncondition noise for RCFG
self.reference_unet(
ref_image_latent,
torch.zeros(size=(1,), dtype=int, device=device),
encoder_hidden_states=uncond_encoder_hidden_states,
return_dict=False,
)
reference_control_reader.update(reference_control_writer)
self.scheduler.set_timesteps(steps_uncond, device=device)
timesteps = self.scheduler.timesteps
with self.progress_bar(total=steps_uncond) as progress_bar:
for i, t in enumerate(timesteps):
_, self.noise_pred_uncond_prev = self.solve_one_step(
t,
latents,
uncond_encoder_hidden_states,
pose_latent,
global_context,
cfg,
delta,
do_classifier_free_guidance,
do_double_pass,
i+1 == steps_uncond, #do_rcfg,
self.noise_pred_uncond_prev,
device,
extra_step_kwargs,
)
progress_bar.update()
reference_control_writer.clear()
reference_control_reader.clear()
#if hasattr(self.scheduler, "_step_index"):
# self.scheduler._step_index = 0
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)
# Prepare timesteps
self.scheduler.set_timesteps(steps, device=device)
timesteps = self.scheduler.timesteps
# denoising loop
num_warmup_steps = len(timesteps) - steps * self.scheduler.order
with self.progress_bar(total=steps) as progress_bar:
for i, t in enumerate(timesteps):
noise_pred = torch.zeros(
(
latents.shape[0] * (2 if do_classifier_free_guidance 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,
latents, self.noise_pred_uncond_prev = self.solve_one_step(
t,
latents,
encoder_hidden_states,
pose_latent,
global_context,
cfg,
delta,
do_classifier_free_guidance,
do_double_pass,
do_rcfg,
self.noise_pred_uncond_prev,
device,
extra_step_kwargs,
)
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_classifier_free_guidance 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_fea[:, :, c] for c in context]
).repeat(2 if do_classifier_free_guidance 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
# perform guidance
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = (noise_pred / counter).chunk(2)
noise_pred = noise_pred_uncond + cfg * (
noise_pred_text - noise_pred_uncond
)
latents = self.scheduler.step(
noise_pred, t, latents, **extra_step_kwargs
).prev_sample
# update progress bar and callback
if i == len(timesteps) - 1 or (
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
@@ -308,5 +501,8 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
if interpolation_factor > 0:
latents = self.interpolate_latents(latents, interpolation_factor, device)
reference_control_reader.clear()
reference_control_writer.clear()
return latents
src/models/main_diffuser_not_working.py
+497
View File
@@ -0,0 +1,497 @@
# 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 torch
from diffusers import DiffusionPipeline
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 .mutual_self_attention import ReferenceAttentionControl
from ..utils.util import get_tensor_interpolation_method, get_context_scheduler
class CFGType(Enum):
rcfg_self_negative = 1
rcfg_one_time_negative = 2
full_double_pass = 3
class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
def __init__(
self,
reference_unet,
denoising_unet,
scheduler: Union[
DDIMScheduler,
LCMScheduler,
PNDMScheduler,
LMSDiscreteScheduler,
EulerDiscreteScheduler,
EulerAncestralDiscreteScheduler,
DPMSolverMultistepScheduler,
],
):
super().__init__()
self.register_modules(
reference_unet=reference_unet,
denoising_unet=denoising_unet,
scheduler=scheduler,
)
self.reference_unet = reference_unet
self.denoising_unet = denoising_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 calculate_residual_noise(self, t, latents, pred_original_sample):
# compute the residual empty/negative condition noise sample
alpha_prod_t = self.scheduler.alphas_cumprod[t].to(latents.device)
beta_prod_t = 1 - alpha_prod_t
if self.scheduler.config.prediction_type == "epsilon":
# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
# rearranged from formular: pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5)
noise_pred_uncond_prev = (latents - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
elif self.scheduler.config.prediction_type == "v_prediction":
# v_prediction: (see section 2.4 of [Imagen Video](https://imagen.research.google/video/paper.pdf)
# rearranged from formular: pred_original_sample = (alpha_prod_t ** (0.5) * latents - (beta_prod_t ** (0.5) * noise_pred
noise_pred_uncond_prev = (alpha_prod_t ** (0.5) * latents - pred_original_sample) / beta_prod_t ** (0.5)
return noise_pred_uncond_prev
def perform_guidance(
self,
noise_pred,
counter,
cfg,
delta,
do_double_pass,
noise_pred_uncond_prev=None,
):
# perform guidance
if do_double_pass: # CFG
noise_pred_uncond, noise_pred_cond = (noise_pred / counter).chunk(2)
elif noise_pred_uncond_prev is not None: # RCFG
noise_pred_cond = noise_pred / counter
noise_pred_uncond = noise_pred_uncond_prev
else: # Direct
noise_pred_cond = noise_pred_uncond = noise_pred / counter
cfg = 0.0
# RCFG of formula (7) from https://arxiv.org/pdf/2312.12491.pdf
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,
):
noise_pred = torch.zeros(
(
latents.shape[0] * (2 if do_double_pass else 1),
*latents.shape[1:],
),
device=latents.device,
dtype=latents.dtype,
)
counter = torch.zeros(
(1, 1, latents.shape[2], 1, 1),
device=latents.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(latents.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.denoising_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,
do_rcfg,
noise_pred_uncond_prev,
extra_step_kwargs,
):
noise_pred, counter = self.denoise_all_one_step(
t,
latents,
encoder_hidden_states,
pose_latent,
global_context,
do_double_pass,
)
noise_pred = self.perform_guidance(
noise_pred,
counter,
cfg,
delta,
do_double_pass,
noise_pred_uncond_prev,
)
scheduler_output = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs)
if do_rcfg:
noise_pred_uncond_prev = self.calculate_residual_noise(t, latents, scheduler_output.pred_original_sample)
return scheduler_output.prev_sample, noise_pred_uncond_prev
@torch.no_grad()
def __call__(
self,
steps,
cfg,
delta,
cfg_type,
steps_uncond,
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,
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
callback_steps: Optional[int] = 1,
):
do_classifier_free_guidance = cfg > 0
do_double_pass = do_classifier_free_guidance and cfg_type == CFGType.full_double_pass
do_rcfg = do_classifier_free_guidance and (cfg_type == CFGType.rcfg_self_negative or cfg_type == CFGType.rcfg_one_time_negative)
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.denoising_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,
)
# Prepare confitional embeds
if do_classifier_free_guidance:
uncond_encoder_hidden_states = torch.zeros_like(encoder_hidden_states)
if cfg_type == CFGType.full_double_pass:
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.denoising_unet,
do_classifier_free_guidance=do_double_pass,
mode="read",
fusion_blocks="full",
)
if do_rcfg:
# predict initial uncondition noise for RCFG
self.reference_unet(
ref_image_latent,
torch.zeros(size=(1,), dtype=int, device=device),
encoder_hidden_states=uncond_encoder_hidden_states,
return_dict=False,
)
reference_control_reader.update(reference_control_writer)
self.scheduler.set_timesteps(steps_uncond, device=device)
timesteps = self.scheduler.timesteps
with self.progress_bar(total=steps_uncond) as progress_bar:
for i, t in enumerate(timesteps):
_, self.noise_pred_uncond_prev = self.solve_one_step(
t,
latents,
encoder_hidden_states,
pose_latent,
global_context,
cfg,
delta,
do_double_pass,
i+1 == steps_uncond,
self.noise_pred_uncond_prev,
extra_step_kwargs,
)
progress_bar.update()
reference_control_writer.clear()
reference_control_reader.clear()
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)
# Prepare timesteps
self.scheduler.set_timesteps(steps, device=device)
timesteps = self.scheduler.timesteps
# denoising loop
num_warmup_steps = len(timesteps) - steps * self.scheduler.order
with self.progress_bar(total=steps) as progress_bar:
for i, t in enumerate(timesteps):
latents, self.noise_pred_uncond_prev = self.solve_one_step(
t,
latents,
encoder_hidden_states,
pose_latent,
global_context,
cfg,
delta,
do_double_pass,
do_rcfg,
self.noise_pred_uncond_prev,
extra_step_kwargs,
)
# update progress bar and callback
if i == len(timesteps) - 1 or (
(i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
):
progress_bar.update()
if callback is not None and i % callback_steps == 0:
step_idx = i // getattr(self.scheduler, "order", 1)
callback(step_idx, t, latents)
if interpolation_factor > 0:
latents = self.interpolate_latents(latents, interpolation_factor, device)
reference_control_reader.clear()
reference_control_writer.clear()
return latents