[Admin maintenance] Support new ZeroGPU hardware

README.md

@@ -4,7 +4,7 @@ emoji: 🌟
 colorFrom: purple
 colorTo: red
 sdk: gradio
-sdk_version: 5.0.1
+sdk_version: 5.49.1
 app_file: app.py
 pinned: false
 short_description: Video Super-Resolution with Text-to-Video Model

app.py

@@ -1,6 +1,18 @@
-import spaces
 import os
+# Disable torch's NVML-based CUDA check — NVML isn't accessible in ZeroGPU
+# sandbox and torch >=2.10 hits an INTERNAL ASSERT (CUDACachingAllocator
+# NVML_SUCCESS == r) when the caching allocator tries to enumerate devices.
+os.environ.setdefault("PYTORCH_NVML_BASED_CUDA_CHECK", "0")
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
+
+import spaces
 import gradio as gr
+
+# Step 8: weights_only flipped to True in torch 2.6+, old checkpoint needs False.
+import torch
+_orig_load = torch.load
+torch.load = lambda *a, **k: _orig_load(*a, **{**k, "weights_only": k.get("weights_only", False)})
+
 from video_super_resolution.scripts.inference_sr import STAR_sr
 
 # Example video and prompt pairs
@@ -11,7 +23,7 @@ examples = [
 ]
 
 # Define a GPU-decorated function for enhancement
-@spaces.GPU()
+@spaces.GPU(duration=100)
 def enhance_with_gpu(input_video, input_text, upscale, max_chunk_len, chunk_size):
     """在每次调用时创建新的 STAR_sr 实例，确保参数正确传递"""
     star = STAR_sr(

requirements.txt

@@ -1,16 +1,14 @@
-torch==2.0.1
-torchvision==0.15.2
-torchaudio==2.0.2
 opencv-python==4.10.0.84
 easydict==1.13
 einops==0.8.0
 open-clip-torch==2.20.0
-xformers==0.0.21
 fairscale==0.4.13
 torchsde==0.2.6
-pytorch-lightning==2.0.1
-diffusers==0.30.0
-huggingface_hub==0.23.3
-gradio==4.41.0
-numpy==1.24
-tqdm
+pytorch-lightning
+diffusers>=0.30
+huggingface_hub
+gradio==5.49.1
+numpy
+tqdm
+soundfile
+requests

video_to_video/modules/unet_v2v.py

@@ -7,8 +7,9 @@ from abc import abstractmethod
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import xformers
-import xformers.ops
+# xformers removed: prebuilt wheel has no kernel that runs on Blackwell sm_120.
+# All memory_efficient_attention call sites in this file have been replaced
+# with torch.nn.functional.scaled_dot_product_attention (fused FA backend).
 from einops import rearrange
 from fairscale.nn.checkpoint import checkpoint_wrapper
 from timm.models.vision_transformer import Mlp
@@ -162,36 +163,19 @@ class MemoryEfficientCrossAttention(nn.Module):
         v = self.to_v(context)
 
         b, _, _ = q.shape
-        q, k, v = map(
-            lambda t: t.unsqueeze(3).reshape(b, t.shape[
-                1], self.heads, self.dim_head).permute(0, 2, 1, 3).reshape(
-                    b * self.heads, t.shape[1], self.dim_head).contiguous(),
-            (q, k, v),
-        )
-
-        # actually compute the attention, what we cannot get enough of.
-        if q.shape[0] > self.max_bs:
-            q_list = torch.chunk(q, q.shape[0] // self.max_bs, dim=0)
-            k_list = torch.chunk(k, k.shape[0] // self.max_bs, dim=0)
-            v_list = torch.chunk(v, v.shape[0] // self.max_bs, dim=0)
-            out_list = []
-            for q_1, k_1, v_1 in zip(q_list, k_list, v_list):
-                out = xformers.ops.memory_efficient_attention(
-                    q_1, k_1, v_1, attn_bias=None, op=self.attention_op)
-                out_list.append(out)
-            out = torch.cat(out_list, dim=0)
-        else:
-            out = xformers.ops.memory_efficient_attention(
-                q, k, v, attn_bias=None, op=self.attention_op)
+        # Reshape to (B, H, M, K) for torch SDPA (Flash/Mem-efficient backends).
+        # xformers MEA prebuilt wheel has no kernel that runs on Blackwell sm_120;
+        # SDPA gives the same fused FA path natively in torch on this hardware.
+        q = q.reshape(b, q.shape[1], self.heads, self.dim_head).permute(0, 2, 1, 3)
+        k = k.reshape(b, k.shape[1], self.heads, self.dim_head).permute(0, 2, 1, 3)
+        v = v.reshape(b, v.shape[1], self.heads, self.dim_head).permute(0, 2, 1, 3)
 
         if exists(mask):
             raise NotImplementedError
-        out = (
-            out.unsqueeze(0).reshape(
-                b, self.heads, out.shape[1],
-                self.dim_head).permute(0, 2, 1,
-                                       3).reshape(b, out.shape[1],
-                                                  self.heads * self.dim_head))
+
+        out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+        # (B, H, M, K) -> (B, M, H*K)
+        out = out.permute(0, 2, 1, 3).reshape(b, out.shape[2], self.heads * self.dim_head)
         return self.to_out(out)
 
 

video_to_video/video_to_video_model.py

@@ -14,19 +14,25 @@ from video_to_video.diffusion.schedules_sdedit import noise_schedule
 from video_to_video.utils.logger import get_logger
 
 from diffusers import AutoencoderKLTemporalDecoder
-import requests
+from huggingface_hub import hf_hub_download
+import shutil
 
 def download_model(url, model_path):
-    if not os.path.exists(os.path.join(model_path, 'model.pt')):
-        print(f"Model not found at {model_path}, downloading...")
-        response = requests.get(url, stream=True)
-        with open(os.path.join(model_path, 'model.pt'), 'wb') as f:
-            for chunk in response.iter_content(chunk_size=1024):
-                if chunk:
-                    f.write(chunk)
-        print(f"Model downloaded to {model_path}")
-    else:
+    target = os.path.join(model_path, 'model.pt')
+    if os.path.exists(target) and os.path.getsize(target) > 1_000_000_000:
         print(f"Model found at {model_path}, skipping download.")
+        return
+    # Remove partial/corrupt file if present
+    if os.path.exists(target):
+        os.remove(target)
+    os.makedirs(model_path, exist_ok=True)
+    print(f"Model not found at {model_path}, downloading via hf_hub_download...")
+    cached = hf_hub_download(
+        repo_id="SherryX/STAR",
+        filename="I2VGen-XL-based/heavy_deg.pt",
+    )
+    shutil.copy(cached, target)
+    print(f"Model downloaded to {target}")
 
 
 logger = get_logger()