init

app.py

@@ -2,17 +2,14 @@ import os
 import os.path as osp
 import sys
 import tempfile
+from uuid import uuid4
 
 import gradio as gr
-import librosa
 import soundfile
 import torch
 import torch.nn.functional as F
-import torchaudio
 from huggingface_hub import snapshot_download
-from moviepy import VideoFileClip
-from pydub import AudioSegment
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, AutoTokenizer, pipeline
+from transformers import AutoTokenizer
 
 from src.internvl.eval import load_video
 from src.moviedubber.infer.utils_infer import (
@@ -22,7 +19,7 @@ from src.moviedubber.infer.utils_infer import (
     sway_sampling_coef,
 )
 from src.moviedubber.infer.video_preprocess import VideoFeatureExtractor
-from src.moviedubber.infer_with_mmlm_result import concat_movie_with_audio, get_spk_emb, load_models
+from src.moviedubber.infer_with_mmlm_result import get_spk_emb, get_video_duration, load_models, merge_video_audio
 from src.moviedubber.model.utils import convert_char_to_pinyin
 
 
@@ -32,25 +29,6 @@ sys.path.append("src/third_party/BigVGAN")
 from InternVL.internvl_chat.internvl.model.internvl_chat.modeling_internvl_chat import InternVLChatModel  # type: ignore
 
 
-def load_asr_model(model_id="openai/whisper-large-v3-turbo"):
-    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-    model = AutoModelForSpeechSeq2Seq.from_pretrained(
-        model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
-    ).to(device)
-
-    processor = AutoProcessor.from_pretrained(model_id)
-
-    pipe = pipeline(
-        "automatic-speech-recognition",
-        model=model,
-        tokenizer=processor.tokenizer,
-        feature_extractor=processor.feature_extractor,
-        torch_dtype=torch_dtype,
-        device=device,
-    )
-    return pipe
-
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 repo_local_path = snapshot_download(repo_id="woak-oa/DeepDubber-V1")
@@ -69,7 +47,6 @@ generation_config = dict(max_new_tokens=1024, do_sample=False)
 
 
 ema_model, vocoder, ort_session = load_models(repo_local_path, device=device)
-asr_pipe = load_asr_model()
 
 videofeature_extractor = VideoFeatureExtractor(device=device)
 
@@ -106,60 +83,26 @@ def deepdubber(video_path: str, subtitle_text: str, audio_path: str = None) -> s
     gen_clip = videofeature_extractor.extract_features(video_path)
     gen_text = subtitle_text
 
-    clip = VideoFileClip(video_path)
-    gen_audio_len = int(clip.duration * 24000 // 256)
+    v_dur = get_video_duration(video_path)
+    gen_audio_len = int(v_dur * 24000 // 256)
 
     gen_clip = gen_clip.unsqueeze(0).to(device=device, dtype=torch.float32).transpose(1, 2)
     gen_clip = F.interpolate(gen_clip, size=(gen_audio_len,), mode="linear", align_corners=False).transpose(1, 2)
 
-    ref_audio_len = None
     if audio_path is not None:
-        print("reference audio is not None, dubbing with reference audio")
-
-        if audio_path.endswith(".mp3"):
-            audio = AudioSegment.from_mp3(audio_path)
-
-            wav_file = audio_path.replace(".mp3", ".wav")
-            audio.export(wav_file, format="wav")
-        else:
-            wav_file = audio_path
-
-        ref_text = asr_pipe(librosa.load(wav_file, sr=16000)[0], generate_kwargs={"language": "english"})["text"]
-        ref_text = ref_text.replace("\n", " ").replace("\r", " ")
-        print(f"Reference text: {ref_text}")
-
-        spk_emb = get_spk_emb(wav_file, ort_session)
+        spk_emb = get_spk_emb(audio_path, ort_session)
         spk_emb = torch.tensor(spk_emb).to(device=device, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
-
-        audio_data, sr = torchaudio.load(wav_file)
-        resampler = torchaudio.transforms.Resample(sr, 24000)
-        if sr != 24000:
-            audio_data = resampler(audio_data)
-
-        if audio_data.shape[0] > 1:
-            audio_data = torch.mean(audio_data, dim=0, keepdim=True)
-
-        audio_data = audio_data.to(device)
-
-        ref_audio_len = int(audio_data.shape[-1] // 256)
-        ref_clip = torch.zeros((1, ref_audio_len, 768)).to(device=device)
-
-        gen_clip = torch.cat((gen_clip, ref_clip), dim=1)
-
-        gen_audio_len = ref_audio_len + gen_audio_len
-
-        gen_text = ref_text + " " + gen_text
-
     else:
-        spk_emb = torch.zeros((1, 1, 192)).to(device=device)
-        audio_data = torch.zeros((1, gen_audio_len, 100)).to(device=device)
+        spk_emb = torch.zeros(1, 1, 256).to(device=device, dtype=torch.float32)
 
     gen_text_batches = chunk_text(gen_text, max_chars=1024)
     final_text_list = convert_char_to_pinyin(gen_text_batches)
 
+    cond = torch.zeros(1, gen_audio_len, 100).to(device)
+
     with torch.inference_mode():
         generated, _ = ema_model.sample(
-            cond=audio_data,
+            cond=cond,
             text=final_text_list,
             clip=gen_clip,
             spk_emb=spk_emb,
@@ -167,14 +110,11 @@ def deepdubber(video_path: str, subtitle_text: str, audio_path: str = None) -> s
             steps=nfe_step,
             cfg_strength=cfg_strength,
             sway_sampling_coef=sway_sampling_coef,
-            no_ref_audio=False,
+            no_ref_audio=True,
         )
 
         generated = generated.to(torch.float32)
 
-        if ref_audio_len is not None:
-            generated = generated[:, ref_audio_len:, :]
-
         generated_mel_spec = generated.permute(0, 2, 1)
         generated_wave = vocoder(generated_mel_spec)
 
@@ -185,7 +125,10 @@ def deepdubber(video_path: str, subtitle_text: str, audio_path: str = None) -> s
         temp_wav_path = temp_wav_file.name
         soundfile.write(temp_wav_path, generated_wave, samplerate=24000)
 
-    concated_video = concat_movie_with_audio(temp_wav_path, video_path, ".")
+    video_out_path = os.path.join(out_dir, f"dubbed_video_{uuid4[:6]}.mp4")
+    concated_video = merge_video_audio(
+        video_path, temp_wav_path, video_out_path, 0, soundfile.info(temp_wav_path).duration
+    )
 
     # Ensure the temporary file is deleted after use
     os.remove(temp_wav_path)
@@ -194,7 +137,9 @@ def deepdubber(video_path: str, subtitle_text: str, audio_path: str = None) -> s
     return response, concated_video
 
 
-def process_video_dubbing(video_path: str, subtitle_text: str, audio_path: str = None) -> str:
+def process_video_dubbing(
+    video_path: str, subtitle_text: str, audio_path: str = None, caption_input: str = None
+) -> str:
     try:
         if not os.path.exists(video_path):
             raise ValueError("Video file does not exist")
@@ -227,6 +172,7 @@ def create_ui():
                 label="Enter the subtitle", placeholder="Enter the subtitle to be dubbed...", lines=5
             )
             audio_input = gr.Audio(label="Upload speech prompt (Optional)", type="filepath")
+            # caption_input = gr.Textbox(label="Enter the description of Video (Optional)", lines=1)
 
         process_btn = gr.Button("Start Dubbing")
 
@@ -240,16 +186,19 @@ def create_ui():
                 "datasets/CoTMovieDubbing/demo/v01input.mp4",
                 "it isn't simply a question of creating a robot who can love",
                 "datasets/CoTMovieDubbing/demo/speech_prompt_01.mp3",
+                # "datasets/CoTMovieDubbing/demo/speech_prompt_01.mp3",
             ],
             [
                 "datasets/CoTMovieDubbing/demo/v02input.mp4",
                 "Me, I'd be happy with one who's not... fixed.",
                 "datasets/CoTMovieDubbing/demo/speech_prompt_02.mp3",
+                # "datasets/CoTMovieDubbing/demo/speech_prompt_02.mp3",
             ],
             [
                 "datasets/CoTMovieDubbing/demo/v03input.mp4",
                 "Man, Papi. What am I gonna do?",
                 "datasets/CoTMovieDubbing/demo/speech_prompt_03.mp3",
+                # "datasets/CoTMovieDubbing/demo/speech_prompt_02.mp3",
             ],
         ]
 
@@ -259,6 +208,7 @@ def create_ui():
             outputs=[output_response, output_video],
         )
 
+        # gr.Examples(examples=examples, inputs=[video_input, subtitle_input, audio_input, caption_input])
         gr.Examples(examples=examples, inputs=[video_input, subtitle_input, audio_input])
 
     return app

src/moviedubber/infer/utils_infer.py

@@ -89,48 +89,11 @@ def load_vocoder(local_path, device=device):
 # load model checkpoint for inference
 
 
-def load_checkpoint(model, ckpt_path, device: str, dtype=None, use_ema=True):
-    if dtype is None:
-        dtype = (
-            torch.float16
-            if "cuda" in device
-            and torch.cuda.get_device_properties(device).major >= 6
-            and not torch.cuda.get_device_name().endswith("[ZLUDA]")
-            else torch.float32
-        )
+def load_checkpoint(model, ckpt_path, device: str, dtype=torch.float32):
     model = model.to(dtype)
 
-    ckpt_type = ckpt_path.split(".")[-1]
-    if ckpt_type == "safetensors":
-        from safetensors.torch import load_file
-
-        checkpoint = load_file(ckpt_path, device=device)
-    else:
-        checkpoint = torch.load(ckpt_path, map_location=device, weights_only=True)
-
-    if use_ema:
-        if ckpt_type == "safetensors":
-            checkpoint = {"ema_model_state_dict": checkpoint}
-        checkpoint["model_state_dict"] = {
-            k.replace("ema_model.", ""): v
-            for k, v in checkpoint["ema_model_state_dict"].items()
-            if k not in ["initted", "step"]
-        }
-
-        # patch for backward compatibility, 305e3ea
-        for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
-            if key in checkpoint["model_state_dict"]:
-                del checkpoint["model_state_dict"][key]
-
-        state_dict_result = model.load_state_dict(checkpoint["model_state_dict"], strict=False)
-        if state_dict_result.unexpected_keys:
-            print("\nUnexpected keys in state_dict:", state_dict_result.unexpected_keys)
-        if state_dict_result.missing_keys:
-            print("\nMissing keys in state_dict:", state_dict_result.missing_keys)
-    else:
-        if ckpt_type == "safetensors":
-            checkpoint = {"model_state_dict": checkpoint}
-        model.load_state_dict(checkpoint["model_state_dict"], strict=True)
+    checkpoint = torch.load(ckpt_path, map_location=device, weights_only=True)
+    model.load_state_dict(checkpoint, strict=True)
 
     del checkpoint
     torch.cuda.empty_cache()
@@ -149,7 +112,6 @@ def load_model(
     mel_spec_type=mel_spec_type,
     vocab_file="",
     ode_method=ode_method,
-    use_ema=True,
     device=device,
 ):
     tokenizer = "custom"
@@ -181,7 +143,7 @@ def load_model(
     ).to(device)
 
     dtype = torch.float32 if mel_spec_type == "bigvgan" else None
-    model = load_checkpoint(model, ckpt_path, device, dtype=dtype, use_ema=use_ema)
+    model = load_checkpoint(model, ckpt_path, device, dtype=dtype)
 
     return model
 

src/moviedubber/infer_with_mmlm_result.py

@@ -1,9 +1,10 @@
 import os
+import subprocess as sp
 
+import cv2
 import onnxruntime
 import torchaudio
 import torchaudio.compliance.kaldi as kaldi
-from moviepy import AudioFileClip, VideoFileClip
 from omegaconf import OmegaConf
 
 from src.moviedubber.infer.utils_infer import (
@@ -13,31 +14,50 @@ from src.moviedubber.infer.utils_infer import (
 from src.moviedubber.model import ControlNetDiT, DiT
 
 
-def concat_movie_with_audio(wav, video_path, out_dir):
-    with (
-        AudioFileClip(str(wav)) as audio_clip,
-        VideoFileClip(str(video_path)) as video_clip,
-    ):
-        duration = min(video_clip.duration, audio_clip.duration)
-
-        video_subclip = video_clip.subclipped(0, duration)
-        audio_subclip = audio_clip.subclipped(0, duration)
-
-        final_video = video_subclip.with_audio(audio_subclip)
-
-        output_path = wav.replace(".wav", ".mp4")
-
-        final_video.write_videofile(
-            str(output_path),
-            codec="libx264",
-            audio_codec="mp3",
-            fps=25,
-            logger=None,
-            threads=1,
-            temp_audiofile_path=out_dir,
-        )
-
-    return output_path
+def get_video_duration(video_path):
+    cap = cv2.VideoCapture(video_path)
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    duration = total_frames / fps
+    return duration
+
+
+def merge_video_audio(video_path, audio_path, output_path, start_time, duration):
+    command = [
+        "ffmpeg",
+        "-y",
+        "-ss",
+        str(start_time),
+        "-t",
+        str(duration),
+        "-i",
+        video_path,
+        "-i",
+        audio_path,
+        "-c:v",
+        "copy",
+        "-c:a",
+        "aac",
+        "-map",
+        "0:v:0",
+        "-map",
+        "1:a:0",
+        "-shortest",
+        "-strict",
+        "experimental",
+        output_path,
+    ]
+
+    try:
+        sp.run(command, check=True, stdout=sp.DEVNULL, stderr=sp.DEVNULL, stdin=sp.DEVNULL)
+        print(f"Successfully merged audio and video into {output_path}")
+        return output_path
+    except sp.CalledProcessError as e:
+        print(f"Error merging audio and video: {e}")
+        return None
 
 
 def get_spk_emb(audio_path, ort_session):