Update app.py

app.py

@@ -1,3 +1,4 @@
+import torch
 from typing import List
 from pydantic import BaseModel
 from lama_cleaner.server import main
@@ -16,7 +17,7 @@ class FakeArgs(BaseModel):
     sd_run_local: bool = False
     local_files_only: bool = False
     cpu_offload: bool = False
-    device: str = "cuda"
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
     gui: bool = False
     gui_size: List[int] = [1000, 1000]
     input: str = ''
@@ -27,12 +28,12 @@ class FakeArgs(BaseModel):
     enable_xformers: bool = False
     enable_interactive_seg: bool = True
     interactive_seg_model: str = "vit_b"
-    interactive_seg_device: str = "cuda"
+    interactive_seg_device: str = "cuda" if torch.cuda.is_available() else "cpu"
     enable_remove_bg: bool = False
     enable_anime_seg: bool = False
     enable_realesrgan: bool = False
     enable_gfpgan: bool = False
-    gfpgan_device: str = "cuda"
+    gfpgan_device: str = "cuda" if torch.cuda.is_available() else "cpu"
     enable_restoreformer: bool = False
     enable_gif: bool = False
     quality: int = 95
@@ -40,4 +41,11 @@ class FakeArgs(BaseModel):
     output_dir: str = None
 
 if __name__ == "__main__":
-    main(FakeArgs())
+    if torch.cuda.is_available():
+        print("CUDA is available. Using GPU.")
+    else:
+        print("CUDA is not available. Using CPU.")
+        
+    args = FakeArgs()
+    print(f"Using device: {args.device}")
+    main(args)