update

app.py

@@ -21,6 +21,12 @@ from einops import rearrange
 import math
 
 
+if torch.cuda.is_available():
+  device = "cuda"
+elif torch.backends.mps.is_available():
+  device = "mps"
+else:
+  device = "cpu"
 
 #### Description ####
 title = r"""<h1 align="center">CustomNet: Object Customization with Variable-Viewpoints in Text-to-Image Diffusion Models</h1>"""
@@ -65,7 +71,7 @@ If you have any questions, please feel free to reach me out at <b>yuanzy22@mails
 negtive_prompt = ""
 
 # load model
-device = torch.device("cuda")
+device = torch.device(device)
 preprocess_model = load_preprocess_model()
 config = OmegaConf.load("configs/config_customnet.yaml") 
 model = instantiate_from_config(config.model)
@@ -299,4 +305,4 @@ if __name__=="__main__":
     parser.add_argument("--port", type=int, default=12345)
     args = parser.parse_args()
 
-    main(args)
+    main(args)