File size: 3,152 Bytes
35698c9
ad8f099
 
35698c9
 
ad8f099
 
 
 
 
35698c9
ad8f099
 
 
35698c9
ad8f099
 
 
35698c9
ad8f099
 
35698c9
ad8f099
 
35698c9
ad8f099
 
 
 
 
 
35698c9
ad8f099
 
 
 
 
 
 
 
35698c9
 
ad8f099
 
35698c9
 
 
 
 
ad8f099
 
35698c9
ad8f099
35698c9
 
 
 
 
ad8f099
35698c9
 
 
 
 
 
 
 
 
 
 
 
ad8f099
35698c9
ad8f099
35698c9
 
 
ad8f099
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
35698c9
 
ad8f099
 
 
 
 
 
 
35698c9
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import os
import sys
import types
import uuid

# --------------------------------------------------------------------
# Fix environment issues
# --------------------------------------------------------------------
# Fix OMP_NUM_THREADS error
os.environ["OMP_NUM_THREADS"] = "1"

# Dummy google.colab for wyn_wav2lip
google = types.ModuleType("google")
colab = types.ModuleType("google.colab")

class _DummyDrive:
    def mount(self, *args, **kwargs):
        print("google.colab.drive.mount() called (dummy).")

colab.drive = _DummyDrive()
google.colab = colab

sys.modules["google"] = google
sys.modules["google.colab"] = colab

# --------------------------------------------------------------------
# Imports
# --------------------------------------------------------------------
import gradio as gr
from wyn_wav2lip.wav2lip import Wav2Lip
import soundfile as sf

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MEDIA_DIR = os.path.join(BASE_DIR, "media")
os.makedirs(MEDIA_DIR, exist_ok=True)

print("Initialising Wav2Lip...")
wav2lip = Wav2Lip()
wav2lip.setup()
print("Wav2Lip ready.")


def run_wav2lip(image_path: str, audio_path: str) -> str:
    """Run Wav2Lip and return absolute path to created video file."""
    existing_mp4 = {
        f for f in os.listdir(MEDIA_DIR)
        if f.lower().endswith(".mp4")
    }

    old_cwd = os.getcwd()
    os.chdir(MEDIA_DIR)
    try:
        # Wav2Lip expects filenames relative to current working dir
        wav2lip.run(
            video_file=os.path.basename(image_path),
            vocal_file=os.path.basename(audio_path),
        )
    finally:
        os.chdir(old_cwd)

    new_mp4 = [
        f for f in os.listdir(MEDIA_DIR)
        if f.lower().endswith(".mp4") and f not in existing_mp4
    ]
    if not new_mp4:
        mp4_candidates = [
            os.path.join(MEDIA_DIR, f)
            for f in os.listdir(MEDIA_DIR)
            if f.lower().endswith(".mp4")
        ]
        if not mp4_candidates:
            raise RuntimeError("No MP4 created by Wav2Lip.")
        return max(mp4_candidates, key=os.path.getmtime)

    return os.path.join(MEDIA_DIR, new_mp4[0])


def lipsync_func(image, audio):
    """
    image: PIL image (from gr.Image)
    audio: (sr, data) tuple (from gr.Audio(type="numpy"))
    """
    if image is None or audio is None:
        return None

    # Save image
    img_id = uuid.uuid4().hex
    image_path = os.path.join(MEDIA_DIR, f"{img_id}.png")
    image.save(image_path)

    # Save audio
    aud_id = uuid.uuid4().hex
    audio_path = os.path.join(MEDIA_DIR, f"{aud_id}.wav")
    sr, data = audio
    sf.write(audio_path, data, sr)

    # Run Wav2Lip
    video_path = run_wav2lip(image_path, audio_path)
    return video_path


demo = gr.Interface(
    fn=lipsync_func,
    inputs=[
        gr.Image(type="pil", label="Teacher image"),
        gr.Audio(type="numpy", label="Teacher audio (.wav)")
    ],
    outputs=gr.Video(label="Lip-synced video"),
    title="Wav2Lip Lipsync Service",
    description="Upload a static teacher image and a WAV audio. The model will generate a talking video."
)

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