probe

README.md

@@ -1,51 +1,14 @@
 ---
-title: Bernini Diffusers v2
+title: Bernini Diffusers V2 Demo
 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: Reference-to-video with ByteDance Bernini-Diffusers-v2
-python_version: "3.12"
-startup_duration_timeout: 1h
-models:
-  - ByteDance/Bernini-Diffusers-v2
+short_description: Probe
+startup_duration_timeout: 3h
 ---
-
-# 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,454 +1,36 @@
-"""Bernini-Diffusers-v2 — reference-to-video (subject-to-video) demo.
+import os, shutil, subprocess, gradio as gr, spaces, torch
 
-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
+def sh(c):
     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
+        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):
     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)
+    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()

requirements.txt

@@ -1,31 +0,0 @@
-# --- 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