From 3c7bd59d33b2fb488a2495de9538f9429fb607e9 Mon Sep 17 00:00:00 2001 From: MrForExample <62230687+MrForExample@users.noreply.github.com> Date: Sun, 21 Jan 2024 11:51:43 +0100 Subject: [PATCH 1/4] Tested Multi Uncondition Steps for RCFG, No improvements visible --- _test.py | 20 +- nodes.py | 55 +-- src/models/main_diffuser.py | 342 ++++++++++++---- src/models/main_diffuser_not_working.py | 497 ++++++++++++++++++++++++ 4 files changed, 791 insertions(+), 123 deletions(-) create mode 100644 src/models/main_diffuser_not_working.py diff --git a/_test.py b/_test.py index be444e6..9b8e6f0 100644 --- a/_test.py +++ b/_test.py @@ -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) - \ No newline at end of file +mc = MyClass() +mc.my_function() +if hasattr(mc, 'z'): + print('MyClass has x') \ No newline at end of file diff --git a/nodes.py b/nodes.py index d959b44..092bd79 100644 --- a/nodes.py +++ b/nodes.py @@ -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: diff --git a/src/models/main_diffuser.py b/src/models/main_diffuser.py index 018aeee..d43451c 100644 --- a/src/models/main_diffuser.py +++ b/src/models/main_diffuser.py @@ -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 \ No newline at end of file diff --git a/src/models/main_diffuser_not_working.py b/src/models/main_diffuser_not_working.py new file mode 100644 index 0000000..92a0c9c --- /dev/null +++ b/src/models/main_diffuser_not_working.py @@ -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 \ No newline at end of file From 34dcb9ecb8ed07877568def04ac8d7eef967e4a0 Mon Sep 17 00:00:00 2001 From: MrForExample <62230687+MrForExample@users.noreply.github.com> Date: Sun, 21 Jan 2024 14:40:42 +0100 Subject: [PATCH 2/4] Tested direct sample first (20 steps) then double pass (5 steps), result lack of details --- nodes.py | 6 +- src/models/main_diffuser.py | 274 +++++++-------- src/models/main_diffuser_worked.py | 513 +++++++++++++++++++++++++++++ 3 files changed, 653 insertions(+), 140 deletions(-) create mode 100644 src/models/main_diffuser_worked.py diff --git a/nodes.py b/nodes.py index 092bd79..b0fe065 100644 --- a/nodes.py +++ b/nodes.py @@ -58,7 +58,7 @@ class Animate_Anyone_Sampler: "cfg": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}), "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}), + "steps_rcfg": ("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}), @@ -100,7 +100,7 @@ class Animate_Anyone_Sampler: seed, steps, cfg, - steps_rcfg_uncondition, + steps_rcfg, delta, cfg_type, context_frames, @@ -160,7 +160,7 @@ class Animate_Anyone_Sampler: cfg, delta, CFGType[cfg_type], - steps_rcfg_uncondition, + steps_rcfg, ref_image_latent, pose_latent, encoder_hidden_states, diff --git a/src/models/main_diffuser.py b/src/models/main_diffuser.py index d43451c..3e99369 100644 --- a/src/models/main_diffuser.py +++ b/src/models/main_diffuser.py @@ -172,30 +172,44 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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_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) + 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 @@ -260,7 +274,6 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): global_context, cfg, delta, - do_classifier_free_guidance, do_double_pass, do_rcfg, noise_pred_uncond_prev, @@ -282,30 +295,93 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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) + noise_pred_uncond_prev = self.calculate_residual_noise(t, latents, scheduler_output.pred_original_sample) return scheduler_output.prev_sample, noise_pred_uncond_prev + + def denoise_loop( + self, + ref_image_latent, + loop_steps, + latents, + encoder_hidden_states, + pose_latent, + global_context, + cfg, + delta, + do_double_pass, + do_rcfg, + interpolation_factor, + device, + extra_step_kwargs, + ): + # Prepare confitional embeds + uncond_encoder_hidden_states = torch.zeros_like(encoder_hidden_states) + if do_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", + ) + + 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(loop_steps, device=device) + timesteps = self.scheduler.timesteps + + with self.progress_bar(total=loop_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, + device, + extra_step_kwargs, + ) + + progress_bar.update() + + 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__( @@ -314,7 +390,7 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): cfg, delta, cfg_type, - steps_uncond, + steps_rcfg, ref_image_latent, pose_latent, encoder_hidden_states, @@ -326,8 +402,6 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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 @@ -396,113 +470,39 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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), + # single pass denoising loop + + latents = self.denoise_loop( + ref_image_latent=ref_image_latent, + loop_steps=steps_rcfg, + latents=latents, encoder_hidden_states=encoder_hidden_states, - return_dict=False, + pose_latent=pose_latent, + global_context=global_context, + cfg=cfg, + delta=delta, + do_double_pass=False, + do_rcfg=False, + interpolation_factor=interpolation_factor, + device=device, + extra_step_kwargs=extra_step_kwargs, ) - 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_classifier_free_guidance, - do_double_pass, - do_rcfg, - self.noise_pred_uncond_prev, - device, - 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() + # double pass fine-tune + latents = self.denoise_loop( + ref_image_latent=ref_image_latent, + loop_steps=steps-steps_rcfg, + latents=latents, + encoder_hidden_states=encoder_hidden_states, + pose_latent=pose_latent, + global_context=global_context, + cfg=cfg, + delta=delta, + do_double_pass=True, + do_rcfg=False, + interpolation_factor=interpolation_factor, + device=device, + extra_step_kwargs=extra_step_kwargs, + ) return latents \ No newline at end of file diff --git a/src/models/main_diffuser_worked.py b/src/models/main_diffuser_worked.py new file mode 100644 index 0000000..db028f4 --- /dev/null +++ b/src/models/main_diffuser_worked.py @@ -0,0 +1,513 @@ +# 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 perform_guidance( + self, + noise_pred, + counter, + cfg, + delta, + do_classifier_free_guidance, + do_double_pass, + 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) + + # compute the residual empty/negative condition noise sample + if do_rcfg: + pred_original_sample = scheduler_output.pred_original_sample + + alpha_prod_t = self.scheduler.alphas_cumprod[t].to(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 scheduler_output.prev_sample, noise_pred_uncond_prev + + @torch.no_grad() + def __call__( + self, + steps, + cfg, + delta, + cfg_type, + steps_rcfg, + 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_rcfg, device=device) + timesteps = self.scheduler.timesteps + + with self.progress_bar(total=steps_rcfg) 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_classifier_free_guidance, + do_double_pass, + i+1 == steps_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): + + 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, + ) + + # 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() + + + + # double pass fine-tune + + + return latents \ No newline at end of file From 27faf22160569622ca343bedad1c89574c02c758 Mon Sep 17 00:00:00 2001 From: MrForExample <62230687+MrForExample@users.noreply.github.com> Date: Sun, 21 Jan 2024 15:33:38 +0100 Subject: [PATCH 3/4] Tested sequential double pass, better but still does not produce ideal result Using two denoising_unet forward pass, one with image embeds condition one with zero condition, instead of double batch forward with both conditions (double pass). Result are still 2 times faster than original double pass, but apparently models architecture using attention mechanism cross both batch, that's why the result is not ideal --- src/models/main_diffuser.py | 101 +++++++++++++----------------------- 1 file changed, 35 insertions(+), 66 deletions(-) diff --git a/src/models/main_diffuser.py b/src/models/main_diffuser.py index 3e99369..7d4222d 100644 --- a/src/models/main_diffuser.py +++ b/src/models/main_diffuser.py @@ -189,43 +189,27 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): return noise_pred_uncond_prev - def perform_guidance( - self, - noise_pred, - counter, - cfg, - delta, - do_double_pass, - noise_pred_uncond_prev=None, - ): - - # perform CFG/RCFG 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, + uncond_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[0], + *latents.shape[1:], + ), + device=device, + dtype=latents.dtype, + ) + noise_pred_uncond = torch.zeros( + ( + latents.shape[0], *latents.shape[1:], ), device=device, @@ -242,7 +226,6 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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 @@ -250,7 +233,7 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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, @@ -258,46 +241,51 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): pose_cond_fea=latent_pose_input, # torch.Size([2, 320, 24, 96, 64]) return_dict=False, )[0] + pred_uncond = self.denoising_unet( + latent_model_input, # torch.Size([2, 4, 24, 96, 64]) + t, + encoder_hidden_states=uncond_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 + noise_pred_uncond[:, :, c] = noise_pred_uncond[:, :, c] + pred_uncond counter[:, :, c] = counter[:, :, c] + 1 - return noise_pred, counter + return noise_pred, noise_pred_uncond, counter def solve_one_step( self, t, latents, encoder_hidden_states, + uncond_encoder_hidden_states, pose_latent, global_context, cfg, delta, - do_double_pass, do_rcfg, noise_pred_uncond_prev, device, extra_step_kwargs, ): - noise_pred, counter = self.denoise_all_one_step( + noise_pred, noise_pred_uncond, counter = self.denoise_all_one_step( t, latents, encoder_hidden_states, + uncond_encoder_hidden_states, pose_latent, global_context, - do_double_pass, device, ) - noise_pred = self.perform_guidance( - noise_pred, - counter, - cfg, - delta, - do_double_pass, - noise_pred_uncond_prev, - ) + + # CFG + noise_pred_cond = noise_pred / counter + noise_pred_uncond = noise_pred_uncond / counter + noise_pred = delta * noise_pred_uncond + cfg * (noise_pred_cond - delta * noise_pred_uncond) scheduler_output = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs) @@ -324,10 +312,6 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): ): # Prepare confitional embeds uncond_encoder_hidden_states = torch.zeros_like(encoder_hidden_states) - if do_double_pass: - encoder_hidden_states = torch.cat( - [uncond_encoder_hidden_states, encoder_hidden_states], dim=0 - ) # Prepare reference attention control reference_control_writer = ReferenceAttentionControl( @@ -346,7 +330,10 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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, + #encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states=torch.cat( + [uncond_encoder_hidden_states, encoder_hidden_states], dim=0 + ) if do_double_pass else encoder_hidden_states, return_dict=False, ) reference_control_reader.update(reference_control_writer) @@ -362,11 +349,11 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): t, latents, encoder_hidden_states, + uncond_encoder_hidden_states, pose_latent, global_context, cfg, delta, - do_double_pass, do_rcfg, self.noise_pred_uncond_prev, device, @@ -470,11 +457,10 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): generator=generator, ) - # single pass denoising loop - + # sequential double pass latents = self.denoise_loop( ref_image_latent=ref_image_latent, - loop_steps=steps_rcfg, + loop_steps=steps, latents=latents, encoder_hidden_states=encoder_hidden_states, pose_latent=pose_latent, @@ -487,22 +473,5 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): device=device, extra_step_kwargs=extra_step_kwargs, ) - - # double pass fine-tune - latents = self.denoise_loop( - ref_image_latent=ref_image_latent, - loop_steps=steps-steps_rcfg, - latents=latents, - encoder_hidden_states=encoder_hidden_states, - pose_latent=pose_latent, - global_context=global_context, - cfg=cfg, - delta=delta, - do_double_pass=True, - do_rcfg=False, - interpolation_factor=interpolation_factor, - device=device, - extra_step_kwargs=extra_step_kwargs, - ) return latents \ No newline at end of file From 5c7a6a3f9b2f490163aefff0b0223e1518a22373 Mon Sep 17 00:00:00 2001 From: MrForExample <62230687+MrForExample@users.noreply.github.com> Date: Sun, 21 Jan 2024 17:56:58 +0100 Subject: [PATCH 4/4] Update After RCFG experiment - Update readme to log the RCFG experiment - Refactor code - Add delta parameter on Animate_Anyone_Sampler node - Add ComfyUI Sampler progress bar --- README.md | 15 +- .../MooreAnimateAnyone_DDIM.json | 541 +++++++++--------- _test.py | 20 +- nodes.py | 10 +- src/models/main_diffuser.py | 162 +++--- src/models/main_model_patcher.py | 13 - 6 files changed, 363 insertions(+), 398 deletions(-) delete mode 100644 src/models/main_model_patcher.py diff --git a/README.md b/README.md index 0d0fc21..eae3023 100644 --- a/README.md +++ b/README.md @@ -35,12 +35,15 @@ - I tried to break it down into as many modules as possible, so the workflow in ComfyUI would closely resemble the original pipeline from AnimateAnyone paper:
![_Example_Workflow\_Other_Imgs\AA_pipeline.png](_Example_Workflow/_Other_Imgs/AA_pipeline.png) -## Will Do Next -- Train a LCM Lora for denoise unet (**Estimated speed up: 5X**) -- Convert Model using [stable-fast](https://github.com/chengzeyi/stable-fast) (**Estimated speed up: 2X**) -- Implement the compoents (Residual CFG) proposed in [StreamDiffusion](https://github.com/cumulo-autumn/StreamDiffusion?tab=readme-ov-file) (**Estimated speed up: 2X**) -- Incorporate the implementation & Pre-trained Models from [Open-AnimateAnyone](https://github.com/guoqincode/Open-AnimateAnyone) & [AnimateAnyone](https://github.com/HumanAIGC/AnimateAnyone) once they released -- Training a new Model using better dataset to improve results quality (Optional, we'll see if there is any need for me to do it ;) +## Roadmap +- [x] Implement the compoents (Residual CFG) proposed in [StreamDiffusion](https://github.com/cumulo-autumn/StreamDiffusion?tab=readme-ov-file) (**Estimated speed up: 2X**) + - **Result:** + Generated result is not good enough when using DDIM Scheduler togather with RCFG, even though it speed up the generating process by about 4X.
+ In StreamDiffusion, RCFG works with LCM, could also be the case here, so keep it in another branch for now. +- [ ] Convert Model using [stable-fast](https://github.com/chengzeyi/stable-fast) (**Estimated speed up: 2X**) +- [ ] Train a LCM Lora for denoise unet (**Estimated speed up: 5X**) +- [ ] Incorporate the implementation & Pre-trained Models from [Open-AnimateAnyone](https://github.com/guoqincode/Open-AnimateAnyone) & [AnimateAnyone](https://github.com/HumanAIGC/AnimateAnyone) once they released +- [ ] Training a new Model using better dataset to improve results quality (Optional, we'll see if there is any need for me to do it ;) - Continuous research, always moving towards something better & faster🚀 ## Install (Will add it to ComfyUI Manager Soon!) diff --git a/_Example_Workflow/MooreAnimateAnyone_DDIM.json b/_Example_Workflow/MooreAnimateAnyone_DDIM.json index ef5c353..b52a618 100644 --- a/_Example_Workflow/MooreAnimateAnyone_DDIM.json +++ b/_Example_Workflow/MooreAnimateAnyone_DDIM.json @@ -1,6 +1,6 @@ { - "last_node_id": 26, - "last_link_id": 32, + "last_node_id": 31, + "last_link_id": 56, "nodes": [ { "id": 15, @@ -56,7 +56,7 @@ { "name": "samples", "type": "LATENT", - "link": 23 + "link": 56 }, { "name": "vae", @@ -79,121 +79,6 @@ "Node name for S&R": "VAEDecode" } }, - { - "id": 11, - "type": "[ComfyUI-3D] Load UNet2D ConditionModel", - "pos": [ - 1435, - 195 - ], - "size": { - "0": 407.1247863769531, - "1": 86.651123046875 - }, - "flags": {}, - "order": 0, - "mode": 0, - "outputs": [ - { - "name": "unet2d", - "type": "UNET2D", - "links": [ - 18 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "[ComfyUI-3D] Load UNet2D ConditionModel" - }, - "widgets_values": [ - "./pretrained_weights/stable-diffusion-v1-5/unet/", - "./pretrained_weights/reference_unet.pth" - ] - }, - { - "id": 12, - "type": "[ComfyUI-3D] Load UNet3D ConditionModel", - "pos": [ - 1471, - 358 - ], - "size": { - "0": 409.384033203125, - "1": 106 - }, - "flags": {}, - "order": 1, - "mode": 0, - "outputs": [ - { - "name": "unet3d", - "type": "UNET3D", - "links": [ - 19 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "[ComfyUI-3D] Load UNet3D ConditionModel" - }, - "widgets_values": [ - "./pretrained_weights/stable-diffusion-v1-5/unet/", - "./pretrained_weights/denoising_unet.pth", - "./pretrained_weights/motion_module.pth" - ] - }, - { - "id": 2, - "type": "VAEEncode", - "pos": [ - 1542, - 88 - ], - "size": { - "0": 210, - "1": 46 - }, - "flags": {}, - "order": 12, - "mode": 0, - "inputs": [ - { - "name": "pixels", - "type": "IMAGE", - "link": 24 - }, - { - "name": "vae", - "type": "VAE", - "link": 1 - } - ], - "outputs": [ - { - "name": "LATENT", - "type": "LATENT", - "links": [ - 20 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "VAEEncode", - "ttNbgOverride": { - "color": "#223", - "bgcolor": "#335", - "groupcolor": "#88A" - } - }, - "color": "#223", - "bgcolor": "#335" - }, { "id": 7, "type": "[ComfyUI-3D] Load Pose Guider", @@ -206,7 +91,7 @@ "1": 58 }, "flags": {}, - "order": 2, + "order": 0, "mode": 0, "outputs": [ { @@ -226,47 +111,6 @@ "./pretrained_weights/pose_guider.pth" ] }, - { - "id": 8, - "type": "[ComfyUI-3D] Pose Guider Encode", - "pos": [ - 1668, - 704 - ], - "size": { - "0": 292.20001220703125, - "1": 46 - }, - "flags": {}, - "order": 14, - "mode": 0, - "inputs": [ - { - "name": "pose_guider", - "type": "POSE_GUIDER", - "link": 5 - }, - { - "name": "pose_images", - "type": "IMAGE", - "link": 27 - } - ], - "outputs": [ - { - "name": "pose_latent", - "type": "POSE_LATENT", - "links": [ - 22 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "[ComfyUI-3D] Pose Guider Encode" - } - }, { "id": 6, "type": "CLIPVisionLoader", @@ -279,7 +123,7 @@ "1": 58 }, "flags": {}, - "order": 3, + "order": 1, "mode": 0, "outputs": [ { @@ -299,47 +143,6 @@ "pytorch_model.bin" ] }, - { - "id": 5, - "type": "CLIPVisionEncode", - "pos": [ - 1544, - 592 - ], - "size": { - "0": 380.4000244140625, - "1": 46 - }, - "flags": {}, - "order": 13, - "mode": 0, - "inputs": [ - { - "name": "clip_vision", - "type": "CLIP_VISION", - "link": 3 - }, - { - "name": "image", - "type": "IMAGE", - "link": 25 - } - ], - "outputs": [ - { - "name": "CLIP_VISION_OUTPUT", - "type": "CLIP_VISION_OUTPUT", - "links": [ - 21 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "CLIPVisionEncode" - } - }, { "id": 4, "type": "LoadImage", @@ -352,7 +155,7 @@ "1": 314 }, "flags": {}, - "order": 4, + "order": 2, "mode": 0, "outputs": [ { @@ -391,7 +194,7 @@ "1": 58 }, "flags": {}, - "order": 5, + "order": 3, "mode": 0, "outputs": [ { @@ -426,7 +229,7 @@ "1": 58 }, "flags": {}, - "order": 6, + "order": 4, "mode": 0, "outputs": [ { @@ -461,7 +264,7 @@ "1": 60.091087341308594 }, "flags": {}, - "order": 7, + "order": 5, "mode": 0, "outputs": [ { @@ -655,7 +458,7 @@ "hidden": false, "paused": false, "params": { - "filename": "video_00019.mp4", + "filename": "video_00087.mp4", "subfolder": "AnimateAnyone", "type": "output", "format": "video/h265-mp4" @@ -664,15 +467,212 @@ } }, { - "id": 21, - "type": "[ComfyUI-3D] Animate Anyone Sampler", + "id": 11, + "type": "[ComfyUI-3D] Load UNet2D ConditionModel", "pos": [ - 2159, - 48 + 1435, + 195 ], "size": { - "0": 328.86407470703125, - "1": 594 + "0": 407.1247863769531, + "1": 86.651123046875 + }, + "flags": {}, + "order": 6, + "mode": 0, + "outputs": [ + { + "name": "unet2d", + "type": "UNET2D", + "links": [ + 51 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Load UNet2D ConditionModel" + }, + "widgets_values": [ + "./pretrained_weights/stable-diffusion-v1-5/unet/", + "./pretrained_weights/reference_unet.pth" + ] + }, + { + "id": 12, + "type": "[ComfyUI-3D] Load UNet3D ConditionModel", + "pos": [ + 1471, + 358 + ], + "size": { + "0": 409.384033203125, + "1": 106 + }, + "flags": {}, + "order": 7, + "mode": 0, + "outputs": [ + { + "name": "unet3d", + "type": "UNET3D", + "links": [ + 52 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Load UNet3D ConditionModel" + }, + "widgets_values": [ + "./pretrained_weights/stable-diffusion-v1-5/unet/", + "./pretrained_weights/denoising_unet.pth", + "./pretrained_weights/motion_module.pth" + ] + }, + { + "id": 2, + "type": "VAEEncode", + "pos": [ + 1542, + 88 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "pixels", + "type": "IMAGE", + "link": 24 + }, + { + "name": "vae", + "type": "VAE", + "link": 1 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 53 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEEncode", + "ttNbgOverride": { + "color": "#223", + "bgcolor": "#335", + "groupcolor": "#88A" + } + }, + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 5, + "type": "CLIPVisionEncode", + "pos": [ + 1544, + 592 + ], + "size": { + "0": 380.4000244140625, + "1": 46 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "clip_vision", + "type": "CLIP_VISION", + "link": 3 + }, + { + "name": "image", + "type": "IMAGE", + "link": 25 + } + ], + "outputs": [ + { + "name": "CLIP_VISION_OUTPUT", + "type": "CLIP_VISION_OUTPUT", + "links": [ + 54 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPVisionEncode" + } + }, + { + "id": 8, + "type": "[ComfyUI-3D] Pose Guider Encode", + "pos": [ + 1668, + 704 + ], + "size": { + "0": 292.20001220703125, + "1": 46 + }, + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "pose_guider", + "type": "POSE_GUIDER", + "link": 5 + }, + { + "name": "pose_images", + "type": "IMAGE", + "link": 27 + } + ], + "outputs": [ + { + "name": "pose_latent", + "type": "POSE_LATENT", + "links": [ + 55 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "[ComfyUI-3D] Pose Guider Encode" + } + }, + { + "id": 31, + "type": "[ComfyUI-3D] Animate Anyone Sampler", + "pos": [ + 2091, + 47 + ], + "size": { + "0": 315, + "1": 618 }, "flags": {}, "order": 15, @@ -681,27 +681,27 @@ { "name": "reference_unet", "type": "UNET2D", - "link": 18 + "link": 51 }, { "name": "denoising_unet", "type": "UNET3D", - "link": 19 + "link": 52 }, { "name": "ref_image_latent", "type": "LATENT", - "link": 20 + "link": 53 }, { "name": "clip_image_embeds", "type": "CLIP_VISION_OUTPUT", - "link": 21 + "link": 54 }, { "name": "pose_latent", "type": "POSE_LATENT", - "link": 22 + "link": 55 } ], "outputs": [ @@ -709,7 +709,7 @@ "name": "latent", "type": "LATENT", "links": [ - 23 + 56 ], "shape": 3, "slot_index": 0 @@ -723,22 +723,23 @@ "fixed", 20, 3.5, + 1, 24, 1, 4, 1, 1, "DDIM", - 0.0008500000000000001, + 0.00085, 0.012, "linear", "v_prediction", "trailing", 1, false, - "cammy_white_v1.safetensors", + true, false, - "lcm-lora-sdv1-5.safetensors" + "cammy_white_v1.safetensors" ] }, { @@ -856,54 +857,6 @@ 0, "IMAGE" ], - [ - 18, - 11, - 0, - 21, - 0, - "UNET2D" - ], - [ - 19, - 12, - 0, - 21, - 1, - "UNET3D" - ], - [ - 20, - 2, - 0, - 21, - 2, - "LATENT" - ], - [ - 21, - 5, - 0, - 21, - 3, - "CLIP_VISION_OUTPUT" - ], - [ - 22, - 8, - 0, - 21, - 4, - "POSE_LATENT" - ], - [ - 23, - 21, - 0, - 14, - 0, - "LATENT" - ], [ 24, 22, @@ -975,6 +928,54 @@ 23, 2, "INT" + ], + [ + 51, + 11, + 0, + 31, + 0, + "UNET2D" + ], + [ + 52, + 12, + 0, + 31, + 1, + "UNET3D" + ], + [ + 53, + 2, + 0, + 31, + 2, + "LATENT" + ], + [ + 54, + 5, + 0, + 31, + 3, + "CLIP_VISION_OUTPUT" + ], + [ + 55, + 8, + 0, + 31, + 4, + "POSE_LATENT" + ], + [ + 56, + 31, + 0, + 14, + 0, + "LATENT" ] ], "groups": [], diff --git a/_test.py b/_test.py index 9b8e6f0..be444e6 100644 --- a/_test.py +++ b/_test.py @@ -1,9 +1,13 @@ -class MyClass: - def my_function(self): - self.x = 123 - pass +import inspect +from diffusers.schedulers import ( + DDIMScheduler, + DPMSolverMultistepScheduler, + LCMScheduler, + EulerDiscreteScheduler, + EulerAncestralDiscreteScheduler, + LMSDiscreteScheduler, + PNDMScheduler, +) -mc = MyClass() -mc.my_function() -if hasattr(mc, 'z'): - print('MyClass has x') \ No newline at end of file +print("clip_samples" in inspect.signature(DDIMScheduler.__init__).parameters) + \ No newline at end of file diff --git a/nodes.py b/nodes.py index b0fe065..28f5a40 100644 --- a/nodes.py +++ b/nodes.py @@ -18,12 +18,10 @@ from diffusers.schedulers import ( PNDMScheduler, ) -from .src.models.mutual_self_attention import ReferenceAttentionControl 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, CFGType +from .src.models.main_diffuser import AADiffusion 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") @@ -57,8 +55,6 @@ class Animate_Anyone_Sampler: "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}), "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": ("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}), @@ -100,9 +96,7 @@ class Animate_Anyone_Sampler: seed, steps, cfg, - steps_rcfg, delta, - cfg_type, context_frames, context_stride, context_overlap, @@ -159,8 +153,6 @@ class Animate_Anyone_Sampler: steps, cfg, delta, - CFGType[cfg_type], - steps_rcfg, ref_image_latent, pose_latent, encoder_hidden_states, diff --git a/src/models/main_diffuser.py b/src/models/main_diffuser.py index 7d4222d..038e9c8 100644 --- a/src/models/main_diffuser.py +++ b/src/models/main_diffuser.py @@ -4,8 +4,10 @@ import math from enum import Enum from typing import Callable, List, Optional, Union +import comfy.utils + import torch -from diffusers import DiffusionPipeline +from diffusers import DiffusionPipeline, StableDiffusionPipeline from diffusers.loaders import LoraLoaderMixin from diffusers.schedulers import ( DDIMScheduler, @@ -20,17 +22,12 @@ 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, + unet, scheduler: Union[ DDIMScheduler, LCMScheduler, @@ -44,11 +41,11 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): super().__init__() self.register_modules( reference_unet=reference_unet, - denoising_unet=denoising_unet, + unet=unet, scheduler=scheduler, ) self.reference_unet = reference_unet - self.denoising_unet = denoising_unet + self.unet = unet self.scheduler = scheduler self.noise_pred_uncond_prev = None @@ -172,44 +169,38 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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 + def perform_guidance( + self, + noise_pred, + counter, + cfg, + delta, + do_double_pass, + ): - 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 + # 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, - uncond_encoder_hidden_states, pose_latent, global_context, + do_double_pass, device, ): noise_pred = torch.zeros( ( - latents.shape[0], - *latents.shape[1:], - ), - device=device, - dtype=latents.dtype, - ) - noise_pred_uncond = torch.zeros( - ( - latents.shape[0], + latents.shape[0] * (2 if do_double_pass else 1), *latents.shape[1:], ), device=device, @@ -226,6 +217,7 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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 @@ -233,66 +225,56 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): b, c, f, h, w = latent_model_input.shape latent_pose_input = torch.cat( [pose_latent[:, :, c] for c in context] - ) - pred = self.denoising_unet( + ).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] - pred_uncond = self.denoising_unet( - latent_model_input, # torch.Size([2, 4, 24, 96, 64]) - t, - encoder_hidden_states=uncond_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 - noise_pred_uncond[:, :, c] = noise_pred_uncond[:, :, c] + pred_uncond counter[:, :, c] = counter[:, :, c] + 1 - return noise_pred, noise_pred_uncond, counter + return noise_pred, counter def solve_one_step( self, t, latents, encoder_hidden_states, - uncond_encoder_hidden_states, pose_latent, global_context, cfg, delta, - do_rcfg, - noise_pred_uncond_prev, + do_double_pass, device, extra_step_kwargs, ): - noise_pred, noise_pred_uncond, counter = self.denoise_all_one_step( - t, - latents, - encoder_hidden_states, - uncond_encoder_hidden_states, - pose_latent, - global_context, - device, + 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, ) - # CFG - noise_pred_cond = noise_pred / counter - noise_pred_uncond = noise_pred_uncond / counter - noise_pred = delta * noise_pred_uncond + cfg * (noise_pred_cond - delta * noise_pred_uncond) + 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) - - 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 + return scheduler_output.prev_sample def denoise_loop( self, @@ -305,13 +287,16 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): cfg, delta, do_double_pass, - do_rcfg, interpolation_factor, device, extra_step_kwargs, ): # Prepare confitional embeds - uncond_encoder_hidden_states = torch.zeros_like(encoder_hidden_states) + 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( @@ -321,7 +306,7 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): fusion_blocks="full", ) reference_control_reader = ReferenceAttentionControl( - self.denoising_unet, + self.unet, do_classifier_free_guidance=do_double_pass, mode="read", fusion_blocks="full", @@ -330,10 +315,7 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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, - encoder_hidden_states=torch.cat( - [uncond_encoder_hidden_states, encoder_hidden_states], dim=0 - ) if do_double_pass else encoder_hidden_states, + encoder_hidden_states=encoder_hidden_states, return_dict=False, ) reference_control_reader.update(reference_control_writer) @@ -342,25 +324,26 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): 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.noise_pred_uncond_prev = self.solve_one_step( - t, - latents, - encoder_hidden_states, - uncond_encoder_hidden_states, - pose_latent, - global_context, - cfg, - delta, - do_rcfg, - self.noise_pred_uncond_prev, - device, - extra_step_kwargs, + 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) @@ -376,8 +359,6 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): steps, cfg, delta, - cfg_type, - steps_rcfg, ref_image_latent, pose_latent, encoder_hidden_states, @@ -391,9 +372,7 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): interpolation_factor=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) + do_double_pass = cfg > 0 device = pose_latent.device generator=torch.manual_seed(seed) @@ -448,7 +427,7 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): # Prepare initial latents latents = self.prepare_latents( batch_size=pose_latent.shape[0], - num_channels_latents=self.denoising_unet.config.in_channels, + num_channels_latents=self.unet.config.in_channels, video_length=video_length, height=pose_latent.shape[-2], width=pose_latent.shape[-1], @@ -467,8 +446,7 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin): global_context=global_context, cfg=cfg, delta=delta, - do_double_pass=False, - do_rcfg=False, + do_double_pass=do_double_pass, interpolation_factor=interpolation_factor, device=device, extra_step_kwargs=extra_step_kwargs, diff --git a/src/models/main_model_patcher.py b/src/models/main_model_patcher.py deleted file mode 100644 index 4ba4731..0000000 --- a/src/models/main_model_patcher.py +++ /dev/null @@ -1,13 +0,0 @@ -import torch - -from comfy.model_patcher import ModelPatcher - -from .unet_3d import UNet3DConditionModel - -class AnimateAnyoneModelPatcher(ModelPatcher): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self.model: UNet3DConditionModel = self.model - - \ No newline at end of file