import os import subprocess import traceback import sys import types import gradio as gr import numpy as np import soundfile as sf import noisereduce as nr import spaces # === PATCH: fix basicsr/GFPGAN import error on newer torchvision === # torchvision >=0.17 removed `torchvision.transforms.functional_tensor` # (renamed to `functional`), but basicsr's degradations.py (used by GFPGAN, # used by SadTalker's face enhancer) still imports the old path. This shim # creates that missing module so the import chain succeeds. Must run BEFORE # importing SadTalker. import torchvision.transforms.functional as _F _shim = types.ModuleType("torchvision.transforms.functional_tensor") _shim.rgb_to_grayscale = _F.rgb_to_grayscale sys.modules["torchvision.transforms.functional_tensor"] = _shim # === END PATCH === import uuid from src.gradio_demo import SadTalker CHECKPOINT_DIR = "checkpoints" sad_talker = SadTalker(checkpoint_path=CHECKPOINT_DIR, config_path="src/config", lazy_load=True) MIN_VALID_MP4_BYTES = 10_000 def get_duration(path): try: result = subprocess.run( [ "ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", path, ], stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, ) return float(result.stdout.strip()) except Exception: return -1.0 def clean_audio(audio_path, trim_seconds): """ Cleaning pipeline for a single uploaded audio clip: - trims ONLY leading/trailing silence (mid-audio pauses/breaths are kept) - loudness-normalizes - reduces background noise - if trim_seconds > 0, hard-trims the final result to that length (e.g. 10/15/20s). trim_seconds = 0 means "no trim, keep full clip". """ if not audio_path: raise gr.Error("Pehle audio upload karein.") os.makedirs("temp", exist_ok=True) ffmpeg_out = "temp/ffmpeg_stage.wav" trimmed_out = "temp/trimmed_stage.wav" out_path = "temp/cleaned_audio.wav" try: subprocess.check_call( [ "ffmpeg", "-y", "-i", audio_path, "-af", "silenceremove=start_periods=1:start_threshold=-45dB:start_silence=0.1," "areverse," "silenceremove=start_periods=1:start_threshold=-45dB:start_silence=0.1," "areverse," "loudnorm=I=-16:TP=-1.5:LRA=11", "-ac", "1", "-ar", "16000", ffmpeg_out, ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) except subprocess.CalledProcessError as e: raise gr.Error(f"Audio clean karte waqt error aayi: {e}") trim_source = ffmpeg_out if trim_seconds and trim_seconds > 0: try: subprocess.check_call( ["ffmpeg", "-y", "-i", ffmpeg_out, "-t", str(trim_seconds), trimmed_out], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) trim_source = trimmed_out except subprocess.CalledProcessError: trim_source = ffmpeg_out # if trim fails, fall back to un-trimmed cleaned audio try: data, sr = sf.read(trim_source) reduced = nr.reduce_noise(y=data.astype(np.float32), sr=sr) sf.write(out_path, reduced, sr) except Exception: out_path = trim_source final_dur = get_duration(out_path) print(f"[INFO] cleaned audio duration: {final_dur:.2f}s (trim={trim_seconds or 'none'}s)") return out_path, gr.update(interactive=True) def estimate_gpu_duration(avatar_image, cleaned_audio, enhance_face, still_mode, trim_seconds, framing, progress=None): """ Uses the ACTUAL final cleaned audio duration (via ffprobe) to size the GPU time request accurately. """ base_overhead = 40 # model load / warmup, roughly fixed audio_len = get_duration(cleaned_audio) if cleaned_audio else 10 if audio_len <= 0: audio_len = float(trim_seconds) if trim_seconds else 10 per_second_cost = 9 if enhance_face else 4.5 estimated = (base_overhead + audio_len * per_second_cost) * 1.25 return int(min(max(estimated, 50), 280)) def _remux_faststart(src_path): """ Fully re-encode (not just stream-copy) into a clean, browser-safe mp4: H.264 + yuv420p + AAC audio + moov atom at the front (+faststart). A plain `-c copy` remux only moves metadata around — if SadTalker's raw output itself has a broken/variable frame timestamp table (which is the actual root cause of "0:00 / NaN:NaN" even after a copy-remux), copy-remuxing does NOT fix it. A full re-encode rebuilds the frame timing table from scratch, which reliably fixes it. Costs a few extra seconds of CPU but is far more robust. """ fixed_path = os.path.join( os.path.dirname(src_path), f"fixed_{os.path.basename(src_path)}" ) try: subprocess.check_call( [ "ffmpeg", "-y", "-i", src_path, "-c:v", "libx264", "-pix_fmt", "yuv420p", "-preset", "veryfast", "-c:a", "aac", "-b:a", "128k", "-movflags", "+faststart", fixed_path, ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) if os.path.isfile(fixed_path) and os.path.getsize(fixed_path) > 0: return fixed_path except subprocess.CalledProcessError as e: print(f"[WARN] re-encode fix failed: {e}") return src_path # fall back to original if re-encode itself failed def _generate_once(avatar_image, cleaned_audio, enhance_face, still_mode, framing): """Single SadTalker generation attempt. Returns a validated mp4 path or raises.""" preprocess_mode = framing # "full" (keeps body/shoulders in frame) or "crop" (tight face close-up) result_path = sad_talker.test( source_image=avatar_image, driven_audio=cleaned_audio, preprocess=preprocess_mode, still_mode=still_mode, use_enhancer=enhance_face, batch_size=1, size=256, pose_style=0, exp_scale=1.0, result_dir="results", ) if result_path and os.path.isdir(result_path): mp4_candidates = [ os.path.join(result_path, f) for f in os.listdir(result_path) if f.lower().endswith(".mp4") ] result_path = max(mp4_candidates, key=os.path.getmtime) if mp4_candidates else None if not result_path or not os.path.isfile(result_path): raise RuntimeError("Video file generate nahi hui (path invalid).") file_size = os.path.getsize(result_path) print(f"[INFO] raw result video: {result_path} ({file_size} bytes)") if file_size < MIN_VALID_MP4_BYTES: raise RuntimeError("Generated video corrupt/incomplete hai (GPU time khatam ho gaya lagta hai).") # Fix the moov atom placement so the player can read duration correctly. fixed_path = _remux_faststart(result_path) # Validate the ACTUAL playable duration via ffprobe — this is what # catches the "file exists and looks big enough but plays as # NaN:NaN" case that a plain file-size check misses. duration = get_duration(fixed_path) print(f"[INFO] final video: {fixed_path} duration={duration:.2f}s") if duration <= 0.2: raise RuntimeError( "Video file bani to hai lekin uska duration/metadata corrupt hai (NaN:NaN jaisa issue)." ) # SadTalker names its output like "Demo##cleaned_audio.mp4" (image name + # "##" + audio name). The "#" character is a URL fragment separator in # browsers, so when this filename is served to the