[Admin maintenance] Support new ZeroGPU hardware

README.md

@@ -4,7 +4,7 @@ emoji: 🐨
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 4.36.0
+sdk_version: 5.49.1
 app_file: app.py
 pinned: false
 license: other

app.py

@@ -1,13 +1,142 @@
 import os
-import torch
 import sys
-import spaces #fixme
+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
-import random
 from configs.infer_config import get_parser
-from huggingface_hub import hf_hub_download
+
+
+# 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'],
@@ -27,6 +156,7 @@ img_examples = [
 
 max_seed = 2 ** 31
 
+
 def download_model():
     REPO_ID = 'Drexubery/ViewCrafter_25'
     filename_list = ['model.ckpt']
@@ -34,29 +164,16 @@ def download_model():
         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() #fixme
-parser = get_parser() # infer_config.py
-opts = parser.parse_args() # default device: 'cuda:0'
+
+
+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)
-test_tensor = torch.Tensor([0]).cuda()
-opts.device = str(test_tensor.device)
-opts.config = './configs/inference_pvd_1024_gradio.yaml' #fixme
-# opts.config = './configs/inference_pvd_1024_local.yaml' #fixme
-
-# install pytorch3d # fixme
-pyt_version_str=torch.__version__.split("+")[0].replace(".", "")
-version_str="".join([
-    f"py3{sys.version_info.minor}_cu",
-    torch.version.cuda.replace(".",""),
-    f"_pyt{pyt_version_str}"
-])
-print(version_str)
-os.system(f"{sys.executable} -m pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html")
-os.system("mkdir -p checkpoints/ && wget https://download.europe.naverlabs.com/ComputerVision/DUSt3R/DUSt3R_ViTLarge_BaseDecoder_512_dpt.pth -P checkpoints/")
-print(f'>>> System info: {version_str}')
+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
@@ -86,7 +203,7 @@ def show_traj(mode):
 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) # fixme
+    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'>\
@@ -106,10 +223,6 @@ def viewcrafter_demo(opts):
 
         with gr.Row():
             with gr.Column():
-                # # step 1: input an image
-                # gr.Markdown("---\n## Step 1: Input an Image, selet an elevation angle and a center_scale factor", show_label=False, visible=True)
-                # gr.Markdown("<div align='left' style='font-size:18px;color: #000000'>1. Estimate an elevation angle  that represents the angle at which the image was taken; a value bigger than 0 indicates a top-down view, and it doesn't need to be precise. <br>2. The origin of the world coordinate system is by default defined at the point cloud corresponding to the center pixel of the input image. You can adjust the position of the origin by modifying center_scale; a value smaller than 1 brings the origin closer to you.</div>")
-
                 with gr.Column():
                     i2v_input_image = gr.Image(label="Input Image",elem_id="input_img")
                     with gr.Row():
@@ -120,11 +233,11 @@ def viewcrafter_demo(opts):
                         left = gr.Button(value = "Left")
                         right = gr.Button(value = "Right")
                         up = gr.Button(value = "Up")
-                    with gr.Row(): 
-                        down = gr.Button(value = "Down")                              
+                    with gr.Row():
+                        down = gr.Button(value = "Down")
                         zin = gr.Button(value = "Zoom in")
                         zout = gr.Button(value = "Zoom out")
-                    with gr.Row(): 
+                    with gr.Row():
                         custom = gr.Button(value = "Customize")
                         reset = gr.Button(value = "Reset")
 
@@ -136,25 +249,21 @@ def viewcrafter_demo(opts):
                                 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],
-                                    )     
+                                    )
 
 
-            # step 3 - Generate video
             with gr.Column():
-                # gr.Markdown("---\n## Step 3: Generate video", show_label=False, visible=True)
-                # gr.Markdown("<div align='left' style='font-size:18px;color: #000000'> You can reduce the sampling steps for faster inference; try different random seed if the result is not satisfying. </div>")
                 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_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,],
-            # examples_per_page=6
-        )            
+        )
 
 
 
@@ -201,6 +310,4 @@ def viewcrafter_demo(opts):
 
 viewcrafter_iface = viewcrafter_demo(opts)
 viewcrafter_iface.queue(max_size=10)
-viewcrafter_iface.launch() #fixme
-# viewcrafter_iface.launch(server_name='11.220.92.96', server_port=80, max_threads=10,debug=True)
-
+viewcrafter_iface.launch()

lvdm/modules/encoders/condition.py

@@ -214,9 +214,16 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder):
     def encode_with_transformer(self, text):
         x = self.model.token_embedding(text)  # [batch_size, n_ctx, d_model]
         x = x + self.model.positional_embedding
