Update app.py

app.py

@@ -7,17 +7,14 @@ import sys
 import json
 import warnings
 from datetime import datetime
-
 import gradio as gr
 warnings.filterwarnings('ignore')
-
 import random
-
 import torch
 import torch.distributed as dist
 from PIL import Image
 import subprocess
-
+import spaces  # Add ZeroGPU import
 import wan
 from wan.configs import SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS
 from wan.utils.utils import cache_image, cache_video, str2bool
@@ -25,7 +22,6 @@ from wan.utils.multitalk_utils import save_video_ffmpeg
 from kokoro import KPipeline
 from transformers import Wav2Vec2FeatureExtractor
 from src.audio_analysis.wav2vec2 import Wav2Vec2Model
-
 import librosa
 import pyloudnorm as pyln
 import numpy as np
@@ -33,10 +29,583 @@ from einops import rearrange
 import soundfile as sf
 import re
 
+# Global variables for model management
+wav2vec_feature_extractor = None
+audio_encoder = None
+wan_i2v = None
+
 def _validate_args(args):
     # Basic check
     assert args.ckpt_dir is not None, "Please specify the checkpoint directory."
     assert args.task in WAN_CONFIGS, f"Unsupport task: {args.task}"
+    
+    # The default sampling steps are 40 for image-to-video tasks and 50 for text-to-video tasks.
+    if args.sample_steps is None:
+        args.sample_steps = 40
+        
+    if args.sample_shift is None:
+        if args.size == 'infinitetalk-480':
+            args.sample_shift = 7
+        elif args.size == 'infinitetalk-720':
+            args.sample_shift = 11
+        else:
+            raise NotImplementedError(f'Not supported size')
+            
+    args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(0, 99999999)
+    
+    # Size check
+    assert args.size in SUPPORTED_SIZES[args.task], f"Unsupport size {args.size} for task {args.task}, supported sizes are: {', '.join(SUPPORTED_SIZES[args.task])}"
+
+def get_default_args():
+    """Get default arguments for ZeroGPU environment"""
+    class Args:
+        def __init__(self):
+            self.task = "infinitetalk-14B"
+            self.size = "infinitetalk-480"
+            self.frame_num = 81
+            self.ckpt_dir = './weights/Wan2.1-I2V-14B-480P'
+            self.quant_dir = None
+            self.infinitetalk_dir = 'weights/InfiniteTalk/single/infinitetalk.safetensors'
+            self.wav2vec_dir = './weights/chinese-wav2vec2-base'
+            self.dit_path = None
+            self.lora_dir = None
+            self.lora_scale = [1.2]
+            self.offload_model = True  # Always offload for ZeroGPU
+            self.ulysses_size = 1
+            self.ring_size = 1
+            self.t5_fsdp = False
+            self.t5_cpu = True  # Use CPU for T5 to save GPU memory
+            self.dit_fsdp = False
+            self.save_file = None
+            self.audio_save_dir = 'save_audio/gradio'
+            self.base_seed = 42
+            self.input_json = 'examples.json'
+            self.motion_frame = 9
+            self.mode = "streaming"
+            self.sample_steps = 8  # Reduced for faster inference
+            self.sample_shift = 7
+            self.sample_text_guide_scale = 5.0
+            self.sample_audio_guide_scale = 4.0
+            self.num_persistent_param_in_dit = None
+            self.use_teacache = False
+            self.teacache_thresh = 0.2
+            self.use_apg = False
+            self.apg_momentum = -0.75
+            self.apg_norm_threshold = 55
+            self.color_correction_strength = 1.0
+            self.quant = None
+    
+    args = Args()
+    _validate_args(args)
+    return args
+
+@spaces.GPU(duration=120)  # Allocate GPU for 2 minutes
+def initialize_models():
+    """Initialize models on GPU"""
+    global wav2vec_feature_extractor, audio_encoder, wan_i2v
+    
+    args = get_default_args()
+    
+    if wav2vec_feature_extractor is None or audio_encoder is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        # Initialize audio models
+        audio_encoder = Wav2Vec2Model.from_pretrained(args.wav2vec_dir, local_files_only=True).to(device)
+        audio_encoder.feature_extractor._freeze_parameters()
+        wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(args.wav2vec_dir, local_files_only=True)
+        
+        os.makedirs(args.audio_save_dir, exist_ok=True)
+        
+        # Initialize WAN pipeline
+        cfg = WAN_CONFIGS[args.task]
+        wan_i2v = wan.InfiniteTalkPipeline(
+            config=cfg,
+            checkpoint_dir=args.ckpt_dir,
+            quant_dir=args.quant_dir,
+            device_id=0,
+            rank=0,
+            t5_fsdp=args.t5_fsdp,
+            dit_fsdp=args.dit_fsdp,
+            use_usp=False,
+            t5_cpu=args.t5_cpu,
+            lora_dir=args.lora_dir,
+            lora_scales=args.lora_scale,
+            quant=args.quant,
+            dit_path=args.dit_path,
+            infinitetalk_dir=args.infinitetalk_dir
+        )
+        
+        if args.num_persistent_param_in_dit is not None:
+            wan_i2v.vram_management = True
+            wan_i2v.enable_vram_management(
+                num_persistent_param_in_dit=args.num_persistent_param_in_dit
+            )
+    
+    return "Models initialized successfully!"
+
+def loudness_norm(audio_array, sr=16000, lufs=-23):
+    meter = pyln.Meter(sr)
+    loudness = meter.integrated_loudness(audio_array)
+    if abs(loudness) > 100:
+        return audio_array
+    normalized_audio = pyln.normalize.loudness(audio_array, loudness, lufs)
+    return normalized_audio
+
+def audio_prepare_multi(left_path, right_path, audio_type, sample_rate=16000):
+    if not (left_path=='None' or right_path=='None'):
+        human_speech_array1 = audio_prepare_single(left_path)
+        human_speech_array2 = audio_prepare_single(right_path)
+    elif left_path=='None':
+        human_speech_array2 = audio_prepare_single(right_path)
+        human_speech_array1 = np.zeros(human_speech_array2.shape[0])
+    elif right_path=='None':
+        human_speech_array1 = audio_prepare_single(left_path)
+        human_speech_array2 = np.zeros(human_speech_array1.shape[0])
+    
+    if audio_type=='para':
+        new_human_speech1 = human_speech_array1
+        new_human_speech2 = human_speech_array2
+    elif audio_type=='add':
+        new_human_speech1 = np.concatenate([human_speech_array1[: human_speech_array1.shape[0]], np.zeros(human_speech_array2.shape[0])])
+        new_human_speech2 = np.concatenate([np.zeros(human_speech_array1.shape[0]), human_speech_array2[:human_speech_array2.shape[0]]])
+    
+    sum_human_speechs = new_human_speech1 + new_human_speech2
+    return new_human_speech1, new_human_speech2, sum_human_speechs
+
+def get_embedding(speech_array, wav2vec_feature_extractor, audio_encoder, sr=16000, device='cuda'):
+    audio_duration = len(speech_array) / sr
+    video_length = audio_duration * 25  # Assume the video fps is 25
+    
+    # wav2vec_feature_extractor
+    audio_feature = np.squeeze(
+        wav2vec_feature_extractor(speech_array, sampling_rate=sr).input_values
+    )
+    audio_feature = torch.from_numpy(audio_feature).float().to(device=device)
+    audio_feature = audio_feature.unsqueeze(0)
+    
+    # audio encoder
+    with torch.no_grad():
+        embeddings = audio_encoder(audio_feature, seq_len=int(video_length), output_hidden_states=True)
+    
+    if len(embeddings) == 0:
+        print("Fail to extract audio embedding")
+        return None
+    
+    audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0)
+    audio_emb = rearrange(audio_emb, "b s d -> s b d")
+    audio_emb = audio_emb.cpu().detach()
+    return audio_emb
+
+def extract_audio_from_video(filename, sample_rate):
+    raw_audio_path = filename.split('/')[-1].split('.')[0]+'.wav'
+    ffmpeg_command = [
+        "ffmpeg",
+        "-y",
+        "-i",
+        str(filename),
+        "-vn",
+        "-acodec",
+        "pcm_s16le",
+        "-ar",
+        "16000",
+        "-ac",
+        "2",
+        str(raw_audio_path),
+    ]
+    subprocess.run(ffmpeg_command, check=True)
+    human_speech_array, sr = librosa.load(raw_audio_path, sr=sample_rate)
+    human_speech_array = loudness_norm(human_speech_array, sr)
+    os.remove(raw_audio_path)
+    return human_speech_array
+
+def audio_prepare_single(audio_path, sample_rate=16000):
+    ext = os.path.splitext(audio_path)[1].lower()
+    if ext in ['.mp4', '.mov', '.avi', '.mkv']:
+        human_speech_array = extract_audio_from_video(audio_path, sample_rate)
+        return human_speech_array
+    else:
+        human_speech_array, sr = librosa.load(audio_path, sr=sample_rate)
+        human_speech_array = loudness_norm(human_speech_array, sr)
+        return human_speech_array
+
+def process_tts_single(text, save_dir, voice1):
+    s1_sentences = []
+    pipeline = KPipeline(lang_code='a', repo_id='weights/Kokoro-82M')
+    voice_tensor = torch.load(voice1, weights_only=True)
+    generator = pipeline(
+        text, voice=voice_tensor,
+        speed=1, split_pattern=r'\n+'
+    )
+    audios = []
+    for i, (gs, ps, audio) in enumerate(generator):
+        audios.append(audio)
+    audios = torch.concat(audios, dim=0)
+    s1_sentences.append(audios)
+    s1_sentences = torch.concat(s1_sentences, dim=0)
+    save_path1 =f'{save_dir}/s1.wav'
+    sf.write(save_path1, s1_sentences, 24000)
+    s1, _ = librosa.load(save_path1, sr=16000)
+    return s1, save_path1
+
+def process_tts_multi(text, save_dir, voice1, voice2):
+    pattern = r'\(s(\d+)\)\s*(.*?)(?=\s*\(s\d+\)|$)'
+    matches = re.findall(pattern, text, re.DOTALL)
+    s1_sentences = []
+    s2_sentences = []
+    pipeline = KPipeline(lang_code='a', repo_id='weights/Kokoro-82M')
+    
+    for idx, (speaker, content) in enumerate(matches):
+        if speaker == '1':
+            voice_tensor = torch.load(voice1, weights_only=True)
+            generator = pipeline(
+                content, voice=voice_tensor,
+                speed=1, split_pattern=r'\n+'
+            )
+            audios = []
+            for i, (gs, ps, audio) in enumerate(generator):
+                audios.append(audio)
+            audios = torch.concat(audios, dim=0)
+            s1_sentences.append(audios)
+            s2_sentences.append(torch.zeros_like(audios))
+        elif speaker == '2':
+            voice_tensor = torch.load(voice2, weights_only=True)
+            generator = pipeline(
+                content, voice=voice_tensor,
+                speed=1, split_pattern=r'\n+'
+            )
+            audios = []
+            for i, (gs, ps, audio) in enumerate(generator):
+                audios.append(audio)
+            audios = torch.concat(audios, dim=0)
+            s2_sentences.append(audios)
+            s1_sentences.append(torch.zeros_like(audios))
+    
+    s1_sentences = torch.concat(s1_sentences, dim=0)
+    s2_sentences = torch.concat(s2_sentences, dim=0)
+    sum_sentences = s1_sentences + s2_sentences
+    
+    save_path1 =f'{save_dir}/s1.wav'
+    save_path2 =f'{save_dir}/s2.wav'
+    save_path_sum = f'{save_dir}/sum.wav'
+    
+    sf.write(save_path1, s1_sentences, 24000)
+    sf.write(save_path2, s2_sentences, 24000)
+    sf.write(save_path_sum, sum_sentences, 24000)
+    
+    s1, _ = librosa.load(save_path1, sr=16000)
+    s2, _ = librosa.load(save_path2, sr=16000)
+    
+    return s1, s2, save_path_sum
+
+@spaces.GPU(duration=300)  # Allocate GPU for 5 minutes for video generation
+def generate_video(img2vid_image, vid2vid_vid, task_mode, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2,
+                  sd_steps, seed, text_guide_scale, audio_guide_scale, mode_selector, tts_text, resolution_select, human1_voice, human2_voice):
+    
+    global wav2vec_feature_extractor, audio_encoder, wan_i2v
+    args = get_default_args()
+    
+    # Ensure models are initialized
+    if wan_i2v is None:
+        initialize_models()
+    
+    # Move models to GPU if needed
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if audio_encoder is not None:
+        audio_encoder = audio_encoder.to(device)
+    
+    input_data = {}
+    input_data["prompt"] = img2vid_prompt
+    
+    if task_mode=='VideoDubbing':
+        input_data["cond_video"] = vid2vid_vid
+    else:
+        input_data["cond_video"] = img2vid_image
+    
+    person = {}
+    if mode_selector == "Single Person(Local File)":
+        person['person1'] = img2vid_audio_1
+    elif mode_selector == "Single Person(TTS)":
+        tts_audio = {}
+        tts_audio['text'] = tts_text
+        tts_audio['human1_voice'] = human1_voice
+        input_data["tts_audio"] = tts_audio
+    elif mode_selector == "Multi Person(Local File, audio add)":
+        person['person1'] = img2vid_audio_1
+        person['person2'] = img2vid_audio_2
+        input_data["audio_type"] = 'add'
+    elif mode_selector == "Multi Person(Local File, audio parallel)":
+        person['person1'] = img2vid_audio_1
+        person['person2'] = img2vid_audio_2
+        input_data["audio_type"] = 'para'
+    else:
+        tts_audio = {}
+        tts_audio['text'] = tts_text
+        tts_audio['human1_voice'] = human1_voice
+        tts_audio['human2_voice'] = human2_voice
+        input_data["tts_audio"] = tts_audio
+    
+    input_data["cond_audio"] = person
+    
+    # Process audio
+    if 'Local File' in mode_selector:
+        if len(input_data['cond_audio'])==2:
+            new_human_speech1, new_human_speech2, sum_human_speechs = audio_prepare_multi(input_data['cond_audio']['person1'], input_data['cond_audio']['person2'], input_data['audio_type'])
+            audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder, device=device)
+            audio_embedding_2 = get_embedding(new_human_speech2, wav2vec_feature_extractor, audio_encoder, device=device)
+            emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+            emb2_path = os.path.join(args.audio_save_dir, '2.pt')
+            sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
+            sf.write(sum_audio, sum_human_speechs, 16000)
+            torch.save(audio_embedding_1, emb1_path)
+            torch.save(audio_embedding_2, emb2_path)
+            input_data['cond_audio']['person1'] = emb1_path
+            input_data['cond_audio']['person2'] = emb2_path
+            input_data['video_audio'] = sum_audio
+        elif len(input_data['cond_audio'])==1:
+            human_speech = audio_prepare_single(input_data['cond_audio']['person1'])
+            audio_embedding = get_embedding(human_speech, wav2vec_feature_extractor, audio_encoder, device=device)
+            emb_path = os.path.join(args.audio_save_dir, '1.pt')
+            sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
+            sf.write(sum_audio, human_speech, 16000)
+            torch.save(audio_embedding, emb_path)
+            input_data['cond_audio']['person1'] = emb_path
+            input_data['video_audio'] = sum_audio
+    elif 'TTS' in mode_selector:
+        if 'human2_voice' not in input_data['tts_audio'].keys():
+            new_human_speech1, sum_audio = process_tts_single(input_data['tts_audio']['text'], args.audio_save_dir, input_data['tts_audio']['human1_voice'])
+            audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder, device=device)
+            emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+            torch.save(audio_embedding_1, emb1_path)
+            input_data['cond_audio']['person1'] = emb1_path
+            input_data['video_audio'] = sum_audio
+        else:
+            new_human_speech1, new_human_speech2, sum_audio = process_tts_multi(input_data['tts_audio']['text'], args.audio_save_dir, input_data['tts_audio']['human1_voice'], input_data['tts_audio']['human2_voice'])
+            audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder, device=device)
+            audio_embedding_2 = get_embedding(new_human_speech2, wav2vec_feature_extractor, audio_encoder, device=device)
+            emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+            emb2_path = os.path.join(args.audio_save_dir, '2.pt')
+            torch.save(audio_embedding_1, emb1_path)
+            torch.save(audio_embedding_2, emb2_path)
+            input_data['cond_audio']['person1'] = emb1_path
+            input_data['cond_audio']['person2'] = emb2_path
+            input_data['video_audio'] = sum_audio
+    
+    logging.info("Generating video ...")
+    
+    # Generate video
+    video = wan_i2v.generate_infinitetalk(
+        input_data,
+        size_buckget=resolution_select,
+        motion_frame=args.motion_frame,
+        frame_num=args.frame_num,
+        shift=args.sample_shift,
+        sampling_steps=sd_steps,
+        text_guide_scale=text_guide_scale,
+        audio_guide_scale=audio_guide_scale,
+        seed=seed,
+        n_prompt=n_prompt,
+        offload_model=args.offload_model,
+        max_frames_num=args.frame_num if args.mode == 'clip' else 1000,
+        color_correction_strength=args.color_correction_strength,
+        extra_args=args,
+    )
+    
+    # Save video
+    if args.save_file is None:
+        formatted_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+        formatted_prompt = input_data['prompt'].replace(" ", "_").replace("/", "_")[:50]
+        args.save_file = f"{args.task}_{args.size.replace('*','x')}_{formatted_prompt}_{formatted_time}"
+    
+    logging.info(f"Saving generated video to {args.save_file}.mp4")
+    save_video_ffmpeg(video, args.save_file, [input_data['video_audio']], high_quality_save=False)
+    logging.info("Finished.")
+    
+    # Clear GPU memory
+    torch.cuda.empty_cache()
+    
+    return args.save_file + '.mp4'
+
+def toggle_audio_mode(mode):
+    if 'TTS' in mode:
+        return [
+            gr.Audio(visible=False, interactive=False),
+            gr.Audio(visible=False, interactive=False),
+            gr.Textbox(visible=True, interactive=True)
+        ]
+    elif 'Single' in mode:
+        return [
+            gr.Audio(visible=True, interactive=True),
+            gr.Audio(visible=False, interactive=False),
+            gr.Textbox(visible=False, interactive=False)
+        ]
+    else:
+        return [
+            gr.Audio(visible=True, interactive=True),
+            gr.Audio(visible=True, interactive=True),
+            gr.Textbox(visible=False, interactive=False)
+        ]
+
+def show_upload(mode):
+    if mode == "SingleImageDriven":
+        return gr.update(visible=True), gr.update(visible=False)
+    else:
+        return gr.update(visible=False), gr.update(visible=True)
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("""
+    <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
+        MeiGen-InfiniteTalk (ZeroGPU)
+    </div>
+    <div style="text-align: center; font-size: 16px; font-weight: normal; margin-bottom: 20px;">
+        InfiniteTalk: Audio-driven Video Generation for Spare-Frame Video Dubbing.
+    </div>
+    <div style="text-align: center; color: orange; margin-bottom: 20px;">
+        ⚠️ This is optimized for ZeroGPU. Processing may take several minutes due to model loading and GPU allocation.
+    </div>
+    """)
+    
+    # Initialize models button
+    with gr.Row():
+        init_button = gr.Button("Initialize Models (Click first!)", variant="primary")
+        init_status = gr.Textbox(label="Status", interactive=False)
+    
+    init_button.click(
+        fn=initialize_models,
+        outputs=[init_status]
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            task_mode = gr.Radio(
+                choices=["SingleImageDriven", "VideoDubbing"],
+                label="Choose SingleImageDriven task or VideoDubbing task",
+                value="VideoDubbing"
+            )
+            
+            vid2vid_vid = gr.Video(
+                label="Upload Input Video",
+                visible=True)
+            
+            img2vid_image = gr.Image(
+                type="filepath",
+                label="Upload Input Image", 
+                elem_id="image_upload",
+                visible=False
+            )
+            
+            img2vid_prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="Describe the video you want to generate",
+            )
+            
+            task_mode.change(
+                fn=show_upload,
+                inputs=task_mode,
+                outputs=[img2vid_image, vid2vid_vid]
+            )
+            
+            with gr.Accordion("Audio Options", open=True):
+                mode_selector = gr.Radio(
+                    choices=["Single Person(Local File)", "Single Person(TTS)", "Multi Person(Local File, audio add)", "Multi Person(Local File, audio parallel)", "Multi Person(TTS)"],
+                    label="Select person and audio mode.",
+                    value="Single Person(Local File)"
+                )
+                
+                resolution_select = gr.Radio(
+                    choices=["infinitetalk-480", "infinitetalk-720"],
+                    label="Select resolution.",
+                    value="infinitetalk-480"
+                )
+                
+                img2vid_audio_1 = gr.Audio(label="Conditioning Audio for speaker 1", type="filepath", visible=True)
+                img2vid_audio_2 = gr.Audio(label="Conditioning Audio for speaker 2", type="filepath", visible=False)
+                
+                tts_text = gr.Textbox(
+                    label="Text for TTS",
+                    placeholder="Refer to the format in the examples",
+                    visible=False,
+                    interactive=False
+                )
+                
+                mode_selector.change(
+                    fn=toggle_audio_mode,
+                    inputs=mode_selector,
+                    outputs=[img2vid_audio_1, img2vid_audio_2, tts_text]
+                )
+            
+            with gr.Accordion("Advanced Options", open=False):
+                with gr.Row():
+                    sd_steps = gr.Slider(
+                        label="Diffusion steps",
+                        minimum=1,
+                        maximum=50,  # Reduced max for ZeroGPU
+                        value=8,
+                        step=1)
+                    
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=-1,
+                        maximum=2147483647,
+                        step=1,
+                        value=42)
+                
+                with gr.Row():
+                    text_guide_scale = gr.Slider(
+                        label="Text Guide scale",
+                        minimum=0,
+                        maximum=20,
+                        value=1.0,
+                        step=1)
+                    
+                    audio_guide_scale = gr.Slider(
+                        label="Audio Guide scale",
+                        minimum=0,
+                        maximum=20,
+                        value=2.0,
+                        step=1)
+                
+                with gr.Row():
+                    human1_voice = gr.Textbox(
+                        label="Voice for the left person",
+                        value="weights/Kokoro-82M/voices/am_adam.pt",
+                    )
+                    
+                    human2_voice = gr.Textbox(
+                        label="Voice for right person",
+                        value="weights/Kokoro-82M/voices/af_heart.pt"
+                    )
+                
+                n_prompt = gr.Textbox(
+                    label="Negative Prompt",
+                    placeholder="Describe the negative prompt you want to add",
+                    value="bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+                )
+            
+            run_i2v_button = gr.Button("Generate Video", variant="primary")
+        
+        with gr.Column(scale=2):
+            result_gallery = gr.Video(
+                label='Generated Video', interactive=False, height=600,)
+            
+            gr.Examples(
+                examples = [
+                    ['SingleImageDriven', 'examples/single/ref_image.png', None, "A woman is passionately singing into a professional microphone in a recording studio.", "Single Person(Local File)", "examples/single/1.wav", None, None],
+                    ['VideoDubbing', None, 'examples/single/ref_video.mp4', "A man is talking", "Single Person(Local File)", "examples/single/1.wav", None, None],
+                ],
+                inputs = [task_mode, img2vid_image, vid2vid_vid, img2vid_prompt, mode_selector, img2vid_audio_1, img2vid_audio_2, tts_text],
+            )
+    
+    run_i2v_button.click(
+        fn=generate_video,
+        inputs=[img2vid_image, vid2vid_vid, task_mode, img2vid_prompt, n_prompt, img2vid_audio_1, img2vid_audio_2,sd_steps, seed, text_guide_scale, audio_guide_scale, mode_selector, tts_text, resolution_select, human1_voice, human2_voice],
+        outputs=[result_gallery],
+    )
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", debug=True)    # Basic check
+    assert args.ckpt_dir is not None, "Please specify the checkpoint directory."
+    assert args.task in WAN_CONFIGS, f"Unsupport task: {args.task}"
 
     # The default sampling steps are 40 for image-to-video tasks and 50 for text-to-video tasks.
     if args.sample_steps is None: