Update app.py

app.py

@@ -1,134 +1,105 @@
 import os
 import uuid
 import subprocess
-import glob
-
 import gradio as gr
 from PIL import Image
+import torch
+import soundfile as sf
+import numpy as np
 
-# -------------------------
-#  Download/prepare SadTalker
-# -------------------------
-from download_sadtalker_models import ensure_sadtalker
-ensure_sadtalker()
-
-# -------------------------
-#  Bark (latest API)
-# -------------------------
-from bark import SAMPLE_RATE, generate_audio, preload_models
-
-# Load Bark SMALL (change to "large" later for upgrade)
-preload_models(model_type="small")
+from transformers import AutoProcessor, AutoModelForTextToWaveform
 
+# -------------------------------
+# Setup output folder
+# -------------------------------
 RESULTS_DIR = "results"
 os.makedirs(RESULTS_DIR, exist_ok=True)
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# -------------------------------
+# Load Parler-TTS
+# -------------------------------
+model_name = "facebook/parler-tts-mini-en"  # HuggingFace-native, stable
+processor = AutoProcessor.from_pretrained(model_name)
+model = AutoModelForTextToWaveform.from_pretrained(model_name).to(device)
 
-# -----------------------------------------
-# Available Bark Male Voice Presets
-# -----------------------------------------
-BARK_MALE_VOICES = [
-    "male_voice_1",
-    "male_voice_2",
-    "angry_male",
-    "male_broadcast",
-    "male_baritone",
-    "us_male_0",
-    "us_male_1",
-    "male_host",
-    "old_male",
-    "rough_male",
-    "male_voice_young",
-    "announcer"
-]
-
-
-# -------------------------
-# Generate Bark audio
-# -------------------------
+# -------------------------------
+# Parler male voices (all)
+# -------------------------------
+PARLER_MALE_VOICES = processor.speakers["male"]  # all male speakers
+
+
+# -------------------------------
+# TTS function
+# -------------------------------
 def generate_tts(script: str, speaker: str):
-    """Generate a WAV file using Bark."""
-    audio_path = os.path.join(RESULTS_DIR, f"audio_{uuid.uuid4().hex}.wav")
+    inputs = processor(text=script, speaker=speaker, return_tensors="pt").to(device)
 
-    # Bark’s updated API
-    audio_array = generate_audio(
-        text=script,
-        speaker=speaker
-    )
+    with torch.no_grad():
+        audio_values = model.generate(**inputs)
 
-    import soundfile as sf
-    sf.write(audio_path, audio_array, SAMPLE_RATE)
+    audio = audio_values.cpu().numpy().squeeze()
+    sample_rate = model.config.sampling_rate
 
-    return audio_path
+    out_path = os.path.join(RESULTS_DIR, f"audio_{uuid.uuid4().hex}.wav")
+    sf.write(out_path, audio, sample_rate)
+    return out_path
 
 
-# -------------------------
-# Run SadTalker
-# -------------------------
-def run_sadtalker(image: Image.Image, audio_path: str):
-    """Run SadTalker to generate a talking-head video."""
+# -------------------------------
+# Wav2Lip function
+# -------------------------------
+def run_wav2lip(image: Image.Image, audio_path: str):
     img_path = os.path.join(RESULTS_DIR, f"torch_{uuid.uuid4().hex}.png")
-    image.save(img_path)
+    video_out = os.path.join(RESULTS_DIR, f"torch_out_{uuid.uuid4().hex}.mp4")
 
-    sadtalker_results = os.path.join("SadTalker", "results")
-    os.makedirs(sadtalker_results, exist_ok=True)
+    image.save(img_path)
 
     cmd = [
         "python", "inference.py",
-        "--driven_audio", os.path.abspath(audio_path),
-        "--source_image", os.path.abspath(img_path),
-        "--result_dir", os.path.abspath(sadtalker_results),
-        "--preprocess", "full",
-        "--still"
+        "--face", img_path,
+        "--audio", audio_path,
+        "--outfile", video_out
     ]
 
