mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-09 12:47:07 +08:00
406 lines
16 KiB
Python
406 lines
16 KiB
Python
import json
|
|
import math
|
|
|
|
from comfy_extras.nodes_custom_sampler import Noise_RandomNoise
|
|
import latent_preview
|
|
from .modules.video_model import inject_model
|
|
import folder_paths
|
|
import node_helpers
|
|
import torch
|
|
import comfy
|
|
|
|
class MD_LoadVideoModel:
|
|
"""
|
|
Loads the DIT model for video generation.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"chkpt_name": (folder_paths.get_filename_list("checkpoints"), {
|
|
"default": "genesis-dit-video-2b.safetensors",
|
|
"tooltip": "The name of the checkpoint (model) to load."
|
|
}),
|
|
"clip_name": (folder_paths.get_filename_list("text_encoders"), {
|
|
"default": "t5xxl_fp16.safetensors",
|
|
"tooltip": "The name of the clip (model) to load."
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
RETURN_NAMES = ("model", "clip", "vae")
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "MemeDeck"
|
|
|
|
def load_model(self, chkpt_name, clip_name):
|
|
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", chkpt_name)
|
|
out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
|
model = out[0]
|
|
vae = out[2]
|
|
|
|
clip_path = folder_paths.get_full_path_or_raise("text_encoders", clip_name)
|
|
clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=8)
|
|
|
|
# modify model
|
|
model.model.diffusion_model = inject_model(model.model.diffusion_model)
|
|
|
|
return (model, clip, vae, )
|
|
|
|
class LatentGuide(torch.nn.Module):
|
|
def __init__(self, latent: torch.Tensor, index) -> None:
|
|
super().__init__()
|
|
self.index = index
|
|
self.register_buffer('latent', latent)
|
|
|
|
class MD_ImgToVideo:
|
|
"""
|
|
Sets the conditioning and dimensions for video generation.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"vae": ("VAE",),
|
|
"image": ("IMAGE",),
|
|
"width": ("INT", {
|
|
"default": 832,
|
|
"description": "The width of the video."
|
|
}),
|
|
"height": ("INT", {
|
|
"default": 832,
|
|
"description": "The height of the video."
|
|
}),
|
|
"length": ("INT", {
|
|
"default": 97,
|
|
"description": "The length of the video."
|
|
}),
|
|
"fps": ("INT", {
|
|
"default": 24,
|
|
"description": "The fps of the video."
|
|
}),
|
|
# LATENT GUIDE INPUTS
|
|
"add_latent_guide_index": ("INT", {
|
|
"default": 0,
|
|
"description": "The index of the latent to add to the guide."
|
|
}),
|
|
"add_latent_guide_insert": ("BOOLEAN", {
|
|
"default": False,
|
|
"description": "Whether to add the latent to the guide."
|
|
}),
|
|
# SCHEDULER INPUTS
|
|
"steps": ("INT", {
|
|
"default": 40,
|
|
"description": "Number of steps to generate the video."
|
|
}),
|
|
"max_shift": ("FLOAT", {
|
|
"default": 1.5,
|
|
"step": 0.01,
|
|
"description": "The maximum shift of the video."
|
|
}),
|
|
"base_shift": ("FLOAT", {
|
|
"default": 0.95,
|
|
"step": 0.01,
|
|
"description": "The base shift of the video."
|
|
}),
|
|
"stretch": ("BOOLEAN", {
|
|
"default": True,
|
|
"description": "Stretch the sigmas to be in the range [terminal, 1]."
|
|
}),
|
|
"terminal": ("FLOAT", {
|
|
"default": 0.1,
|
|
"step": 0.01,
|
|
"description": "The terminal values of the sigmas after stretching."
|
|
}),
|
|
# ATTENTION OVERRIDE INPUTS
|
|
"attention_override": ("STRING", {
|
|
"default": 14,
|
|
"description": "The amount of attention to override the model with."
|
|
}),
|
|
"attention_adjustment_scale": ("FLOAT", {
|
|
"default": 1.0,
|
|
"description": "The scale of the attention adjustment."
|
|
}),
|
|
"attention_adjustment_rescale": ("FLOAT", {
|
|
"default": 0.5,
|
|
"description": "The scale of the attention adjustment."
|
|
}),
|
|
"attention_adjustment_cfg": ("FLOAT", {
|
|
"default": 3.0,
|
|
"description": "The scale of the attention adjustment."
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING", "SIGMAS", "LATENT", "STRING")
|
|
RETURN_NAMES = ("model", "positive", "negative", "sigmas", "latent", "img2vid_metadata")
|
|
FUNCTION = "img_to_video"
|
|
CATEGORY = "MemeDeck"
|
|
|
|
def img_to_video(self, model, positive, negative, vae, image, width, height, length, fps, add_latent_guide_index, add_latent_guide_insert, steps, max_shift, base_shift, stretch, terminal, attention_override, attention_adjustment_scale, attention_adjustment_rescale, attention_adjustment_cfg):
|
|
batch_size = 1
|
|
pixels = comfy.utils.common_upscale(image.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
|
|
encode_pixels = pixels[:, :, :, :3]
|
|
t = vae.encode(encode_pixels)
|
|
positive = node_helpers.conditioning_set_values(positive, {"guiding_latent": t})
|
|
negative = node_helpers.conditioning_set_values(negative, {"guiding_latent": t})
|
|
|
|
latent = torch.zeros([batch_size, 128, ((length - 1) // 8) + 1, height // 32, width // 32], device=comfy.model_management.intermediate_device())
|
|
latent[:, :, :t.shape[2]] = t
|
|
latent_samples = {"samples": latent}
|
|
|
|
positive = node_helpers.conditioning_set_values(positive, {"frame_rate": fps})
|
|
negative = node_helpers.conditioning_set_values(negative, {"frame_rate": fps})
|
|
|
|
# 2. add latent guide
|
|
model, latent_updated = self.add_latent_guide(model, latent_samples, latent_samples, add_latent_guide_index, add_latent_guide_insert)
|
|
|
|
# 3. apply attention override
|
|
attn_override_layers = self.attention_override(attention_override)
|
|
model = self.apply_attention_override(model, attention_adjustment_scale, attention_adjustment_rescale, attention_adjustment_cfg, attn_override_layers)
|
|
|
|
# 5. configure scheduler
|
|
sigmas = self.get_sigmas(steps, max_shift, base_shift, stretch, terminal, latent_updated)
|
|
|
|
# all parameters starting with width, height, fps, crf, etc
|
|
img2vid_metadata = {
|
|
"width": width,
|
|
"height": height,
|
|
"length": length,
|
|
"fps": fps,
|
|
"steps": steps,
|
|
"max_shift": max_shift,
|
|
"base_shift": base_shift,
|
|
"stretch": stretch,
|
|
"terminal": terminal,
|
|
"attention_override": attention_override,
|
|
"attention_adjustment_scale": attention_adjustment_scale,
|
|
"attention_adjustment_rescale": attention_adjustment_rescale,
|
|
"attention_adjustment_cfg": attention_adjustment_cfg,
|
|
}
|
|
|
|
json_img2vid_metadata = json.dumps(img2vid_metadata)
|
|
return (model, positive, negative, sigmas, latent_updated, json_img2vid_metadata)
|
|
|
|
# -----------------------------
|
|
# Attention functions
|
|
# -----------------------------
|
|
# 1. Add latent guide
|
|
def add_latent_guide(self, model, latent, image_latent, index, insert):
|
|
image_latent = image_latent['samples']
|
|
latent = latent['samples'].clone()
|
|
|
|
# # Convert negative index to positive
|
|
# if insert:
|
|
# index = max(0, min(index, latent.shape[2])) # Clamp index
|
|
# latent = torch.cat([
|
|
# latent[:,:,:index],
|
|
# image_latent[:,:,0:1],
|
|
# latent[:,:,index:]
|
|
# ], dim=2)
|
|
# else:
|
|
# latent[:,:,index] = image_latent[:,:,0]
|
|
if insert:
|
|
# Handle insertion
|
|
if index == 0:
|
|
# Insert at beginning
|
|
latent = torch.cat([image_latent[:,:,0:1], latent], dim=2)
|
|
elif index >= latent.shape[2] or index < 0:
|
|
# Append to end
|
|
latent = torch.cat([latent, image_latent[:,:,0:1]], dim=2)
|
|
else:
|
|
# Insert in middle
|
|
latent = torch.cat([
|
|
latent[:,:,:index],
|
|
image_latent[:,:,0:1],
|
|
latent[:,:,index:]
|
|
], dim=2)
|
|
else:
|
|
# Original replacement behavior
|
|
latent[:,:,index] = image_latent[:,:,0]
|
|
|
|
model = model.clone()
|
|
guiding_latent = LatentGuide(image_latent, index)
|
|
model.set_model_patch(guiding_latent, 'guiding_latents')
|
|
|
|
return (model, {"samples": latent},)
|
|
|
|
# 2. Apply attention override
|
|
def is_integer(self, string):
|
|
try:
|
|
int(string)
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
def attention_override(self, layers: str = "14"):
|
|
try:
|
|
return set(map(int, layers.split(',')))
|
|
except ValueError:
|
|
return set()
|
|
|
|
# layers_map = set([])
|
|
# return set(map(int, layers.split(',')))
|
|
# for block in layers.split(','):
|
|
# block = block.strip()
|
|
# if self.is_integer(block):
|
|
# layers_map.add(block)
|
|
|
|
# return layers_map
|
|
|
|
def apply_attention_override(self, model, scale, rescale, cfg, attention_override: set):
|
|
m = model.clone()
|
|
|
|
def pag_fn(q, k,v, heads, attn_precision=None, transformer_options=None):
|
|
return v
|
|
|
|
def post_cfg_function(args):
|
|
model = args["model"]
|
|
|
|
cond_pred = args["cond_denoised"]
|
|
uncond_pred = args["uncond_denoised"]
|
|
|
|
len_conds = 1 if args.get('uncond', None) is None else 2
|
|
|
|
cond = args["cond"]
|
|
sigma = args["sigma"]
|
|
model_options = args["model_options"].copy()
|
|
x = args["input"]
|
|
|
|
if scale == 0:
|
|
if len_conds == 1:
|
|
return cond_pred
|
|
return uncond_pred + (cond_pred - uncond_pred)
|
|
|
|
for block_idx in attention_override:
|
|
model_options = comfy.model_patcher.set_model_options_patch_replace(model_options, pag_fn, f"layer", "self_attn", int(block_idx))
|
|
|
|
(perturbed,) = comfy.samplers.calc_cond_batch(model, [cond], x, sigma, model_options)
|
|
|
|
output = uncond_pred + cfg * (cond_pred - uncond_pred) \
|
|
+ scale * (cond_pred - perturbed)
|
|
if rescale > 0:
|
|
factor = cond_pred.std() / output.std()
|
|
factor = rescale * factor + (1 - rescale)
|
|
output = output * factor
|
|
|
|
return output
|
|
|
|
|
|
m.set_model_sampler_post_cfg_function(post_cfg_function)
|
|
|
|
return m
|
|
|
|
# -----------------------------
|
|
# Scheduler
|
|
# -----------------------------
|
|
def get_sigmas(self, steps, max_shift, base_shift, stretch, terminal, latent=None):
|
|
if latent is None:
|
|
tokens = 4096
|
|
else:
|
|
tokens = math.prod(latent["samples"].shape[2:])
|
|
|
|
sigmas = torch.linspace(1.0, 0.0, steps + 1)
|
|
|
|
x1 = 1024
|
|
x2 = 4096
|
|
mm = (max_shift - base_shift) / (x2 - x1)
|
|
b = base_shift - mm * x1
|
|
sigma_shift = (tokens) * mm + b
|
|
|
|
power = 1
|
|
sigmas = torch.where(
|
|
sigmas != 0,
|
|
math.exp(sigma_shift) / (math.exp(sigma_shift) + (1 / sigmas - 1) ** power),
|
|
0,
|
|
)
|
|
|
|
# Stretch sigmas so that its final value matches the given terminal value.
|
|
if stretch:
|
|
non_zero_mask = sigmas != 0
|
|
non_zero_sigmas = sigmas[non_zero_mask]
|
|
one_minus_z = 1.0 - non_zero_sigmas
|
|
scale_factor = one_minus_z[-1] / (1.0 - terminal)
|
|
stretched = 1.0 - (one_minus_z / scale_factor)
|
|
sigmas[non_zero_mask] = stretched
|
|
|
|
return sigmas
|
|
|
|
|
|
KSAMPLER_NAMES = ["euler", "ddim", "euler_ancestral", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_cfg_pp", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
|
|
"lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_2s_ancestral_cfg_pp", "dpmpp_sde", "dpmpp_sde_gpu",
|
|
"dpmpp_2m", "dpmpp_2m_cfg_pp", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
|
|
"ipndm", "ipndm_v", "deis"]
|
|
|
|
class MD_VideoSampler:
|
|
"""
|
|
Samples the video.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"sigmas": ("SIGMAS",),
|
|
"latent_image": ("LATENT",),
|
|
"sampler": (KSAMPLER_NAMES, ),
|
|
"noise_seed": ("INT", {
|
|
"default": 42,
|
|
"description": "The seed of the noise."
|
|
}),
|
|
"cfg": ("FLOAT", {
|
|
"default": 5.0,
|
|
"min": 0.0,
|
|
"max": 30.0,
|
|
"step": 0.01,
|
|
"description": "The cfg of the video."
|
|
}),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("LATENT", "LATENT", "STRING")
|
|
RETURN_NAMES = ("output", "denoised_output", "img2vid_metadata")
|
|
FUNCTION = "video_sampler"
|
|
CATEGORY = "MemeDeck"
|
|
|
|
def video_sampler(self, model, positive, negative, sigmas, latent_image, sampler, noise_seed, cfg):
|
|
latent = latent_image
|
|
latent_image = latent["samples"]
|
|
latent = latent.copy()
|
|
latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image)
|
|
latent["samples"] = latent_image
|
|
|
|
sampler_name = sampler
|
|
noise = Noise_RandomNoise(noise_seed).generate_noise(latent)
|
|
sampler = comfy.samplers.sampler_object(sampler)
|
|
|
|
noise_mask = None
|
|
if "noise_mask" in latent:
|
|
noise_mask = latent["noise_mask"]
|
|
|
|
x0_output = {}
|
|
callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
|
|
|
|
disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
|
|
samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
|
|
|
|
out = latent.copy()
|
|
out["samples"] = samples
|
|
if "x0" in x0_output:
|
|
out_denoised = latent.copy()
|
|
out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
|
|
else:
|
|
out_denoised = out
|
|
|
|
sampler_metadata = {
|
|
"sampler": sampler_name,
|
|
"noise_seed": noise_seed,
|
|
"cfg": cfg,
|
|
}
|
|
|
|
json_sampler_metadata = json.dumps(sampler_metadata)
|
|
return (out, out_denoised, json_sampler_metadata) |