app.py

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()