-    subprocess.run(cmd, cwd="SadTalker", check=True)
-
-    mp4_files = glob.glob(os.path.join(sadtalker_results, "**", "*.mp4"), recursive=True)
-    if not mp4_files:
-        raise RuntimeError("SadTalker produced no output video.")
-
-    latest = max(mp4_files, key=os.path.getmtime)
-
-    out_path = os.path.join(RESULTS_DIR, f"torch_out_{uuid.uuid4().hex}.mp4")
-    subprocess.run(["cp", latest, out_path])
-
-    return out_path
+    subprocess.run(cmd, cwd="Wav2Lip", check=True)
+    return video_out
 
 
-# -------------------------
-# Pipeline
-# -------------------------
+# -------------------------------
+# Full pipeline
+# -------------------------------
 def pipeline(script, voice, image):
     if not script.strip():
         raise gr.Error("Script is empty.")
     if image is None:
-        raise gr.Error("No image uploaded.")
+        raise gr.Error("Upload your Torch PNG first.")
 
     audio = generate_tts(script, voice)
-    video = run_sadtalker(image, audio)
+    video = run_wav2lip(image, audio)
     return video
 
 
-# -------------------------
+# -------------------------------
 # Gradio UI
-# -------------------------
+# -------------------------------
 def build_ui():
     with gr.Blocks() as demo:
-        gr.Markdown("## 🔥 Torch Em — Bark TTS + SadTalker Lipsync")
+        gr.Markdown("## 🔥 Torch Em — Parler-TTS + Wav2Lip (Stable)")
 
         with gr.Row():
             with gr.Column():
                 script = gr.Textbox(
                     label="Script",
-                    lines=4,
-                    placeholder="Type any dialogue here…"
+                    lines=3,
+                    placeholder="Enter 2–3 second intro line…"
                 )
 
                 voice = gr.Dropdown(
-                    label="Male Voice Preset",
-                    choices=BARK_MALE_VOICES,
-                    value="male_voice_1"
+                    label="Voice",
+                    choices=PARLER_MALE_VOICES,
+                    value=PARLER_MALE_VOICES[0]
                 )
 
                 image = gr.Image(
@@ -136,15 +107,15 @@ def build_ui():
                     type="pil"
                 )
 
-                generate_btn = gr.Button("Generate Video")
+                btn = gr.Button("Generate Video")
 
             with gr.Column():
-                output_video = gr.Video(label="Output")
+                output = gr.Video(label="Output Video")
 
-        generate_btn.click(
+        btn.click(
             pipeline,
             inputs=[script, voice, image],
-            outputs=output_video
+            outputs=output
         )
 
     return demo
@@ -153,9 +124,9 @@ def build_ui():
 demo = build_ui()
 
 
-# -------------------------
-# API endpoint for n8n
-# -------------------------
+# -------------------------------
+# FastAPI endpoint for n8n
+# -------------------------------
 from fastapi import FastAPI, UploadFile, Form
 from fastapi.responses import FileResponse
 import uvicorn
@@ -168,19 +139,18 @@ async def generate_api(
     voice: str = Form(...),
     image: UploadFile = Form(...)
 ):
-    img_path = f"tmp_{uuid.uuid4().hex}.png"
-    with open(img_path, "wb") as f:
+    tmp_img = f"tmp_{uuid.uuid4().hex}.png"
+    with open(tmp_img, "wb") as f:
         f.write(await image.read())
 
-    pil_img = Image.open(img_path).convert("RGB")
+    pil_img = Image.open(tmp_img).convert("RGB")
 
     audio = generate_tts(script, voice)
-    video = run_sadtalker(pil_img, audio)
+    video = run_wav2lip(pil_img, audio)
 
     return FileResponse(video)
 
 
-# Mount Gradio under FastAPI
 app = gr.mount_gradio_app(api, demo, path="/")