app.py

"""Wav2Lip Free API v3 - using official Rudrabha model code.

Simplest possible wrapper around the proven Wav2Lip inference pipeline.
"""
import os, shutil, subprocess, sys, warnings
from pathlib import Path

warnings.filterwarnings("ignore")
sys.path.insert(0, str(Path(__file__).parent))

WORK = Path("/data/wav2lip_app")
HF_TOKEN = os.environ.get("HF_TOKEN") or None


def ensure_setup():
    from huggingface_hub import hf_hub_download
    code_dir = WORK / "Wav2Lip"
    code_dir.mkdir(parents=True, exist_ok=True)
    target = code_dir / "wav2lip.pth"
    if target.exists() and target.stat().st_size > 100_000_000:
        print(f"[setup] cached: {target} ({target.stat().st_size//1024//1024}MB)")
    else:
        print("[setup] downloading wav2lip.pth from Nekochu/Wav2Lip (~436MB, first run only)...")
        d = hf_hub_download(repo_id="Nekochu/Wav2Lip", filename="wav2lip.pth",
                           local_dir=str(code_dir), token=HF_TOKEN)
        sp = Path(d)
        if str(sp) != str(target):
            shutil.copy2(sp, target)
        print(f"  [setup] downloaded: {target}")
    return code_dir


def detect_face_mediapipe(image):
    import mediapipe as mp
    import numpy as np
    mp_face = mp.solutions.face_detection
    fd = mp_face.FaceDetection(min_detection_confidence=0.5)
    res = fd.process(image)
    if not res.detections:
        return None
    d = res.detections[0].location_data.relative_bounding_box
    h, w = image.shape[:2]
    x1 = max(0, int(d.xmin * w))
    y1 = max(0, int(d.ymin * h))
    x2 = min(w, int((d.xmin + d.width) * w))
    y2 = min(h, int((d.ymin + d.height) * h))
    return (x1, y1, x2, y2)


def load_audio_mel(audio_path):
    """Generate mel spectrogram chunks matching Wav2Lip requirements."""
    import librosa
    import numpy as np
    wav, _ = librosa.load(audio_path, sr=16000)
    wav = np.concatenate([np.zeros(6400), wav])
    mel = librosa.feature.melspectrogram(y=wav, sr=16000, n_fft=800,
                                          win_length=800, hop_length=200, n_mels=80)
    mel = 20 * np.log10(np.maximum(1e-5, mel))
    mel = np.maximum(mel, mel.max() - 8)
    mel = (mel + 4) / 4
    mel_chunks = []
    i = 0
    while i + 16 <= mel.shape[1]:
        mel_chunks.append(mel[:, i:i+16])
        i += 5
    if not mel_chunks:
        mel_chunks = [np.pad(mel, ((0,0),(0,16-mel.shape[1])), mode='constant')]
    return np.array(mel_chunks)