-        x = x.permute(1, 0, 2)  # NLD -> LND
+        # open-clip-torch >= 2.20 creates nn.MultiheadAttention with batch_first=True
+        # by default; skip the NLD <-> LND permutes when that's the case.
+        batch_first = getattr(
+            self.model.transformer.resblocks[0].attn, "batch_first", False
+        )
+        if not batch_first:
+            x = x.permute(1, 0, 2)  # NLD -> LND
         x = self.text_transformer_forward(x, attn_mask=self.model.attn_mask)
-        x = x.permute(1, 0, 2)  # LND -> NLD
+        if not batch_first:
+            x = x.permute(1, 0, 2)  # LND -> NLD
         x = self.model.ln_final(x)
         return x
 
@@ -343,7 +350,8 @@ class FrozenOpenCLIPImageEmbedderV2(AbstractEncoder):
         x = self.preprocess(x)
 
         # to patches - whether to use dual patchnorm - https://arxiv.org/abs/2302.01327v1
-        if self.model.visual.input_patchnorm:
+        # open-clip-torch >= 2.20 removed the input_patchnorm attribute; default False.
+        if getattr(self.model.visual, "input_patchnorm", False):
             # einops - rearrange(x, 'b c (h p1) (w p2) -> b (h w) (c p1 p2)')
             x = x.reshape(x.shape[0], x.shape[1], self.model.visual.grid_size[0], self.model.visual.patch_size[0], self.model.visual.grid_size[1], self.model.visual.patch_size[1])
             x = x.permute(0, 2, 4, 1, 3, 5)

requirements.txt

@@ -2,33 +2,29 @@ decord==0.6.0
 einops==0.6.1
 imageio==2.27.0
 imageio-ffmpeg==0.4.8
-torch==2.0.0
 torchvision
 kornia
 matplotlib
 moviepy==1.0.3
-numpy==1.23.5
-open-clip-torch==2.17.1
-opencv-python==4.7.0.72
-Pillow==9.4.0
-pip==23.0.1
-pyglet==1.5.0
-pytorch-lightning==1.8.3
-PyYAML==6.0
-roma==1.5.0
-setuptools==65.6.3
-scikit-image==0.20.0
-scikit-learn==1.2.2
-scipy==1.9.1
-timm==0.6.13
-tqdm==4.65.0
-transformers==4.25.1
+numpy<2
+open-clip-torch
+opencv-python
+Pillow
+pyglet
+pytorch-lightning
+PyYAML
+roma
+scikit-image
+scikit-learn
+scipy
+timm
+tqdm
+transformers
 trimesh==4.4.3
 omegaconf==2.3.0
-triton
 av
 xformers
 gradio
 iopath
 fvcore
-# pytorch3d @ git+https://github.com/facebookresearch/pytorch3d.git
+huggingface_hub

utils/pvd_utils.py

@@ -44,7 +44,16 @@ def save_video(data,images_path,folder=None):
         images = [np.array(Image.open(os.path.join(folder_name,path))) for folder_name,path in zip(folder,data)]
         stacked_images = np.stack(images, axis=0)
         tensor_data = torch.from_numpy(stacked_images).to(torch.uint8)
-    torchvision.io.write_video(images_path, tensor_data, fps=8, video_codec='h264', options={'crf': '10'})
+    # torchvision >= 0.22 removed io.write_video; fall back to imageio.
+    if hasattr(torchvision.io, "write_video"):
+        torchvision.io.write_video(images_path, tensor_data, fps=8, video_codec='h264', options={'crf': '10'})
+    else:
+        import imageio
+        frames_np = tensor_data.numpy() if hasattr(tensor_data, 'numpy') else np.asarray(tensor_data)
+        writer = imageio.get_writer(images_path, fps=8, codec='libx264', quality=8)
+        for frame in frames_np:
+            writer.append_data(frame)
+        writer.close()
 
 def get_input_dict(img_tensor,idx,dtype = torch.float32):
 
@@ -510,7 +519,11 @@ def visualizer_frame(camera_poses, highlight_index):
     fig.canvas.draw()
     width, height = fig.canvas.get_width_height()
     
-    img = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8').reshape(height, width, 3)
+    # matplotlib >= 3.10 removed tostring_rgb(); use buffer_rgba() and drop alpha.
+    if hasattr(fig.canvas, "tostring_rgb"):
+        img = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8').reshape(height, width, 3)
+    else:
+        img = np.asarray(fig.canvas.buffer_rgba())[..., :3].copy()
     # new_width = int(width * 0.6)
     # start_x = (width - new_width) // 2 + new_width // 5
     # end_x = start_x + new_width