app.py

import os
import sys
import subprocess

# ---------------------------------------------------------------------------
# Blackwell ZeroGPU shim — env + heavy CUDA-extension build
# ---------------------------------------------------------------------------
# pytorch3d publishes prebuilt wheels only for torch <= 2.4. On the new
# Blackwell ZeroGPU stack (torch 2.10/2.11, CUDA 13) the old offline-install
# block in this file simply fails. Build pytorch3d from source the first time
# we get a GPU, the same way the Blackwell playbook recommends for
# nvdiffrast / diff_gaussian_rasterization.

import spaces  # MUST come before torch / CUDA-touching imports
import torch

# torch.load weights_only flipped in 2.6 — old ckpts (DUSt3R/dynamicrafter) use
# argparse Namespaces / numpy scalars that the new default refuses to unpickle.
_orig_torch_load = torch.load
def _patched_torch_load(*args, **kwargs):
    kwargs.setdefault("weights_only", False)
    return _orig_torch_load(*args, **kwargs)
torch.load = _patched_torch_load

# xformers on the Blackwell ZeroGPU wheel ships without CUDA-built ops for
# sm_120: FA3 needs cap <= 9.0, Cutlass needs cap <= 9.0 too. Every op raises
# `NotImplementedError`. Replace `xformers.ops.memory_efficient_attention` with
# a torch-native SDPA equivalent so existing call-sites keep working.
def _mea_sdpa(q, k, v, attn_bias=None, p=0.0, scale=None, op=None):
    # xformers convention: q/k/v shaped (B, M, H, K) or (B*H, M, K).
    # The lvdm code path passes 3D (B*H, M, K). Convert to (B*H, 1, M, K)
    # for F.scaled_dot_product_attention.
    import torch.nn.functional as F
    is_3d = (q.ndim == 3)
    if is_3d:
        q = q.unsqueeze(1); k = k.unsqueeze(1); v = v.unsqueeze(1)
    elif q.ndim == 4:
        # (B, M, H, K) -> (B, H, M, K)
        q = q.transpose(1, 2); k = k.transpose(1, 2); v = v.transpose(1, 2)
    mask = None
    if attn_bias is not None and hasattr(attn_bias, "to_tensor"):
        mask = attn_bias.to_tensor()
    elif attn_bias is not None and torch.is_tensor(attn_bias):
        mask = attn_bias
    out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=p, scale=scale)
    if is_3d:
        out = out.squeeze(1)
    else:
        out = out.transpose(1, 2)
    return out

try:
    import xformers.ops as _xops  # noqa: E402
    _xops.memory_efficient_attention = _mea_sdpa
except Exception as _e:
    print(f"[xformers shim] could not patch xformers.ops: {_e}")


CUDA_HOME = "/cuda-image/usr/local/cuda-13.0"


@spaces.GPU(duration=600)
def _first_gpu_setup():
    """Build pytorch3d from source against the live torch on first GPU acquire."""
    try:
        import pytorch3d  # noqa: F401
        return
    except ImportError:
        pass

    import tempfile
    patch_dir = tempfile.mkdtemp(prefix="torch_cuda_patch_")
    with open(os.path.join(patch_dir, "sitecustomize.py"), "w") as f:
        f.write(
            "try:\n"
            "    import torch.utils.cpp_extension as _c\n"
            "    _c._check_cuda_version = lambda *a, **k: None\n"
            "except Exception:\n"
            "    pass\n"
        )

    env = os.environ.copy()
    env["CUDA_HOME"] = CUDA_HOME
    env["CUDA_PATH"] = CUDA_HOME
    env["PATH"] = os.path.join(CUDA_HOME, "bin") + os.pathsep + env.get("PATH", "")
    env["PYTHONPATH"] = patch_dir + os.pathsep + env.get("PYTHONPATH", "")
    env["TORCH_CUDA_ARCH_LIST"] = "12.0"
    env["FORCE_CUDA"] = "1"
    # CUDA 13 changed default symbol visibility; pytorch3d's pulsar renderer
    # needs the old behaviour or it fails to link with "undefined reference".
    # https://github.com/facebookresearch/pytorch3d/issues/2011
    env["NVCC_FLAGS"] = "-static-global-template-stub=false"

    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "--no-deps",
         "setuptools", "wheel", "ninja", "packaging"],
    )
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install",
         "--no-build-isolation", "--no-deps",
         "git+https://github.com/facebookresearch/pytorch3d.git@stable"],
        env=env,
    )


# Pre-download the DUSt3R checkpoint at module scope (CPU-only). Using
# hf_hub_download lets us avoid a fresh wget every boot.
from huggingface_hub import hf_hub_download

