Spaces:
Sleeping
Sleeping
| import os | |
| import sys | |
| import math | |
| import uuid | |
| import shutil | |
| import glob | |
| import tempfile | |
| import subprocess | |
| import json | |
| import time | |
| import warnings | |
| from pathlib import Path | |
| from typing import Optional, Tuple | |
| import asyncio | |
| import librosa | |
| import soundfile as sf | |
| from pydub import AudioSegment | |
| import cv2 | |
| import gradio as gr | |
| import torch | |
| # Setup paths for Hugging Face Spaces | |
| BASE = Path("/tmp") if os.path.exists("/tmp") else Path.cwd() | |
| WORK = BASE / "ai_avatar_work" | |
| OUT = BASE / "ai_avatar_out" | |
| WORK.mkdir(exist_ok=True, parents=True) | |
| OUT.mkdir(exist_ok=True, parents=True) | |
| # Setup SadTalker | |
| SADTALKER_DIR = BASE / "SadTalker" | |
| def setup_sadtalker(): | |
| """Setup SadTalker if not already available.""" | |
| if not SADTALKER_DIR.exists(): | |
| print("Setting up SadTalker...") | |
| try: | |
| # Clone SadTalker | |
| subprocess.run([ | |
| "git", "clone", "https://github.com/OpenTalker/SadTalker.git", | |
| str(SADTALKER_DIR) | |
| ], check=True, capture_output=True, text=True) | |
| # Install requirements | |
| requirements_path = SADTALKER_DIR / "requirements.txt" | |
| if requirements_path.exists(): | |
| subprocess.run([ | |
| sys.executable, "-m", "pip", "install", "-r", str(requirements_path) | |
| ], check=True, capture_output=True, text=True) | |
| # Download models | |
| download_script = SADTALKER_DIR / "scripts" / "download_models.sh" | |
| if download_script.exists(): | |
| subprocess.run([ | |
| "bash", str(download_script) | |
| ], cwd=str(SADTALKER_DIR), check=True, capture_output=True, text=True) | |
| print("✅ SadTalker setup complete!") | |
| except subprocess.CalledProcessError as e: | |
| print(f"❌ SadTalker setup failed: {e}") | |
| print(f"stdout: {e.stdout}") | |
| print(f"stderr: {e.stderr}") | |
| return False | |
| return True | |
| # Initialize SadTalker on startup | |
| setup_sadtalker() | |
| # -------------------- Configuration -------------------- | |
| class AgentConfig: | |
| def __init__(self, | |
| language="en", | |
| grab_frame_at=1.0, | |
| expr_scale=1.0, | |
| pose_scale=1.0, | |
| fps=25): | |
| self.language = language | |
| self.grab_frame_at = grab_frame_at | |
| self.expr_scale = float(expr_scale) | |
| self.pose_scale = float(pose_scale) | |
| self.fps = int(fps) | |
| class AgentLogs: | |
| def __init__(self): | |
| self.lines = [] | |
| def log(self, msg): | |
| print(msg) | |
| self.lines.append(msg) | |
| # -------------------- Utility Functions -------------------- | |
| def run_cmd(cmd: list, check=True): | |
| print("▶", " ".join(cmd)) | |
| p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) | |
| print(p.stdout) | |
| if check and p.returncode != 0: | |
| raise RuntimeError(f"Command failed: {' '.join(cmd)}") | |
| return p | |
| def extract_audio_to_wav(video_path: str, wav_out: str, sr: int = 16000): | |
| """Extract mono WAV audio from video.""" | |
| cmd = [ | |
| "ffmpeg", "-y", "-i", video_path, "-ac", "1", "-ar", str(sr), wav_out | |
| ] | |
| run_cmd(cmd) | |
| def cut_audio_segment(in_wav: str, out_wav: str, start_s: float, dur_s: float): | |
| """Cut segment from audio using ffmpeg.""" | |
| cmd = ["ffmpeg", "-y", "-ss", f"{start_s:.3f}", "-t", f"{dur_s:.3f}", | |
| "-i", in_wav, "-acodec", "pcm_s16le", out_wav] | |
| run_cmd(cmd) | |
| def ensure_exact_duration(in_wav: str, out_wav: str, target_sec: float = 20.0): | |
| """Trim or pad with silence to exactly target_sec.""" | |
| audio = AudioSegment.from_file(in_wav) | |
| cur = len(audio) / 1000.0 | |
| if cur > target_sec: | |
| audio = audio[:int(target_sec * 1000)] | |
| elif cur < target_sec: | |
| silence = AudioSegment.silent(duration=int((target_sec - cur) * 1000)) | |
| audio = audio + silence | |
| audio.export(out_wav, format="wav") | |
| def grab_frame_from_video(video_path: str, out_img: str, at_sec: float = 1.0): | |
| """Extract a single frame from a video at time t.""" | |
| cap = cv2.VideoCapture(video_path) | |
| fps = cap.get(cv2.CAP_PROP_FPS) or 25 | |
| frame_no = int(at_sec * fps) | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, frame_no) | |
| ok, frame = cap.read() | |
| if not ok: | |
| cap.set(cv2.CAP_PROP_POS_FRAMES, 0) | |
| ok, frame = cap.read() | |
| cap.release() | |
| if not ok: | |
| raise RuntimeError("Could not extract frame from video.") | |
| cv2.imwrite(out_img, frame) | |
| # -------------------- Speech VAD -------------------- | |
| def find_voice_reference_chunk(wav_path: str, sr: int = 16000, target_dur: float = 6.0) -> Tuple[float, float]: | |
| """Use Silero VAD to find a ~target_dur chunk with speech.""" | |
| try: | |
| wav, file_sr = librosa.load(wav_path, sr=sr, mono=True) | |
| vad_model, utils = torch.hub.load(repo_or_dir='snakers4/silero-vad', | |
| model='silero_vad', force_reload=False) | |
| (get_speech_timestamps, _, read_audio, _, collect_chunks) = utils | |
| speech_timestamps = get_speech_timestamps( | |
| torch.tensor(wav, dtype=torch.float32), vad_model, sampling_rate=sr | |
| ) | |
| if not speech_timestamps: | |
| total = len(wav) / sr | |
| mid_start = max(0.0, (total / 2) - (target_dur / 2)) | |
| return mid_start, min(target_dur, total - mid_start) | |
| best = max(speech_timestamps, key=lambda x: x['end'] - x['start']) | |
| cand_dur = (best['end'] - best['start']) / sr | |
| if cand_dur >= target_dur: | |
| start = best['start'] / sr | |
| return start, target_dur | |
| start = best['start'] / sr | |
| dur = cand_dur | |
| for s in speech_timestamps: | |
| if s is best: | |
| continue | |
| if abs(s['start'] / sr - (start + dur)) < 0.1: | |
| dur += (s['end'] - s['start']) / sr | |
| if dur >= target_dur: | |
| break | |
| return start, min(target_dur, len(wav) / sr - start) | |
| except Exception as e: | |
| print(f"VAD failed: {e}, using fallback") | |
| return 1.0, min(6.0, target_dur) | |
| # -------------------- Script Generation -------------------- | |
| def generate_20s_script(user_text: str = None, language: str = "en") -> str: | |
| """Generate or return user-provided script text.""" | |
| if user_text and user_text.strip(): | |
| return user_text.strip() | |
| # For Hugging Face deployment, we'll use a simple fallback | |
| # You can add API keys as Hugging Face Secrets if needed | |
| fallback_scripts = { | |
| "en": "Hello! I'm your AI avatar. I'm here to demonstrate lifelike speech, natural lip-sync, and subtle head movement. This short clip shows voice cloning from a brief sample, then animates a still image for a realistic talking experience. Thanks for watching!", | |
| "es": "¡Hola! Soy tu avatar de IA. Estoy aquí para demostrar un habla realista, sincronización labial natural y movimientos sutiles de cabeza. Este breve clip muestra clonación de voz a partir de una muestra breve.", | |
| "fr": "Bonjour! Je suis votre avatar IA. Je suis ici pour démontrer une parole réaliste, une synchronisation labiale naturelle et des mouvements subtils de la tête. Ce court clip montre le clonage vocal.", | |
| } | |
| return fallback_scripts.get(language, fallback_scripts["en"]) | |
| # -------------------- TTS -------------------- | |
| def tts_20s_voice_clone(script_text: str, ref_wav: str, out_wav: str, language: str = "en") -> str: | |
| """TTS with voice cloning fallbacks.""" | |
| tmp = str(WORK / f"tts_raw_{uuid.uuid4().hex}.wav") | |
| try: | |
| from TTS.api import TTS | |
| xtts = TTS("tts_models/multilingual/multi-dataset/xtts_v2") | |
| xtts.tts_to_file(text=script_text, speaker_wav=ref_wav, | |
| language=language, file_path=tmp) | |
| print("✓ XTTS v2 voice clone success") | |
| except Exception as e: | |
| print("XTTS v2 clone failed:", e) | |
| try: | |
| from TTS.api import TTS | |
| tts = TTS("tts_models/en/ljspeech/tacotron2-DDC") | |
| tts.tts_to_file(text=script_text, file_path=tmp) | |
| print("✓ XTTS fallback (non-clone) success") | |
| except Exception as e2: | |
| print("XTTS non-clone failed:", e2) | |
| try: | |
| import edge_tts | |
| async def gen_edge(): | |
| communicate = edge_tts.Communicate(script_text, voice="en-US-JennyNeural") | |
| await communicate.save(tmp) | |
| asyncio.run(gen_edge()) | |
| print("✓ edge-tts fallback success") | |
| except Exception as e3: | |
| raise RuntimeError(f"All TTS fallbacks failed: {e3}") | |
| ensure_exact_duration(tmp, out_wav, 20.0) | |
| return out_wav | |
| # -------------------- SadTalker with Fallback -------------------- | |
| def run_sadtalker(source_img: str, driven_wav: str, out_dir: str, | |
| expr_scale: float = 1.0, pose_scale: float = 1.0, fps: int = 25) -> str: | |
| """Call SadTalker inference with fallback.""" | |
| if not SADTALKER_DIR.exists(): | |
| if not setup_sadtalker(): | |
| return create_static_video_fallback(source_img, driven_wav, out_dir, fps) | |
| out_dir = str(Path(out_dir)) | |
| os.makedirs(out_dir, exist_ok=True) | |
| inference_script = SADTALKER_DIR / "inference.py" | |
| if not inference_script.exists(): | |
| print("❌ SadTalker inference script not found, using fallback") | |
| return create_static_video_fallback(source_img, driven_wav, out_dir, fps) | |
| try: | |
| args = [ | |
| sys.executable, str(inference_script), | |
| "--driven_audio", driven_wav, | |
| "--source_image", source_img, | |
| "--preprocess", "full", | |
| "--still", | |
| "--enhancer", "gfpgan", | |
| "--expression_scale", str(expr_scale), | |
| "--pose_scale", str(pose_scale), | |
| "--result_dir", out_dir, | |
| "--fps", str(fps), | |
| ] | |
| # Change to SadTalker directory for execution | |
| original_cwd = os.getcwd() | |
| try: | |
| os.chdir(str(SADTALKER_DIR)) | |
| run_cmd(args) | |
| finally: | |
| os.chdir(original_cwd) | |
| mp4s = sorted(glob.glob(os.path.join(out_dir, "**", "*.mp4"), recursive=True), | |
| key=os.path.getmtime) | |
| if not mp4s: | |
| print("❌ SadTalker produced no output, using fallback") | |
| return create_static_video_fallback(source_img, driven_wav, out_dir, fps) | |
| return mp4s[-1] | |
| except Exception as e: | |
| print(f"❌ SadTalker failed: {e}, using fallback") | |
| return create_static_video_fallback(source_img, driven_wav, out_dir, fps) | |
| def create_static_video_fallback(source_img: str, driven_wav: str, out_dir: str, fps: int = 25) -> str: | |
| """Create a static video with the image and audio as fallback.""" | |
| output_path = os.path.join(out_dir, "fallback_output.mp4") | |
| # Get audio duration | |
| audio = AudioSegment.from_file(driven_wav) | |
| duration = len(audio) / 1000.0 # Convert to seconds | |
| # Create video with static image and audio | |
| cmd = [ | |
| "ffmpeg", "-y", | |
| "-loop", "1", "-i", source_img, | |
| "-i", driven_wav, | |
| "-c:v", "libx264", "-tune", "stillimage", "-c:a", "aac", | |
| "-b:a", "192k", "-pix_fmt", "yuv420p", | |
| "-shortest", "-r", str(fps), | |
| "-t", str(duration), | |
| output_path | |
| ] | |
| try: | |
| run_cmd(cmd) | |
| print(f"✅ Created fallback static video: {output_path}") | |
| return output_path | |
| except Exception as e: | |
| raise RuntimeError(f"Even fallback video creation failed: {e}") | |
| # -------------------- Final Muxing -------------------- | |
| def mux_audio_video(video_path: str, audio_wav: str, final_mp4: str, fps: int = 25): | |
| """Replace video audio with our exact 20s wav.""" | |
| cmd = [ | |
| "ffmpeg", "-y", "-i", video_path, | |
| "-t", "20.0", | |
| "-r", str(fps), | |
| "-i", audio_wav, | |
| "-map", "0:v:0", "-map", "1:a:0", | |
| "-c:v", "libx264", "-pix_fmt", "yuv420p", | |
| "-c:a", "aac", "-shortest", final_mp4 | |
| ] | |
| run_cmd(cmd) | |
| # -------------------- Main Agent Function -------------------- | |
| def run_agent(video_path: str, | |
| maybe_image_path: Optional[str], | |
| user_script_text: Optional[str], | |
| cfg: AgentConfig): | |
| """Main agent orchestrator function.""" | |
| logs = AgentLogs() | |
| try: | |
| # Check SadTalker setup first | |
| logs.log("Checking SadTalker setup...") | |
| if not SADTALKER_DIR.exists(): | |
| logs.log("Setting up SadTalker (first run may take a few minutes)...") | |
| if not setup_sadtalker(): | |
| logs.log("⚠️ SadTalker setup failed, will use static video fallback") | |
| else: | |
| logs.log("✅ SadTalker ready") | |
| video_path = str(video_path) | |
| vid_name = Path(video_path).stem | |
| session = WORK / f"run_{uuid.uuid4().hex[:8]}_{vid_name}" | |
| session.mkdir(parents=True, exist_ok=True) | |
| full_audio = str(session / "audio.wav") | |
| ref_audio = str(session / "ref.wav") | |
| tts_audio = str(session / "tts_20s.wav") | |
| still_img = str(session / "still.jpg") | |
| sadtalker_out = str(session / "sadtalker_out") | |
| final_mp4 = str(OUT / f"{vid_name}_avatar_20s.mp4") | |
| logs.log("Step 1) Extracting audio...") | |
| extract_audio_to_wav(video_path, full_audio, sr=16000) | |
| logs.log("Step 2) Finding speech reference (~6s) via VAD...") | |
| start, dur = find_voice_reference_chunk(full_audio, sr=16000, target_dur=6.0) | |
| logs.log(f" - ref start: {start:.2f}s, dur: {dur:.2f}s") | |
| cut_audio_segment(full_audio, ref_audio, start, dur) | |
| logs.log("Step 3) Generating 20s script text...") | |
| script_text = generate_20s_script(user_script_text, cfg.language) | |
| logs.log(" - script preview: " + (script_text[:140] + ("..." if len(script_text) > 140 else ""))) | |
| logs.log("Step 4) TTS voice cloning to 20s...") | |
| tts_20s_voice_clone(script_text, ref_audio, tts_audio, language=cfg.language) | |
| logs.log(f" - TTS saved: {tts_audio}") | |
| logs.log("Step 5) Prepare still image...") | |
| if maybe_image_path and str(maybe_image_path).strip(): | |
| shutil.copy(maybe_image_path, still_img) | |
| logs.log(" - Using provided still image.") | |
| else: | |
| grab_frame_from_video(video_path, still_img, at_sec=cfg.grab_frame_at) | |
| logs.log(f" - Grabbed frame at {cfg.grab_frame_at}s from video.") | |
| logs.log("Step 6) Run SadTalker animation (or fallback)...") | |
| raw_video = run_sadtalker(still_img, tts_audio, sadtalker_out, | |
| expr_scale=cfg.expr_scale, | |
| pose_scale=cfg.pose_scale, | |
| fps=cfg.fps) | |
| logs.log(f" - Video output: {raw_video}") | |
| logs.log("Step 7) Mux final MP4 (20s, audio + avatar)...") | |
| mux_audio_video(raw_video, tts_audio, final_mp4, fps=cfg.fps) | |
| logs.log(f"✅ DONE: {final_mp4}") | |
| return final_mp4, "\n".join(logs.lines) | |
| except Exception as e: | |
| logs.log(f"❌ ERROR: {e}") | |
| return "", "\n".join(logs.lines) | |
| # -------------------- Gradio Interface -------------------- | |
| def ui_run(video, image, script_text, language, grab_t, expr, pose, fps): | |
| if video is None: | |
| return None, "Please upload a ~30s video." | |
| cfg = AgentConfig(language=language, grab_frame_at=grab_t, | |
| expr_scale=expr, pose_scale=pose, fps=int(fps)) | |
| out, logs = run_agent(video, image, script_text, cfg) | |
| return out if out else None, logs | |
| # Create Gradio interface | |
| with gr.Blocks(title="AI Avatar Agentic Flow", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # 🎬🧠 AI Avatar Agentic Flow — Voice Clone + SadTalker (20s MP4) | |
| Upload a short (~30s) video to clone the voice, then generate a 20-second talking avatar with lip-sync and head movements. | |
| **Features:** | |
| - Automatic voice reference detection using VAD | |
| - Voice cloning with TTS fallbacks (XTTS v2 → edge-tts) | |
| - Animated avatar with SadTalker | |
| - 20-second output with perfect audio sync | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| video = gr.Video( | |
| label="Upload ~30s Video (used for voice reference + optional frame)", | |
| height=300 | |
| ) | |
| image = gr.Image( | |
| type="filepath", | |
| label="Optional Still Image (else a frame is grabbed from the video)", | |
| height=300 | |
| ) | |
| with gr.Column(): | |
| script_text = gr.Textbox( | |
| lines=4, | |
| label="Optional 20s Script Text (leave blank to auto-generate)", | |
| placeholder="Enter your custom script here..." | |
| ) | |
| with gr.Row(): | |
| language = gr.Dropdown( | |
| choices=["en", "es", "fr", "de", "it", "hi", "ur"], | |
| value="en", | |
| label="Language" | |
| ) | |
| grab_t = gr.Slider( | |
| 0.0, 5.0, value=1.0, step=0.1, | |
| label="Grab Frame at (sec) if no image" | |
| ) | |
| with gr.Row(): | |
| expr = gr.Slider( | |
| 0.5, 2.0, value=1.0, step=0.1, | |
| label="Expression Scale" | |
| ) | |
| pose = gr.Slider( | |
| 0.5, 2.0, value=1.0, step=0.1, | |
| label="Pose Scale" | |
| ) | |
| fps = gr.Slider( | |
| 15, 30, value=25, step=1, | |
| label="FPS" | |
| ) | |
| run_btn = gr.Button("🚀 Run Agent", variant="primary", size="lg") | |
| with gr.Row(): | |
| out_video = gr.Video(label="Final 20s MP4", height=400) | |
| logs = gr.Textbox(label="Logs", lines=20, max_lines=30) | |
| run_btn.click( | |
| ui_run, | |
| inputs=[video, image, script_text, language, grab_t, expr, pose, fps], | |
| outputs=[out_video, logs] | |
| ) | |
| gr.Markdown(""" | |
| ## 🧰 Tips & Troubleshooting | |
| - **Processing Time:** First run may take longer due to model downloads | |
| - **Audio Length:** Output is enforced to exactly 20 seconds | |
| - **Voice Reference:** Auto-finds ~6s speech chunk using Silero VAD | |
| - **Language Support:** XTTS v2 supports multiple languages | |
| - **Fallbacks:** Script generation and TTS have multiple fallback options | |
| """) | |
| if __name__ == "__main__": | |
| demo.launch() |