Update app.py

app.py

@@ -1,47 +1,52 @@
-import gradio as gr
-from ultralytics import YOLO
-import cv2
-import numpy as np
+import spaces
+import datetime
 import os
-import requests
+import subprocess
 import torch
+import gradio as gr
 
-# Ensure the model file is in the correct location
-model_path = "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
-if not os.path.exists(model_path):
-    # Download the model file if it doesn't exist
-    model_url = "https://huggingface.co/DILHTWD/documentlayoutsegmentation_YOLOv8_ondoclaynet/resolve/main/yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
-    response = requests.get(model_url)
-    with open(model_path, "wb") as f:
-        f.write(response.content)
-
-# Load the document segmentation model on CPU
-device = torch.device('cpu')
-docseg_model = YOLO(model_path).to(device)
-
-def process_image(image):
-    # Convert image to the format YOLO model expects
-    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    results = docseg_model(image)
-
-    # Extract annotated image from results
-    annotated_img = results[0].plot()
-    annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
-
-    # Prepare detected areas and labels as text output
-    detected_areas_labels = "\n".join(
-        [f"{box.label}: {box.conf:.2f}" for box in results[0].boxes]
+CUSTOM_CSS = """
+#output_box textarea {
+    font-family: IBM Plex Mono, ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
+}
+"""
+
+zero = torch.Tensor([0]).cuda()
+print(zero.device) # <-- 'cpu' 🤔
+
+@spaces.GPU
+def run_gpu() -> str:
+    print(zero.device) # <-- 'cuda:0' 🤗
+    output: str = ""
+    try:
+        output = subprocess.check_output(["nvidia-smi"], text=True)
+    except FileNotFoundError:
+        output = "nvidia-smi failed"
+    comment = (
+        datetime.datetime.now().replace(microsecond=0).isoformat().replace("T", " ")
     )
+    return f"# {comment}\n\n{output}"
+
+def run(check: bool) -> str:
+    if check:
+        return run_gpu()
+    else:
+        comment = (
+            datetime.datetime.now().replace(microsecond=0).isoformat().replace("T", " ")
+        )
+        return f"# {comment}\n\nThis is running on CPU\n\nClick on 'Run on GPU' below to move to GPU instantly and run nvidia-smi"
+
+output = gr.Textbox(
+    label="Command Output", max_lines=32, elem_id="output_box", value=run(False)
+)
 
-    return annotated_img, detected_areas_labels
+with gr.Blocks(css=CUSTOM_CSS) as demo:
+    gr.Markdown("#### `zero-gpu`: how to run on serverless GPU for free on Spaces 🔥")
 
-# Define the Gradio interface
-interface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(type="pil", label="Annotated Image"),
-             gr.Textbox(label="Detected Areas and Labels")]
-)
+    output.render()
+
+    check = gr.Checkbox(label="Run on GPU")
+
+    check.change(run, inputs=[check], outputs=output, every=1)
 
-if __name__ == "__main__":
-    interface.launch()
+demo.queue().launch(show_api=False)