app.py

import os

# Allocator config for the video-DiT transient spikes (see zerogpu known-errors).
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  MUST be before torch / any CUDA-touching import

import sys
import shutil
import subprocess
import tempfile
import random
from pathlib import Path

import torch  # noqa: E402
import gradio as gr  # noqa: E402
from huggingface_hub import hf_hub_download, snapshot_download  # noqa: E402

# ----------------------------------------------------------------------------
# 0. Vendor ComfyUI (the reference implementation lives entirely in ComfyUI's
#    Wan S2V nodes + the Bernini custom node). We drive its node `execute`/
#    FUNCTION methods directly ("ComfyUI as a library") rather than running the
#    server, so everything happens in-process and works with the ZeroGPU hijack.
# ----------------------------------------------------------------------------
COMFY_COMMIT = "04a30fb375a6c1312365fbbbdd0a7d0669212e92"
COMFY_DIR = Path(__file__).parent / "ComfyUI"

if not COMFY_DIR.exists():
    subprocess.run(
        ["git", "clone", "https://github.com/comfyanonymous/ComfyUI.git", str(COMFY_DIR)],
        check=True,
    )
    subprocess.run(["git", "checkout", COMFY_COMMIT], cwd=str(COMFY_DIR), check=True)

sys.path.insert(0, str(COMFY_DIR))

# ----------------------------------------------------------------------------
# 1. Candidate + dependency model IDs
# ----------------------------------------------------------------------------
BERNINI_REPO = "rzgar/Bernini-R-S2V"
LORA_REPO = "rzgar/Bernini-R-LightX2V-4step-loras"
COMFY_WAN22 = "Comfy-Org/Wan_2.2_ComfyUI_Repackaged"
COMFY_WAN21 = "Comfy-Org/Wan_2.1_ComfyUI_repackaged"

MODELS_ROOT = Path(os.environ.get("COMFY_MODELS", "/data/models"))
try:
    MODELS_ROOT.mkdir(parents=True, exist_ok=True)
except Exception:
    MODELS_ROOT = Path(tempfile.gettempdir()) / "comfy_models"
    MODELS_ROOT.mkdir(parents=True, exist_ok=True)

DIRS = {
    "diffusion_models": MODELS_ROOT / "diffusion_models",
    "text_encoders": MODELS_ROOT / "text_encoders",
    "vae": MODELS_ROOT / "vae",
    "audio_encoders": MODELS_ROOT / "audio_encoders",
    "loras": MODELS_ROOT / "loras",
}
for d in DIRS.values():
    d.mkdir(parents=True, exist_ok=True)


def _fetch(repo, filename, dest_dir, out_name=None):
    out = DIRS[dest_dir] / (out_name or Path(filename).name)
    p = hf_hub_download(repo_id=repo, filename=filename)
    if out.exists() or out.is_symlink():
        out.unlink()
    # symlink into ComfyUI's model dir (avoid copying tens of GB)
    try:
        out.symlink_to(p)
    except OSError:
        shutil.copy(p, out)
    return out


print("Downloading model weights (this happens once) ...", flush=True)
HIGH = _fetch(BERNINI_REPO, "Bernini-R-S2V-FP8/wan2.2_bernini_r_high_noise_fp8_scaled_s2v.safetensors",
              "diffusion_models", "bernini_high_fp8_s2v.safetensors")
LOW = _fetch(BERNINI_REPO, "Bernini-R-S2V-FP8/wan2.2_bernini_r_low_noise_fp8_scaled_s2v.safetensors",
             "diffusion_models", "bernini_low_fp8_s2v.safetensors")
WAV2VEC = _fetch(BERNINI_REPO, "audio_encoders/wav2vec2_large_english_fp16.safetensors",
                 "audio_encoders")
UMT5 = _fetch(COMFY_WAN22, "split_files/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors",
              "text_encoders")
VAE_FILE = _fetch(COMFY_WAN21, "split_files/vae/wan_2.1_vae.safetensors", "vae")
LORA_HIGH = _fetch(LORA_REPO, "Bernini-R_LightX2V_high_noise.safetensors", "loras")
LORA_LOW = _fetch(LORA_REPO, "Bernini-R_LightX2V_low_noise.safetensors", "loras")
print("All weights present.", flush=True)

# ----------------------------------------------------------------------------
# 2. Point ComfyUI's folder_paths at our model dirs, then import nodes.
# ----------------------------------------------------------------------------
import folder_paths  # noqa: E402

