[Release] cpu support

demo_img.py

@@ -9,7 +9,7 @@ from networks.amtl import Model as AMTL
 from networks.amtg import Model as AMTG
 from utils import img2tensor, tensor2img, InputPadder
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cpu' if torch.cuda.is_available() else 'cpu')
 model_dict = {
     'AMT-S': AMTS, 'AMT-L': AMTL, 'AMT-G': AMTG
 }
@@ -18,7 +18,6 @@ def img2vid(model_type, img0, img1, frame_ratio, iters):
     model = model_dict[model_type]()
     model.to(device)
     ckpt_path = hf_hub_download(repo_id='lalala125/AMT', filename=f'{model_type.lower()}.pth')
-    print(model_type)
     ckpt = torch.load(ckpt_path, map_location=torch.device('cpu'))
     model.load_state_dict(ckpt['state_dict'])
     model.eval()

demo_vid.py

@@ -19,7 +19,6 @@ def vid2vid(model_type, video, iters):
     model = model_dict[model_type]()
     model.to(device)
     ckpt_path = hf_hub_download(repo_id='lalala125/AMT', filename=f'{model_type.lower()}.pth')
-    print(model_type)
     ckpt = torch.load(ckpt_path)
     model.load_state_dict(ckpt['state_dict'], map_location=torch.device('cpu'))
     model.eval()