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