Restore working r2v app (revert probe)

README.md

@@ -1,14 +1,51 @@
 ---
-title: Bernini Diffusers V2 Demo
+title: Bernini Diffusers v2
 emoji: 🗿
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
 sdk_version: 6.15.0
-python_version: "3.10"
 app_file: app.py
 pinned: false
 license: apache-2.0
-short_description: Probe
-startup_duration_timeout: 3h
+short_description: Reference-to-video with ByteDance Bernini-Diffusers-v2
+python_version: "3.12"
+startup_duration_timeout: 2h
+models:
+  - ByteDance/Bernini-Diffusers-v2
 ---
+
+# Bernini-Diffusers-v2 — reference-to-video
+
+Give it a handful of **reference images** (a subject, an outfit, a prop, a scene…) and a prompt
+that points at them as `image0`, `image1`, … Bernini's Qwen2.5-VL planner reads the references
+together with the instruction and *plans* a target visual embedding with a flow-matching head; the
+Wan2.2-A14B MoE renderer (two 14 B DiTs, high-noise + low-noise) turns that plan into a video.
+
+- Model: [`ByteDance/Bernini-Diffusers-v2`](https://huggingface.co/ByteDance/Bernini-Diffusers-v2)
+- Code: [`bytedance/Bernini`](https://github.com/bytedance/Bernini)
+
+## What this Space runs
+
+The `r2v` (reference-to-video) task, matching the authors' `scripts/bernini_v2/run_r2v.sh`
+one-for-one: `guidance_mode=vae_txt_vit_wapg`, `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`, `flow_shift=5.0`, `max_image_size=842`, 16 fps, and the same system / negative
+prompt.
+
+The only deviation is the default clip length and step count (33 frames / 16 steps instead of
+81 / 40), so a generation fits inside a single ZeroGPU slot — both are sliders under
+**Advanced settings**. At the defaults a video takes about 4 minutes.
+
+The released checkpoint is fp32 (~180 GB); it is loaded in bf16, which is the dtype the reference
+pipeline computes in anyway.
+
+## Credits
+
+The `bernini/` package and the `veomni/` subset shipped alongside `app.py` are vendored from
+[`bytedance/Bernini`](https://github.com/bytedance/Bernini) and
+[`ByteDance-Seed/VeOmni`](https://github.com/ByteDance-Seed/VeOmni) (v0.1.11), both Apache-2.0,
+because both declare `requires-python` ranges that exclude this runtime.
+
+The example reference images in `examples/` are the authors' own r2v test case assets from
+`bytedance/Bernini` (Apache-2.0).

app.py

@@ -1,36 +1,454 @@
-import os, shutil, subprocess, gradio as gr, spaces, torch
+"""Bernini-Diffusers-v2 — reference-to-video (subject-to-video) demo.
 
-def sh(c):
+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:
-        return subprocess.run(c, shell=True, capture_output=True, text=True, timeout=60).stdout
-    except Exception as e:
-        return repr(e)
-
-INFO = []
-INFO.append("== df -h ==\n" + sh("df -h"))
-INFO.append("== free -g ==\n" + sh("free -g"))
-INFO.append("== nproc ==\n" + sh("nproc"))
-INFO.append("== HOME/cwd ==\n" + os.path.expanduser("~") + " " + os.getcwd())
-INFO.append("== env HF ==\n" + "\n".join(f"{k}={v}" for k, v in os.environ.items() if "HF" in k or "CACHE" in k or "TMP" in k))
-print("\n\n".join(INFO), flush=True)
-
-@spaces.GPU(duration=60)
-def probe(size_label):
+        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 = []
-    out.append(sh("nvidia-smi"))
-    free, total = torch.cuda.mem_get_info()
-    out.append(f"torch VRAM free={free/2**30:.1f}GiB total={total/2**30:.1f}GiB")
-    out.append(torch.cuda.get_device_name(0))
-    out.append("torch " + torch.__version__)
-    out.append(sh("df -h"))
-    out.append(sh("free -g"))
-    return "\n\n".join(out)
-
-with gr.Blocks(theme=gr.themes.Citrus()) as demo:
-    gr.Markdown("probe")
-    b = gr.Button("probe")
-    t = gr.Textbox(lines=30)
-    gr.Markdown("\n\n".join(INFO))
-    b.click(probe, gr.Textbox(value="x", visible=False), t)
-
-demo.queue().launch()
+    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)

requirements.txt

@@ -0,0 +1,31 @@
+# --- core stack (versions from the Bernini repo's own requirements.txt) ---
+transformers==4.57.3
+diffusers==0.35.2
+accelerate
+safetensors
+torchvision
+einops
+numpy
+Pillow
+tqdm
+ftfy
+scipy
+sentencepiece
+packaging
+psutil
+hf_transfer
+
+# --- video / image I/O ---
+decord
+imageio
+imageio-ffmpeg
+
+# --- veomni deps that its inference-side modules actually touch ---
+# (the veomni package itself is vendored in ./veomni, Apache-2.0, v0.1.11,
+#  because pip-installing it drags in datasets<=2.21.0 / torchdata / wandb)
+
+# --- FlashAttention 2 ---
+# bernini/models/modeling_qwen2_5_vl.py raises at import time unless flash_attn
+# is importable, and the MLLM's vision tower asks for flash_attention_2.
+# sm_120 (Blackwell) prebuilt wheel, cp312 / torch 2.11:
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl