From fa383bd37efdcbb2e8c1b73d693d50615656fc92 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 | 49 +-
nodes.py | 12 +-
src/models/main_diffuser.py | 296 +++++++---
src/models/main_model_patcher.py | 13 -
6 files changed, 547 insertions(+), 379 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 be444e6..9153f41 100644
--- a/_test.py
+++ b/_test.py
@@ -1,13 +1,38 @@
-import inspect
-from diffusers.schedulers import (
- DDIMScheduler,
- DPMSolverMultistepScheduler,
- LCMScheduler,
- EulerDiscreteScheduler,
- EulerAncestralDiscreteScheduler,
- LMSDiscreteScheduler,
- PNDMScheduler,
-)
+import time
+import torch
+from diffusers import (StableDiffusionPipeline,
+ EulerAncestralDiscreteScheduler)
+from sfast.compilers.diffusion_pipeline_compiler import (compile,
+ CompilationConfig)
-print("clip_samples" in inspect.signature(DDIMScheduler.__init__).parameters)
-
\ No newline at end of file
+def load_model():
+ model = StableDiffusionPipeline.from_pretrained(
+ 'runwayml/stable-diffusion-v1-5',
+ torch_dtype=torch.float16)
+
+ model.scheduler = EulerAncestralDiscreteScheduler.from_config(
+ model.scheduler.config)
+ model.safety_checker = None
+ model.to(torch.device('cuda'))
+ return model
+
+model = load_model()
+
+config = CompilationConfig.Default()
+# xformers and Triton are suggested for achieving best performance.
+try:
+ import xformers
+ config.enable_xformers = True
+except ImportError:
+ print('xformers not installed, skip')
+try:
+ import triton
+ config.enable_triton = True
+except ImportError:
+ print('Triton not installed, skip')
+# CUDA Graph is suggested for small batch sizes and small resolutions to reduce CPU overhead.
+# But it can increase the amount of GPU memory used.
+# For StableVideoDiffusionPipeline it is not needed.
+config.enable_cuda_graph = True
+
+model = compile(model, config)
\ No newline at end of file
diff --git a/nodes.py b/nodes.py
index d959b44..8effeba 100644
--- a/nodes.py
+++ b/nodes.py
@@ -18,11 +18,9 @@ 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
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "./ComfyUI-AnimateAnyone-Evolved")
@@ -55,7 +53,8 @@ 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}),
+ "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 +94,8 @@ class Animate_Anyone_Sampler:
pose_latent,
seed,
steps,
- cfg,
+ cfg,
+ delta,
context_frames,
context_stride,
context_overlap,
@@ -179,7 +179,9 @@ class Animate_Anyone_Sampler:
samples = diffuser(
steps,
- cfg,
+ cfg,
+ delta,
+ ref_image_latent,
pose_latent,
encoder_hidden_states,
seed,
diff --git a/src/models/main_diffuser.py b/src/models/main_diffuser.py
index 018aeee..a588814 100644
--- a/src/models/main_diffuser.py
+++ b/src/models/main_diffuser.py
@@ -3,8 +3,10 @@ import inspect
import math
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,
@@ -18,11 +20,11 @@ from diffusers.schedulers import (
from diffusers.utils.torch_utils import randn_tensor
from ..utils.util import get_tensor_interpolation_method, get_context_scheduler
-
class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
def __init__(
self,
+ reference_unet,
unet,
scheduler: Union[
DDIMScheduler,
@@ -36,9 +38,11 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
):
super().__init__()
self.register_modules(
+ reference_unet=reference_unet,
unet=unet,
scheduler=scheduler,
)
+ self.reference_unet = reference_unet
self.unet = unet
self.scheduler = scheduler
@@ -160,13 +164,199 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
new_index += 1
return new_latents
+
+ def perform_guidance(
+ self,
+ noise_pred,
+ counter,
+ cfg,
+ delta,
+ do_double_pass,
+ ):
+
+ # perform CFG
+ if do_double_pass:
+ noise_pred_uncond, noise_pred_cond = (noise_pred / counter).chunk(2)
+ else:
+ noise_pred_cond = noise_pred_uncond = noise_pred / counter
+ cfg = 0.0
+ noise_pred = delta * noise_pred_uncond + cfg * (noise_pred_cond - delta * noise_pred_uncond)
+
+ return noise_pred
+
+ def denoise_all_one_step(
+ self,
+ t,
+ latents,
+ encoder_hidden_states,
+ pose_latent,
+ global_context,
+ do_double_pass,
+ device,
+ ):
+ noise_pred = torch.zeros(
+ (
+ latents.shape[0] * (2 if do_double_pass else 1),
+ *latents.shape[1:],
+ ),
+ device=device,
+ dtype=latents.dtype,
+ )
+ counter = torch.zeros(
+ (1, 1, latents.shape[2], 1, 1),
+ device=device,
+ dtype=latents.dtype,
+ )
+
+ for context in global_context:
+ # 3.1 expand the latents if we are doing classifier free guidance
+ latent_model_input = (
+ torch.cat([latents[:, :, c] for c in context]) # Concatenates the context in the same batch
+ .to(device)
+ .repeat(2 if do_double_pass else 1, 1, 1, 1, 1)
+ )
+ latent_model_input = self.scheduler.scale_model_input(
+ latent_model_input, t
+ )
+ b, c, f, h, w = latent_model_input.shape
+ latent_pose_input = torch.cat(
+ [pose_latent[:, :, c] for c in context]
+ ).repeat(2 if do_double_pass else 1, 1, 1, 1, 1)
+ pred = self.unet(
+ latent_model_input, # torch.Size([2, 4, 24, 96, 64])
+ t,
+ encoder_hidden_states=encoder_hidden_states[:b],
+ pose_cond_fea=latent_pose_input, # torch.Size([2, 320, 24, 96, 64])
+ return_dict=False,
+ )[0]
+
+ for j, c in enumerate(context):
+ noise_pred[:, :, c] = noise_pred[:, :, c] + pred
+ counter[:, :, c] = counter[:, :, c] + 1
+
+ return noise_pred, counter
+
+ def solve_one_step(
+ self,
+ t,
+ latents,
+ encoder_hidden_states,
+ pose_latent,
+ global_context,
+ cfg,
+ delta,
+ do_double_pass,
+ device,
+ extra_step_kwargs,
+ ):
+
+ noise_pred, counter = self.denoise_all_one_step(
+ t=t,
+ latents=latents,
+ encoder_hidden_states=encoder_hidden_states,
+ pose_latent=pose_latent,
+ global_context=global_context,
+ do_double_pass=do_double_pass,
+ device=device,
+ )
+
+ noise_pred = self.perform_guidance(
+ noise_pred=noise_pred,
+ counter=counter,
+ cfg=cfg,
+ delta=delta,
+ do_double_pass=do_double_pass,
+ )
+
+ scheduler_output = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs)
+
+ return scheduler_output.prev_sample
+
+ def denoise_loop(
+ self,
+ ref_image_latent,
+ loop_steps,
+ latents,
+ encoder_hidden_states,
+ pose_latent,
+ global_context,
+ cfg,
+ delta,
+ do_double_pass,
+ interpolation_factor,
+ device,
+ extra_step_kwargs,
+ ):
+ # Prepare confitional embeds
+ if do_double_pass:
+ uncond_encoder_hidden_states = torch.zeros_like(encoder_hidden_states)
+ encoder_hidden_states=torch.cat(
+ [uncond_encoder_hidden_states, encoder_hidden_states], dim=0
+ )
+
+ # Prepare reference attention control
+ reference_control_writer = ReferenceAttentionControl(
+ self.reference_unet,
+ do_classifier_free_guidance=do_double_pass,
+ mode="write",
+ fusion_blocks="full",
+ )
+ reference_control_reader = ReferenceAttentionControl(
+ self.unet,
+ do_classifier_free_guidance=do_double_pass,
+ mode="read",
+ fusion_blocks="full",
+ )
+
+ self.reference_unet(
+ ref_image_latent.repeat(2 if do_double_pass else 1, 1, 1, 1),
+ torch.zeros(size=(1,), dtype=int, device=device),
+ encoder_hidden_states=encoder_hidden_states,
+ return_dict=False,
+ )
+ reference_control_reader.update(reference_control_writer)
+
+ # Prepare timesteps
+ self.scheduler.set_timesteps(loop_steps, device=device)
+ timesteps = self.scheduler.timesteps
+
+ comfy_pbar = comfy.utils.ProgressBar(loop_steps)
+
+ with self.progress_bar(total=loop_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+
+ latents = self.solve_one_step(
+ t=t,
+ latents=latents,
+ encoder_hidden_states=encoder_hidden_states,
+ pose_latent=pose_latent,
+ global_context=global_context,
+ cfg=cfg,
+ delta=delta,
+ do_double_pass=do_double_pass,
+ device=device,
+ extra_step_kwargs=extra_step_kwargs,
+ )
+
+ progress_bar.update()
+ comfy_pbar.update_absolute(i + 1)
+
+ if interpolation_factor > 0:
+ latents = self.interpolate_latents(latents, interpolation_factor, device)
+
+ reference_control_reader.clear()
+ reference_control_writer.clear()
+
+ return latents
@torch.no_grad()
def __call__(
self,
steps,
cfg,
- pose_fea,
+ delta,
+ ref_image_latent,
+ pose_latent,
encoder_hidden_states,
seed = 999999999,
eta: float = 0.0,
@@ -180,8 +370,9 @@ class AADiffusion(DiffusionPipeline, LoraLoaderMixin):
callback_steps: Optional[int] = 1,
):
- device = pose_fea.device
-
+ do_double_pass = cfg > 0
+
+ device = pose_latent.device
generator=torch.manual_seed(seed)
latents = self.prepare_latents(
batch_size=pose_fea.shape[0],
@@ -239,74 +430,33 @@ 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
- # 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,
- )
-
- 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
- ):
- 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)
+ # Prepare initial latents
+ latents = self.prepare_latents(
+ batch_size=pose_latent.shape[0],
+ num_channels_latents=self.unet.config.in_channels,
+ video_length=video_length,
+ height=pose_latent.shape[-2],
+ width=pose_latent.shape[-1],
+ dtype=encoder_hidden_states.dtype,
+ device=device,
+ generator=generator,
+ )
+
+ # sequential double pass
+ latents = self.denoise_loop(
+ ref_image_latent=ref_image_latent,
+ loop_steps=steps,
+ latents=latents,
+ encoder_hidden_states=encoder_hidden_states,
+ pose_latent=pose_latent,
+ global_context=global_context,
+ cfg=cfg,
+ delta=delta,
+ do_double_pass=do_double_pass,
+ interpolation_factor=interpolation_factor,
+ device=device,
+ extra_step_kwargs=extra_step_kwargs,
+ )
return latents
\ No newline at end of file
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