NameError

app.py

@@ -1,4 +1,3 @@
-
 import os
 import io
 import cv2
@@ -11,6 +10,7 @@ from functools import lru_cache
 from huggingface_hub import hf_hub_download, snapshot_download
 from torchvision.transforms.functional import normalize
 import glob
+import traceback
 
 
 from restormerRFR_arch import RestormerRFR
@@ -83,39 +83,41 @@ def get_model_and_device():
     return model, device
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=240)
 def restore_image(pil_img: Image.Image) -> Image.Image:
     """
     输入一张图片，输出复原后的图片（与 RAM++ RestormerRFR + DINO 特征推理一致）
     """
-    model, device = get_model_and_device()
-    dino_extractor = get_dino_extractor(device)
-
+    try:
+        model, device = get_model_and_device()
+        dino_extractor = get_dino_extractor(device)
 
-    img_bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR).astype(np.float32) / 255.0
-    img = torch.from_numpy(np.transpose(img_bgr[:, :, [2, 1, 0]], (2, 0, 1))).float()  # (3,H,W), RGB
-    img = img.unsqueeze(0).to(device)  # (1,3,H,W)
 
+        img_bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR).astype(np.float32) / 255.0
+        img = torch.from_numpy(np.transpose(img_bgr[:, :, [2, 1, 0]], (2, 0, 1))).float()  # (3,H,W), RGB
+        img = img.unsqueeze(0).to(device)  # (1,3,H,W)
 
-    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
-    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
-    normalize(img, mean, std, inplace=True)
 
-    with torch.no_grad():
-        dino_features = dino_extractor(img)
-        output = model(img, dino_features)
+        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+        std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+        normalize(img, mean, std, inplace=True)
 
+        with torch.no_grad():
+            dino_features = dino_extractor(img)
+            output = model(img, dino_features)
 
-    output = normalize(output, -1 * mean / std, 1 / std)
-    output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()  # (3,H,W)
-    output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))  # (H,W,RGB)
-    output = (output * 255.0).round().astype(np.uint8)
-    out_pil = Image.fromarray(output, mode="RGB")
-    return out_pil
 
+        output = normalize(output, -1 * mean / std, 1 / std)
+        output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()  # (3,H,W)
+        output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))  # (H,W,RGB)
+        output = (output * 255.0).round().astype(np.uint8)
+        out_pil = Image.fromarray(output, mode="RGB")
+        return out_pil
+    except Exception as e:
+        raise gr.Error(f"{e}\n{traceback.format_exc()}")
 
 DESCRIPTION = """
-# RAM++ Demo 
+# RAM++: Robust Representation Learning via  Adaptive Mask for All-in-One Image Restoration
 """
 
 with gr.Blocks(title="RAM++ ZeroGPU Demo") as demo:

dino_feature_extractor.py

@@ -10,8 +10,9 @@ class DinoFeatureModule(nn.Module):
     def __init__(self, model_id: str = "facebook/dinov2-giant"):
         super(DinoFeatureModule, self).__init__()
         dtype = torch.float32
+        self.model_id = model_id
         self.dino = AutoModel.from_pretrained(
-            model_id,
+            self.model_id,
             torch_dtype=dtype
         )
         
@@ -110,7 +111,7 @@ class DinoFeatureModule(nn.Module):
        
         shortest_edge = min(target_h, target_w)
         processor = AutoImageProcessor.from_pretrained(
-            model_id,
+            self.model_id,
             local_files_only=False,
             do_rescale=False,
             do_center_crop=False,