579 lines
16 KiB
Python
579 lines
16 KiB
Python
"""
|
|
@author: Fuou Marinas
|
|
@title: ComfyUI-EbSynth
|
|
@nickname: EbSynth
|
|
@description: Run EbSynth in ComfyUI.
|
|
"""
|
|
|
|
import cv2
|
|
import torch
|
|
|
|
from .constants import EBTYPES, MAX_GUIDES, ONLY_DEFAULT_MODE, ONLY_MODES
|
|
from .Ezsynth.ezsynth.aux_classes import RunConfig
|
|
from .Ezsynth.ezsynth.constants import (
|
|
DEFAULT_EDGE_METHOD,
|
|
DEFAULT_FLOW_MODEL,
|
|
DEFAULT_FLOW_ARCH,
|
|
EDGE_METHODS,
|
|
FLOW_ARCHS,
|
|
FLOW_MODELS,
|
|
EF_RAFT_MODELS,
|
|
FLOW_DIFF_MODEL,
|
|
)
|
|
from .Ezsynth.ezsynth.main_ez import EzsynthBase, ImageSynthBase
|
|
from .utils import (
|
|
batched_tensor_to_cv2_list,
|
|
cv2_img_to_tensor,
|
|
out_video,
|
|
resize_cv2_list,
|
|
)
|
|
|
|
flow_all_models = []
|
|
flow_all_models.extend([f"RAFT_{model}" for model in FLOW_MODELS])
|
|
flow_all_models.extend([f"EfRAFT_{model}" for model in EF_RAFT_MODELS])
|
|
flow_all_models.append(f"FLOWDIFF_{FLOW_DIFF_MODEL}")
|
|
|
|
|
|
def create_guide_arg(i: int):
|
|
return {
|
|
f"src_img_{i}": ("IMAGE",),
|
|
f"tgt_img_{i}": ("IMAGE",),
|
|
f"weight_{i}": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.1},
|
|
),
|
|
}
|
|
|
|
|
|
def build_guide_args(max_guides=MAX_GUIDES):
|
|
required = create_guide_arg(0)
|
|
optional = dict()
|
|
for i in range(1, max_guides):
|
|
optional.update(create_guide_arg(i))
|
|
return {"required": required, "optional": optional}
|
|
|
|
|
|
class ES_Guides7:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return build_guide_args()
|
|
|
|
RETURN_TYPES = (
|
|
"IMAGE",
|
|
"IMAGE",
|
|
"STRING",
|
|
)
|
|
RETURN_NAMES = (
|
|
"src_imgs",
|
|
"tgt_imgs",
|
|
"wgts",
|
|
)
|
|
|
|
FUNCTION = "todo"
|
|
CATEGORY = "EbSynth"
|
|
|
|
def todo(self, **args):
|
|
src_imgs = []
|
|
tgt_imgs = []
|
|
weights = []
|
|
|
|
for i in range(MAX_GUIDES):
|
|
src_key = f"src_img_{i}"
|
|
tgt_key = f"tgt_img_{i}"
|
|
weight_key = f"weight_{i}"
|
|
if src_key in args and tgt_key in args:
|
|
src_imgs.append(args[src_key])
|
|
tgt_imgs.append(args[tgt_key])
|
|
weights.append(args[weight_key])
|
|
|
|
src_tensor = torch.cat(src_imgs, dim=0)
|
|
tgt_tensor = torch.cat(tgt_imgs, dim=0)
|
|
weights_str = serialize_floats(weights)
|
|
|
|
return (src_tensor, tgt_tensor, weights_str)
|
|
|
|
|
|
class ES_Translate:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
required = {
|
|
"style_image": ("IMAGE",),
|
|
"source_image": ("IMAGE",),
|
|
"target_image": ("IMAGE",),
|
|
"weight": (
|
|
"FLOAT",
|
|
{"default": 0.9, "min": 0.1},
|
|
),
|
|
}
|
|
required.update(EBTYPES)
|
|
return {
|
|
"required": required,
|
|
"optional": {
|
|
"src_imgs": ("IMAGE",),
|
|
"tgt_imgs": ("IMAGE",),
|
|
"wgts": ("STRING", {}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"IMAGE",
|
|
"IMAGE",
|
|
)
|
|
RETURN_NAMES = (
|
|
"result_img",
|
|
"error_img",
|
|
)
|
|
FUNCTION = "todo"
|
|
CATEGORY = "EbSynth"
|
|
|
|
def todo(
|
|
self,
|
|
style_image: torch.Tensor,
|
|
source_image: torch.Tensor,
|
|
target_image: torch.Tensor,
|
|
weight: float,
|
|
uniformity: float,
|
|
patch_size: int,
|
|
pyramid_levels: int,
|
|
search_vote_iters: int,
|
|
patch_match_iters: int,
|
|
extra_pass_3x3: bool,
|
|
src_imgs: torch.Tensor | None = None,
|
|
tgt_imgs: torch.Tensor | None = None,
|
|
wgts: str | None = None,
|
|
):
|
|
print(f"{style_image.shape=}")
|
|
print(f"{source_image.shape=}")
|
|
print(f"{target_image.shape=}")
|
|
|
|
style_img = batched_tensor_to_cv2_list(style_image)[0]
|
|
src_img = batched_tensor_to_cv2_list(source_image)[0]
|
|
tgt_img = batched_tensor_to_cv2_list(target_image)[0]
|
|
|
|
guides = []
|
|
|
|
if src_imgs is not None and tgt_imgs is not None and wgts is not None:
|
|
g_srcs = batched_tensor_to_cv2_list(src_imgs)
|
|
g_tgt = batched_tensor_to_cv2_list(tgt_imgs)
|
|
g_wgts = deserialize_floats(wgts)
|
|
for g_src, g_tgt, g_wgt in zip(g_srcs, g_tgt, g_wgts):
|
|
guides.append((g_src, g_tgt, g_wgt))
|
|
|
|
print(wgts)
|
|
|
|
params = {
|
|
"style_img": style_img,
|
|
"src_img": src_img,
|
|
"tgt_img": tgt_img,
|
|
"cfg": RunConfig(
|
|
uniformity=uniformity,
|
|
patchsize=patch_size,
|
|
pyramidlevels=pyramid_levels,
|
|
searchvoteiters=search_vote_iters,
|
|
patchmatchiters=patch_match_iters,
|
|
extrapass3x3=extra_pass_3x3,
|
|
img_wgt=weight,
|
|
),
|
|
}
|
|
|
|
ezsynner = ImageSynthBase(**params)
|
|
result, error = ezsynner.run(guides=guides)
|
|
|
|
result_tensor = cv2_img_to_tensor(result)
|
|
error_tensor = cv2_img_to_tensor(error)
|
|
|
|
print(f"{result_tensor.shape=}")
|
|
|
|
return (
|
|
result_tensor,
|
|
error_tensor,
|
|
)
|
|
|
|
|
|
class ES_VideoTransfer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
required = {
|
|
"source_video": ("IMAGE",),
|
|
"style_images": ("IMAGE",),
|
|
"style_idxes": ("STRING", {"default": "0"}),
|
|
"edge_method": (EDGE_METHODS, {"default": DEFAULT_EDGE_METHOD}),
|
|
"flow_model": (FLOW_MODELS, {"default": DEFAULT_FLOW_MODEL}),
|
|
"only_mode": (ONLY_MODES, {"default": ONLY_DEFAULT_MODE}),
|
|
"do_mask": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
"pre_mask": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
"feather": (
|
|
"INT",
|
|
{"default": 5, "min": 0, "step": 1},
|
|
),
|
|
"style_weight": (
|
|
"FLOAT",
|
|
{"default": 6.0, "min": 0.1},
|
|
),
|
|
"edge_weight": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.1},
|
|
),
|
|
"warp_weight": (
|
|
"FLOAT",
|
|
{"default": 0.5, "min": 0.1},
|
|
),
|
|
"pos_weight": (
|
|
"FLOAT",
|
|
{"default": 2.0, "min": 0.1},
|
|
),
|
|
}
|
|
required.update(EBTYPES)
|
|
required.update(
|
|
{
|
|
"use_gpu_hist_blend": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
"use_lsqr": (
|
|
"BOOLEAN",
|
|
{"default": True},
|
|
),
|
|
"use_poisson_cupy": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
"poisson_maxiter": (
|
|
"INT",
|
|
{"default": 0, "min": 0, "step": 1},
|
|
),
|
|
}
|
|
)
|
|
return {
|
|
"required": required,
|
|
"optional": {
|
|
"source_mask": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"IMAGE",
|
|
"IMAGE",
|
|
)
|
|
RETURN_NAMES = (
|
|
"result_video",
|
|
"error_video",
|
|
)
|
|
FUNCTION = "todo"
|
|
CATEGORY = "EbSynth"
|
|
|
|
def todo(
|
|
self,
|
|
source_video: torch.Tensor,
|
|
style_images: torch.Tensor,
|
|
style_idxes: str,
|
|
edge_method: str,
|
|
flow_model: str,
|
|
only_mode: str,
|
|
# Masking params
|
|
do_mask: bool,
|
|
pre_mask: bool,
|
|
feather: float,
|
|
# Ebsynth guide weights params
|
|
style_weight: float,
|
|
edge_weight: float,
|
|
warp_weight: float,
|
|
pos_weight: float,
|
|
# Ebsynth gen params
|
|
uniformity: float,
|
|
patch_size: int,
|
|
pyramid_levels: int,
|
|
search_vote_iters: int,
|
|
patch_match_iters: int,
|
|
extra_pass_3x3: bool,
|
|
# Blending params
|
|
use_gpu_hist_blend: bool,
|
|
use_lsqr: bool,
|
|
use_poisson_cupy: bool,
|
|
poisson_maxiter: int,
|
|
source_mask: torch.Tensor | None = None,
|
|
):
|
|
print(f"{source_video.shape=}")
|
|
print(f"{style_images.shape=}")
|
|
|
|
img_frs_seq = batched_tensor_to_cv2_list(source_video)
|
|
stl_frs = batched_tensor_to_cv2_list(style_images)
|
|
stl_idxes = sorted(deserialize_integers(style_idxes))
|
|
|
|
if len(stl_frs) != len(stl_idxes):
|
|
raise ValueError(
|
|
f"Style indices mismatch: There are {len(stl_frs)}, but only [{stl_idxes}]"
|
|
)
|
|
|
|
if source_mask is not None:
|
|
print(f"{source_mask.shape=}")
|
|
msk_frs_seq = batched_tensor_to_cv2_list(
|
|
source_mask, color=cv2.COLOR_RGB2GRAY
|
|
)
|
|
else:
|
|
msk_frs_seq = None
|
|
|
|
ezrunner = EzsynthBase(
|
|
style_frs=stl_frs,
|
|
style_idxes=stl_idxes,
|
|
img_frs_seq=img_frs_seq,
|
|
cfg=RunConfig(
|
|
only_mode=only_mode,
|
|
pre_mask=pre_mask,
|
|
feather=feather,
|
|
uniformity=uniformity,
|
|
patchsize=patch_size,
|
|
pyramidlevels=pyramid_levels,
|
|
searchvoteiters=search_vote_iters,
|
|
patchmatchiters=patch_match_iters,
|
|
extrapass3x3=extra_pass_3x3,
|
|
img_wgt=style_weight,
|
|
edg_wgt=edge_weight,
|
|
wrp_wgt=warp_weight,
|
|
pos_wgt=pos_weight,
|
|
use_gpu=use_gpu_hist_blend,
|
|
use_lsqr=use_lsqr,
|
|
use_poisson_cupy=use_poisson_cupy,
|
|
poisson_maxiter=poisson_maxiter if poisson_maxiter > 0 else None,
|
|
),
|
|
edge_method=edge_method,
|
|
raft_flow_model_name=flow_model,
|
|
do_mask=do_mask,
|
|
msk_frs_seq=msk_frs_seq,
|
|
)
|
|
|
|
stylized_frames, err_frames = ezrunner.run_sequences()
|
|
style_tensor = out_video(stylized_frames)
|
|
err_tensor = out_video(err_frames)
|
|
|
|
return (
|
|
style_tensor,
|
|
err_tensor,
|
|
)
|
|
|
|
|
|
class ES_VideoTransferExtra:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
required = {
|
|
"source_video": ("IMAGE",),
|
|
"style_images": ("IMAGE",),
|
|
"style_idxes": ("STRING", {"default": "0"}),
|
|
"edge_method": (EDGE_METHODS, {"default": DEFAULT_EDGE_METHOD}),
|
|
"flow_arch": (FLOW_ARCHS, {"default": DEFAULT_FLOW_ARCH}),
|
|
"flow_model": (flow_all_models, {"default": DEFAULT_FLOW_MODEL}),
|
|
"only_mode": (ONLY_MODES, {"default": ONLY_DEFAULT_MODE}),
|
|
"return_flow": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
"do_mask": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
"pre_mask": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
"feather": (
|
|
"INT",
|
|
{"default": 5, "min": 0, "step": 1},
|
|
),
|
|
"style_weight": (
|
|
"FLOAT",
|
|
{"default": 6.0, "min": 0.1},
|
|
),
|
|
"edge_weight": (
|
|
"FLOAT",
|
|
{"default": 1.0, "min": 0.1},
|
|
),
|
|
"warp_weight": (
|
|
"FLOAT",
|
|
{"default": 0.5, "min": 0.1},
|
|
),
|
|
"pos_weight": (
|
|
"FLOAT",
|
|
{"default": 2.0, "min": 0.1},
|
|
),
|
|
}
|
|
required.update(EBTYPES)
|
|
required.update(
|
|
{
|
|
"use_gpu_hist_blend": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
"use_lsqr": (
|
|
"BOOLEAN",
|
|
{"default": True},
|
|
),
|
|
"use_poisson_cupy": (
|
|
"BOOLEAN",
|
|
{"default": False},
|
|
),
|
|
"poisson_maxiter": (
|
|
"INT",
|
|
{"default": 0, "min": 0, "step": 1},
|
|
),
|
|
}
|
|
)
|
|
return {
|
|
"required": required,
|
|
"optional": {
|
|
"source_mask": ("IMAGE",),
|
|
"source_edge": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"IMAGE",
|
|
"IMAGE",
|
|
"IMAGE",
|
|
)
|
|
RETURN_NAMES = (
|
|
"result_video",
|
|
"error_video",
|
|
"flow_video",
|
|
)
|
|
FUNCTION = "todo"
|
|
CATEGORY = "EbSynth"
|
|
|
|
def todo(
|
|
self,
|
|
source_video: torch.Tensor,
|
|
style_images: torch.Tensor,
|
|
style_idxes: str,
|
|
edge_method: str,
|
|
flow_arch: str,
|
|
flow_model: str,
|
|
only_mode: str,
|
|
return_flow: bool,
|
|
# Masking params
|
|
do_mask: bool,
|
|
pre_mask: bool,
|
|
feather: float,
|
|
# Ebsynth guide weights params
|
|
style_weight: float,
|
|
edge_weight: float,
|
|
warp_weight: float,
|
|
pos_weight: float,
|
|
# Ebsynth gen params
|
|
uniformity: float,
|
|
patch_size: int,
|
|
pyramid_levels: int,
|
|
search_vote_iters: int,
|
|
patch_match_iters: int,
|
|
extra_pass_3x3: bool,
|
|
# Blending params
|
|
use_gpu_hist_blend: bool,
|
|
use_lsqr: bool,
|
|
use_poisson_cupy: bool,
|
|
poisson_maxiter: int,
|
|
source_mask: torch.Tensor | None = None,
|
|
source_edge: torch.Tensor | None = None,
|
|
):
|
|
print(f"{source_video.shape=}")
|
|
print(f"{style_images.shape=}")
|
|
|
|
img_frs_seq = batched_tensor_to_cv2_list(source_video)
|
|
stl_frs = batched_tensor_to_cv2_list(style_images)
|
|
stl_idxes = sorted(deserialize_integers(style_idxes))
|
|
|
|
if len(stl_frs) != len(stl_idxes):
|
|
raise ValueError(
|
|
f"Style indices mismatch: There are {len(stl_frs)}, but only [{stl_idxes}]"
|
|
)
|
|
|
|
if source_mask is not None:
|
|
print(f"{source_mask.shape=}")
|
|
msk_frs_seq = batched_tensor_to_cv2_list(
|
|
source_mask, color=cv2.COLOR_RGB2GRAY
|
|
)
|
|
else:
|
|
msk_frs_seq = None
|
|
|
|
edge_guides = None
|
|
if source_edge is not None:
|
|
print(f"{source_edge.shape=}")
|
|
edge_guides = resize_cv2_list(
|
|
batched_tensor_to_cv2_list(source_edge), img_frs_seq[0]
|
|
)
|
|
do_compute_edge = False
|
|
else:
|
|
do_compute_edge = True
|
|
|
|
flow_model = flow_model[flow_model.index("_") + 1 :]
|
|
|
|
ezrunner = EzsynthBase(
|
|
style_frs=stl_frs,
|
|
style_idxes=stl_idxes,
|
|
img_frs_seq=img_frs_seq,
|
|
cfg=RunConfig(
|
|
only_mode=only_mode,
|
|
pre_mask=pre_mask,
|
|
feather=feather,
|
|
uniformity=uniformity,
|
|
patchsize=patch_size,
|
|
pyramidlevels=pyramid_levels,
|
|
searchvoteiters=search_vote_iters,
|
|
patchmatchiters=patch_match_iters,
|
|
extrapass3x3=extra_pass_3x3,
|
|
img_wgt=style_weight,
|
|
edg_wgt=edge_weight,
|
|
wrp_wgt=warp_weight,
|
|
pos_wgt=pos_weight,
|
|
use_gpu=use_gpu_hist_blend,
|
|
use_lsqr=use_lsqr,
|
|
use_poisson_cupy=use_poisson_cupy,
|
|
poisson_maxiter=poisson_maxiter if poisson_maxiter > 0 else None,
|
|
),
|
|
edge_method=edge_method,
|
|
raft_flow_model_name=flow_model,
|
|
do_mask=do_mask,
|
|
msk_frs_seq=msk_frs_seq,
|
|
flow_arch=flow_arch,
|
|
do_compute_edge=do_compute_edge,
|
|
)
|
|
|
|
if not do_compute_edge and edge_guides is not None:
|
|
ezrunner.edge_guides = edge_guides
|
|
|
|
stylized_frames, err_frames, flow_frames = ezrunner.run_sequences_full(
|
|
return_flow=return_flow
|
|
)
|
|
style_tensor = out_video(stylized_frames)
|
|
err_tensor = out_video(err_frames)
|
|
|
|
if not flow_frames:
|
|
flow_tensor = style_tensor
|
|
else:
|
|
flow_tensor = out_video(flow_frames)
|
|
|
|
return (
|
|
style_tensor,
|
|
err_tensor,
|
|
flow_tensor,
|
|
)
|
|
|
|
|
|
def serialize_floats(lst: list[float]):
|
|
return ",".join(map(str, lst))
|
|
|
|
|
|
def deserialize_floats(floats_str: str):
|
|
return [float(w) for w in floats_str.split(",") if w.strip()]
|
|
|
|
|
|
def serialize_integers(int_list):
|
|
return "_".join(map(str, int_list))
|
|
|
|
|
|
def deserialize_integers(int_string):
|
|
return list(map(int, int_string.split("_")))
|