Update app.py

app.py

@@ -1,7 +1,9 @@
-import os, sys
+import os
+import sys
 import cv2
 import time
-import datetime, pytz
+import datetime
+import pytz
 import gradio as gr
 import torch
 import numpy as np
@@ -34,7 +36,11 @@ def auto_download_if_needed(weight_path):
         os.system(f"mv {filename} pretrained")
 
 def load_model_with_device(loader_func, weight_path, scale, device):
-    # First load the state dict
+    # Check if CUDA is available, otherwise fall back to CPU
+    if device.type == 'cuda' and not torch.cuda.is_available():
+        device = torch.device('cpu')  # Force CPU if CUDA is not available
+    
+    # Load the state dict with map_location to handle CPU if necessary
     state_dict = torch.load(weight_path, map_location=device)
     
     # Initialize the model
@@ -152,4 +158,4 @@ if __name__ == '__main__':
 
         run_btn.click(inference, inputs=[input_image, model_name], outputs=[output_image])
 
-    block.launch()
+    block.launch()