for key, d in DIRS.items():
    folder_paths.add_model_folder_path(key, str(d))

import asyncio  # noqa: E402
import nodes as comfy_nodes  # noqa: E402

# init_extra_nodes is async in this ComfyUI revision.
_res = comfy_nodes.init_extra_nodes(init_custom_nodes=False, init_api_nodes=False)
if asyncio.iscoroutine(_res):
    asyncio.new_event_loop().run_until_complete(_res)

import comfy.model_management as mm  # noqa: E402
import comfy.utils  # noqa: E402
import node_helpers  # noqa: E402
from comfy_extras.nodes_wan import get_audio_embed_bucket_fps, linear_interpolation  # noqa: E402

# ----------------------------------------------------------------------------
# 3. Apply the Bernini S2V model patch (adds context_latents support to the
#    Wan S2V forward). Ported verbatim from the candidate's custom node so the
#    output matches the ComfyUI reference path.
# ----------------------------------------------------------------------------
import inspect  # noqa: E402
import logging  # noqa: E402


def _append_context_latents(self, x, kwargs):
    context_latents = kwargs.get("context_latents", None)
    if context_latents is None:
        return x
    for lat in context_latents:
        cl = self.patch_embedding(lat.float().to(x.device)).to(x.dtype).flatten(2).transpose(1, 2)
        x = torch.cat([x, cl], dim=1)
    return x


def _apply_bernini_patch():
    from comfy.ldm.wan.model import WanModel_S2V, sinusoidal_embedding_1d  # noqa: F401

    try:
        source = inspect.getsource(WanModel_S2V.forward_orig)
    except (OSError, TypeError):
        source = ""
    if "context_latents" in source:
        return
    if getattr(WanModel_S2V.forward_orig, "__wan_bernini_s2v_patch__", False):
        return
    original = WanModel_S2V.forward_orig

    def forward_orig(self, x, t, context, audio_embed=None, reference_latent=None,
                     control_video=None, reference_motion=None, clip_fea=None, freqs=None,
                     transformer_options={}, **kwargs):
        from comfy.ldm.wan.model import sinusoidal_embedding_1d
        if audio_embed is not None:
            num_embeds = x.shape[-3] * 4
            audio_emb_global, audio_emb = self.casual_audio_encoder(audio_embed[:, :, :, :num_embeds])
        else:
            audio_emb = None
            audio_emb_global = None

        bs, _, time, height, width = x.shape
        x = self.patch_embedding(x.float()).to(x.dtype)
        if control_video is not None:
            x = x + self.cond_encoder(control_video)

        if t.ndim == 1:
            t = t.unsqueeze(1).repeat(1, x.shape[2])

        grid_sizes = x.shape[2:]
        x = x.flatten(2).transpose(1, 2)
        seq_len = x.size(1)

        cond_mask_weight = mm.cast_to(self.trainable_cond_mask.weight, dtype=x.dtype,
                                      device=x.device).unsqueeze(1).unsqueeze(1)
        x = x + cond_mask_weight[0]
        x = _append_context_latents(self, x, kwargs)

        if reference_latent is not None:
            ref = self.patch_embedding(reference_latent.float()).to(x.dtype)
            ref = ref.flatten(2).transpose(1, 2)
            freqs_ref = self.rope_encode(reference_latent.shape[-3], reference_latent.shape[-2],
                                         reference_latent.shape[-1], t_start=max(30, time + 9),
                                         device=x.device, dtype=x.dtype)
            ref = ref + cond_mask_weight[1]
            x = torch.cat([x, ref], dim=1)
            freqs = torch.cat([freqs, freqs_ref], dim=1)
            t = torch.cat([t, torch.zeros((t.shape[0], reference_latent.shape[-3]),
                                          device=t.device, dtype=t.dtype)], dim=1)

        if reference_motion is not None:
            motion_encoded, freqs_motion = self.frame_packer(reference_motion, self)
            motion_encoded = motion_encoded + cond_mask_weight[2]
            x = torch.cat([x, motion_encoded], dim=1)
            freqs = torch.cat([freqs, freqs_motion], dim=1)
            t = torch.repeat_interleave(t, 2, dim=1)
            t = torch.cat([t, torch.zeros((t.shape[0], 3), device=t.device, dtype=t.dtype)], dim=1)

        e = self.time_embedding(
            sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x[0].dtype))
        e = e.reshape(t.shape[0], -1, e.shape[-1])
        e0 = self.time_projection(e).unflatten(2, (6, self.dim))

        context = self.text_embedding(context)

        patches_replace = transformer_options.get("patches_replace", {})
        blocks_replace = patches_replace.get("dit", {})
        transformer_options["total_blocks"] = len(self.blocks)
        transformer_options["block_type"] = "double"
        for i, block in enumerate(self.blocks):
            transformer_options["block_index"] = i
            if ("double_block", i) in blocks_replace:
                def block_wrap(args):
                    out = {}
                    out["img"] = block(args["img"], context=args["txt"], e=args["vec"],
                                       freqs=args["pe"], transformer_options=args["transformer_options"])
                    return out
                out = blocks_replace[("double_block", i)](
                    {"img": x, "txt": context, "vec": e0, "pe": freqs,
                     "transformer_options": transformer_options},
                    {"original_block": block_wrap})
                x = out["img"]
            else:
                x = block(x, e=e0, freqs=freqs, context=context, transformer_options=transformer_options)
            if audio_emb is not None:
                x = self.audio_injector(x, i, audio_emb, audio_emb_global, seq_len)
        x = self.head(x, e)
        x = self.unpatchify(x, grid_sizes)
        return x

    forward_orig.__wan_bernini_s2v_patch__ = True
    forward_orig.__wan_bernini_s2v_original__ = original
    WanModel_S2V.forward_orig = forward_orig
    logging.info("Applied Bernini S2V context_latents patch")


