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)