Update core/test_xportrait.py

core/test_xportrait.py

@@ -221,7 +221,7 @@ def x_portrait_data_prep(source_image_path, driving_video_path, device, best_fra
 # You can now use the modified state_dict without the deleted keys
 def load_state_dict(model, ckpt_path, reinit_hint_block=False, strict=True, map_location="cpu"):
     print(f"Loading model state dict from {ckpt_path} ...")
-    state_dict = torch.load(ckpt_path, map_location=map_location)
+    state_dict = torch.load(ckpt_path, map_location=map_location, weights_only=True)
     state_dict = state_dict.get('state_dict', state_dict)
     if reinit_hint_block:
         print("Ignoring hint block parameters from checkpoint!")
@@ -341,7 +341,7 @@ def visualize_mm(args, name, batch_data, infer_model, nSample, local_image_dir,
         
     noise = pre_noise.to(c_cross.device)
 
-    with torch.cuda.amp.autocast(enabled=args.use_fp16, dtype=FP16_DTYPE):
+    with torch.amp.autocast("cuda", enabled=args.use_fp16, dtype=FP16_DTYPE):
         infer_model.to(args.device)
         infer_model.eval()