ComfyUI-WanBerniniS2V_v2/nodes.py

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()