abidlabs's picture
abidlabs HF Staff
Fix ZeroGPU detection: launch Gradio UI instead of FastAPI/uvicorn
752e4f7
Raw
History Blame Contribute Delete
1.37 kB
import cv2
import gradio as gr
import spaces
from fastrtc import Stream, get_twilio_turn_credentials
from gradio.utils import get_space
try:
from demo.object_detection.inference import YOLOv10
except (ImportError, ModuleNotFoundError):
from inference import YOLOv10
# Load the model and place it on CUDA at module level. On ZeroGPU this uses the
# PyTorch CUDA emulation; the real GPU is attached inside @spaces.GPU below.
model = YOLOv10("yolov10n.pt").to("cuda")
@spaces.GPU
def detection(image, conf_threshold=0.3):
new_image = model.detect_objects(image, conf_threshold)
return cv2.resize(new_image, (500, 500))
stream = Stream(
handler=detection,
modality="video",
mode="send-receive",
additional_inputs=[gr.Slider(minimum=0, maximum=1, step=0.01, value=0.3)],
rtc_configuration=get_twilio_turn_credentials() if get_space() else None,
concurrency_limit=2 if get_space() else None,
)
# ZeroGPU only detects @spaces.GPU functions when Gradio is the launched app,
# so we expose and launch the built-in fastrtc UI (a Gradio Blocks) rather than
# mounting the stream on a FastAPI/uvicorn app.
demo = stream.ui
if __name__ == "__main__":
import os
if (mode := os.getenv("MODE")) == "PHONE":
stream.fastphone(host="0.0.0.0", port=7860)
else:
demo.launch(server_name="0.0.0.0", server_port=7860)