app.py

"""Bernini-Diffusers-v2 — reference-to-video (subject-to-video) demo.

Bernini couples a Qwen2.5-VL planner (which reads the reference images and the
instruction, then *plans* a target visual embedding with a flow-matching head)
to a Wan2.2-A14B MoE renderer (two 14B DiTs, high-noise + low-noise).

This Space mirrors the authors' own ``scripts/bernini_v2/run_r2v.sh`` /
``gradio_demo.py`` single-GPU path 1:1 (same guidance mode, omegas, planning
steps, system prompt and negative prompt); only the frame count / step count
defaults are lowered so a generation fits inside a ZeroGPU slot.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

import spaces  # noqa: E402  (must precede torch / CUDA touching imports)

import gc  # noqa: E402
import logging  # noqa: E402
import random  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402

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

logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(name)s: %(message)s")
logging.getLogger("bernini.pipeline").setLevel(logging.INFO)

MODEL_ID = "ByteDance/Bernini-Diffusers-v2"


def _stat(tag):
    import shutil

    du = shutil.disk_usage("/tmp")
    rss = 0
    try:
        with open("/proc/self/status") as f:
            for line in f:
                if line.startswith("VmRSS"):
                    rss = int(line.split()[1]) / 1e6
    except Exception:
        pass
    print(f"[stat] {tag}: rss={rss:.1f}GB disk_used={du.used / 1e9:.1f}GB "
          f"free={du.free / 1e9:.1f}GB", flush=True)


# ---------------------------------------------------------------- weights ---
# The released checkpoint is fp32: `bernini/` alone is 180 GB, which blows past
# the Space's 150 GB disk quota.  So only the small components are materialised
# up-front; the 38 big shards are streamed one at a time, cast to bf16 straight
# into a meta-initialised model, and deleted immediately.  bf16 is the dtype the
# reference pipeline computes in anyway (`BerniniPipeline.weight_dtype`), so
# nothing is lost.  Peak disk for the shard stream is one shard (~5 GB).
#
# `mllm/*.safetensors` is skipped too: config.json sets `scratch_mllm: true`, so
# the MLLM is built from config and filled from the `bernini/` shards.
MODEL_DIR = snapshot_download(
    MODEL_ID,
    allow_patterns=[
        "config.json",
        "transformer_config.json",
        "transformer_2_config.json",
        "scheduler/*",
        "vae/*",
        "t5_text_encoder/*",
        "t5_tokenizer/*",
        "mllm/*.json",
        "mllm/*.txt",
        "mllm/*.model",
    ],
    max_workers=8,
)
_stat("after small snapshot")

# ------------------------------------------------------------------ model ---
import json  # noqa: E402

from accelerate import init_empty_weights  # noqa: E402
from safetensors import safe_open  # noqa: E402

from bernini.models import BerniniConfig, BerniniModel  # noqa: E402
from bernini.pipeline import BerniniPipeline, _localize_bernini_config  # noqa: E402
from diffusers.models import AutoencoderKLWan  # noqa: E402
from transformers import AutoProcessor, AutoTokenizer  # noqa: E402

config = BerniniConfig.from_pretrained(
    MODEL_DIR,
    use_unipc=True,
    use_src_id_rotary_emb=True,
    interpolate_src_id=True,
    max_trained_src_id=5,
)
_localize_bernini_config(config, MODEL_DIR)
config.mllm_attn_implementation = "sdpa"

with init_empty_weights():
    model = BerniniModel(config)
model.eval()
model.requires_grad_(False)
_stat("after meta init")

_index_path = hf_hub_download(MODEL_ID, f"{config.bernini_ckpt_subfolder}/model.safetensors.index.json")
_weight_map = json.load(open(_index_path))["weight_map"]
_shards = sorted(set(_weight_map.values()))
_pending = set(_weight_map)

for _i, _shard in enumerate(_shards, 1):
    _p = hf_hub_download(MODEL_ID, f"{config.bernini_ckpt_subfolder}/{_shard}")
    _sd = {}
    with safe_open(_p, framework="pt", device="cpu") as _f:
        for _k in _f.keys():
            _t = _f.get_tensor(_k)
            _sd[_k] = _t.to(torch.bfloat16) if _t.is_floating_point() else _t
            del _t
    model.load_state_dict(_sd, strict=False, assign=True)
    _pending -= set(_sd)
    del _sd
    for _f2 in {os.path.realpath(_p), _p}:
        try:
            os.remove(_f2)
        except OSError:
            pass
    gc.collect()
    print(f"[load] shard {_i}/{len(_shards)} {_shard}", flush=True)

_stat("after shard stream")
_meta = [n for n, p in model.named_parameters() if p.device.type == "meta"]
if _meta:
    print(f"[load] WARNING {len(_meta)} params still on meta, e.g. {_meta[:8]}", flush=True)
if _pending:
    print(f"[load] WARNING {len(_pending)} checkpoint keys unconsumed, e.g. {sorted(_pending)[:8]}", flush=True)

# transformer_2 is loaded inside diff_dec_low and attached back before sampling
setattr(model.diff_dec, "transformer_2", model.diff_dec_low.transformer_2)

t5_tokenizer = AutoTokenizer.from_pretrained(
    config.t5_tokenizer_path, subfolder=config.t5_tokenizer_subfolder, trust_remote_code=True
)
vit_processor = AutoProcessor.from_pretrained(
    config.processor_config_path,
    subfolder=config.processor_subfolder,
    padding_side="right",
    trust_remote_code=True,
)
vae = AutoencoderKLWan.from_pretrained(
    config.vae_model_path, subfolder=config.vae_subfolder, torch_dtype=torch.float32
)
vae.eval()
vae.requires_grad_(False)

PIPE = BerniniPipeline(config, model, vae, t5_tokenizer, vit_processor, "cuda")

# The two 14B renderer DiTs (~56 GB bf16) live on the GPU for the whole life of
# the Space.  The planner stack (MLLM / connector / vit head / T5 / VAE) is much
# smaller and the reference pipeline moves it on and off the device around its
# own phases, so it is left where that code expects to find it.
model.diff_dec.transformer.to("cuda")
model.diff_dec.transformer_2.to("cuda")
gc.collect()
_stat("after DiTs -> cuda")

# ------------------------------------------------------------------- task ---
# Verbatim from scripts/bernini_v2/run_r2v.sh
SYSTEM_PROMPT = "You are a helpful assistant specialized in subject-to-video generation."
NEG_PROMPT = (
    "vivid tones, overexposed, static, blurry details, subtitles, style, artwork, painting, "
    "image, motionless, overall grayish, worst quality, low quality, JPEG compression artifacts, "
    "ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn face, deformed, disfigured, "
    "malformed limbs, fused fingers, still frame, cluttered background, three legs, "
    "too many people in the background, walking backwards"
)
R2V = dict(
    guidance_mode="vae_txt_vit_wapg",
    max_image_size=842,
    flow_shift=5.0,
    fps=16,
    omega_txt=4.5,
    omega_tgt=1.5,
    omega_img=3.0,
    omega_vid=1.0,
    omega_scale=0.75,
    planning_step=50,
    vit_denoising_step=1,
    vit_txt_cfg=1.2,
    vit_img_cfg=1.0,
    eta=0.5,
    momentum=0.0,
    norm_threshold=(50.0, 50.0, 50.0),
)

RESOLUTIONS = {
    "Landscape · 848×480": (480, 848),
    "Portrait · 480×848": (848, 480),
    "Square · 640×640": (640, 640),
}
MAX_SEED = 2**31 - 1


def _coerce_gallery_paths(gallery_input):
    """gr.Gallery hands back a list of (path, caption) tuples."""
    if not gallery_input:
        return None
    out = []
    for item in gallery_input:
        if isinstance(item, (list, tuple)) and item:
            item = item[0]
        if isinstance(item, str):
            out.append(item)
        elif isinstance(item, dict) and item.get("path"):
            out.append(item["path"])
        elif hasattr(item, "name"):
            out.append(item.name)
    return out or None


def _estimate(*args, **kwargs):
    """Runtime scales with (denoising steps x latent tokens)."""
    try:
        n_images = max(1, len(args[0] or []))
        num_frames = int(args[2])
        steps = int(args[3])
        resolution = args[4]
    except Exception:
        return 420
    height, width = RESOLUTIONS.get(resolution, (480, 848))
    latent_frames = (int(num_frames) - 1) // 4 + 1
    tokens = latent_frames * (height // 16) * (width // 16)
    # Fitted on this Space (33f/848x480/16 steps unless noted):
    #   2 refs, 17f, 8 steps ->  95.1 s
    #   2 refs               -> 231.6 s warm / 254.3 s on a cold slot
    #   5 refs               -> 322.8 s
    # Planning cost scales with the reference count, sampling with steps x latent tokens.
    secs = 15.0 + 22.8 * n_images + 9.7e-4 * steps * tokens
    return int(min(800, max(90, secs * 1.15)))


@spaces.GPU(duration=_estimate, size="xlarge")
def generate(
    reference_images,
    prompt,
    num_frames=33,
    num_inference_steps=16,
    resolution="Landscape · 848×480",
    seed=42,
    randomize_seed=False,
    negative_prompt=NEG_PROMPT,
    omega_txt=4.5,
    omega_img=3.0,
    omega_tgt=1.5,
    omega_scale=0.75,
    progress=gr.Progress(track_tqdm=True),
):
    images = _coerce_gallery_paths(reference_images)
    if not images:
        raise gr.Error("Please add at least one reference image.")
    if len(images) > 8:
        raise gr.Error("Please use at most 8 reference images.")
    if not prompt or not prompt.strip():
        raise gr.Error("Please write a prompt describing the video you want.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    height, width = RESOLUTIONS[resolution]

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    kwargs = dict(R2V)
    kwargs.update(
        omega_txt=float(omega_txt),
        omega_img=float(omega_img),
        omega_tgt=float(omega_tgt),
        omega_scale=float(omega_scale),
    )

    t0 = time.perf_counter()
    PIPE(
        "r2v",
        prompt.strip(),
        images=images,
        neg_prompt=negative_prompt or "",
        system_prompt=SYSTEM_PROMPT,
        num_frames=int(num_frames),
        height=int(height),
        width=int(width),
        num_inference_steps=int(num_inference_steps),
        seed=int(seed),
        output_path=out_path,
        **kwargs,
    )
    elapsed = time.perf_counter() - t0
    torch.cuda.empty_cache()
    print(f"[bernini] generated in {elapsed:.1f}s "
          f"({num_frames}f {width}x{height} {num_inference_steps} steps)", flush=True)
    return out_path, int(seed)


# --------------------------------------------------------------------- UI ---
EX1_PROMPT = (
    "The marble statue from image0, wearing the black T-shirt from image2, the tropical floral "
    "shorts from image3, and the pink cat-ear headphones from image1, sits on the wooden bench in "
    "the beach sunset setting from image4, facing the camera and gently bobbing and swaying to the "
    "music in a medium shot. Generate a video where the marble statue from image0 is the main "
    "subject, with the same muscular stone body, curly sculpted hair, and classical carved "
    "appearance, now humorously dressed in the black short-sleeve T-shirt from image2 with the "
    'white word "bernini" across the chest, the bright blue tropical floral shorts from image3 '
    "with large red, orange, and yellow flowers and green leaves, and the pink over-ear cat-ear "
    "headphones from image1. He is seated on the wooden bench from image4, centered in the frame "
    "and facing directly toward the camera in a medium shot. Keep the environment unchanged from "
    "image4: a seaside promenade with the wooden bench in the foreground, sandy beach and calm "
    "ocean behind it, palm trees rising on the left, and a vivid sunset sky glowing with warm "
    "orange, pink, and purple tones. He begins moving subtly and rhythmically as if listening to "
    "music through the headphones, gently nodding his head, swaying his upper body slightly, and "
    "rocking side to side in a natural music-driven motion, always remaining seated on the bench "
    "and facing the camera."
)
EX2_PROMPT = (
    "Place the male marble sculpture from image0 on the bench in image1, wearing the black T-shirt "
    'from image2 with the word "bernini" across the chest, holding the brown ceramic cup from '
    "image3 and slowly drinking from it with no steam visible, always facing the camera in a fixed "
    "medium shot. Keep the seaside sunset setting from image1 unchanged: the wooden bench centered "
    "on a paved path, palm trees on the left, and the beach, ocean and glowing sun in the "
    "background under a pink and orange sky. He starts seated upright holding the cup near his "
    "torso with a subtle rhythmic sway of the shoulders, then slowly lifts the cup toward his "
    "mouth in a controlled motion, gently tilts it and takes a sip, and finally lowers it while "
    "continuing a soft bobbing motion of the head and torso."
)

EXAMPLES = [
    [
        [
            "examples/source_img0.png",
            "examples/source_img1.png",
            "examples/source_img2.png",
            "examples/source_img3.png",
            "examples/source_img4.png",
        ],
        EX1_PROMPT,
    ],
    [
        [
            "examples/source_img0.png",
            "examples/source_img4.png",
            "examples/source_img2.png",
            "examples/source_img7.png",
        ],
        EX2_PROMPT,
    ],
]

CSS = """
#col-container { margin: 0 auto; max-width: 1100px; }
"""

with gr.Blocks(title="Bernini-Diffusers-v2") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
# Bernini-Diffusers-v2 — reference-to-video

Drop in a few **reference images** (a subject, an outfit, a prop, a scene…), then describe the
video you want while pointing at them as `image0`, `image1`, … Bernini's Qwen2.5-VL planner reads
the references plus your instruction and plans a target visual embedding, which the Wan2.2-A14B
MoE renderer turns into a video.

[model](https://huggingface.co/ByteDance/Bernini-Diffusers-v2) ·
[code](https://github.com/bytedance/Bernini)
"""
        )

        with gr.Row():
            with gr.Column(scale=1):
                reference_images = gr.Gallery(
                    label="Reference images (order matters → image0, image1, …)",
                    file_types=["image"],
                    type="filepath",
                    columns=4,
                    height=240,
                    object_fit="contain",
                    interactive=True,
                    show_label=True,
                )
                prompt = gr.Textbox(
                    label="Prompt",
                    lines=6,
                    placeholder="The statue from image0, wearing the shirt from image1, sits on a "
                    "bench at sunset and gently sways to the music in a medium shot…",
                )
                run_btn = gr.Button("Generate video", variant="primary")
            with gr.Column(scale=1):
                video_out = gr.Video(label="Result", autoplay=True, height=380)
                used_seed = gr.Number(label="Seed used", interactive=False)

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                num_frames = gr.Slider(
                    label="Frames (16 fps)", minimum=17, maximum=49, step=4, value=33
                )
                num_inference_steps = gr.Slider(
                    label="Denoising steps", minimum=8, maximum=24, step=1, value=16
                )
            resolution = gr.Radio(
                label="Resolution",
                choices=list(RESOLUTIONS.keys()),
                value="Landscape · 848×480",
            )
            with gr.Row():
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
            negative_prompt = gr.Textbox(label="Negative prompt", value=NEG_PROMPT, lines=3)
            gr.Markdown("Guidance weights — the defaults are the authors' `run_r2v.sh` values.")
            with gr.Row():
                omega_txt = gr.Slider(label="omega_txt", minimum=1.0, maximum=8.0, step=0.1, value=4.5)
                omega_img = gr.Slider(label="omega_img", minimum=0.0, maximum=8.0, step=0.1, value=3.0)
                omega_tgt = gr.Slider(label="omega_tgt", minimum=0.0, maximum=6.0, step=0.1, value=1.5)
                omega_scale = gr.Slider(label="omega_scale", minimum=0.0, maximum=1.0, step=0.05, value=0.75)

        gr.Markdown(
            "Longer clips and more steps look better but cost more GPU time. The defaults "
            "(33 frames ≈ 2 s at 16 fps, 16 steps) take about 4 minutes; the authors' reference "
            "setting is 81 frames / 40 steps, which does not fit in a single ZeroGPU slot."
        )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[reference_images, prompt],
            outputs=[video_out, used_seed],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
            label="Official Bernini r2v examples",
        )

    inputs = [
        reference_images,
        prompt,
        num_frames,
        num_inference_steps,
        resolution,
        seed,
        randomize_seed,
        negative_prompt,
        omega_txt,
        omega_img,
        omega_tgt,
        omega_scale,
    ]
    run_btn.click(fn=generate, inputs=inputs, outputs=[video_out, used_seed], api_name="generate")

demo.queue(max_size=12).launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)