[Admin maintenance] Support new ZeroGPU hardware

README.md

@@ -4,7 +4,7 @@ emoji: 🐨
 colorFrom: green
 colorTo: indigo
 sdk: gradio
-sdk_version: 4.36.1
+sdk_version: 5.49.1
 app_file: app_recon.py
 pinned: false
 ---

app_recon.py

@@ -2,6 +2,18 @@ import functools
 import os
 import shutil
 import sys
+
+import spaces
+
+# Stub huggingface_hub.cached_download (removed in hub>=0.26) for old diffusers.
+import huggingface_hub
+if not hasattr(huggingface_hub, "cached_download"):
+    def _cached_download_removed(*a, **k):
+        raise NotImplementedError(
+            "huggingface_hub.cached_download was removed in 0.26."
+        )
+    huggingface_hub.cached_download = _cached_download_removed
+
 import git
 
 import gradio as gr
@@ -10,13 +22,39 @@ import torch as torch
 from PIL import Image
 
 print(torch.version.cuda)
-os.system('locate libcusolver.so.11')
+
+# xformers.ops.memory_efficient_attention fails on Blackwell sm_120 with the
+# fp32 inputs used by this pipeline. Reimplement against torch SDPA so it
+# works on the new GPUs without changing any call sites.
+import xformers
+import xformers.ops as _xops
+import torch.nn.functional as _F
+
+def _mea_sdpa(query, key, value, attn_bias=None, p=0.0, scale=None, op=None):
+    # xformers shape (B, M, K) or (B, M, H, K) -> SDPA shape (B, H, M, K)
+    if query.dim() == 3:
+        q = query.unsqueeze(1)
+        k = key.unsqueeze(1)
+        v = value.unsqueeze(1)
+        squeeze = True
+    else:
+        q = query.transpose(1, 2)
+        k = key.transpose(1, 2)
+        v = value.transpose(1, 2)
+        squeeze = False
+    out = _F.scaled_dot_product_attention(
+        q, k, v, attn_mask=attn_bias, dropout_p=p, scale=scale,
+    )
+    if squeeze:
+        return out.squeeze(1)
+    return out.transpose(1, 2)
+
+_xops.memory_efficient_attention = _mea_sdpa
+xformers.ops.memory_efficient_attention = _mea_sdpa
 
 from gradio_imageslider import ImageSlider
 from bilateral_normal_integration.bilateral_normal_integration_cupy import bilateral_normal_integration_function
 
-import spaces
-
 import fire
 
 import argparse

models/unet_2d_condition.py

@@ -38,13 +38,17 @@ from diffusers.models.embeddings import (
     ImageHintTimeEmbedding,
     ImageProjection,
     ImageTimeEmbedding,
-    PositionNet,
     TextImageProjection,
     TextImageTimeEmbedding,
     TextTimeEmbedding,
     TimestepEmbedding,
     Timesteps,
 )
+try:
+    from diffusers.models.embeddings import PositionNet
+except ImportError:
+    # diffusers >= 0.26 renamed PositionNet to GLIGENTextBoundingboxProjection
+    from diffusers.models.embeddings import GLIGENTextBoundingboxProjection as PositionNet
 from diffusers.models.modeling_utils import ModelMixin
 
 from models.unet_2d_blocks import (

requirements.txt

@@ -1,19 +1,15 @@
---extra-index-url https://download.pytorch.org/whl/cu117
-torch == 2.0.1
-torchvision == 0.15.2
-xformers == 0.0.21
-diffusers == 0.25.0
+torchvision
+xformers
+diffusers==0.25.0
 opencv-python
-numpy == 1.23.1
 Pillow
 h5py
 matplotlib==3.7.5
 scikit-image
-gradio==4.43.0
 gradio-imageslider==0.0.16
 trimesh==4.0.5
-tqdm == 4.65.0
-accelerate==0.28.0
+tqdm
+accelerate
 transformers==4.39.1
 imgaug
 GitPython==3.1.40
@@ -23,7 +19,5 @@ datetime
 rembg
 segment_anything
 pyvista
-cupy-cuda11x
-trimesh
+cupy-cuda13x
 onnxruntime
-huggingface-hub==0.25.2