def run_inference(code_dir, face_path, audio_path, output_path):
    import cv2, numpy as np, torch, mediapipe as mp

    sys.path.insert(0, str(code_dir))
    # 修复：模型在 app.py 同目录的 models/ 下，不在 code_dir
    app_dir = Path(__file__).parent
    if str(app_dir) not in sys.path:
        sys.path.insert(0, str(app_dir))
    from models import Wav2Lip
    from models import Wav2Lip as Wav2Lip_class

    # Load model (cached)
    global _MODEL
    if _MODEL is None:
        print("[infer] loading model (first call)...")
        _MODEL = Wav2Lip_class()
        ckpt = torch.load(code_dir / "wav2lip.pth", map_location="cpu", weights_only=False)
        sd = ckpt.get("state_dict", ckpt)
        # Remove DataParallel 'module.' prefix if present
        sd = {k.replace("module.", ""): v for k, v in sd.items()}
        _MODEL.load_state_dict(sd, strict=False)
        _MODEL.eval()
        print("[infer] model loaded")

    # Read face
    if face_path.lower().endswith((".png", ".jpg", ".jpeg")):
        img = cv2.imread(face_path)
        mel_chunks = load_audio_mel(audio_path)
        n_frames = max(125, len(mel_chunks) + 5)
        frames = [img] * n_frames
        fps = 25
    else:
        cap = cv2.VideoCapture(face_path)
        fps = cap.get(cv2.CAP_PROP_FPS) or 25
        frames = []
        while True:
            ret, fr = cap.read()
            if not ret: break
            frames.append(fr)
        cap.release()
        mel_chunks = load_audio_mel(audio_path)
        # Limit frames to mel chunks + 5
        max_frames = len(mel_chunks) + 5
        if len(frames) > max_frames:
            frames = frames[:max_frames]
        while len(frames) < len(mel_chunks) + 5:
            frames.append(frames[-1])

    # Detect face in first frame
    box = detect_face_mediapipe(frames[0])
    if box is None:
        raise RuntimeError("No face detected in input")
    x1, y1, x2, y2 = box
    y2 = min(frames[0].shape[0], y2 + int((y2 - y1) * 0.2))

    IMG_SIZE = 96
    out_frames = []
    BATCH = 4

    for i in range(0, len(mel_chunks), BATCH):
        batch_end = min(i + BATCH, len(mel_chunks))
        actual_batch = batch_end - i
        if actual_batch <= 0:
            break

        # Get mel batch (always BATCH size for tensor consistency)
        if actual_batch < BATCH:
            mel_batch = np.zeros((BATCH, 80, 16), dtype=np.float32)
            mel_batch[:actual_batch] = mel_chunks[i:batch_end]
        else:
            mel_batch = mel_chunks[i:batch_end]

        # Get face batch (RGB float, 0-1)
        face_batch = []
        for j in range(BATCH):
            fi = i + j
            if fi >= len(frames):
                face = frames[-1]
            else:
                face = frames[fi]
            face_crop = face[y1:y2, x1:x2]
            if face_crop.size == 0:
                face_crop = frames[0][y1:y2, x1:x2]
            face_resized = cv2.resize(face_crop, (IMG_SIZE, IMG_SIZE))
            face_resized = face_resized.astype(np.float32) / 255.0
            face_resized = face_resized[..., ::-1].copy()  # BGR->RGB
            face_batch.append(face_resized)

        face_batch = np.array(face_batch)  # (BATCH, 96, 96, 3)
        # Wav2Lip expects 6 channels: [face with bottom-half masked, face]
        img_masked = face_batch.copy()
        img_masked[:, IMG_SIZE//2:] = 0  # zero bottom half (where mouth is)
        face_batch_6ch = np.concatenate((img_masked, face_batch), axis=3)  # (BATCH, 96, 96, 6)
        face_t = torch.from_numpy(face_batch_6ch).float().permute(0, 3, 1, 2)
        mel_t = torch.from_numpy(mel_batch).float().unsqueeze(1)

        with torch.no_grad():
            pred = _MODEL(mel_t, face_t)

        pred = pred.cpu().numpy().transpose(0, 2, 3, 1)
        pred = np.clip(pred * 255, 0, 255).astype(np.uint8)
        pred = pred[..., ::-1].copy()  # RGB->BGR

        for j in range(actual_batch):
            fi = i + j
            out = frames[fi].copy()
            ph, pw = pred[j].shape[:2]
            try:
                if (y2 - y1) != ph or (x2 - x1) != pw:
                    p_resized = cv2.resize(pred[j], (x2 - x1, y2 - y1))
                else:
                    p_resized = pred[j]
                out[y1:y2, x1:x2] = p_resized
                out_frames.append(out)
            except Exception:
                out_frames.append(frames[fi].copy())

    print(f"[infer] generated {len(out_frames)} frames")

    # Write video
    h, w = out_frames[0].shape[:2]
    tmp_out = "/tmp/wav2lip_out.mp4"
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    vw = cv2.VideoWriter(tmp_out, fourcc, 25, (w, h))
    for f in out_frames:
        vw.write(f)
    vw.release()

    # Mux audio
    subprocess.run([
        "ffmpeg", "-y", "-i", tmp_out, "-i", audio_path,
        "-c:v", "libx264", "-c:a", "aac", "-shortest", output_path
    ], capture_output=True)
    if os.path.exists(tmp_out):
        os.remove(tmp_out)
    print(f"[infer] done: {output_path}")


_MODEL = None

# Initialize at startup
print("[init] downloading weights...")
CODE_DIR = ensure_setup()
print(f"[init] ready")


from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import FileResponse
import uvicorn

app = FastAPI(title="Wav2Lip Free API v3", version="3.0")


@app.get("/")
def root():
    return {
        "name": "Wav2Lip Free API v3",
        "license": "Non-commercial only (LRS2 dataset)",
        "endpoints": {"GET /health": "health check", "POST /lipsync": "multipart face+audio -> MP4"},
        "speed": "~30-90s per 1s of input video (CPU free tier)",
    }


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/lipsync")
async def lipsync(face: UploadFile = File(...), audio: UploadFile = File(...)):
    face_path = f"/tmp/in_face_{os.getpid()}.png"
    audio_path = f"/tmp/in_audio_{os.getpid()}.mp3"
    output_path = "/tmp/wav2lip_result.mp4"
    for p in [face_path, audio_path, output_path]:
        if os.path.exists(p):
            os.remove(p)
    with open(face_path, "wb") as f:
        f.write(await face.read())
    with open(audio_path, "wb") as f:
        f.write(await audio.read())
    try:
        run_inference(CODE_DIR, face_path, audio_path, output_path)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        for p in [face_path, audio_path]:
            if os.path.exists(p):
                os.remove(p)
    if not os.path.exists(output_path):
        raise HTTPException(status_code=500, detail="No output produced")
    return FileResponse(output_path, media_type="video/mp4", filename="lipsync.mp4")



if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)