Enable ZeroGPU: add spaces dep and @spaces.GPU decorator

README.md

@@ -7,6 +7,7 @@ sdk: gradio
 sdk_version: 6.15.2
 python_version: "3.12"
 app_file: app.py
+suggested_hardware: zero-a10g
 pinned: false
 license: mit
 ---

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import numpy as np
 import torch
@@ -6,27 +7,26 @@ from transformers import AutoImageProcessor, AutoModel
 
 REPO_NAME = "p1atdev/MangaLineExtraction-hf"
 
-# Load once at startup. trust_remote_code is required: this model ships a
-# custom architecture in its repo.
+# Load on CPU at startup. On ZeroGPU there is no CUDA device until we enter a
+# @spaces.GPU function, so the actual .to("cuda") happens inside extract_lines.
 model = AutoModel.from_pretrained(REPO_NAME, trust_remote_code=True)
 processor = AutoImageProcessor.from_pretrained(REPO_NAME, trust_remote_code=True)
 model.eval()
+model = model.to(dtype=torch.float32)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-# The checkpoint is bf16; cast to float32 on CPU for stable inference.
-model = model.to(device=device, dtype=torch.float32)
 
-
-@torch.no_grad()
+@spaces.GPU
 def extract_lines(image: Image.Image) -> Image.Image:
     if image is None:
         return None
 
     image = image.convert("RGB")
     inputs = processor(image, return_tensors="pt")
-    pixel_values = inputs.pixel_values.to(device=device, dtype=torch.float32)
 
-    outputs = model(pixel_values)
+    with torch.no_grad():
+        model.to("cuda")
+        pixel_values = inputs.pixel_values.to(device="cuda", dtype=torch.float32)
+        outputs = model(pixel_values)
 
     line = outputs.pixel_values[0].cpu().float().numpy()
     line = np.clip(line, 0, 255).astype("uint8")
@@ -41,7 +41,8 @@ demo = gr.Interface(
     description=(
         "Extract clean line art from manga / illustration images using "
         "[p1atdev/MangaLineExtraction-hf](https://huggingface.co/p1atdev/MangaLineExtraction-hf). "
-        "Upload an image and the model returns a grayscale line drawing."
+        "Upload an image and the model returns a grayscale line drawing. "
+        "Runs on ZeroGPU."
     ),
     flagging_mode="never",
 )

requirements.txt

@@ -1,3 +1,4 @@
+spaces
 torch
 transformers
 huggingface_hub<1.0