app.py

# -*- coding: utf-8 -*-
"""
Bernini-R-S2V Talking Head — image + speech audio -> lip-synced talking video.

Bernini-R-S2V (rzgar/Bernini-R-S2V) is a speech-driven video extension of
Bernini-R, built on Wan 2.2 S2V. It adds single-speaker audio conditioning so a
still portrait can be animated to speak/sing in sync with an audio track.

The Bernini finetune ships as ComfyUI single-file checkpoints (dual high/low
noise Wan 2.2 MoE) plus a wav2vec2 audio encoder. There is no diffusers-format
export of those exact weights, and the community diffusers Wan 2.2 S2V pipeline
uses a single-transformer S2V architecture. To give the same
image + audio -> talking-video experience self-hosted on ZeroGPU, this Space
runs the base Wan 2.2 S2V pipeline that Bernini-R-S2V is finetuned from,
via the community WanSpeechToVideoPipeline (diffusers PR #12258).
"""
import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import math
import tempfile

import spaces
import torch
import gradio as gr
from PIL import Image

from diffusers import AutoencoderKLWan, WanSpeechToVideoPipeline
from diffusers.utils import export_to_video, load_audio
try:
    from diffusers.utils import export_to_merged_video_audio
    HAS_MERGE = True
except Exception:
    HAS_MERGE = False
from transformers import Wav2Vec2ForCTC

MODEL_ID = "tolgacangoz/Wan2.2-S2V-14B-Diffusers"

# Wan default negative prompt (Chinese) — steers away from low quality / artifacts.
NEG = ("色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，"
       "低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，"
       "毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走")

DEFAULT_PROMPT = "a person talking to the camera, natural expression, slight head movement"

# Load the pipeline eagerly at module scope so ZeroGPU packs the weights and
# streams them into VRAM on the first @spaces.GPU call (rather than paying the
# load cost inside every request).
print("Loading Wan 2.2 S2V pipeline (base weights for Bernini-R-S2V)...")
audio_encoder = Wav2Vec2ForCTC.from_pretrained(MODEL_ID, subfolder="audio_encoder", dtype=torch.float32)
vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
pipe = WanSpeechToVideoPipeline.from_pretrained(
    MODEL_ID, vae=vae, audio_encoder=audio_encoder, torch_dtype=torch.bfloat16
)
pipe.to("cuda")
print("Pipeline ready.")


def _snap_size(img, target_area, divisor=64, max_side=1280):
    w, h = img.size
    scale = math.sqrt(target_area / (w * h))
    nw = max(divisor, min(max_side, int(round(w * scale / divisor) * divisor)))
    nh = max(divisor, min(max_side, int(round(h * scale / divisor) * divisor)))
    return nh, nw


def _audio_seconds(audio_path):
    try:
        import soundfile as sf
        info = sf.info(audio_path)
        return info.frames / info.samplerate
    except Exception:
        return 6.0


def _estimate_duration(image_path, audio_path, prompt, negative_prompt, steps,
                       guidance, res_area, max_seconds, seed, *args, **kwargs):
    # Video length grows with audio length (chunked at ~5s / chunk). Estimate
    # GPU seconds from the (capped) audio duration and step count.
    try:
        secs = min(float(max_seconds), _audio_seconds(audio_path)) if audio_path else 5.0
    except Exception:
        secs = 5.0
    n_chunks = max(1, math.ceil(secs / 5.0))
    per_chunk = 8.0 * (int(steps) / 20.0) * (1.0 if res_area == "480p" else 2.0)
    return int(min(700, 45 + n_chunks * (12 + per_chunk * 6)))


