| 75 | |
| 76 | |
| 77 | class InferenceAgent: |
| 78 | def __init__(self, opt): |
| 79 | torch.cuda.empty_cache() |
| 80 | self.opt = opt |
| 81 | self.rank = opt.rank |
| 82 | |
| 83 | # Load Model |
| 84 | self.load_model() |
| 85 | self.load_weight(opt.ckpt_path, rank=self.rank) |
| 86 | self.G.to(self.rank) |
| 87 | self.G.eval() |
| 88 | |
| 89 | # Load Data Processor |
| 90 | self.data_processor = DataProcessor(opt) |
| 91 | |
| 92 | def load_model(self) -> None: |
| 93 | self.G = FLOAT(self.opt) |
| 94 | |
| 95 | def load_weight(self, checkpoint_path: str, rank: int) -> None: |
| 96 | state_dict = torch.load(checkpoint_path, map_location='cpu', weights_only=True) |
| 97 | with torch.no_grad(): |
| 98 | for model_name, model_param in self.G.named_parameters(): |
| 99 | if model_name in state_dict: |
| 100 | model_param.copy_(state_dict[model_name].to(rank)) |
| 101 | elif "wav2vec2" in model_name: pass |
| 102 | else: |
| 103 | print(f"! Warning; {model_name} not found in state_dict.") |
| 104 | |
| 105 | del state_dict |
| 106 | |
| 107 | def save_video(self, vid_target_recon: torch.Tensor, video_path: str, audio_path: str) -> str: |
| 108 | with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_video: |
| 109 | temp_filename = temp_video.name |
| 110 | vid = vid_target_recon.permute(0, 2, 3, 1) |
| 111 | vid = vid.detach().clamp(-1, 1).cpu() |
| 112 | vid = ((vid + 1) / 2 * 255).type('torch.ByteTensor') |
| 113 | torchvision.io.write_video(temp_filename, vid, fps=self.opt.fps) |
| 114 | if audio_path is not None: |
| 115 | with open(os.devnull, 'wb') as f: |
| 116 | command = "ffmpeg -i {} -i {} -c:v copy -c:a aac {} -y".format(temp_filename, audio_path, video_path) |
| 117 | subprocess.call(command, shell=True, stdout=f, stderr=f) |
| 118 | if os.path.exists(video_path): |
| 119 | os.remove(temp_filename) |
| 120 | else: |
| 121 | os.rename(temp_filename, video_path) |
| 122 | return video_path |
| 123 | |
| 124 | @torch.no_grad() |
| 125 | def run_inference( |
| 126 | self, |
| 127 | res_video_path: str, |
| 128 | ref_path: str, |
| 129 | audio_path: str, |
| 130 | a_cfg_scale: float = 2.0, |
| 131 | r_cfg_scale: float = 1.0, |
| 132 | e_cfg_scale: float = 1.0, |
| 133 | emo: str = 'S2E', |
| 134 | nfe: int = 10, |