os.makedirs("./checkpoints/", exist_ok=True)
if not os.path.exists("./checkpoints/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth"):
    try:
        hf_hub_download(
            repo_id="naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt",
            filename="DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth",
            local_dir="./checkpoints/",
        )
    except Exception as _e:
        print(f"[dust3r hf_hub_download fallback to wget] {_e}")
        subprocess.check_call([
            "wget", "-q", "-c",
            "https://download.europe.naverlabs.com/ComputerVision/DUSt3R/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth",
            "-P", "checkpoints/",
        ])


# ---------------------------------------------------------------------------
# Original app
# ---------------------------------------------------------------------------
import random
import gradio as gr
from configs.infer_config import get_parser


# Build pytorch3d before any module-scope code tries to import it. We need a
# GPU here because pytorch3d's CUDA kernels link against the CUDA runtime
# headers — run inside @spaces.GPU.
_first_gpu_setup()


traj_examples = [
    ['0 -35; 0 0; 0 -0.1'],
    ['0 -3 -15 -20 -17 -5 0; 0 -2 -5 -10 -8 -5 0 2 5 3 0; 0 0'],
    ['0 3 10 20 17 10 0; 0 -2 -8 -6 0 2 5 3 0; 0 -0.02 -0.09 -0.16 -0.09 0'],
    ['0 30; 0 -1 -5 -4 0 1 5 4 0; 0 -0.2'],
]


img_examples = [
    ['test/images/boy.png',0,1],
    ['test/images/car.jpeg',5,1],
    ['test/images/fruit.jpg',5,1],
    ['test/images/room.png',10,1],
    ['test/images/castle.png',-4,1],
]

max_seed = 2 ** 31


def download_model():
    REPO_ID = 'Drexubery/ViewCrafter_25'
    filename_list = ['model.ckpt']
    for filename in filename_list:
        local_file = os.path.join('./checkpoints/', filename)
        if not os.path.exists(local_file):
            hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir='./checkpoints/', force_download=True)


download_model()
parser = get_parser()
opts = parser.parse_args()
tmp = str(random.randint(10**(5-1), 10**5 - 1))
opts.save_dir = f'./{tmp}'
os.makedirs(opts.save_dir, exist_ok=True)
opts.device = "cuda" if torch.cuda.is_available() else "cpu"
opts.config = './configs/inference_pvd_1024_gradio.yaml'


from viewcrafter import ViewCrafter


CAMERA_MOTION_MODE = ["Basic Camera Trajectory", "Custom Camera Trajectory"]


def show_traj(mode):
    if mode == 'Left':
        return gr.update(value='0 -35; 0 0; 0 0',visible=True),gr.update(visible=False)
    elif mode == 'Right':
        return gr.update(value='0 35; 0 0; 0 0',visible=True),gr.update(visible=False)
    elif mode == 'Up':
        return gr.update(value='0 0; 0 -30; 0 0',visible=True),gr.update(visible=False)
    elif mode == 'Down':
        return gr.update(value='0 0; 0 20; 0 0',visible=True), gr.update(visible=False)
    elif mode == 'Zoom in':
        return gr.update(value='0 0; 0 0; 0 -0.4',visible=True), gr.update(visible=False)
    elif mode == 'Zoom out':
        return gr.update(value='0 0; 0 0; 0 0.4',visible=True), gr.update(visible=False)
    elif mode == 'Customize':
        return gr.update(value='0 0; 0 0; 0 0',visible=True), gr.update(visible=True)
    elif mode == 'Reset':
        return gr.update(value='0 0; 0 0; 0 0',visible=False), gr.update(visible=False)

