Update app.py

app.py

@@ -4,13 +4,13 @@ from PIL import Image
 import torch
 import gradio as gr
 
-
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 model2 = torch.hub.load(
     "AK391/animegan2-pytorch:main",
     "generator",
     pretrained=True,
-    device="cuda",
+    device=DEVICE,  #"cuda",
     progress=False
 )
 
@@ -18,7 +18,7 @@ model2 = torch.hub.load(
 model1 = torch.hub.load("AK391/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v1",  device="cuda")
 face2paint = torch.hub.load(
     'AK391/animegan2-pytorch:main', 'face2paint', 
-    size=512, device="cuda",side_by_side=False
+    size=512, device=DEVICE,side_by_side=False
 )
 def inference(img, ver):
     if ver == 'version 2 (🔺 robustness,🔻 stylization)':