Update

app.py

@@ -2,8 +2,6 @@
 
 from __future__ import annotations
 
-import functools
-import os
 import pathlib
 import sys
 import tarfile
@@ -13,6 +11,7 @@ import gradio as gr
 import huggingface_hub
 import numpy as np
 import PIL.Image
+import spaces
 import torch
 
 sys.path.insert(0, "yolov5_anime")
@@ -23,7 +22,19 @@ from utils.general import non_max_suppression, scale_coords
 
 DESCRIPTION = "# [zymk9/yolov5_anime](https://github.com/zymk9/yolov5_anime)"
 
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+torch.set_grad_enabled(False)
 MODEL_REPO = "public-data/yolov5_anime"
+model_path = huggingface_hub.hf_hub_download(MODEL_REPO, "yolov5x_anime.pth")
+config_path = huggingface_hub.hf_hub_download(MODEL_REPO, "yolov5x.yaml")
+state_dict = torch.load(model_path)
+model = Model(cfg=config_path)
+model.load_state_dict(state_dict)
+if device.type != "cpu":
+    model.half()
+model.to(device)
+model.eval()
 
 
 def load_sample_image_paths() -> list[pathlib.Path]:
@@ -36,24 +47,9 @@ def load_sample_image_paths() -> list[pathlib.Path]:
     return sorted(image_dir.glob("*"))
 
 
-def load_model(device: torch.device) -> torch.nn.Module:
-    torch.set_grad_enabled(False)
-    model_path = huggingface_hub.hf_hub_download(MODEL_REPO, "yolov5x_anime.pth")
-    config_path = huggingface_hub.hf_hub_download(MODEL_REPO, "yolov5x.yaml")
-    state_dict = torch.load(model_path)
-    model = Model(cfg=config_path)
-    model.load_state_dict(state_dict)
-    model.to(device)
-    if device.type != "cpu":
-        model.half()
-    model.eval()
-    return model
-
-
+@spaces.GPU
 @torch.inference_mode()
-def predict(
-    image: PIL.Image.Image, score_threshold: float, iou_threshold: float, device: torch.device, model: torch.nn.Module
-) -> np.ndarray:
+def predict(image: PIL.Image.Image, score_threshold: float, iou_threshold: float) -> np.ndarray:
     orig_image = np.asarray(image)
 
     image = letterbox(orig_image, new_shape=640)[0]
@@ -83,9 +79,6 @@ def predict(
 image_paths = load_sample_image_paths()
 examples = [[path.as_posix(), 0.4, 0.5] for path in image_paths]
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-model = load_model(device)
-fn = functools.partial(predict, device=device, model=model)
 
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
@@ -103,15 +96,14 @@ with gr.Blocks(css="style.css") as demo:
         examples=examples,
         inputs=inputs,
         outputs=result,
-        fn=fn,
-        cache_examples=os.getenv("CACHE_EXAMPLES") == "1",
+        fn=predict,
     )
     run_button.click(
-        fn=fn,
+        fn=predict,
         inputs=inputs,
         outputs=result,
         api_name="predict",
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=15).launch()
+    demo.queue(max_size=20).launch()

requirements.txt

@@ -1,5 +1,6 @@
 gradio==4.31.5
 opencv-python-headless==4.9.0.80
 scipy==1.13.1
+spaces==0.28.3
 torch==2.0.1
 torchvision==0.15.2