def viewcrafter_demo(opts):
    css = """#input_img {max-width: 1024px !important} #output_vid {max-width: 1024px; max-height:576px} #random_button {max-width: 100px !important}"""
    image2video = ViewCrafter(opts, gradio = True)
    image2video.run_both = spaces.GPU(image2video.run_both, duration=290)
    with gr.Blocks(analytics_enabled=False, css=css) as viewcrafter_iface:
        gr.Markdown("<div align='center'> <h1> ViewCrafter: Taming Video Diffusion Models for High-fidelity Novel View Synthesis </span> </h1> \
                      <h2 style='font-weight: 450; font-size: 1rem; margin: 0rem'>\
                        <a href='https://scholar.google.com/citations?user=UOE8-qsAAAAJ&hl=zh-CN'>Wangbo Yu</a>, \
                        <a href='https://doubiiu.github.io/'>Jinbo Xing</a>, <a href=''>Li Yuan</a>, \
                        <a href='https://wbhu.github.io/'>Wenbo Hu</a>, <a href='https://xiaoyu258.github.io/'>Xiaoyu Li</a>,\
                        <a href=''>Zhipeng Huang</a>, <a href='https://scholar.google.com/citations?user=qgdesEcAAAAJ&hl=en/'>Xiangjun Gao</a>,\
                        <a href='https://www.cse.cuhk.edu.hk/~ttwong/myself.html/'>Tien-Tsin Wong</a>,\
                        <a href='https://scholar.google.com/citations?hl=en&user=4oXBp9UAAAAJ&view_op=list_works&sortby=pubdate/'>Ying Shan</a>\
                        <a href=''>Yonghong Tian</a>\
                    </h2> \
                     <a style='font-size:18px;color: #000000' href='https://arxiv.org/abs/2409.02048'> [ArXiv] </a>\
                     <a style='font-size:18px;color: #000000' href='https://drexubery.github.io/ViewCrafter/'> [Project Page] </a>\
                     <a style='font-size:18px;color: #FF5DB0' href='https://github.com/Drexubery/ViewCrafter'> [Github] </a>\
                     <a style='font-size:18px;color: #000000' href='https://www.youtube.com/watch?v=WGIEmu9eXmU'> [Video] </a> </div>")


        with gr.Row():
            with gr.Column():
                with gr.Column():
                    i2v_input_image = gr.Image(label="Input Image",elem_id="input_img")
                    with gr.Row():
                        i2v_elevation = gr.Slider(minimum=-45, maximum=45, step=1, elem_id="elevation", label="elevation", value=5)
                        i2v_center_scale = gr.Slider(minimum=0.1, maximum=2, step=0.1, elem_id="i2v_center_scale", label="center_scale", value=1)
                with  gr.Column():
                    with gr.Row():
                        left = gr.Button(value = "Left")
                        right = gr.Button(value = "Right")
                        up = gr.Button(value = "Up")
                    with gr.Row():
                        down = gr.Button(value = "Down")
                        zin = gr.Button(value = "Zoom in")
                        zout = gr.Button(value = "Zoom out")
                    with gr.Row():
                        custom = gr.Button(value = "Customize")
                        reset = gr.Button(value = "Reset")

                with gr.Column():
                     with gr.Row():
                        with gr.Column():
                            i2v_pose = gr.Text(value = '0 0; 0 0; 0 0', label="Camera trajectory (d_phi sequence; d_theta sequence; d_r sequence)",visible=False)
                            with gr.Column(visible=False) as i2v_egs:
                                gr.Markdown("<div align='left' style='font-size:18px;color: #000000'>Please refer to the <a href='https://github.com/Drexubery/ViewCrafter/blob/main/docs/gradio_tutorial.md' target='_blank'>tutorial</a> for customizing camera trajectory.</div>")
                                gr.Examples(examples=traj_examples,
                                        inputs=[i2v_pose],
                                    )


            with gr.Column():
                i2v_output_video = gr.Video(label="Generated Video",elem_id="output_vid",autoplay=True,show_share_button=True)
                with gr.Row():
                    i2v_steps = gr.Slider(minimum=1, maximum=50, step=1, elem_id="i2v_steps", label="Sampling steps", value=50)
                    i2v_seed = gr.Slider(label='Random seed', minimum=0, maximum=max_seed, step=1, value=0)
                i2v_end_btn = gr.Button("Generate video")
                i2v_traj_video = gr.Video(label="Camera Trajectory",elem_id="traj_vid",autoplay=True,show_share_button=True)


        gr.Examples(examples=img_examples,
            inputs=[i2v_input_image,i2v_elevation, i2v_center_scale,],
        )



        i2v_end_btn.click(inputs=[i2v_input_image, i2v_elevation, i2v_center_scale, i2v_pose, i2v_steps, i2v_seed],
                        outputs=[i2v_output_video,i2v_traj_video],
                        fn = image2video.run_both
        )

        left.click(inputs=[left],
                      outputs=[i2v_pose,i2v_egs],
                      fn = show_traj
                      )
        right.click(inputs=[right],
                      outputs=[i2v_pose,i2v_egs],
                      fn = show_traj
                      )
        up.click(inputs=[up],
                      outputs=[i2v_pose,i2v_egs],
                      fn = show_traj
                      )
        down.click(inputs=[down],
                      outputs=[i2v_pose,i2v_egs],
                      fn = show_traj
                      )
        zin.click(inputs=[zin],
                      outputs=[i2v_pose,i2v_egs],
                      fn = show_traj
                      )
        zout.click(inputs=[zout],
                      outputs=[i2v_pose,i2v_egs],
                      fn = show_traj
                      )
        custom.click(inputs=[custom],
                      outputs=[i2v_pose,i2v_egs],
                      fn = show_traj
                      )
        reset.click(inputs=[reset],
                      outputs=[i2v_pose,i2v_egs],
                      fn = show_traj
                      )

    return viewcrafter_iface


viewcrafter_iface = viewcrafter_demo(opts)
viewcrafter_iface.queue(max_size=10)
viewcrafter_iface.launch()