hanxxing's picture
Duplicate from rzgar/Bernini-R-S2V
832fee5
Raw
History Blame Contribute Delete
4.24 kB
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()