File size: 6,927 Bytes
832fee5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | import torch
from typing_extensions import override
import comfy.model_management
import comfy.utils
import node_helpers
from comfy_api.latest import ComfyExtension, io
from .audio_mask import build_timeline_audio_inject_mask
from .wan_audio import apply_timeline_audio_conditioning, resolve_timeline_segment_ranges
def _resize_long_edge(image, max_size, stride=16):
h, w = image.shape[1], image.shape[2]
scale = min(max_size / max(h, w), 1.0)
nh = max(stride, round(h * scale / stride) * stride)
nw = max(stride, round(w * scale / stride) * stride)
return comfy.utils.common_upscale(image[:, :, :, :3].movedim(-1, 1), nw, nh, "area", "disabled").movedim(1, -1)
def _build_context_latents(vae, width, height, length, source_video=None, reference_video=None, reference_images=None, ref_max_size=848):
context = []
if source_video is not None:
vid = comfy.utils.common_upscale(source_video[:length, :, :, :3].movedim(-1, 1), width, height, "area", "center").movedim(1, -1)
context.append(vae.encode(vid[:, :, :, :3]))
if reference_video is not None:
ref_vid = _resize_long_edge(reference_video[:length], ref_max_size)
context.append(vae.encode(ref_vid[:, :, :, :3]))
if reference_images:
for name in sorted(reference_images):
imgs = reference_images[name]
if imgs is None:
continue
for i in range(imgs.shape[0]):
img = _resize_long_edge(imgs[i:i + 1], ref_max_size)
context.append(vae.encode(img[:, :, :, :3]))
return context
class BerniniS2VConditioningV2(io.ComfyNode):
@classmethod
def define_schema(cls):
return io.Schema(
node_id="BerniniS2VConditioningV2",
display_name="Bernini S2V Conditioning v2",
category="model/conditioning/bernini",
description="Bernini in-context conditioning with masked S2V audio for one or two speakers. Requires a Wan 2.2 S2V grafted Bernini-R model. Paint speaker masks on the output frame; reference_image_0 maps to image0 in prompts.",
inputs=[
io.Conditioning.Input("positive"),
io.Conditioning.Input("negative"),
io.Vae.Input("vae"),
io.Int.Input("width", default=832, min=16, max=8192, step=16),
io.Int.Input("height", default=480, min=16, max=8192, step=16),
io.Int.Input("length", default=81, min=1, max=8192, step=4),
io.Int.Input("batch_size", default=1, min=1, max=4096),
io.AudioEncoderOutput.Input("audio_1"),
io.Mask.Input("mask_1", tooltip="White = speaker 1 lip-sync region on the output frame."),
io.AudioEncoderOutput.Input("audio_2", optional=True),
io.Mask.Input("mask_2", optional=True, tooltip="Required when audio_2 is connected."),
io.Int.Input("speaker_2_start_frame", default=-1, min=-1, max=8192, step=1,
tooltip="-1 auto-starts speaker 2 when speaker 1 audio ends."),
io.Image.Input("source_video", optional=True),
io.Image.Input("reference_video", optional=True),
io.Autogrow.Input("reference_images", optional=True,
template=io.Autogrow.TemplatePrefix(
input=io.Image.Input("reference_image"),
prefix="reference_image_", min=0, max=8)),
io.Int.Input("ref_max_size", default=848, min=16, max=8192, step=16, optional=True),
io.Int.Input("mask_crossfade_frames", default=4, min=0, max=64, step=1,
tooltip="Softens the mask handoff between speakers. 0 = hard cut."),
io.Float.Input("audio_inject_scale", default=1.0, min=0.0, max=10.0, step=0.01),
],
outputs=[
io.Conditioning.Output(display_name="positive"),
io.Conditioning.Output(display_name="negative"),
io.Latent.Output(display_name="latent"),
],
)
@classmethod
def execute(
cls,
positive,
negative,
vae,
width,
height,
length,
batch_size,
audio_1,
mask_1,
audio_2=None,
mask_2=None,
speaker_2_start_frame=-1,
source_video=None,
reference_video=None,
reference_images=None,
ref_max_size=848,
mask_crossfade_frames=4,
audio_inject_scale=1.0,
) -> io.NodeOutput:
if audio_1 is None:
raise ValueError("Bernini S2V Conditioning v2 requires audio_1.")
if mask_1 is None:
raise ValueError("Bernini S2V Conditioning v2 requires mask_1.")
if audio_2 is not None and mask_2 is None:
raise ValueError("mask_2 is required when audio_2 is connected.")
segments = [{"audio_encoder_output": audio_1, "start_frame": 0, "mask_image": mask_1}]
if audio_2 is not None:
segments.append({
"audio_encoder_output": audio_2,
"start_frame": speaker_2_start_frame,
"mask_image": mask_2,
})
latent = torch.zeros(
[batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8],
device=comfy.model_management.intermediate_device())
context = _build_context_latents(vae, width, height, length, source_video, reference_video, reference_images, ref_max_size)
if context:
positive = node_helpers.conditioning_set_values(positive, {"context_latents": context})
negative = node_helpers.conditioning_set_values(negative, {"context_latents": context})
positive, negative = apply_timeline_audio_conditioning(positive, negative, length, segments)
resolved_segments = resolve_timeline_segment_ranges(length, segments)
cond_values = {
"audio_inject_scale": audio_inject_scale,
"audio_inject_mask": build_timeline_audio_inject_mask(
width, height, length, resolved_segments,
crossfade_frames=mask_crossfade_frames,
device=comfy.model_management.intermediate_device(),
),
}
positive = node_helpers.conditioning_set_values(positive, cond_values)
negative_values = dict(cond_values)
negative_values["audio_inject_mask"] = negative_values["audio_inject_mask"] * 0.0
negative = node_helpers.conditioning_set_values(negative, negative_values)
return io.NodeOutput(positive, negative, {"samples": latent})
class WanBerniniS2VV2Extension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [BerniniS2VConditioningV2]
async def comfy_entrypoint() -> WanBerniniS2VV2Extension:
return WanBerniniS2VV2Extension() |