kani1021's picture
Add ZeroGPU decorator
ec94ca0
Raw
History Blame Contribute Delete
693 Bytes
try:
import spaces # HF ZeroGPU 需要先於 torch/transformers 匯入
except ImportError:
spaces = None
import gradio as gr
from transformers import pipeline
classifier = pipeline(
task="image-classification",
model="model",
image_processor="model",
)
gpu = spaces.GPU if spaces else lambda fn: fn
@gpu
def predict(image):
results = classifier(image, top_k=3)
return {item["label"]: item["score"] for item in results}
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil", label="Upload image"),
outputs=gr.Label(num_top_classes=3, label="Prediction"),
title="Ball Sport Classifier",
)
if __name__ == "__main__":
demo.launch()