import os import sys import tempfile import subprocess import numpy as np import cv2 import torch import torchvision import librosa import face_alignment import gradio as gr from PIL import Image import torchvision.transforms as transforms from transformers import Wav2Vec2FeatureExtractor from tqdm import tqdm import random from huggingface_hub import hf_hub_download # 引入 spaces，用于 ZeroGPU 支持 import spaces sys.path.append(os.path.dirname(os.path.abspath(__file__))) # 尝试导入本地模块 try: from generator.FM import FMGenerator from renderer.models import IMTRenderer except ImportError as e: print(f"Import Error: {e}") print("Please ensure 'generator' and 'renderer' folders are in the same directory.") exit(1) # ========================================== # 自动下载模型权重的逻辑 # ========================================== def ensure_checkpoints(): print("Checking model checkpoints...") REPO_ID = "cbsjtu01/IMTalker" REPO_TYPE = "model" files_to_download = [ "renderer.ckpt", "generator.ckpt", "wav2vec2-base-960h/config.json", "wav2vec2-base-960h/pytorch_model.bin", "wav2vec2-base-960h/preprocessor_config.json", "wav2vec2-base-960h/feature_extractor_config.json", ] TARGET_DIR = "checkpoints" os.makedirs(TARGET_DIR, exist_ok=True) for remote_filename in files_to_download: local_file_path = os.path.join(TARGET_DIR, remote_filename) # 检查文件是否存在且大小正常 (大于 1KB) if not os.path.exists(local_file_path) or os.path.getsize(local_file_path) < 1024: print(f"Downloading {remote_filename} to {TARGET_DIR}...") try: hf_hub_download( repo_id=REPO_ID, filename=remote_filename, repo_type=REPO_TYPE, local_dir=TARGET_DIR, local_dir_use_symlinks=False ) except Exception as e: print(f"Failed to download {remote_filename}: {e}") pass else: print(f"File {local_file_path} already exists. Skipping download.") ensure_checkpoints() class AppConfig: def __init__(self): # 关键：在 ZeroGPU 环境启动时，必须先设为 CPU，不能直接占满显存，否则会被杀掉 self.device = "cpu" self.seed = 42 self.fix_noise_seed = False self.renderer_path = "./checkpoints/renderer.ckpt" self.generator_path = "./checkpoints/generator.ckpt" self.wav2vec_model_path = "./checkpoints/wav2vec2-base-960h" self.input_size = 512 self.input_nc = 3 self.fps = 25.0 self.rank = "cuda" self.sampling_rate = 16000 self.audio_marcing = 2 self.wav2vec_sec = 2.0 self.attention_window = 5 self.only_last_features = True self.audio_dropout_prob = 0.1 self.style_dim = 512 self.dim_a = 512 self.dim_h = 512 self.dim_e = 7 self.dim_motion = 32 self.dim_c = 32 self.dim_w = 32 self.fmt_depth = 8 self.num_heads = 8 self.mlp_ratio = 4.0 self.no_learned_pe = False self.num_prev_frames = 10 self.max_grad_norm = 1.0 self.ode_atol = 1e-5 self.ode_rtol = 1e-5 self.nfe = 10 self.torchdiffeq_ode_method = 'euler' self.a_cfg_scale = 3.0 self.swin_res_threshold = 128 self.window_size = 8 self.ref_path = None self.pose_path = None self.gaze_path = None self.aud_path = None self.crop = True self.source_path = None self.driving_path = None class DataProcessor: def __init__(self, opt): self.opt = opt self.fps = opt.fps self.sampling_rate = opt.sampling_rate print(f"Loading Face Alignment (CPU first)...") # 强制使用 CPU 加载 FaceAlignment，避免初始化时占用 GPU self.fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, device='cpu', flip_input=False) print("Loading Wav2Vec2...") local_path = opt.wav2vec_model_path if os.path.exists(local_path) and os.path.exists(os.path.join(local_path, "config.json")): print(f"Loading local wav2vec from {local_path}") self.wav2vec_preprocessor = Wav2Vec2FeatureExtractor.from_pretrained(local_path, local_files_only=True) else: print("Local wav2vec model not found, downloading from 'facebook/wav2vec2-base-960h'...") self.wav2vec_preprocessor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h") self.transform = transforms.Compose([transforms.Resize((512, 512)), transforms.ToTensor()]) def process_img(self, img: Image.Image) -> Image.Image: img_arr = np.array(img) # 处理灰度图和 RGBA 图 if img_arr.ndim == 2: img_arr = cv2.cvtColor(img_arr, cv2.COLOR_GRAY2RGB) elif img_arr.shape[2] == 4: img_arr = cv2.cvtColor(img_arr, cv2.COLOR_RGBA2RGB) h, w = img_arr.shape[:2] try: print("Detecting face on original high-res image...") # 提示一下，因为大图检测会慢一点 bboxes = self.fa.face_detector.detect_from_image(img_arr) if bboxes is None or len(bboxes) == 0: # 如果原图检测失败，可以考虑保留一个回退方案，或者直接报错 print("Face detection failed on original image.") bboxes = None except Exception as e: print(f"Face detection failed: {e}") bboxes = None valid_bboxes = [] if bboxes is not None: valid_bboxes = [(int(x1), int(y1), int(x2), int(y2), score) for (x1, y1, x2, y2, score) in bboxes if score > 0.5] if not valid_bboxes: print("Warning: No face detected. Using center crop.") cx, cy = w // 2, h // 2 # 默认裁剪一个边长为短边一半的正方形 half = min(w, h) // 4 x1_new, x2_new = cx - half, cx + half y1_new, y2_new = cy - half, cy + half else: # 使用置信度最高的人脸 x1, y1, x2, y2, _ = max(valid_bboxes, key=lambda x: x[4]) cx = (x1 + x2) // 2 cy = (y1 + y2) // 2 w_face = x2 - x1 h_face = y2 - y1 # 计算扩边后的正方形边长的一半 # 这里可以调整扩边系数 0.8 来控制裁剪区域大小 half_side = int(max(w_face, h_face) * 0.8) x1_new = cx - half_side y1_new = cy - half_side x2_new = cx + half_side y2_new = cy + half_side # 边界检查和调整，确保不超出图像范围，同时保持正方形 if x1_new < 0: x2_new += (0 - x1_new); x1_new = 0 if y1_new < 0: y2_new += (0 - y1_new); y1_new = 0 if x2_new > w: x1_new -= (x2_new - w); x2_new = w if y2_new > h: y1_new -= (y2_new - h); y2_new = h # 再次确保坐标在有效范围内 x1_new = max(0, x1_new); y1_new = max(0, y1_new); x2_new = min(w, x2_new); y2_new = min(h, y2_new) # 确保裁剪区域是正方形 (以防万一边界调整破坏了正方形) curr_w = x2_new - x1_new curr_h = y2_new - y1_new min_side = min(curr_w, curr_h) # 如果需要严格中心对齐，这里的调整可能需要更复杂一点，简单起见从左上角调整 x2_new = x1_new + min_side y2_new = y1_new + min_side # 在原图上进行裁剪 crop_img = img_arr[int(y1_new):int(y2_new), int(x1_new):int(x2_new)] crop_pil = Image.fromarray(crop_img) resized_pil = crop_pil.resize((self.opt.input_size, self.opt.input_size), Image.LANCZOS) return resized_pil def process_audio(self, path: str) -> torch.Tensor: speech_array, sampling_rate = librosa.load(path, sr=self.sampling_rate) return self.wav2vec_preprocessor(speech_array, sampling_rate=sampling_rate, return_tensors='pt').input_values[0] def crop_video_stable(self, from_mp4_file_path, to_mp4_file_path, expanded_ratio=0.6, skip_per_frame=1): if os.path.exists(to_mp4_file_path): os.remove(to_mp4_file_path) video = cv2.VideoCapture(from_mp4_file_path) index = 0 bboxes_lists = [] width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) print(f"Analyzing video for stable cropping: {from_mp4_file_path}") while video.isOpened(): success = video.grab() if not success: break if index % skip_per_frame == 0: success, frame = video.retrieve() if not success: break h, w = frame.shape[:2] mult = 360.0 / h resized_frame = cv2.resize(frame, dsize=(0, 0), fx=mult, fy=mult, interpolation=cv2.INTER_AREA if mult < 1 else cv2.INTER_CUBIC) try: detected_bboxes = self.fa.face_detector.detect_from_image(resized_frame) except: detected_bboxes = None current_frame_bboxes = [] if detected_bboxes is not None: for d_box in detected_bboxes: bx1, by1, bx2, by2, score = d_box if score > 0.5: current_frame_bboxes.append([int(bx1 / mult), int(by1 / mult), int(bx2 / mult), int(by2 / mult), score]) if len(current_frame_bboxes) > 0: max_bboxes = max(current_frame_bboxes, key=lambda bbox: bbox[2] - bbox[0]) bboxes_lists.append(max_bboxes) index += 1 video.release() x_center_lists, y_center_lists, width_lists, height_lists = [], [], [], [] for bbox in bboxes_lists: x1, y1, x2, y2 = bbox[:4] x_center, y_center = (x1 + x2) / 2, (y1 + y2) / 2 x_center_lists.append(x_center) y_center_lists.append(y_center) width_lists.append(x2 - x1) height_lists.append(y2 - y1) if not (x_center_lists and y_center_lists and width_lists and height_lists): import shutil shutil.copy(from_mp4_file_path, to_mp4_file_path) return x_center = sorted(x_center_lists)[len(x_center_lists) // 2] y_center = sorted(y_center_lists)[len(y_center_lists) // 2] median_width = sorted(width_lists)[len(width_lists) // 2] median_height = sorted(height_lists)[len(height_lists) // 2] expanded_width = int(median_width * (1 + expanded_ratio)) expanded_height = int(median_height * (1 + expanded_ratio)) fixed_cropped_width = min(max(expanded_width, expanded_height), width, height) x1, y1 = int(x_center - fixed_cropped_width / 2), int(y_center - fixed_cropped_width / 2) x1 = max(0, x1); y1 = max(0, y1) if x1 + fixed_cropped_width > width: x1 = width - fixed_cropped_width if y1 + fixed_cropped_width > height: y1 = height - fixed_cropped_width target_size = self.opt.input_size cmd = (f'ffmpeg -i "{from_mp4_file_path}" -filter:v "crop={fixed_cropped_width}:{fixed_cropped_width}:{x1}:{y1},scale={target_size}:{target_size}:flags=lanczos" -c:v libx264 -crf 18 -preset slow -c:a aac -b:a 128k "{to_mp4_file_path}" -y -loglevel error') if os.system(cmd) != 0: import shutil shutil.copy(from_mp4_file_path, to_mp4_file_path) class InferenceAgent: def __init__(self, opt): torch.cuda.empty_cache() self.opt = opt self.device = opt.device self.data_processor = DataProcessor(opt) print("Loading Models...") self.renderer = IMTRenderer(self.opt).to(self.device) self.generator = FMGenerator(self.opt).to(self.device) if not os.path.exists(self.opt.renderer_path) or not os.path.exists(self.opt.generator_path): raise FileNotFoundError("Checkpoints not found even after download attempt.") self._load_ckpt(self.renderer, self.opt.renderer_path, "gen.") self._load_fm_ckpt(self.generator, self.opt.generator_path) self.renderer.eval() self.generator.eval() # 关键：ZeroGPU 需要在函数内部动态将模型移动到 CUDA def to(self, device): if self.device != device: print(f"Moving models to {device}...") self.device = device self.renderer = self.renderer.to(device) self.generator = self.generator.to(device) def _load_ckpt(self, model, path, prefix="gen."): if not os.path.exists(path): print(f"Warning: Checkpoint {path} not found.") return checkpoint = torch.load(path, map_location="cpu") state_dict = checkpoint.get("state_dict", checkpoint) clean_state_dict = {k.replace(prefix, ""): v for k, v in state_dict.items() if k.startswith(prefix)} model.load_state_dict(clean_state_dict, strict=False) def _load_fm_ckpt(self, model, path): if not os.path.exists(path): return checkpoint = torch.load(path, map_location='cpu') state_dict = checkpoint.get('state_dict', checkpoint) if 'model' in state_dict: state_dict = state_dict['model'] prefix = 'model.' clean_dict = {k[len(prefix):]: v for k, v in state_dict.items() if k.startswith(prefix)} with torch.no_grad(): for name, param in model.named_parameters(): if name in clean_dict: param.copy_(clean_dict[name].to(self.device)) def save_video(self, vid_tensor, fps, audio_path=None): with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp: raw_path = tmp.name if vid_tensor.dim() == 4: vid = vid_tensor.permute(0, 2, 3, 1).detach().cpu().numpy() if vid.min() < 0: vid = (vid + 1) / 2 vid = np.clip(vid, 0, 1) vid = (vid * 255).astype(np.uint8) height, width = vid.shape[1], vid.shape[2] fourcc = cv2.VideoWriter_fourcc(*'mp4v') writer = cv2.VideoWriter(raw_path, fourcc, fps, (width, height)) for frame in vid: writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) writer.release() if audio_path: with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_out: final_path = tmp_out.name cmd = f"ffmpeg -y -i {raw_path} -i {audio_path} -c:v copy -c:a aac -shortest {final_path}" subprocess.call(cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) if os.path.exists(raw_path): os.remove(raw_path) return final_path else: return raw_path @torch.no_grad() def run_audio_inference(self, img_pil, aud_path, crop, seed, nfe, cfg_scale): s_pil = self.data_processor.process_img(img_pil) if crop else img_pil.resize((self.opt.input_size, self.opt.input_size)) s_tensor = self.data_processor.transform(s_pil).unsqueeze(0).to(self.device) a_tensor = self.data_processor.process_audio(aud_path).unsqueeze(0).to(self.device) data = {'s': s_tensor, 'a': a_tensor, 'pose': None, 'cam': None, 'gaze': None, 'ref_x': None} f_r, g_r = self.renderer.dense_feature_encoder(s_tensor) t_lat = self.renderer.latent_token_encoder(s_tensor) if isinstance(t_lat, tuple): t_lat = t_lat[0] data['ref_x'] = t_lat torch.manual_seed(seed) sample = self.generator.sample(data, a_cfg_scale=cfg_scale, nfe=nfe, seed=seed) d_hat = [] T = sample.shape[1] ta_r = self.renderer.adapt(t_lat, g_r) m_r = self.renderer.latent_token_decoder(ta_r) for t in range(T): ta_c = self.renderer.adapt(sample[:, t, ...], g_r) m_c = self.renderer.latent_token_decoder(ta_c) out_frame = self.renderer.decode(m_c, m_r, f_r) d_hat.append(out_frame) vid_tensor = torch.stack(d_hat, dim=1).squeeze(0) return self.save_video(vid_tensor, self.opt.fps, aud_path) @torch.no_grad() def run_video_inference(self, source_img_pil, driving_video_path, crop): s_pil = self.data_processor.process_img(source_img_pil) if crop else source_img_pil.resize((self.opt.input_size, self.opt.input_size)) s_tensor = self.data_processor.transform(s_pil).unsqueeze(0).to(self.device) f_r, i_r = self.renderer.app_encode(s_tensor) t_r = self.renderer.mot_encode(s_tensor) ta_r = self.renderer.adapt(t_r, i_r) ma_r = self.renderer.mot_decode(ta_r) final_driving_path = driving_video_path temp_crop_video = None if crop: with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp: temp_crop_video = tmp.name self.data_processor.crop_video_stable(driving_video_path, temp_crop_video) final_driving_path = temp_crop_video cap = cv2.VideoCapture(final_driving_path) fps = cap.get(cv2.CAP_PROP_FPS) vid_results = [] while True: ret, frame = cap.read() if not ret: break frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame_pil = Image.fromarray(frame).resize((self.opt.input_size, self.opt.input_size)) d_tensor = self.data_processor.transform(frame_pil).unsqueeze(0).to(self.device) t_c = self.renderer.mot_encode(d_tensor) ta_c = self.renderer.adapt(t_c, i_r) ma_c = self.renderer.mot_decode(ta_c) out = self.renderer.decode(ma_c, ma_r, f_r) vid_results.append(out.cpu()) cap.release() if temp_crop_video and os.path.exists(temp_crop_video): os.remove(temp_crop_video) if not vid_results: raise Exception("Driving video reading failed.") vid_tensor = torch.cat(vid_results, dim=0) return self.save_video(vid_tensor, fps=fps, audio_path=driving_video_path) print("Initializing Configuration...") cfg = AppConfig() agent = None try: if os.path.exists(cfg.renderer_path) and os.path.exists(cfg.generator_path): agent = InferenceAgent(cfg) else: print("Error: Checkpoints not found. Please upload 'renderer.ckpt' and 'generator.ckpt' via the Files tab.") except Exception as e: print(f"Initialization Error: {e}") import traceback traceback.print_exc() # 添加 @spaces.GPU 装饰器，必须添加！ @spaces.GPU def fn_audio_driven(image, audio, crop, seed, nfe, cfg_scale, progress=gr.Progress()): if agent is None: raise gr.Error("Models not loaded properly. Check logs.") if image is None or audio is None: raise gr.Error("Missing image or audio.") # 动态移动模型到 GPU if torch.cuda.is_available(): agent.to("cuda") img_pil = Image.fromarray(image).convert('RGB') try: return agent.run_audio_inference(img_pil, audio, crop, int(seed), int(nfe), float(cfg_scale)) except Exception as e: raise gr.Error(f"Error: {e}") # 添加 @spaces.GPU 装饰器，必须添加！ @spaces.GPU def fn_video_driven(source_image, driving_video, crop, progress=gr.Progress()): if agent is None: raise gr.Error("Models not loaded properly. Check logs.") if source_image is None or driving_video is None: raise gr.Error("Missing inputs.") # 动态移动模型到 GPU if torch.cuda.is_available(): agent.to("cuda") img_pil = Image.fromarray(source_image).convert('RGB') try: return agent.run_video_inference(img_pil, driving_video, crop) except Exception as e: import traceback traceback.print_exc() raise gr.Error(f"Error: {e}") with gr.Blocks(title="IMTalker Demo") as demo: gr.Markdown("# 🗣️ IMTalker: Efficient Audio-driven Talking Face Generation") # 最佳实践说明 with gr.Accordion("💡 Best Practices (Click to read)", open=False): gr.Markdown(""" To obtain the highest quality generation results, we recommend following these guidelines: 1. **Input Image Composition**: Please ensure the input image features the person's head as the primary subject. Since our model is explicitly trained on facial data, it does not support full-body video generation. * The inference pipeline automatically **crops the input image** to focus on the face by default. * **Note on Resolution**: The model generates video at a fixed resolution of **512×512**. Using extremely high-resolution inputs will result in downscaling, so prioritize facial clarity over raw image dimensions. 2. **Audio Selection**: Our model was trained primarily on **English datasets**. Consequently, we recommend using **English audio** inputs to achieve the best lip-synchronization performance and naturalness. 3. **Background Quality**: We strongly recommend using source images with **solid colored** or **blurred (bokeh)** backgrounds. Complex or highly detailed backgrounds may lead to visual artifacts or jitter in the generated video. """) with gr.Tabs(): # ========================== # Tab 1: Audio Driven # ========================== with gr.TabItem("Audio Driven"): with gr.Row(): with gr.Column(): # 1. 图片输入 a_img = gr.Image(label="Source Image", type="numpy", height=512, width=512) # --- 图片示例 (独立) --- # 请确保 examples 文件夹下有对应的 source_x.png 文件 gr.Examples( examples=[ ["examples/source_1.png"], ["examples/source_2.png"], ["examples/source_3.jpg"], ["examples/source_4.png"], ["examples/source_5.png"], ["examples/source_6.png"], ], inputs=[a_img], label="Example Images", cache_examples=False, ) # 2. 音频输入 a_aud = gr.Audio(label="Driving Audio", type="filepath") # --- 音频示例 (独立) --- # 请确保 examples 文件夹下有对应的 audio_x.wav 文件 gr.Examples( examples=[ ["examples/audio_1.wav"], ["examples/audio_2.wav"], ["examples/audio_3.wav"], ["examples/audio_4.wav"], ["examples/audio_5.wav"], ], inputs=[a_aud], label="Example Audios", cache_examples=False, ) with gr.Accordion("Settings", open=True): a_crop = gr.Checkbox(label="Auto Crop Face", value=False) a_seed = gr.Number(label="Seed", value=42) a_nfe = gr.Slider(5, 50, value=10, step=1, label="Steps (NFE)") a_cfg = gr.Slider(1.0, 5.0, value=2.0, label="CFG Scale") a_btn = gr.Button("Generate (Audio Driven)", variant="primary") with gr.Column(): a_out = gr.Video(label="Result", height=512, width=512) a_btn.click(fn_audio_driven, [a_img, a_aud, a_crop, a_seed, a_nfe, a_cfg], a_out) # ========================== # Tab 2: Video Driven # ========================== with gr.TabItem("Video Driven"): with gr.Row(): with gr.Column(): # 1. 图片输入 v_img = gr.Image(label="Source Image", type="numpy", height=512, width=512) # --- 图片示例 (独立) --- gr.Examples( examples=[ ["examples/source_7.png"], ["examples/source_8.png"], ["examples/source_9.jpg"], ["examples/source_10.png"], ["examples/source_11.png"], ], inputs=[v_img], label="Example Images", cache_examples=False, ) # 2. 视频输入 v_vid = gr.Video(label="Driving Video", sources=["upload"], height=512, width=512) # --- 视频示例 (独立) --- # 请确保 examples 文件夹下有对应的 driving_x.mp4 文件 gr.Examples( examples=[ ["examples/driving_1.mp4"], ["examples/driving_2.mp4"], ["examples/driving_3.mp4"], ["examples/driving_4.mp4"], ["examples/driving_5.mp4"], ], inputs=[v_vid], label="Example Videos", cache_examples=False, ) v_crop = gr.Checkbox(label="Auto Crop (Both Source & Driving)", value=True) v_btn = gr.Button("Generate (Video Driven)", variant="primary") with gr.Column(): v_out = gr.Video(label="Result", height=512, width=512) v_btn.click(fn_video_driven, [v_img, v_vid, v_crop], v_out) if __name__ == "__main__": demo.queue().launch()