@spaces.GPU(duration=_estimate_duration)
def generate(image_path, audio_path, prompt=DEFAULT_PROMPT, negative_prompt=NEG,
             steps=20, guidance=4.5, res_area="480p", max_seconds=6, seed=42,
             progress=gr.Progress(track_tqdm=True)):
    """Generate a lip-synced talking-head video from a portrait image and speech audio.

    Args:
        image_path: Path to the reference portrait image (a clear front-facing face works best).
        audio_path: Path to the speech audio file (mono, clear speech gives the best lip-sync).
        prompt: Text prompt describing the scene / motion.
        negative_prompt: Things to avoid in the generation.
        steps: Number of denoising steps (higher = slower, more detailed).
        guidance: Classifier-free guidance scale.
        res_area: Target resolution, "480p" or "720p".
        max_seconds: Cap on how many seconds of audio to animate (controls video length / runtime).
        seed: Random seed for reproducibility.

    Returns:
        Path to the generated MP4 video with the audio track muxed in.
    """
    if not image_path:
        raise gr.Error("Please provide a portrait image first.")
    if not audio_path:
        raise gr.Error("Please provide a speech audio file first.")

    image = Image.open(image_path).convert("RGB")
    audio, sampling_rate = load_audio(audio_path)

    # Trim overly long audio to keep generation within the GPU time budget.
    max_seconds = float(max_seconds)
    max_samples = int(max_seconds * sampling_rate)
    if audio.shape[0] > max_samples:
        audio = audio[:max_samples]

    target = 480 * 832 if res_area == "480p" else 720 * 1280
    h, w = _snap_size(image, target_area=target)
    gen = torch.Generator(device="cuda").manual_seed(int(seed))

    frames = pipe(
        image=image,
        audio=audio,
        sampling_rate=sampling_rate,
        prompt=(prompt or DEFAULT_PROMPT),
        negative_prompt=(negative_prompt or NEG),
        height=h,
        width=w,
        num_frames_per_chunk=80,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        generator=gen,
    ).frames[0]

    out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    export_to_video(frames, out, fps=16)
    if HAS_MERGE:
        try:
            export_to_merged_video_audio(out, audio_path)
        except Exception as e:
            print("merge audio failed:", e)
    return out


CSS = """
#col-container { max-width: 1150px; 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
            Animate a **portrait + speech audio** into a **lip-synced talking video**.

            Powered by [Bernini-R-S2V](https://huggingface.co/rzgar/Bernini-R-S2V) — a speech-driven
            video extension of Bernini-R built on **Wan 2.2 S2V**. Give it a clear front-facing face
            and a mono speech clip; keep audio to ~3–10 s for the best sync and fastest results.
            """
        )
        with gr.Row():
            with gr.Column():
                image = gr.Image(type="filepath", label="Portrait image (clear front-facing face)")
                audio = gr.Audio(type="filepath", label="Speech audio (mono, clear speech)")
                prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT)
                with gr.Accordion("Advanced settings", open=False):
                    negative_prompt = gr.Textbox(label="Negative prompt", value=NEG, lines=2)
                    res_area = gr.Radio(["480p", "720p"], value="480p",
                                        label="Resolution (720p is slower)")
                    steps = gr.Slider(8, 40, value=20, step=1, label="Steps")
                    guidance = gr.Slider(1.0, 8.0, value=4.5, step=0.5, label="Guidance scale")
                    max_seconds = gr.Slider(2, 12, value=6, step=1,
                                            label="Max audio seconds to animate")
                    seed = gr.Number(value=42, label="Seed", precision=0)
                btn = gr.Button("Generate talking video", variant="primary")
            with gr.Column():
                out_video = gr.Video(label="Result")

        gr.Examples(
            examples=[
                ["examples/man.jpg", "examples/speech_short.wav"],
                ["examples/woman.jpg", "examples/speech_med.wav"],
                ["examples/elder.jpg", "examples/speech_short.wav"],
            ],
            inputs=[image, audio],
            outputs=out_video,
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    btn.click(
        generate,
        inputs=[image, audio, prompt, negative_prompt, steps, guidance, res_area, max_seconds, seed],
        outputs=out_video,
        api_name="generate",
    )

if __name__ == "__main__":
    demo.queue().launch(mcp_server=True)