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

Update After RCFG experiment

Browse Source 
- Update readme to log the RCFG experiment
- Refactor code
- Add delta parameter on Animate_Anyone_Sampler node
- Add ComfyUI Sampler progress bar
This commit is contained in:
MrForExample
2024-01-21 17:56:58 +01:00
parent c2dd4665a7
commit fa383bd37e
6 changed files with 547 additions and 379 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+9 -6
View File
@@ -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:
<br>![_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.<br>
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!)
_Example_Workflow/MooreAnimateAnyone_DDIM.json
+271 -270
View File
@@ -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": [],
_test.py
+37 -12
View File
@@ -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)
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)
nodes.py
+7 -5
View File
@@ -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,
src/models/main_diffuser.py
+223 -73
View File
@@ -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
src/models/main_model_patcher.py
-13
View File
@@ -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