Browse Source

added cpu support to img2img.py

SpandexWizard 4 năm trước cách đây
mục cha
commit
d3f1ebcbd7
1 tập tin đã thay đổi với 2 bổ sung và 1 xóa
  1. 2 1
      scripts/img2img.py

+ 2 - 1
scripts/img2img.py

@@ -40,7 +40,8 @@ def load_model_from_config(config, ckpt, verbose=False):
         print("unexpected keys:")
         print("unexpected keys:")
         print(u)
         print(u)
 
 
-    model.cuda()
+    if torch.cuda.is_available():
+        model.cuda()
     model.eval()
     model.eval()
     return model
     return model