# ----------------------------------------------------------------------------
# 4. Load models at module scope (ZeroGPU eager pattern).
# ----------------------------------------------------------------------------
NODES = comfy_nodes.NODE_CLASS_MAPPINGS

print("Loading models ...", flush=True)
_apply_bernini_patch()

clip = NODES["CLIPLoader"]().load_clip("umt5_xxl_fp8_e4m3fn_scaled.safetensors", "wan", "default")[0]
vae = NODES["VAELoader"]().load_vae("wan_2.1_vae.safetensors")[0]
audio_encoder = NODES["AudioEncoderLoader"].execute("wav2vec2_large_english_fp16.safetensors").result[0]

_model_high_base = NODES["UNETLoader"]().load_unet("bernini_high_fp8_s2v.safetensors", "default")[0]
_model_low_base = NODES["UNETLoader"]().load_unet("bernini_low_fp8_s2v.safetensors", "default")[0]

# Fuse the LightX2V 4-step distillation LoRAs (this is what makes 4-step,
# sub-minute inference feasible and matches the reference workflow).
lora_node = NODES["LoraLoaderModelOnly"]()
model_high = lora_node.load_lora_model_only(_model_high_base, "Bernini-R_LightX2V_high_noise.safetensors", 1.0)[0]
model_low = lora_node.load_lora_model_only(_model_low_base, "Bernini-R_LightX2V_low_noise.safetensors", 1.0)[0]

# ModelSamplingSD3 shift = 8 (from the reference workflow).
msd3 = NODES["ModelSamplingSD3"]()
model_high = msd3.patch(model_high, 8.0)[0]
model_low = msd3.patch(model_low, 8.0)[0]

print("Models loaded.", flush=True)

_reencode = NODES["CLIPTextEncode"]()
_ksampler = NODES["KSamplerAdvanced"]()
_vaedecode = NODES["VAEDecode"]()
_vaedecode_tiled = NODES["VAEDecodeTiled"]()

DEFAULT_NEG = ("Vivid color tone, overexposed, static, unclear details, subtitles, style, artwork, "
               "painting, image, still, overall grayish, worst quality, low quality, leftover JPEG "
               "compression artifacts, ugly, incomplete, missing parts, extra fingers, poorly drawn "
               "hands, poorly drawn face, disfigured, malformed body parts, fused fingers, a completely "
               "motionless image, messy background, three legs, many people in the background, walking backward.")

VIDEO_FPS = 16  # generation fps (matches the S2V audio bucketing)


# ----------------------------------------------------------------------------
# 5. Helpers: load image / audio into ComfyUI tensor formats.
# ----------------------------------------------------------------------------
def _load_image_tensor(path):
    from PIL import Image, ImageOps
    import numpy as np
    img = Image.open(path)
    img = ImageOps.exif_transpose(img).convert("RGB")
    arr = np.array(img).astype("float32") / 255.0
    return torch.from_numpy(arr)[None,]  # [1, H, W, 3]


