| 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 .wan_audio import apply_wan_s2v_audio_conditioning |
|
|
|
|
| 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) |
|
|
|
|
| class BerniniS2VConditioning(io.ComfyNode): |
| """Bernini in-context conditioning plus Wan 2.2 S2V audio.""" |
|
|
| @classmethod |
| def define_schema(cls): |
| return io.Schema( |
| node_id="BerniniS2VConditioning", |
| display_name="Bernini S2V Conditioning", |
| category="model/conditioning/bernini", |
| description="Bernini in-context video/image conditioning with optional S2V audio. Requires a Wan2.2 S2V grafted diffusion model.", |
| 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_encoder_output", optional=True), |
| 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), |
| ], |
| 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_encoder_output=None, source_video=None, reference_video=None, reference_images=None, ref_max_size=848) -> io.NodeOutput: |
| positive, negative, _ = apply_wan_s2v_audio_conditioning(positive, negative, length, audio_encoder_output=audio_encoder_output) |
|
|
| latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device()) |
|
|
| 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])) |
|
|
| if context: |
| positive = node_helpers.conditioning_set_values(positive, {"context_latents": context}) |
| negative = node_helpers.conditioning_set_values(negative, {"context_latents": context}) |
|
|
| return io.NodeOutput(positive, negative, {"samples": latent}) |
|
|
|
|
| class WanBerniniS2VExtension(ComfyExtension): |
| @override |
| async def get_node_list(self) -> list[type[io.ComfyNode]]: |
| return [BerniniS2VConditioning] |
|
|
|
|
| async def comfy_entrypoint() -> WanBerniniS2VExtension: |
| return WanBerniniS2VExtension() |