564 lines
23 KiB
Python
564 lines
23 KiB
Python
import os
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
import gc
|
|
|
|
import folder_paths
|
|
import comfy.model_management as mm
|
|
import comfy.utils
|
|
|
|
try:
|
|
import diffusers.models.activations
|
|
def patch_geglu_inplace():
|
|
"""Patch GEGLU with inplace multiplication to save GPU memory."""
|
|
def forward(self, hidden_states):
|
|
hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1)
|
|
return hidden_states.mul_(self.gelu(gate))
|
|
diffusers.models.activations.GEGLU.forward = forward
|
|
except:
|
|
pass
|
|
|
|
from .pipeline.pipeline_stable_video_diffusion_controlnext import StableVideoDiffusionPipelineControlNeXt, tensor2vid
|
|
|
|
from .models.controlnext_vid_svd import ControlNeXtSDVModel
|
|
from .models.unet_spatio_temporal_condition_controlnext import UNetSpatioTemporalConditionControlNeXtModel
|
|
from .utils.scheduling_euler_discrete_karras_fix import EulerDiscreteScheduler as EulerDiscreteSchedulerKarras
|
|
from diffusers.schedulers import EulerDiscreteScheduler
|
|
|
|
from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor
|
|
from diffusers import AutoencoderKLTemporalDecoder
|
|
|
|
script_directory = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
|
from contextlib import nullcontext
|
|
try:
|
|
from accelerate import init_empty_weights
|
|
from accelerate.utils import set_module_tensor_to_device
|
|
is_accelerate_available = True
|
|
except:
|
|
is_accelerate_available = False
|
|
pass
|
|
|
|
import logging
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
log = logging.getLogger(__name__)
|
|
|
|
def loglinear_interp(t_steps, num_steps):
|
|
"""
|
|
Performs log-linear interpolation of a given array of decreasing numbers.
|
|
"""
|
|
xs = np.linspace(0, 1, len(t_steps))
|
|
ys = np.log(t_steps[::-1])
|
|
|
|
new_xs = np.linspace(0, 1, num_steps)
|
|
new_ys = np.interp(new_xs, xs, ys)
|
|
|
|
interped_ys = np.exp(new_ys)[::-1].copy()
|
|
return interped_ys
|
|
|
|
class DownloadAndLoadControlNeXt:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
|
|
"precision": (
|
|
[
|
|
'fp32',
|
|
'fp16',
|
|
'bf16',
|
|
], {
|
|
"default": 'fp16'
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CONTROLNEXT_PIPE",)
|
|
RETURN_NAMES = ("controlnext_pipeline",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "ControlNeXtWrapper"
|
|
|
|
def loadmodel(self, precision):
|
|
device = mm.get_torch_device()
|
|
mm.soft_empty_cache()
|
|
dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
|
|
|
|
pbar = comfy.utils.ProgressBar(5)
|
|
|
|
download_path = os.path.join(folder_paths.models_dir, "diffusers", "controlnext")
|
|
unet_model_path = os.path.join(download_path, "controlnext-svd_v2-unet-fp16.safetensors")
|
|
contolnet_model_path = os.path.join(download_path, "controlnext-svd_v2-controlnet-fp16.safetensors")
|
|
|
|
if not os.path.exists(unet_model_path):
|
|
log.info(f"Downloading model to: {unet_model_path}")
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(repo_id="Kijai/ControlNeXt-SVD-V2-Comfy",
|
|
ignore_patterns=["*converted*"],
|
|
local_dir=download_path,
|
|
local_dir_use_symlinks=False)
|
|
|
|
log.info(f"Loading model from: {unet_model_path}")
|
|
pbar.update(1)
|
|
|
|
svd_path = os.path.join(folder_paths.models_dir, "diffusers", "stable-video-diffusion-img2vid-xt-1-1")
|
|
if not os.path.exists(svd_path):
|
|
log.info(f"Downloading SVD model to: {svd_path}")
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(repo_id="vdo/stable-video-diffusion-img2vid-xt-1-1",
|
|
allow_patterns=[f"*.json", "*fp16*"],
|
|
ignore_patterns=["*unet*"],
|
|
local_dir=svd_path,
|
|
local_dir_use_symlinks=False)
|
|
pbar.update(1)
|
|
|
|
svd_path = os.path.join(folder_paths.models_dir, "diffusers", "stable-video-diffusion-img2vid-xt-1-1")
|
|
|
|
unet_config = UNetSpatioTemporalConditionControlNeXtModel.load_config(os.path.join(script_directory, "configs", "unet_config.json"))
|
|
log.info("Loading UNET")
|
|
with (init_empty_weights() if is_accelerate_available else nullcontext()):
|
|
self.unet = UNetSpatioTemporalConditionControlNeXtModel.from_config(unet_config)
|
|
sd = comfy.utils.load_torch_file(os.path.join(unet_model_path))
|
|
if is_accelerate_available:
|
|
for key in sd:
|
|
set_module_tensor_to_device(self.unet, key, dtype=dtype, device=device, value=sd[key])
|
|
else:
|
|
self.unet.load_state_dict(sd, strict=False)
|
|
del sd
|
|
pbar.update(1)
|
|
|
|
log.info("Loading VAE")
|
|
self.vae = AutoencoderKLTemporalDecoder.from_pretrained(svd_path, subfolder="vae", variant="fp16", low_cpu_mem_usage=True).to(dtype).to(device).eval()
|
|
|
|
log.info("Loading IMAGE_ENCODER")
|
|
self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(svd_path, subfolder="image_encoder", variant="fp16", low_cpu_mem_usage=True).to(dtype).to(device).eval()
|
|
pbar.update(1)
|
|
self.noise_scheduler = EulerDiscreteScheduler.from_pretrained(svd_path, subfolder="scheduler")
|
|
self.feature_extractor = CLIPImageProcessor.from_pretrained(svd_path, subfolder="feature_extractor")
|
|
|
|
log.info("Loading ControlNeXt")
|
|
self.controlnext = ControlNeXtSDVModel()
|
|
self.controlnext.load_state_dict(comfy.utils.load_torch_file(os.path.join(contolnet_model_path)))
|
|
self.controlnext = self.controlnext.to(dtype).to(device).eval()
|
|
|
|
pipeline = StableVideoDiffusionPipelineControlNeXt(
|
|
vae = self.vae,
|
|
image_encoder = self.image_encoder,
|
|
unet = self.unet,
|
|
scheduler = self.noise_scheduler,
|
|
feature_extractor = self.feature_extractor,
|
|
controlnext=self.controlnext,
|
|
)
|
|
|
|
controlnextsvd_model = {
|
|
'pipeline': pipeline,
|
|
'dtype': dtype,
|
|
}
|
|
pbar.update(1)
|
|
return (controlnextsvd_model,)
|
|
|
|
|
|
|
|
class ControlNextDiffusersScheduler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"scheduler": (
|
|
[
|
|
'EulerDiscreteScheduler',
|
|
'EulerDiscreteSchedulerKarras',
|
|
'EulerDiscreteScheduler_AYS',
|
|
],
|
|
),
|
|
"sigma_min": ("FLOAT", {"default": 0.002, "min": 0.0, "max": 700.0, "step": 0.001}),
|
|
"sigma_max": ("FLOAT", {"default": 700.0, "min": 0.0, "max": 700.0, "step": 0.001}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("DIFFUSERS_SCHEDULER",)
|
|
RETURN_NAMES = ("scheduler",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "ControlNeXtSVD"
|
|
|
|
def loadmodel(self, scheduler, sigma_min, sigma_max):
|
|
|
|
scheduler_config = {
|
|
"beta_end": 0.012,
|
|
"beta_schedule": "scaled_linear",
|
|
"beta_start": 0.00085,
|
|
"clip_sample": False,
|
|
"interpolation_type": "linear",
|
|
"num_train_timesteps": 1000,
|
|
"prediction_type": "v_prediction",
|
|
"set_alpha_to_one": False,
|
|
"sigma_max": sigma_max,
|
|
"sigma_min": sigma_min,
|
|
"skip_prk_steps": True,
|
|
"steps_offset": 1,
|
|
"timestep_spacing": "leading",
|
|
"timestep_type": "continuous",
|
|
"trained_betas": None,
|
|
"use_karras_sigmas": False
|
|
}
|
|
if scheduler == 'EulerDiscreteScheduler':
|
|
noise_scheduler = EulerDiscreteScheduler.from_config(scheduler_config)
|
|
sigmas = None
|
|
elif scheduler == 'EulerDiscreteScheduler_AYS':
|
|
noise_scheduler = EulerDiscreteScheduler.from_config(scheduler_config)
|
|
sigmas = [700.00, 54.5, 15.886, 7.977, 4.248, 1.789, 0.981, 0.403, 0.173, 0.034, 0.002]
|
|
elif scheduler == 'EulerDiscreteSchedulerKarras':
|
|
scheduler_config['use_karras_sigmas'] = True
|
|
noise_scheduler = EulerDiscreteSchedulerKarras.from_config(scheduler_config)
|
|
sigmas = None
|
|
|
|
scheduler_options = {
|
|
"noise_scheduler": noise_scheduler,
|
|
"sigmas": sigmas,
|
|
}
|
|
|
|
return (scheduler_options,)
|
|
|
|
class ControlNextSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"controlnext_pipeline": ("CONTROLNEXT_PIPE",),
|
|
"ref_image": ("IMAGE",),
|
|
"pose_images": ("IMAGE",),
|
|
"steps": ("INT", {"default": 25, "min": 1, "max": 200, "step": 1}),
|
|
"motion_bucket_id": ("INT", {"default": 127, "min": 0, "max": 1000, "step": 1}),
|
|
"cfg_min": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 20.0, "step": 0.01}),
|
|
"cfg_max": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 20.0, "step": 0.01}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"fps": ("INT", {"default": 7, "min": 2, "max": 100, "step": 1}),
|
|
"controlnext_cond_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
"noise_aug_strength": ("FLOAT", {"default": 0.02, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
"context_size": ("INT", {"default": 24, "min": 1, "max": 128, "step": 1}),
|
|
"context_overlap": ("INT", {"default": 6, "min": 1, "max": 128, "step": 1}),
|
|
"keep_model_loaded": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"optional_scheduler": ("DIFFUSERS_SCHEDULER",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("samples",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "ControlNeXtSVD"
|
|
|
|
def process(self, controlnext_pipeline, ref_image, pose_images, cfg_min, cfg_max, controlnext_cond_scale, motion_bucket_id, steps, seed, noise_aug_strength, fps, keep_model_loaded,
|
|
context_size, context_overlap, optional_scheduler=None):
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
mm.unload_all_models()
|
|
mm.soft_empty_cache()
|
|
dtype = controlnext_pipeline['dtype']
|
|
pipeline = controlnext_pipeline['pipeline']
|
|
|
|
original_scheduler = pipeline.scheduler
|
|
|
|
if optional_scheduler is not None:
|
|
log.info(f"Using optional scheduler: {optional_scheduler['noise_scheduler']}")
|
|
pipeline.scheduler = optional_scheduler['noise_scheduler']
|
|
sigmas = optional_scheduler['sigmas']
|
|
|
|
if sigmas is not None and (steps + 1) != len(sigmas):
|
|
sigmas = loglinear_interp(sigmas, steps + 1)
|
|
sigmas = sigmas[-(steps + 1):]
|
|
sigmas[-1] = 0
|
|
log.info(f"Using timesteps: {sigmas}")
|
|
else:
|
|
pipeline.scheduler = original_scheduler
|
|
sigmas = None
|
|
|
|
B, H, W, C = pose_images.shape
|
|
|
|
assert B >= context_size, "The number of poses must be greater than the context size"
|
|
|
|
ref_image = ref_image.permute(0, 3, 1, 2)
|
|
pose_images = pose_images.permute(0, 3, 1, 2)
|
|
pose_images = pose_images * 2 - 1
|
|
|
|
ref_image = ref_image.to(device).to(dtype)
|
|
pose_images = pose_images.to(device).to(dtype)
|
|
|
|
generator = torch.Generator(device=device)
|
|
generator.manual_seed(seed)
|
|
|
|
frames = pipeline(
|
|
ref_image,
|
|
pose_images,
|
|
num_frames=B,
|
|
frames_per_batch=context_size,
|
|
overlap=context_overlap,
|
|
motion_bucket_id=motion_bucket_id,
|
|
min_guidance_scale=cfg_min,
|
|
max_guidance_scale=cfg_max,
|
|
controlnext_cond_scale=controlnext_cond_scale,
|
|
height=H,
|
|
width=W,
|
|
fps=fps,
|
|
noise_aug_strength=noise_aug_strength,
|
|
num_inference_steps=steps,
|
|
generator=generator,
|
|
sigmas = sigmas,
|
|
decode_chunk_size=2,
|
|
output_type="latent",
|
|
return_dict="false",
|
|
#device=device,
|
|
).frames
|
|
|
|
if not keep_model_loaded:
|
|
pipeline.unet.to(offload_device)
|
|
pipeline.vae.to(offload_device)
|
|
mm.soft_empty_cache()
|
|
gc.collect()
|
|
|
|
return {"samples": frames},
|
|
|
|
class ControlNextDecode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"controlnext_pipeline": ("CONTROLNEXT_PIPE",),
|
|
"samples": ("LATENT",),
|
|
"decode_chunk_size": ("INT", {"default": 4, "min": 1, "max": 200, "step": 1})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "ControlNeXtSVD"
|
|
|
|
def process(self, controlnext_pipeline, samples, decode_chunk_size):
|
|
mm.soft_empty_cache()
|
|
|
|
pipeline = controlnext_pipeline['pipeline']
|
|
num_frames = samples['samples'].shape[0]
|
|
try:
|
|
frames = pipeline.decode_latents(samples['samples'], num_frames, decode_chunk_size)
|
|
except:
|
|
frames = pipeline.decode_latents(samples['samples'], num_frames, 1)
|
|
frames = tensor2vid(frames, pipeline.image_processor, output_type="pt")
|
|
|
|
frames = frames.squeeze(1)[1:].permute(0, 2, 3, 1).cpu().float()
|
|
|
|
return frames,
|
|
|
|
class ControlNextGetPoses:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"ref_image": ("IMAGE",),
|
|
"pose_images": ("IMAGE",),
|
|
"include_body": ("BOOLEAN", {"default": True}),
|
|
"include_hand": ("BOOLEAN", {"default": True}),
|
|
"include_face": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE",)
|
|
RETURN_NAMES = ("poses_with_ref", "pose_images")
|
|
FUNCTION = "process"
|
|
CATEGORY = "ControlNextWrapper"
|
|
|
|
def process(self, ref_image, pose_images, include_body, include_hand, include_face):
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
from .dwpose.util import draw_pose
|
|
from .dwpose.dwpose_detector import DWposeDetector
|
|
|
|
assert ref_image.shape[1:3] == pose_images.shape[1:3], "ref_image and pose_images must have the same resolution"
|
|
|
|
#yolo_model = "yolox_l.onnx"
|
|
#dw_pose_model = "dw-ll_ucoco_384.onnx"
|
|
dw_pose_model = "dw-ll_ucoco_384_bs5.torchscript.pt"
|
|
yolo_model = "yolox_l.torchscript.pt"
|
|
|
|
model_base_path = os.path.join(script_directory, "models", "DWPose")
|
|
|
|
model_det=os.path.join(model_base_path, yolo_model)
|
|
model_pose=os.path.join(model_base_path, dw_pose_model)
|
|
|
|
if not os.path.exists(model_det):
|
|
log.info(f"Downloading yolo model to: {model_base_path}")
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(repo_id="hr16/yolox-onnx",
|
|
allow_patterns=[f"*{yolo_model}*"],
|
|
local_dir=model_base_path,
|
|
local_dir_use_symlinks=False)
|
|
|
|
if not os.path.exists(model_pose):
|
|
log.info(f"Downloading dwpose model to: {model_base_path}")
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(repo_id="hr16/DWPose-TorchScript-BatchSize5",
|
|
allow_patterns=[f"*{dw_pose_model}*"],
|
|
local_dir=model_base_path,
|
|
local_dir_use_symlinks=False)
|
|
|
|
model_det=os.path.join(model_base_path, yolo_model)
|
|
model_pose=os.path.join(model_base_path, dw_pose_model)
|
|
|
|
if not hasattr(self, "det") or not hasattr(self, "pose"):
|
|
self.det = torch.jit.load(model_det)
|
|
self.pose = torch.jit.load(model_pose)
|
|
|
|
self.dwprocessor = DWposeDetector(
|
|
model_det=self.det,
|
|
model_pose=self.pose)
|
|
|
|
ref_image = ref_image.squeeze(0).cpu().numpy() * 255
|
|
|
|
self.det = self.det.to(device)
|
|
self.pose = self.pose.to(device)
|
|
|
|
# select ref-keypoint from reference pose for pose rescale
|
|
ref_pose = self.dwprocessor(ref_image)
|
|
#ref_keypoint_id = [0, 1, 2, 5, 8, 11, 14, 15, 16, 17]
|
|
ref_keypoint_id = [0, 1, 2, 5, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
|
|
ref_keypoint_id = [i for i in ref_keypoint_id \
|
|
#if ref_pose['bodies']['score'].shape[0] > 0 and ref_pose['bodies']['score'][0][i] > 0.3]
|
|
if len(ref_pose['bodies']['subset']) > 0 and ref_pose['bodies']['subset'][0][i] >= .0]
|
|
ref_body = ref_pose['bodies']['candidate'][ref_keypoint_id]
|
|
|
|
height, width, _ = ref_image.shape
|
|
pose_images_np = pose_images.cpu().numpy() * 255
|
|
|
|
# read input video
|
|
pbar = comfy.utils.ProgressBar(len(pose_images_np))
|
|
detected_poses_np_list = []
|
|
for img_np in pose_images_np:
|
|
detected_poses_np_list.append(self.dwprocessor(img_np))
|
|
pbar.update(1)
|
|
|
|
self.det = self.det.to(offload_device)
|
|
self.pose = self.pose.to(offload_device)
|
|
|
|
detected_bodies = np.stack(
|
|
[p['bodies']['candidate'] for p in detected_poses_np_list if p['bodies']['candidate'].shape[0] == 18])[:,
|
|
ref_keypoint_id]
|
|
# compute linear-rescale params
|
|
ay, by = np.polyfit(detected_bodies[:, :, 1].flatten(), np.tile(ref_body[:, 1], len(detected_bodies)), 1)
|
|
fh, fw, _ = pose_images_np[0].shape
|
|
ax = ay / (fh / fw / height * width)
|
|
bx = np.mean(np.tile(ref_body[:, 0], len(detected_bodies)) - detected_bodies[:, :, 0].flatten() * ax)
|
|
a = np.array([ax, ay])
|
|
b = np.array([bx, by])
|
|
output_pose = []
|
|
# pose rescale
|
|
for detected_pose in detected_poses_np_list:
|
|
if include_body:
|
|
detected_pose['bodies']['candidate'] = detected_pose['bodies']['candidate'] * a + b
|
|
if include_hand:
|
|
detected_pose['faces'] = detected_pose['faces'] * a + b
|
|
if include_face:
|
|
detected_pose['hands'] = detected_pose['hands'] * a + b
|
|
im = draw_pose(detected_pose, height, width, include_body=include_body, include_hand=include_hand, include_face=include_face)
|
|
output_pose.append(np.array(im))
|
|
|
|
output_pose_tensors = [torch.tensor(np.array(im)) for im in output_pose]
|
|
output_tensor = torch.stack(output_pose_tensors) / 255
|
|
|
|
ref_pose_img = draw_pose(ref_pose, height, width, include_body=include_body, include_hand=include_hand, include_face=include_face)
|
|
ref_pose_tensor = torch.tensor(np.array(ref_pose_img)) / 255
|
|
output_tensor = torch.cat((ref_pose_tensor.unsqueeze(0), output_tensor))
|
|
output_tensor = output_tensor.permute(0, 2, 3, 1).cpu().float()
|
|
|
|
return output_tensor, output_tensor[1:]
|
|
|
|
|
|
class ControlNextSVDApply:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"model": ("MODEL",),
|
|
"pose_images": ("IMAGE",),
|
|
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
|
|
"blocks": ("STRING",{"default": "3"}),
|
|
"input_block_patch_after_skip": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
RETURN_NAMES = ("model", )
|
|
FUNCTION = "patch"
|
|
CATEGORY = "ControlNeXtSVD"
|
|
|
|
def patch(self, model, pose_images, strength, blocks, input_block_patch_after_skip):
|
|
|
|
device = mm.get_torch_device()
|
|
dtype = mm.unet_dtype()
|
|
|
|
B, H, W, C = pose_images.shape
|
|
|
|
pose_images = pose_images.clone()
|
|
pose_images = pose_images.permute(0, 3, 1, 2).unsqueeze(0)
|
|
#pose_images = pose_images * 2 - 1
|
|
#pose_images = pose_images.to(device).to(dtype)
|
|
|
|
if not hasattr(self, 'controlnext'):
|
|
self.controlnext = ControlNeXtSDVModel()
|
|
self.controlnext.load_state_dict(comfy.utils.load_torch_file(os.path.join(script_directory, 'models', 'controlnext-svd_v2-controlnet-fp16.safetensors')))
|
|
self.controlnext = self.controlnext.to(dtype).to(device).eval()
|
|
|
|
block_list = [int(x) for x in blocks.split(',')] #for testing, blocks 0-3 possible to apply to, 3 after skip so far best
|
|
|
|
def input_block_patch(h, transformer_options):
|
|
if transformer_options['block'][1] in block_list and 0 in transformer_options["cond_or_uncond"]:
|
|
|
|
sigma = transformer_options["sigmas"][0]
|
|
|
|
log_sigma = sigma.log()
|
|
min_log_sigma = torch.tensor(0.0002).log()
|
|
max_log_sigma = torch.tensor(700).log() #can I get these from the model?
|
|
normalized_log_sigma = (log_sigma - min_log_sigma) / (max_log_sigma - min_log_sigma)
|
|
|
|
#AnimateDiff-Evolved context windowing, is this method slower than it should be?
|
|
if "ad_params" in transformer_options and transformer_options["ad_params"]['sub_idxs'] is not None:
|
|
sub_idxs = transformer_options['ad_params']['sub_idxs']
|
|
controlnext_input = pose_images[:,sub_idxs].to(h.dtype).to(h.device).contiguous()
|
|
|
|
controlnext_input[:, 0, ...] = pose_images[:, 0, ...]
|
|
else:
|
|
controlnext_input = pose_images.to(h.dtype).to(h.device)
|
|
|
|
#print("controlnext_input shape: ", controlnext_input.shape)
|
|
#print("h shape: ", h.shape)
|
|
|
|
conditional_controls = self.controlnext(controlnext_input, normalized_log_sigma)['output']
|
|
|
|
mean_latents, std_latents = torch.mean(h, dim=(1, 2, 3), keepdim=True), torch.std(h, dim=(1, 2, 3), keepdim=True)
|
|
mean_control, std_control = torch.mean(conditional_controls, dim=(1, 2, 3), keepdim=True), torch.std(conditional_controls, dim=(1, 2, 3), keepdim=True)
|
|
conditional_controls = (conditional_controls - mean_control) * (std_latents / (std_control + 1e-5)) + mean_latents
|
|
conditional_controls = F.adaptive_avg_pool2d(conditional_controls, h.shape[-2:])
|
|
|
|
h = h + conditional_controls * 0.2 * strength
|
|
|
|
return h
|
|
model_clone = model.clone()
|
|
if not input_block_patch_after_skip:
|
|
model_clone.set_model_input_block_patch(input_block_patch)
|
|
else:
|
|
model_clone.set_model_input_block_patch_after_skip(input_block_patch)
|
|
|
|
return (model_clone, )
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DownloadAndLoadControlNeXt": DownloadAndLoadControlNeXt,
|
|
"ControlNextSampler": ControlNextSampler,
|
|
"ControlNextDecode": ControlNextDecode,
|
|
"ControlNextGetPoses": ControlNextGetPoses,
|
|
"ControlNextDiffusersScheduler": ControlNextDiffusersScheduler,
|
|
"ControlNextSVDApply": ControlNextSVDApply
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DownloadAndLoadControlNeXt": "(Down)Load ControlNeXt",
|
|
"ControlNextSampler": "ControlNext Sampler",
|
|
"ControlNextDecode": "ControlNext Decode",
|
|
"ControlNextGetPoses": "ControlNext GetPoses",
|
|
"ControlNextDiffusersScheduler": "ControlNext Diffusers Scheduler",
|
|
"ControlNextSVDApply": "ControlNext SVD Apply"
|
|
}
|