Update app.py

app.py

@@ -1,26 +1,18 @@
 import gradio as gr
 from ultralytics import YOLO
-import os
+import spaces  # Import the `spaces` library
 
-# Load pre-trained YOLOv8 models
-docseg_model1 = YOLO("yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt")
-# docseg_model2 = YOLO("path/to/your/second/model.pt")  # Replace with your second model's path
+# Load pre-trained YOLOv8 model
+model = YOLO("yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt")
 
-# Available models
-MODELS = {
-    "DocLayNet YOLOv8": docseg_model1,
-    # "Your Second Model": docseg_model2  # Uncomment and add more as needed
-}
-
-def process_image(image, model_name):
+# Decorate the `process_image` function with `@spaces.GPU`
+@spaces.GPU(duration=60)  # Optional: Set the duration if needed
+def process_image(image):
     try:
-        # Select the model
-        model = MODELS[model_name]
-
         # Process the image
         results = model(source=image, save=False, show_labels=True, show_conf=True, show_boxes=True)
         result = results[0]
-        
+
         # Extract the annotated image and the labels/confidence scores
         annotated_image = result.plot()
         detected_areas_labels = "\n".join(
@@ -30,17 +22,12 @@ def process_image(image, model_name):
         return annotated_image, detected_areas_labels
     except Exception as e:
         return None, f"Error processing image: {e}"
-    
-
 
 # Create the Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Document Segmentation Demo")
-    
+    gr.Markdown("# Document Segmentation Demo (ZeroGPU)")  
     # Input Components
-    with gr.Row():
-        input_image = gr.Image(type="pil", label="Upload Image")
-        model_dropdown = gr.Dropdown(list(MODELS.keys()), label="Select Model", value=list(MODELS.keys())[0])
+    input_image = gr.Image(type="pil", label="Upload Image")
 
     # Output Components
     output_image = gr.Image(type="pil", label="Annotated Image")
@@ -48,7 +35,7 @@ with gr.Blocks() as demo:
 
     # Button to trigger inference
     btn = gr.Button("Run Document Segmentation")
-    btn.click(fn=process_image, inputs=[input_image, model_dropdown], outputs=[output_image, output_text])
+    btn.click(fn=process_image, inputs=input_image, outputs=[output_image, output_text])
 
 # Launch the demo
-demo.launch()
+demo.queue(max_size=1).launch()  # Queue to handle concurrent requests