223 lines
20 KiB
Python
223 lines
20 KiB
Python
from comfy_api.latest import io
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
from typing import Dict, List
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
import folder_paths
|
|
from comfy.model_patcher import ModelPatcher
|
|
from .ad_settings import AnimateDiffSettings, AdjustGroup, AdjustPE, AdjustWeight
|
|
from .context import ContextOptionsGroup, ContextOptions, ContextSchedules
|
|
from .logger import logger
|
|
from .utils_model import Folders, BetaSchedules, get_available_motion_models
|
|
from .utils_motion import ADKeyframeGroup
|
|
from .motion_lora import MotionLoraList
|
|
from .model_injection import ModelPatcherHelper, InjectionParams, MotionModelGroup, get_mm_attachment, load_motion_module_gen1
|
|
from .sampling import outer_sample_wrapper, sliding_calc_cond_batch
|
|
from .sample_settings import SampleSettings
|
|
|
|
class AnimateDiffLoaderDEPR(io.ComfyNode):
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(node_id='AnimateDiffLoaderV1', display_name='🚫AnimateDiff Loader [DEPRECATED] 🎭🅐🅓', category='', inputs=[io.Model.Input('model'), io.Custom('LATENT').Input('latents'), io.Combo.Input('model_name', options=get_available_motion_models()), io.Boolean.Input('unlimited_area_hack', default=False), io.Combo.Input('beta_schedule', options=['sqrt_linear (AnimateDiff)', 'use existing', 'autoselect', 'linear (AnimateDiff-SDXL)', 'linear (HotshotXL/default)', 'avg(sqrt_linear,linear)', 'lcm avg(sqrt_linear,linear)', 'lcm', 'lcm[100_ots]', 'lcm >> sqrt_linear', 'sqrt', 'cosine', 'squaredcos_cap_v2'])], outputs=[io.Model.Output('MODEL'), io.Custom('LATENT').Output('LATENT')], is_deprecated=True)
|
|
|
|
@classmethod
|
|
def execute(cls, model: ModelPatcher, latents: Dict[str, torch.Tensor], model_name: str, unlimited_area_hack: bool, beta_schedule: str):
|
|
motion_model = load_motion_module_gen1(model_name, model)
|
|
init_frames_len = len(latents['samples'])
|
|
params = InjectionParams(unlimited_area_hack=unlimited_area_hack, apply_v2_properly=False)
|
|
model = model.clone()
|
|
helper = ModelPatcherHelper(model)
|
|
helper.set_all_properties(outer_sampler_wrapper=outer_sample_wrapper, calc_cond_batch_wrapper=sliding_calc_cond_batch, params=params, motion_models=MotionModelGroup(motion_model))
|
|
if beta_schedule == BetaSchedules.AUTOSELECT and (not model.motion_models.is_empty()):
|
|
beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
|
|
new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
|
|
if new_model_sampling is not None:
|
|
model.add_object_patch('model_sampling', new_model_sampling)
|
|
del motion_model
|
|
return io.NodeOutput(model, latents)
|
|
|
|
class AnimateDiffLoaderAdvancedDEPR(io.ComfyNode):
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(node_id='ADE_AnimateDiffLoaderV1Advanced', display_name='🚫AnimateDiff Loader (Advanced) [DEPRECATED] 🎭🅐🅓', category='', inputs=[io.Model.Input('model'), io.Custom('LATENT').Input('latents'), io.Combo.Input('model_name', options=get_available_motion_models()), io.Boolean.Input('unlimited_area_hack', default=False), io.Int.Input('context_length', default=16, max=1000, min=0), io.Int.Input('context_stride', default=1, max=1000, min=1), io.Int.Input('context_overlap', default=4, max=1000, min=0), io.Combo.Input('context_schedule', options=['uniform']), io.Boolean.Input('closed_loop', default=False), io.Combo.Input('beta_schedule', options=['sqrt_linear (AnimateDiff)', 'use existing', 'autoselect', 'linear (AnimateDiff-SDXL)', 'linear (HotshotXL/default)', 'avg(sqrt_linear,linear)', 'lcm avg(sqrt_linear,linear)', 'lcm', 'lcm[100_ots]', 'lcm >> sqrt_linear', 'sqrt', 'cosine', 'squaredcos_cap_v2'])], outputs=[io.Model.Output('MODEL'), io.Custom('LATENT').Output('LATENT')], is_deprecated=True)
|
|
|
|
@classmethod
|
|
def execute(cls, model: ModelPatcher, latents: Dict[str, torch.Tensor], model_name: str, unlimited_area_hack: bool, context_length: int, context_stride: int, context_overlap: int, context_schedule: str, closed_loop: bool, beta_schedule: str):
|
|
motion_model = load_motion_module_gen1(model_name, model)
|
|
init_frames_len = len(latents['samples'])
|
|
params = InjectionParams(unlimited_area_hack=unlimited_area_hack, apply_v2_properly=False)
|
|
context_group = ContextOptionsGroup()
|
|
context_group.add(ContextOptions(context_length=context_length, context_stride=context_stride, context_overlap=context_overlap, context_schedule=context_schedule, closed_loop=closed_loop))
|
|
params.set_context(context_options=context_group)
|
|
model = model.clone()
|
|
helper = ModelPatcherHelper(model)
|
|
helper.set_all_properties(outer_sampler_wrapper=outer_sample_wrapper, calc_cond_batch_wrapper=sliding_calc_cond_batch, params=params, motion_models=MotionModelGroup(motion_model))
|
|
if beta_schedule == BetaSchedules.AUTOSELECT and (not model.motion_models.is_empty()):
|
|
beta_schedule = model.motion_models[0].model.get_best_beta_schedule(log=True)
|
|
new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
|
|
if new_model_sampling is not None:
|
|
model.add_object_patch('model_sampling', new_model_sampling)
|
|
del motion_model
|
|
return io.NodeOutput(model, latents)
|
|
|
|
class LegacyAnimateDiffLoaderWithContextDEPR(io.ComfyNode):
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(node_id='ADE_AnimateDiffLoaderWithContext', display_name='AnimateDiff Loader [Legacy] 🎭🅐🅓①', category='Animate Diff 🎭🅐🅓/① Gen1 nodes ①', inputs=[io.Model.Input('model'), io.Combo.Input('model_name', options=get_available_motion_models()), io.Combo.Input('beta_schedule', options=['autoselect', 'use existing', 'sqrt_linear (AnimateDiff)', 'linear (AnimateDiff-SDXL)', 'linear (HotshotXL/default)', 'avg(sqrt_linear,linear)', 'lcm avg(sqrt_linear,linear)', 'lcm', 'lcm[100_ots]', 'lcm >> sqrt_linear', 'sqrt', 'cosine', 'squaredcos_cap_v2'], default='autoselect'), io.Custom('CONTEXT_OPTIONS').Input('context_options', optional=True), io.Custom('MOTION_LORA').Input('motion_lora', optional=True), io.Custom('AD_SETTINGS').Input('ad_settings', optional=True), io.Custom('SAMPLE_SETTINGS').Input('sample_settings', optional=True), io.Float.Input('motion_scale', optional=True, default=1.0, min=0.0, step=0.001), io.Boolean.Input('apply_v2_models_properly', optional=True, default=True), io.Custom('AD_KEYFRAMES').Input('ad_keyframes', optional=True)], outputs=[io.Model.Output('MODEL')], is_deprecated=True)
|
|
|
|
@classmethod
|
|
def execute(cls, model: ModelPatcher, model_name: str, beta_schedule: str, context_options: ContextOptionsGroup=None, motion_lora: MotionLoraList=None, ad_settings: AnimateDiffSettings=None, motion_model_settings: AnimateDiffSettings=None, sample_settings: SampleSettings=None, motion_scale: float=1.0, apply_v2_models_properly: bool=False, ad_keyframes: ADKeyframeGroup=None):
|
|
if ad_settings is not None:
|
|
motion_model_settings = ad_settings
|
|
motion_model = load_motion_module_gen1(model_name, model, motion_lora=motion_lora, motion_model_settings=motion_model_settings)
|
|
params = InjectionParams(unlimited_area_hack=False, apply_v2_properly=apply_v2_models_properly)
|
|
if context_options:
|
|
params.set_context(context_options)
|
|
if not motion_model_settings:
|
|
motion_model_settings = AnimateDiffSettings()
|
|
motion_model_settings.attn_scale = motion_scale
|
|
params.set_motion_model_settings(motion_model_settings)
|
|
attachment = get_mm_attachment(motion_model)
|
|
if params.motion_model_settings.mask_attn_scale is not None:
|
|
attachment.scale_multival = params.motion_model_settings.mask_attn_scale * params.motion_model_settings.attn_scale
|
|
else:
|
|
attachment.scale_multival = params.motion_model_settings.attn_scale
|
|
attachment.keyframes = ad_keyframes.clone() if ad_keyframes else ADKeyframeGroup()
|
|
model = model.clone()
|
|
helper = ModelPatcherHelper(model)
|
|
helper.set_all_properties(outer_sampler_wrapper=outer_sample_wrapper, calc_cond_batch_wrapper=sliding_calc_cond_batch, params=params, sample_settings=sample_settings, motion_models=MotionModelGroup(motion_model))
|
|
sample_settings = helper.get_sample_settings()
|
|
if sample_settings.custom_cfg is not None:
|
|
logger.info('[Sample Settings] custom_cfg is set; will override any KSampler cfg values or patches.')
|
|
if sample_settings.sigma_schedule is not None:
|
|
logger.info('[Sample Settings] sigma_schedule is set; will override beta_schedule.')
|
|
model.add_object_patch('model_sampling', sample_settings.sigma_schedule.clone().model_sampling)
|
|
else:
|
|
if beta_schedule == BetaSchedules.AUTOSELECT and helper.get_motion_models():
|
|
beta_schedule = helper.get_motion_models()[0].model.get_best_beta_schedule(log=True)
|
|
new_model_sampling = BetaSchedules.to_model_sampling(beta_schedule, model)
|
|
if new_model_sampling is not None:
|
|
model.add_object_patch('model_sampling', new_model_sampling)
|
|
del motion_model
|
|
return io.NodeOutput(model)
|
|
|
|
class AnimateDiffCombineDEPR(io.ComfyNode):
|
|
|
|
@classmethod
|
|
def get_formats(cls):
|
|
ffmpeg_path = shutil.which('ffmpeg')
|
|
if ffmpeg_path is not None:
|
|
return ['image/gif', 'image/webp'] + ['video/' + x[:-5] for x in folder_paths.get_filename_list(Folders.VIDEO_FORMATS)]
|
|
cls.ffmpeg_warning_already_shown = True
|
|
return ['image/gif', 'image/webp']
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(node_id='ADE_AnimateDiffCombine', display_name='🚫AnimateDiff Combine [DEPRECATED, Use Video Combine (VHS) Instead!] 🎭🅐🅓', category='', inputs=[io.Image.Input('images'), io.Int.Input('frame_rate', default=8, max=24, min=1, step=1), io.Int.Input('loop_count', default=0, max=100, min=0, step=1), io.String.Input('filename_prefix', default='AnimateDiff'), io.Combo.Input('format', options=cls.get_formats()), io.Boolean.Input('pingpong', default=False), io.Boolean.Input('save_image', default=True)], outputs=[io.Custom('GIF').Output('GIF')], is_deprecated=True, is_output_node=True, hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo])
|
|
ffmpeg_warning_already_shown = False
|
|
|
|
@classmethod
|
|
def execute(cls, images, frame_rate: int, loop_count: int, filename_prefix='AnimateDiff', format='image/gif', pingpong=False, save_image=True, prompt=None, extra_pnginfo=None):
|
|
prompt = cls.hidden.prompt
|
|
extra_pnginfo = cls.hidden.extra_pnginfo
|
|
logger.warning('Do not use AnimateDiff Combine node, it is deprecated. Use Video Combine node from ComfyUI-VideoHelperSuite instead. Video nodes from VideoHelperSuite are actively maintained, more feature-rich, and also automatically attempts to get ffmpeg.')
|
|
frames: List[Image.Image] = []
|
|
for image in images:
|
|
img = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
|
|
frames.append(img)
|
|
output_dir = folder_paths.get_output_directory() if save_image else folder_paths.get_temp_directory()
|
|
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(filename_prefix, output_dir)
|
|
metadata = PngInfo()
|
|
if prompt is not None:
|
|
metadata.add_text('prompt', json.dumps(prompt))
|
|
if extra_pnginfo is not None:
|
|
for x in extra_pnginfo:
|
|
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
|
file = f'{filename}_{counter:05}_.png'
|
|
file_path = os.path.join(full_output_folder, file)
|
|
frames[0].save(file_path, pnginfo=metadata, compress_level=4)
|
|
if pingpong:
|
|
frames = frames + frames[-2:0:-1]
|
|
format_type, format_ext = format.split('/')
|
|
file = f'{filename}_{counter:05}_.{format_ext}'
|
|
file_path = os.path.join(full_output_folder, file)
|
|
if format_type == 'image':
|
|
frames[0].save(file_path, format=format_ext.upper(), save_all=True, append_images=frames[1:], duration=round(1000 / frame_rate), loop=loop_count, compress_level=4)
|
|
else:
|
|
ffmpeg_path = shutil.which('ffmpeg')
|
|
if ffmpeg_path is None:
|
|
raise ProcessLookupError('Could not find ffmpeg')
|
|
video_format_path = folder_paths.get_full_path('video_formats', format_ext + '.json')
|
|
with open(video_format_path, 'r') as stream:
|
|
video_format = json.load(stream)
|
|
file = f"{filename}_{counter:05}_.{video_format['extension']}"
|
|
file_path = os.path.join(full_output_folder, file)
|
|
dimensions = f'{frames[0].width}x{frames[0].height}'
|
|
args = [ffmpeg_path, '-v', 'error', '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-s', dimensions, '-r', str(frame_rate), '-i', '-'] + video_format['main_pass'] + [file_path]
|
|
env = os.environ.copy()
|
|
if 'environment' in video_format:
|
|
env.update(video_format['environment'])
|
|
with subprocess.Popen(args, stdin=subprocess.PIPE, env=env) as proc:
|
|
for frame in frames:
|
|
proc.stdin.write(frame.tobytes())
|
|
previews = [{'filename': file, 'subfolder': subfolder, 'type': 'output' if save_image else 'temp', 'format': format}]
|
|
return io.NodeOutput.from_dict({'ui': {'images': previews, 'animated': (True,)}})
|
|
|
|
class AnimateDiffModelSettingsDEPR(io.ComfyNode):
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(node_id='ADE_AnimateDiffModelSettings_Release', display_name='🚫[DEPR] Motion Model Settings 🎭🅐🅓①', category='', inputs=[io.Float.Input('min_motion_scale', default=1.0, min=0.0, step=0.001), io.Float.Input('max_motion_scale', default=1.0, min=0.0, step=0.001), io.Mask.Input('mask_motion_scale', optional=True)], outputs=[io.Custom('AD_SETTINGS').Output('AD_SETTINGS')], is_deprecated=True)
|
|
|
|
@classmethod
|
|
def execute(cls, mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
|
|
motion_model_settings = AnimateDiffSettings(mask_attn_scale=mask_motion_scale, mask_attn_scale_min=min_motion_scale, mask_attn_scale_max=max_motion_scale)
|
|
return io.NodeOutput(motion_model_settings)
|
|
|
|
class AnimateDiffModelSettingsSimpleDEPR(io.ComfyNode):
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(node_id='ADE_AnimateDiffModelSettingsSimple', display_name='🚫[DEPR] Motion Model Settings (Simple) 🎭🅐🅓①', category='', inputs=[io.Int.Input('motion_pe_stretch', default=0, min=0, step=1), io.Mask.Input('mask_motion_scale', optional=True), io.Float.Input('min_motion_scale', optional=True, default=1.0, min=0.0, step=0.001), io.Float.Input('max_motion_scale', optional=True, default=1.0, min=0.0, step=0.001)], outputs=[io.Custom('AD_SETTINGS').Output('AD_SETTINGS')], is_deprecated=True)
|
|
|
|
@classmethod
|
|
def execute(cls, motion_pe_stretch: int, mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
|
|
adjust_pe = AdjustGroup(AdjustPE(motion_pe_stretch=motion_pe_stretch))
|
|
motion_model_settings = AnimateDiffSettings(adjust_pe=adjust_pe, mask_attn_scale=mask_motion_scale, mask_attn_scale_min=min_motion_scale, mask_attn_scale_max=max_motion_scale)
|
|
return io.NodeOutput(motion_model_settings)
|
|
|
|
class AnimateDiffModelSettingsAdvancedDEPR(io.ComfyNode):
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(node_id='ADE_AnimateDiffModelSettings', display_name='🚫[DEPR] Motion Model Settings (Advanced) 🎭🅐🅓①', category='', inputs=[io.Float.Input('pe_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Float.Input('attn_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Float.Input('other_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Int.Input('motion_pe_stretch', default=0, min=0, step=1), io.Int.Input('cap_initial_pe_length', default=0, min=0, step=1), io.Int.Input('interpolate_pe_to_length', default=0, min=0, step=1), io.Int.Input('initial_pe_idx_offset', default=0, min=0, step=1), io.Int.Input('final_pe_idx_offset', default=0, min=0, step=1), io.Mask.Input('mask_motion_scale', optional=True), io.Float.Input('min_motion_scale', optional=True, default=1.0, min=0.0, step=0.001), io.Float.Input('max_motion_scale', optional=True, default=1.0, min=0.0, step=0.001)], outputs=[io.Custom('AD_SETTINGS').Output('AD_SETTINGS')], is_deprecated=True)
|
|
|
|
@classmethod
|
|
def execute(cls, pe_strength: float, attn_strength: float, other_strength: float, motion_pe_stretch: int, cap_initial_pe_length: int, interpolate_pe_to_length: int, initial_pe_idx_offset: int, final_pe_idx_offset: int, mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
|
|
adjust_pe = AdjustGroup(AdjustPE(motion_pe_stretch=motion_pe_stretch, cap_initial_pe_length=cap_initial_pe_length, interpolate_pe_to_length=interpolate_pe_to_length, initial_pe_idx_offset=initial_pe_idx_offset, final_pe_idx_offset=final_pe_idx_offset))
|
|
adjust_weight = AdjustGroup(AdjustWeight(pe_MULT=pe_strength, attn_MULT=attn_strength, other_MULT=other_strength))
|
|
motion_model_settings = AnimateDiffSettings(adjust_pe=adjust_pe, adjust_weight=adjust_weight, mask_attn_scale=mask_motion_scale, mask_attn_scale_min=min_motion_scale, mask_attn_scale_max=max_motion_scale)
|
|
return io.NodeOutput(motion_model_settings)
|
|
|
|
class AnimateDiffModelSettingsAdvancedAttnStrengthsDEPR(io.ComfyNode):
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(node_id='ADE_AnimateDiffModelSettingsAdvancedAttnStrengths', display_name='🚫[DEPR] Motion Model Settings (Adv. Attn) 🎭🅐🅓①', category='', inputs=[io.Float.Input('pe_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Float.Input('attn_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Float.Input('attn_q_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Float.Input('attn_k_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Float.Input('attn_v_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Float.Input('attn_out_weight_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Float.Input('attn_out_bias_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Float.Input('other_strength', default=1.0, max=10.0, min=0.0, step=0.0001), io.Int.Input('motion_pe_stretch', default=0, min=0, step=1), io.Int.Input('cap_initial_pe_length', default=0, min=0, step=1), io.Int.Input('interpolate_pe_to_length', default=0, min=0, step=1), io.Int.Input('initial_pe_idx_offset', default=0, min=0, step=1), io.Int.Input('final_pe_idx_offset', default=0, min=0, step=1), io.Mask.Input('mask_motion_scale', optional=True), io.Float.Input('min_motion_scale', optional=True, default=1.0, min=0.0, step=0.001), io.Float.Input('max_motion_scale', optional=True, default=1.0, min=0.0, step=0.001)], outputs=[io.Custom('AD_SETTINGS').Output('AD_SETTINGS')], is_deprecated=True)
|
|
|
|
@classmethod
|
|
def execute(cls, pe_strength: float, attn_strength: float, attn_q_strength: float, attn_k_strength: float, attn_v_strength: float, attn_out_weight_strength: float, attn_out_bias_strength: float, other_strength: float, motion_pe_stretch: int, cap_initial_pe_length: int, interpolate_pe_to_length: int, initial_pe_idx_offset: int, final_pe_idx_offset: int, mask_motion_scale: torch.Tensor=None, min_motion_scale: float=1.0, max_motion_scale: float=1.0):
|
|
adjust_pe = AdjustGroup(AdjustPE(motion_pe_stretch=motion_pe_stretch, cap_initial_pe_length=cap_initial_pe_length, interpolate_pe_to_length=interpolate_pe_to_length, initial_pe_idx_offset=initial_pe_idx_offset, final_pe_idx_offset=final_pe_idx_offset))
|
|
adjust_weight = AdjustGroup(AdjustWeight(pe_MULT=pe_strength, attn_MULT=attn_strength, attn_q_MULT=attn_q_strength, attn_k_MULT=attn_k_strength, attn_v_MULT=attn_v_strength, attn_out_weight_MULT=attn_out_weight_strength, attn_out_bias_MULT=attn_out_bias_strength, other_MULT=other_strength))
|
|
motion_model_settings = AnimateDiffSettings(adjust_pe=adjust_pe, adjust_weight=adjust_weight, mask_attn_scale=mask_motion_scale, mask_attn_scale_min=min_motion_scale, mask_attn_scale_max=max_motion_scale)
|
|
return io.NodeOutput(motion_model_settings)
|