156 lines
4.5 KiB
Python
156 lines
4.5 KiB
Python
from .keysync_wrapper.download_utils import download_or_check_models
|
|
from .keysync_wrapper.model_utils import load_keysync_model_config
|
|
from .keysync_wrapper.runner import run_keysync
|
|
|
|
|
|
DEFAULT_MODELS_DIR = "D:/ComfyUI/models/keysync"
|
|
DEFAULT_KEYSYNC_REPO_PATH = "D:/ComfyUI/custom_nodes/ComfyUI-Keysync-wrapper/keysync"
|
|
|
|
|
|
class KeySyncDownloadModels:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"download": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("models_dir",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "video/lipsync"
|
|
|
|
def run(self, download):
|
|
resolved_models_dir = download_or_check_models(
|
|
models_dir=DEFAULT_MODELS_DIR,
|
|
repo_id="toninio19/keysync",
|
|
download=download,
|
|
)
|
|
return (resolved_models_dir,)
|
|
|
|
|
|
class KeySyncModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {}}
|
|
|
|
RETURN_TYPES = ("KEYSYNC_MODEL",)
|
|
RETURN_NAMES = ("keysync_model",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "video/lipsync"
|
|
|
|
def load(self):
|
|
model = load_keysync_model_config(
|
|
keysync_repo_path=DEFAULT_KEYSYNC_REPO_PATH,
|
|
models_dir=DEFAULT_MODELS_DIR,
|
|
keyframes_ckpt="checkpoints/keyframe_dub.pt",
|
|
interpolation_ckpt="checkpoints/interpolation_dub.pt",
|
|
)
|
|
return (model,)
|
|
|
|
|
|
class KeySyncLipSync:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"keysync_model": ("KEYSYNC_MODEL",),
|
|
"images": ("IMAGE",),
|
|
"audio": ("AUDIO",),
|
|
"input_fps": ("INT", {"default": 25, "min": 1, "max": 120, "step": 1}),
|
|
"compute_until": ("INT", {"default": 30, "min": 1, "max": 300, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "AUDIO", "INT", "STRING")
|
|
RETURN_NAMES = ("images", "audio", "frame_rate", "output_video_path")
|
|
FUNCTION = "run"
|
|
CATEGORY = "video/lipsync"
|
|
|
|
def run(
|
|
self,
|
|
keysync_model,
|
|
images,
|
|
audio,
|
|
input_fps,
|
|
compute_until,
|
|
):
|
|
return run_keysync(
|
|
keysync_model=keysync_model,
|
|
images=images,
|
|
audio=audio,
|
|
input_fps=input_fps,
|
|
compute_until=compute_until,
|
|
fix_occlusion=False,
|
|
position="None",
|
|
start_frame=0,
|
|
force_25fps_16k_audio=True,
|
|
trim_audio_to_frames=True,
|
|
keep_intermediate_outputs=True,
|
|
subprocess_timeout_sec=0,
|
|
use_subprocess=True,
|
|
)
|
|
|
|
|
|
class KeySyncLipSyncAdvanced(KeySyncLipSync):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"keysync_model": ("KEYSYNC_MODEL",),
|
|
"images": ("IMAGE",),
|
|
"audio": ("AUDIO",),
|
|
"input_fps": ("INT", {"default": 25, "min": 1, "max": 120, "step": 1}),
|
|
"compute_until": ("INT", {"default": 30, "min": 1, "max": 300, "step": 1}),
|
|
"fix_occlusion": ("BOOLEAN", {"default": False}),
|
|
"position": ("STRING", {"default": "None", "multiline": False}),
|
|
"start_frame": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}),
|
|
}
|
|
}
|
|
|
|
FUNCTION = "run_advanced"
|
|
CATEGORY = "video/lipsync"
|
|
|
|
def run_advanced(
|
|
self,
|
|
keysync_model,
|
|
images,
|
|
audio,
|
|
input_fps,
|
|
compute_until,
|
|
fix_occlusion,
|
|
position,
|
|
start_frame,
|
|
):
|
|
return run_keysync(
|
|
keysync_model=keysync_model,
|
|
images=images,
|
|
audio=audio,
|
|
input_fps=input_fps,
|
|
compute_until=compute_until,
|
|
fix_occlusion=fix_occlusion,
|
|
position=position,
|
|
start_frame=start_frame,
|
|
force_25fps_16k_audio=True,
|
|
trim_audio_to_frames=True,
|
|
keep_intermediate_outputs=True,
|
|
subprocess_timeout_sec=0,
|
|
use_subprocess=True,
|
|
)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"KeySyncDownloadModels": KeySyncDownloadModels,
|
|
"KeySyncModelLoader": KeySyncModelLoader,
|
|
"KeySyncLipSync": KeySyncLipSync,
|
|
"KeySyncLipSyncAdvanced": KeySyncLipSyncAdvanced,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"KeySyncDownloadModels": "KeySync Download Models",
|
|
"KeySyncModelLoader": "KeySync Model Loader",
|
|
"KeySyncLipSync": "KeySync LipSync",
|
|
"KeySyncLipSyncAdvanced": "KeySync LipSync Advanced",
|
|
}
|