Update app.py

app.py

@@ -6,12 +6,13 @@ import os
 import requests
 import torch
 import huggingface_hub
-
-# Initialize ZeroGPU 
-zero_gpu_is_available = huggingface_hub.utils.is_google_colab() or huggingface_hub.utils.is_notebook()
-if zero_gpu_is_available:
-    from accelerate import Accelerator
-    accelerator = Accelerator()
+from accelerate import Accelerator
+from huggingface_hub import notebook_login  # Added this for HF login
+from huggingface_hub.utils import HfHubHTTPError  # Added this to catch HF login errors
+# Initialize Hugging Face Hub login
+notebook_login()
+# Initialize Accelerator
+accelerator = Accelerator()
 
 
 # Load the model file
@@ -19,25 +20,28 @@ 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)
-
+    try:
+        response = requests.get(model_url)
+        with open(model_path, "wb") as f:
+            f.write(response.content)
+    except HfHubHTTPError as e:
+        if e.response.status_code == 401:
+            print("Authentication error. Please login to Hugging Face Hub.")
+        else:
+            raise e
 # Load the document segmentation model
 docseg_model = YOLO(model_path) 
 
-if zero_gpu_is_available:
-    docseg_model.to(accelerator.device)  # Put the model on the accelerator's device.
 
+docseg_model = accelerator.prepare(docseg_model)
 
 def process_image(image):
     try:
         # Convert image to the format YOLO model expects
         image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
-        # If Zero GPU, move image to accelerator
-        if zero_gpu_is_available:
-            image = torch.from_numpy(image).to(accelerator.device)
+        # Move image to accelerator
+        image = torch.from_numpy(image).to(accelerator.device)
 
         results = docseg_model.predict(image)
         result = results[0]  # Get the first (and usually only) result