wamikabro/TimeCapsuleAI

1import os2import sys3import subprocess4import tempfile5from datetime import datetime6 7# --- Ensure persistent cache dirs on HF Spaces ---8os.makedirs("/data", exist_ok=True)9os.environ.setdefault("TORCH_HOME", "/data/torch")10os.environ.setdefault("BARK_CACHE_DIR", "/data/bark_cache")11 12# --- Clone SadTalker into persistent /data once (won't re-clone on restarts) ---13SADTALKER_DIR = "/data/SadTalker"14if not os.path.exists(SADTALKER_DIR):15    print("📥 Cloning SadTalker into /data ... (only once)")16    subprocess.run(["git", "clone", "https://github.com/OpenTalker/SadTalker.git", SADTALKER_DIR], check=True)17else:18    print("✅ SadTalker already present in /data")19 20# Add SadTalker to path21sys.path.append(SADTALKER_DIR)22 23# --- Imports that depend on installed packages ---24try:25    import torch26    import soundfile as sf27    from bark import generate_audio, preload_models28    # SadTalker import depends on the SadTalker repo layout29    from SadTalker.inference import SadTalker30except Exception as e:31    # Helpful error to surface missing deps during build/logs32    raise RuntimeError(33        "Missing dependency or failed import. Make sure requirements.txt contains the listed libraries "34        "and the Space has finished building. Underlying error: " + str(e)35    )36 37# --- Preload / init (may be heavy, especially on first run) ---38print("🔊 Preloading Bark models (this may download models)...")39preload_models()40 41device = "cuda" if torch.cuda.is_available() else "cpu"42print(f"🎥 Initializing SadTalker on {device} ...")43sad_talker = SadTalker(device=device)44 45# --- Main function: generate audio with Bark, animate with SadTalker, return mp4 path ---46def generate_talking_video(image, text):47    if image is None or not text or text.strip() == "":48        return "⚠️ Please upload an image and enter text.", None49 50    with tempfile.TemporaryDirectory(dir="/data") as tmpdir:51        try:52            # 1) Bark TTS -> WAV53            audio_path = os.path.join(tmpdir, "speech.wav")54            audio_arr = generate_audio(text)  # returns numpy array or similar55            sf.write(audio_path, audio_arr, 22050)56            print("✅ Bark audio saved:", audio_path)57 58            # 2) Create output dir and run SadTalker test/inference59            output_dir = os.path.join(tmpdir, "output")60            os.makedirs(output_dir, exist_ok=True)61 62            # SadTalker expects paths for source_image and driven_audio (image is filepath from Gradio)63            # If Gradio passes a PIL image path, image should be a filepath (we used gr.Image(type="filepath"))64            sad_talker.test(65                source_image=image,66                driven_audio=audio_path,67                result_dir=output_dir,68                enhancer="gfpgan"  # optional, may require extra dependencies69            )70            print("✅ SadTalker finished")71 72            # 3) Find generated .mp473            video_file = None74            for root, _, files in os.walk(output_dir):75                for f in files:76                    if f.endswith(".mp4"):77                        video_file = os.path.join(root, f)78                        break79                if video_file:80                    break81 82            if not video_file:83                return "❌ Failed to generate video — check logs.", None84 85            # Move/copy video to /data so it persists and can be downloaded later86            ts = datetime.now().strftime("%Y%m%d_%H%M%S")87            dst = f"/data/timecapsule_result_{ts}.mp4"88            subprocess.run(["cp", video_file, dst], check=True)89            return "✅ Video generated successfully!", dst90 91        except Exception as e:92            return f"❌ Error during generation: {str(e)}", None93 94 95# --- Gradio UI ---96import gradio as gr97 98title = "🕰️ TimeCapsule AI – Talking Historical Figures"99description = (100    "Upload a historical image and enter text. The system uses Bark (TTS) and SadTalker (animation). "101    "First build may take long; GPU is strongly recommended for real results."102)103 104demo = gr.Interface(105    fn=generate_talking_video,106    inputs=[107        gr.Image(type="filepath", label="Upload Image (photo/painting)"),108        gr.Textbox(label="What should they say?", placeholder="Enter speech text here...")109    ],110    outputs=[111        gr.Textbox(label="Status"),112        gr.Video(label="Preview & Download")113    ],114    title=title,115    description=description,116    allow_flagging="never"117)118 119if __name__ == "__main__":120    demo.launch()121