def _load_audio(path, max_seconds):
    # Use soundfile directly: newer torchaudio.load dispatches to torchcodec
    # (which isn't installed on the Space) and raises ImportError.
    import soundfile as sf
    import numpy as np
    data, sr = sf.read(path, dtype="float32", always_2d=True)  # (frames, channels)
    waveform = torch.from_numpy(np.ascontiguousarray(data.T))   # (channels, frames)
    if max_seconds:
        waveform = waveform[:, : int(sr * max_seconds)]
    return {"waveform": waveform.unsqueeze(0), "sample_rate": sr}


def _bernini_s2v_conditioning(positive, negative, width, height, length,
                              audio_encoder_output, ref_image, ref_max_size=848):
    """Port of BerniniS2VConditioning.execute for the single-reference-image case."""
    latent_t = ((length - 1) // 4) + 1
    # audio conditioning
    if audio_encoder_output is not None:
        feat = torch.cat(audio_encoder_output["encoded_audio_all_layers"])
        video_rate = 30
        fps = 16
        feat = linear_interpolation(feat, input_fps=50, output_fps=video_rate)
        batch_frames = latent_t * 4
        audio_embed_bucket, _ = get_audio_embed_bucket_fps(feat, fps=fps, batch_frames=batch_frames,
                                                           m=0, video_rate=video_rate)
        audio_embed_bucket = audio_embed_bucket.unsqueeze(0)
        if len(audio_embed_bucket.shape) == 3:
            audio_embed_bucket = audio_embed_bucket.permute(0, 2, 1)
        elif len(audio_embed_bucket.shape) == 4:
            audio_embed_bucket = audio_embed_bucket.permute(0, 2, 3, 1)
        audio_embed_bucket = audio_embed_bucket[:, :, :, 0:batch_frames]
        if audio_embed_bucket.shape[3] > 0:
            positive = node_helpers.conditioning_set_values(positive, {"audio_embed": audio_embed_bucket})
            negative = node_helpers.conditioning_set_values(negative, {"audio_embed": audio_embed_bucket * 0.0})

    latent = torch.zeros([1, 16, latent_t, height // 8, width // 8],
                         device=mm.intermediate_device())

    context = []
    if ref_image is not None:
        h, w = ref_image.shape[1], ref_image.shape[2]
        scale = min(ref_max_size / max(h, w), 1.0)
        stride = 16
        nh = max(stride, round(h * scale / stride) * stride)
        nw = max(stride, round(w * scale / stride) * stride)
        img = comfy.utils.common_upscale(ref_image[:, :, :, :3].movedim(-1, 1), nw, nh,
                                         "area", "disabled").movedim(1, -1)
        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 positive, negative, {"samples": latent}


def _estimate(image, audio, seconds, *args, **kwargs):
    seconds = int(seconds) if seconds else 5
    return min(300, 90 + seconds * 22)


@spaces.GPU(duration=_estimate)
def generate(image, audio, seconds=5, prompt="", width=832, height=480,
             seed=0, randomize_seed=True,
             progress=gr.Progress(track_tqdm=True)):
    """Generate a lip-synced talking-head video from a reference image and speech audio.

    Args:
        image: reference image (a portrait / character to animate).
        audio: speech audio file to drive the lip-sync and motion.
        seconds: length of the output video in seconds (max ~10).
        prompt: optional text prompt describing the scene / action.
        width: output width in pixels.
        height: output height in pixels.
        seed: RNG seed.
        randomize_seed: pick a fresh random seed each run.
    """
    if image is None or audio is None:
        raise gr.Error("Please provide both a reference image and an audio file.")

    if randomize_seed:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)

    seconds = max(1, min(10, int(seconds)))
    # length must be 4k+1 frames; 16 fps generation.
    length = int(round(seconds * VIDEO_FPS))
    length = ((length - 1) // 4) * 4 + 1
    width = int(width) // 16 * 16
    height = int(height) // 16 * 16

    # Run the whole graph under inference_mode, matching ComfyUI's own
    # execution context (its samplers create inference-mode tensors; decoding
    # outside that context otherwise raises "Inplace update to inference
    # tensor outside InferenceMode is not allowed").
    with torch.inference_mode():
        ref_img = _load_image_tensor(image)
        audio_dict = _load_audio(audio, seconds)

        aenc = NODES["AudioEncoderEncode"].execute(audio_encoder, audio_dict).result[0]

        positive = _reencode.encode(clip, prompt or "")[0]
        negative = _reencode.encode(clip, DEFAULT_NEG)[0]

        positive, negative, latent = _bernini_s2v_conditioning(
            positive, negative, width, height, length, aenc, ref_img)

        # Two-expert MoE sampling: high-noise (steps 0-2) then low-noise (steps 2-4).
        # 4 total steps, cfg=1 (LightX2V distilled), dpmpp_2m_sde / sgm_uniform.
        stage1 = _ksampler.sample(
            model_high, "enable", seed, 4, 1.0, "dpmpp_2m_sde", "sgm_uniform",
            positive, negative, latent, start_at_step=0, end_at_step=2,
            return_with_leftover_noise="enable")[0]

        stage2 = _ksampler.sample(
            model_low, "disable", seed, 4, 1.0, "dpmpp_2m_sde", "sgm_uniform",
            positive, negative, stage1, start_at_step=2, end_at_step=4,
            return_with_leftover_noise="disable")[0]

        # Tiled decode keeps the video-VAE peak VRAM bounded (avoids allocator
        # fragmentation / NVML assert on the transient decode spike).
        images = _vaedecode_tiled.decode(
            vae, stage2, tile_size=256, overlap=64,
            temporal_size=16, temporal_overlap=8,
        )[0]  # [T, H, W, 3]
        images = images.detach().clone()

    # write to mp4 with audio, at generation fps
    out_path = _write_video(images, audio_dict, VIDEO_FPS)
    return out_path, seed


def _write_video(images, audio_dict, fps):
    import numpy as np
    import av

    frames = (images.cpu().numpy() * 255.0).clip(0, 255).astype("uint8")  # [T,H,W,3]
    T, H, W, _ = frames.shape

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    container = av.open(out_path, mode="w")
    vstream = container.add_stream("libx264", rate=fps)
    vstream.width = W
    vstream.height = H
    vstream.pix_fmt = "yuv420p"
    vstream.options = {"crf": "19"}

    # audio stream (resample waveform to fit)
    wav = audio_dict["waveform"][0]  # [C, N]
    sr = audio_dict["sample_rate"]
    if wav.shape[0] > 1:
        wav = wav.mean(dim=0, keepdim=True)
    astream = container.add_stream("aac", rate=sr)
    astream.layout = "mono"

    for i in range(T):
        frame = av.VideoFrame.from_ndarray(frames[i], format="rgb24")
        for packet in vstream.encode(frame):
            container.mux(packet)
    for packet in vstream.encode():
        container.mux(packet)

    audio_np = wav.numpy().astype("float32")
    aframe = av.AudioFrame.from_ndarray(audio_np, format="fltp", layout="mono")
    aframe.sample_rate = sr
    for packet in astream.encode(aframe):
        container.mux(packet)
    for packet in astream.encode():
        container.mux(packet)

    container.close()
    return out_path


# ----------------------------------------------------------------------------
# 6. UI
# ----------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🎙️ Bernini-R-S2V — Talking Head\n"
            "Speech-driven video generation on **[rzgar/Bernini-R-S2V]"
            "(https://huggingface.co/rzgar/Bernini-R-S2V)** (Wan2.2 S2V grafted onto Bernini-R). "
            "Upload a portrait and a speech clip to get a lip-synced talking-head video. "
            "4-step LightX2V distilled sampling for fast inference on ZeroGPU."
        )
        with gr.Row():
            with gr.Column():
                image = gr.Image(label="Reference image (portrait)", type="filepath", sources=["upload"])
                audio = gr.Audio(label="Speech audio (mono, clear speech)", type="filepath", sources=["upload"])
                seconds = gr.Slider(1, 10, value=5, step=1, label="Video length (seconds)")
                prompt = gr.Textbox(label="Prompt (optional)", placeholder="a person talking, cinematic")
                run = gr.Button("Generate", variant="primary")
            with gr.Column():
                out_video = gr.Video(label="Result")
                used_seed = gr.Number(label="Seed used", interactive=False)

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                width = gr.Slider(320, 1280, value=832, step=16, label="Width")
                height = gr.Slider(320, 1280, value=480, step=16, label="Height")
            with gr.Row():
                seed = gr.Number(label="Seed", value=0, precision=0)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

        run.click(
            fn=generate,
            inputs=[image, audio, seconds, prompt, width, height, seed, randomize_seed],
            outputs=[out_video, used_seed],
            api_name="generate",
        )

        gr.Examples(
            examples=[
                ["examples/demo_face.jpg", "examples/scream_mono.wav", 5,
                 "a person expressing intense emotion, cinematic"],
            ],
            inputs=[image, audio, seconds, prompt],
            outputs=[out_video, used_seed],
            fn=generate,
            cache_examples=False,
            run_on_click=True,
        )

demo.launch(mcp_server=True)