app.py

# Patch for basicsr torchvision compatibility
import sys
import torchvision.transforms.functional as TF
sys.modules['torchvision.transforms.functional_tensor'] = TF

import os
import subprocess
import spaces
import gradio as gr

# Patch load_file_from_url to avoid PermissionError in HF spaces (read-only site-packages)
import os
os.makedirs("gfpgan/weights", exist_ok=True)

try:
    import facexlib.utils
    orig_facex_load = facexlib.utils.load_file_from_url
    def patched_facex_load(*args, **kwargs):
        kwargs["model_dir"] = "gfpgan/weights"
        kwargs["save_dir"] = "gfpgan/weights"
        if len(args) > 1:
            args = list(args)
            args[1] = "gfpgan/weights"
            args = tuple(args)
        try:
            return orig_facex_load(*args, **kwargs)
        except TypeError as e:
            if "save_dir" in str(e) and "save_dir" in kwargs:
                kwargs.pop("save_dir")
                return orig_facex_load(*args, **kwargs)
            raise
    facexlib.utils.load_file_from_url = patched_facex_load
except Exception:
    pass

try:
    import basicsr.utils.download_util
    orig_basicsr_load = basicsr.utils.download_util.load_file_from_url
    def patched_basicsr_load(*args, **kwargs):
        kwargs["model_dir"] = "gfpgan/weights"
        kwargs["save_dir"] = "gfpgan/weights"
        if len(args) > 1:
            args = list(args)
            args[1] = "gfpgan/weights"
            args = tuple(args)
        try:
            return orig_basicsr_load(*args, **kwargs)
        except TypeError as e:
            if "save_dir" in str(e) and "save_dir" in kwargs:
                kwargs.pop("save_dir")
                return orig_basicsr_load(*args, **kwargs)
            raise
    basicsr.utils.download_util.load_file_from_url = patched_basicsr_load
except Exception:
    pass

def download_models():
    if not os.path.exists("checkpoints/mapping_00109-model.pth.tar"):
        print("Downloading SadTalker checkpoints...")
        os.makedirs("checkpoints", exist_ok=True)
        
        models = [
            ("https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/mapping_00109-model.pth.tar", "checkpoints/mapping_00109-model.pth.tar"),
            ("https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/mapping_00229-model.pth.tar", "checkpoints/mapping_00229-model.pth.tar"),
            ("https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/SadTalker_V0.0.2_256.safetensors", "checkpoints/SadTalker_V0.0.2_256.safetensors"),
            ("https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/SadTalker_V0.0.2_512.safetensors", "checkpoints/SadTalker_V0.0.2_512.safetensors")
        ]
        
        gfpgan_models = [
            ("https://github.com/xinntao/facexlib/releases/download/v0.1.0/alignment_WFLW_4HG.pth", "gfpgan/weights/alignment_WFLW_4HG.pth"),
            ("https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth", "gfpgan/weights/detection_Resnet50_Final.pth"),
            ("https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth", "gfpgan/weights/GFPGANv1.4.pth"),
            ("https://github.com/xinntao/facexlib/releases/download/v0.2.2/parsing_parsenet.pth", "gfpgan/weights/parsing_parsenet.pth"),
            ("https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth", "gfpgan/weights/RealESRGAN_x2plus.pth")
        ]
        
        for url, path in models:
            if not os.path.exists(path):
                subprocess.run(["wget", "-nc", url, "-O", path])
                
        os.makedirs("gfpgan/weights", exist_ok=True)
        for url, path in gfpgan_models:
            if not os.path.exists(path):
                subprocess.run(["wget", "-nc", url, "-O", path])
                
        print("Downloads complete.")

download_models()

from src.gradio_demo import SadTalker  

# Initialize with lazy_load=True so weights only move to GPU during the wrapped function call
sad_talker = SadTalker(checkpoint_path='checkpoints', config_path='src/config', lazy_load=True)

import traceback

@spaces.GPU(duration=120)
def generate_video(source_image, driven_audio, preprocess_type, is_still_mode, enhancer, batch_size, size_of_image, pose_style):
    try:
        if not source_image or not driven_audio:
            return None, "Please provide both a source image and an audio file."
        vid_path = sad_talker.test(source_image, driven_audio, preprocess_type, is_still_mode, enhancer, batch_size, size_of_image, pose_style)
        import shutil, tempfile
        out_dir = tempfile.mkdtemp()
        safe_vid_path = os.path.join(out_dir, "output_avatar.mp4")
        shutil.copy(vid_path, safe_vid_path)
        return safe_vid_path, "Success!"
    except Exception as e:
        import traceback
        err = traceback.format_exc()
        print(f"ERROR: {err}")
        return None, f"Error: {str(e)}\n\nTraceback:\n{err}"

# Sleek Web UI
css = """
#header { text-align: center; margin-bottom: 20px; }
#header h1 { font-family: 'Inter', sans-serif; font-weight: 800; color: #E0E7FF; font-size: 2.5rem; }
#header p { font-family: 'Inter', sans-serif; color: #9CA3AF; font-size: 1.1rem; }
.generate-btn { background: linear-gradient(135deg, #6366f1 0%, #8b5cf6 100%); color: white !important; font-weight: bold; border: none !important; transition: transform 0.2s; }
.generate-btn:hover { transform: translateY(-2px); }
"""

with gr.Blocks(css=css, theme=gr.themes.Monochrome()) as demo:
    with gr.Column(elem_id="header"):
        gr.Markdown("# 🗣️ Real-Live Talking Avatar")
        gr.Markdown("Convert any portrait photo into a highly realistic speaking avatar using audio. Powered by **SadTalker** & **ZeroGPU**.")
        
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Inputs")
            source_image = gr.Image(label="Source Image (Face)", type="filepath")
            driven_audio = gr.Audio(label="Driving Audio (Speech)", type="filepath")
            
            with gr.Accordion("Advanced Settings", open=False):
                preprocess_type = gr.Radio(['crop', 'resize', 'full', 'extcrop', 'extfull'], value='crop', label='Preprocess')
                is_still_mode = gr.Checkbox(label="Still Mode (Fewer head movements, best with 'full' preprocess)")
                enhancer = gr.Checkbox(label="GFPGAN Face Enhancement (Sharper face)", value=True)
                size_of_image = gr.Radio([256, 512], value=256, label='Face Model Resolution')
                pose_style = gr.Slider(minimum=0, maximum=46, step=1, label="Pose Style Index", value=0)
                batch_size = gr.Slider(label="Batch Size", step=1, maximum=10, value=2)
                
            submit_btn = gr.Button("🎬 GENERATE AVATAR (ZEROGPU)", elem_classes="generate-btn")
            
        with gr.Column(scale=1):
            gr.Markdown("### 2. Output")
            gen_video = gr.Video(label="Generated Avatar")
            error_log = gr.Textbox(label="Error Logs (if any)", lines=5)
            
    submit_btn.click(
        fn=generate_video,
        inputs=[
            source_image, driven_audio, preprocess_type, 
            is_still_mode, enhancer, batch_size, 
            size_of_image, pose_style
        ],
        outputs=[gen_video, error_log]
    )

if __name__ == "__main__":
    demo.launch(show_error=True)