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] 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:
-## 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