Update app.py

app.py

@@ -20,13 +20,17 @@ import datetime
 import random
 import math
 import subprocess
+from moviepy.editor import VideoFileClip
 from huggingface_hub import snapshot_download
-import requests
 import shutil
 
-# --- 全域設定 ---
+try:
+    from audio_separator.separator import Separator
+except:
+    print("Unable to use vocal separation feature. Please install audio-separator[gpu].")
+
 if torch.cuda.is_available():
-    device = "cuda"
+    device = "cuda" 
     if torch.cuda.get_device_capability()[0] >= 8:
         dtype = torch.bfloat16
     else:
@@ -36,236 +40,304 @@ else:
     dtype = torch.float32
 
 def filter_kwargs(cls, kwargs):
-    """過濾掉不屬於類別建構函式的關鍵字參數"""
     import inspect
     sig = inspect.signature(cls.__init__)
     valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
     filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
     return filtered_kwargs
 
-# <<< 核心修正：加入 Git LFS 指標檔的偵測與處理 >>>
-def is_lfs_pointer(file_path):
-    """檢查一個檔案是否為 Git LFS 指標檔。"""
-    try:
-        # 指標檔通常很小 (< 2KB)
-        if os.path.getsize(file_path) > 2048:
-            return False
-        with open(file_path, 'r') as f:
-            first_line = f.readline().strip()
-            # 指標檔的第一行通常是 'version https://git-lfs.github.com/spec/v1'
-            if 'git-lfs' in first_line:
-                return True
-    except (OSError, UnicodeDecodeError):
-        # 如果檔案無法讀取或不是文字檔，那它就不是指標檔
-        return False
-    return False
-
-def resolve_path(user_path, repo_root):
+def load_transformer_model(model_version):
     """
-    以正確的優先級解析檔案路徑，並處理 Git LFS 指標檔問題。
-    """
-    # 檢查本地路徑是否存在
-    if os.path.exists(user_path):
-        # 檢查它是否為一個無效的 LFS 指標檔
-        if is_lfs_pointer(user_path):
-            print(f"警告：本地檔案 '{user_path}' 是一個 Git LFS 指標檔。將嘗試從 Hugging Face 快取中尋找完整檔案。")
-            # 如果是指標檔，則忽略它，並在下一步從 HF 快取中尋找
-        else:
-            # 如果是個正常檔案，直接使用
-            print(f"找到本地檔案： {os.path.abspath(user_path)}")
-            return os.path.abspath(user_path)
+    根据选择的模型版本加载对应的transformer模型
     
-    # 如果本地檔案不存在或是 LFS 指標檔，則從 HF 快取目錄中尋找
-    potential_repo_path = os.path.join(repo_root, user_path)
-    if os.path.exists(potential_repo_path):
-        print(f"在 Hugging Face 快取目錄中找到檔案： {potential_repo_path}")
-        return potential_repo_path
-        
-    return None
-# <<< 修正結束 >>>
-
-def setup_models(repo_root, model_version):
-    """載入所有必要的模型和設定"""
-    pretrained_model_name_or_path = os.path.join(repo_root, "Wan2.1-Fun-V1.1-1.3B-InP")
-    pretrained_wav2vec_path = os.path.join(repo_root, "wav2vec2-base-960h")
-
-    config_relative_path = "deepspeed_config/wan2.1/wan_civitai.yaml"
-    config_path = resolve_path(config_relative_path, repo_root)
-    if not config_path:
-        raise FileNotFoundError(f"設定檔 '{config_relative_path}' 在當前目錄或 HF 快取中都找不到。")
-
-    print(f"正在從 {config_path} 載入設定...")
-    config = OmegaConf.load(config_path)
-    sampler_name = "Flow"
-    
-    print("正在載入 Tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(os.path.join(pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')))
-    
-    print("正在載入 Text Encoder...")
-    text_encoder = WanT5EncoderModel.from_pretrained(
-        os.path.join(pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
-        additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
-        low_cpu_mem_usage=True,
-        torch_dtype=dtype,
-    ).eval()
-    
-    print("正在載入 VAE...")
-    vae = AutoencoderKLWan.from_pretrained(
-        os.path.join(pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')),
-        additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
-    )
-    
-    print("正在載入 Wav2Vec...")
-    wav2vec_processor = Wav2Vec2Processor.from_pretrained(pretrained_wav2vec_path)
-    wav2vec = Wav2Vec2Model.from_pretrained(pretrained_wav2vec_path).to("cpu")
+    Args:
+        model_version (str): 模型版本，"square" 或 "rec_vec"
     
-    print("正在載入 CLIP Image Encoder...")
-    clip_image_encoder = CLIPModel.from_pretrained(os.path.join(pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder'))).eval()
+    Returns:
+        WanTransformer3DFantasyModel: 加载的transformer模型
+    """
+    global transformer3d
     
-    print("正在載入 Transformer 3D 基礎模型...")
-    transformer3d = WanTransformer3DFantasyModel.from_pretrained(
-        os.path.join(pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
-        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
-        low_cpu_mem_usage=False,
-        torch_dtype=dtype,
-    )
-
     if model_version == "square":
         transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", "transformer3d-square.pt")
-    else: # rec_vec
+    elif model_version == "rec_vec":
         transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", "transformer3d-rec-vec.pt")
-
+    else:
+        # 默认使用square版本
+        transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", "transformer3d-square.pt")
+    
+    print(f"正在加载模型: {transformer_path}")
+    
     if os.path.exists(transformer_path):
-        print(f"正在從 {transformer_path} 載入 StableAvatar 權重...")
         state_dict = torch.load(transformer_path, map_location="cpu")
         state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
         m, u = transformer3d.load_state_dict(state_dict, strict=False)
-        print(f"StableAvatar 權重載入成功。 Missing keys: {len(m)}; Unexpected keys: {len(u)}")
+        print(f"模型加载成功: {transformer_path}")
+        print(f"Missing keys: {len(m)}; Unexpected keys: {len(u)}")
+        return transformer3d
     else:
-        raise FileNotFoundError(f"找不到 StableAvatar 權重檔案：{transformer_path}。請確保模型已完整下載。")
-
-    scheduler_class = { "Flow": FlowMatchEulerDiscreteScheduler }[sampler_name]
-    scheduler = scheduler_class(**filter_kwargs(scheduler_class, OmegaConf.to_container(config['scheduler_kwargs'])))
-    
-    print("正在建立 Pipeline...")
-    pipeline = WanI2VTalkingInferenceLongPipeline(
-        tokenizer=tokenizer, text_encoder=text_encoder, vae=vae,
-        transformer=transformer3d, clip_image_encoder=clip_image_encoder,
-        scheduler=scheduler, wav2vec_processor=wav2vec_processor, wav2vec=wav2vec,
-    )
-    
-    return pipeline, transformer3d, vae
+        print(f"错误：模型文件不存在: {transformer_path}")
+        return None
 
-def run_inference(
-    pipeline, transformer3d, vae, image_path, audio_path, prompt,
-    negative_prompt, seed, output_filename, gpu_memory_mode="model_cpu_offload",
-    width=512, height=512, num_inference_steps=50, fps=25, **kwargs
+def generate_video(
+    GPU_memory_mode="model_cpu_offload",
+    teacache_threshold=0.0,
+    num_skip_start_steps=5,
+    image_path=None,
+    audio_path=None,
+    prompt="",
+    negative_prompt="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
+    width=512,
+    height=512,
+    guidance_scale=6.0,
+    num_inference_steps=50,
+    text_guide_scale=3.0,
+    audio_guide_scale=5.0,
+    motion_frame=25,
+    fps=25,
+    overlap_window_length=10,
+    seed_param=42,
+    overlapping_weight_scheme="uniform",
 ):
-    """執行推理以生成影片。"""
-    if seed < 0:
+    global pipeline, transformer3d
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    
+    if seed_param < 0:
         seed = random.randint(0, np.iinfo(np.int32).max)
-    print(f"使用的種子： {seed}")
+    else:
+        seed = seed_param
 
-    if gpu_memory_mode == "sequential_cpu_offload":
+    print(f"使用种子: {seed}")
+    print(f"输入图片: {image_path}")
+    print(f"输入音频: {audio_path}")
+    print(f"提示词: {prompt}")
+
+    if GPU_memory_mode == "sequential_cpu_offload":
+        replace_parameters_by_name(transformer3d, ["modulation", ], device=device)
+        transformer3d.freqs = transformer3d.freqs.to(device=device)
         pipeline.enable_sequential_cpu_offload(device=device)
-    elif gpu_memory_mode == "model_cpu_offload":
+    elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+        convert_model_weight_to_float8(transformer3d, exclude_module_name=["modulation", ])
+        convert_weight_dtype_wrapper(transformer3d, dtype)
+        pipeline.enable_model_cpu_offload(device=device)
+    elif GPU_memory_mode == "model_cpu_offload":
         pipeline.enable_model_cpu_offload(device=device)
     else:
         pipeline.to(device=device)
+        
+    if teacache_threshold > 0:
+        coefficients = get_teacache_coefficients(pretrained_model_name_or_path)
+        pipeline.transformer.enable_teacache(
+            coefficients,
+            num_inference_steps,
+            teacache_threshold,
+            num_skip_start_steps=num_skip_start_steps,
+        )
 
     with torch.no_grad():
-        print("正在準備輸入資料...")
-        video_length = 81
+        video_length = int((clip_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if clip_sample_n_frames != 1 else 1
         input_video, input_video_mask, clip_image = get_image_to_video_latent(image_path, None, video_length=video_length, sample_size=[height, width])
-        
         sr = 16000
-        vocal_input, _ = librosa.load(audio_path, sr=sr)
+        vocal_input, sample_rate = librosa.load(audio_path, sr=sr)
         
-        print("Pipeline 執行中... 這可能需要一些時間。")
+        print("开始生成视频...")
         sample = pipeline(
-            prompt, num_frames=video_length, negative_prompt=negative_prompt,
-            width=width, height=height, guidance_scale=6.0,
-            generator=torch.Generator().manual_seed(seed), num_inference_steps=num_inference_steps,
-            video=input_video, mask_video=input_video_mask, clip_image=clip_image,
-            text_guide_scale=3.0, audio_guide_scale=5.0, vocal_input_values=vocal_input,
-            motion_frame=25, fps=fps, sr=sr, cond_file_path=image_path,
-            overlap_window_length=10, seed=seed, overlapping_weight_scheme="uniform",
+            prompt,
+            num_frames=video_length,
+            negative_prompt=negative_prompt,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale,
+            generator=torch.Generator().manual_seed(seed),
+            num_inference_steps=num_inference_steps,
+            video=input_video,
+            mask_video=input_video_mask,
+            clip_image=clip_image,
+            text_guide_scale=text_guide_scale,
+            audio_guide_scale=audio_guide_scale,
+            vocal_input_values=vocal_input,
+            motion_frame=motion_frame,
+            fps=fps,
+            sr=sr,
+            cond_file_path=image_path,
+            overlap_window_length=overlap_window_length,
+            seed=seed,
+            overlapping_weight_scheme=overlapping_weight_scheme,
         ).videos
         
-        print("正在儲存影片...")
         os.makedirs("outputs", exist_ok=True)
-        video_path = os.path.join("outputs", f"{output_filename}.mp4")
+        video_path = os.path.join("outputs", f"{timestamp}.mp4")
         save_videos_grid(sample, video_path, fps=fps)
+        output_video_with_audio = os.path.join("outputs", f"{timestamp}_audio.mp4")
         
-        output_video_with_audio = os.path.join("outputs", f"{output_filename}_audio.mp4")
-        
-        print("正在將音訊合併到影片中...")
+        print("合并音频到视频...")
         subprocess.run([
-            "ffmpeg", "-y", "-loglevel", "quiet", "-i", video_path, "-i", audio_path,
-            "-c:v", "copy", "-c:a", "aac", "-strict", "experimental",
+            "ffmpeg", "-y", "-loglevel", "quiet", "-i", video_path, "-i", audio_path, 
+            "-c:v", "copy", "-c:a", "aac", "-strict", "experimental", 
             output_video_with_audio
         ], check=True)
         
-        os.remove(video_path)
-        
-    print(f"✅ 生成完成！影片已儲存至： {output_video_with_audio}")
+    print(f"生成完成! 输出文件: {output_video_with_audio}")
     return output_video_with_audio, seed
 
+def parse_args():
+    parser = argparse.ArgumentParser(description="StableAvatar Video Generation")
+    parser.add_argument("--prompt", type=str, default="", help="文本提示词")
+    parser.add_argument("--input_image", type=str, required=True, help="输入图片路径或URL")
+    parser.add_argument("--input_audio", type=str, required=True, help="输入音频路径或URL")
+    parser.add_argument("--seed", type=int, default=42, help="随机种子，-1为随机")
+    parser.add_argument("--negative_prompt", type=str, 
+                       default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
+                       help="负面提示词")
+    parser.add_argument("--width", type=int, default=512, help="视频宽度")
+    parser.add_argument("--height", type=int, default=512, help="视频高度")
+    parser.add_argument("--guidance_scale", type=float, default=6.0, help="引导权重")
+    parser.add_argument("--num_inference_steps", type=int, default=50, help="推理步数")
+    parser.add_argument("--text_guide_scale", type=float, default=3.0, help="文本引导权重")
+    parser.add_argument("--audio_guide_scale", type=float, default=5.0, help="音频引导权重")
+    parser.add_argument("--motion_frame", type=int, default=25, help="运动帧数")
+    parser.add_argument("--fps", type=int, default=25, help="视频帧率")
+    parser.add_argument("--overlap_window_length", type=int, default=10, help="重叠窗口长度")
+    parser.add_argument("--overlapping_weight_scheme", type=str, default="uniform", 
+                       choices=["uniform", "log"], help="重叠权重方案")
+    parser.add_argument("--GPU_memory_mode", type=str, default="model_cpu_offload",
+                       choices=["Normal", "model_cpu_offload", "model_cpu_offload_and_qfloat8", "sequential_cpu_offload"],
+                       help="GPU内存模式")
+    parser.add_argument("--teacache_threshold", type=float, default=0.0, help="TeaCache阈值")
+    parser.add_argument("--num_skip_start_steps", type=int, default=5, help="跳过开始步数")
+    parser.add_argument("--model_version", type=str, default="square", 
+                       choices=["square", "rec_vec"], help="模型版本")
+    
+    return parser.parse_args()
+
+def download_file(url, local_path):
+    """下载远程文件到本地"""
+    import urllib.request
+    print(f"正在下载 {url} 到 {local_path}")
+    urllib.request.urlretrieve(url, local_path)
+    print(f"下载完成: {local_path}")
+    return local_path
+
 def main():
-    parser = argparse.ArgumentParser(description="StableAvatar 命令列推理工具")
-    parser.add_argument('--prompt', type=str, default="a beautiful woman is talking, masterpiece, best quality", help='正面提示詞')
-    parser.add_argument('--input_image', type=str, default="example_case/case-6/reference.png", help='輸入圖片的路徑')
-    parser.add_argument('--input_audio', type=str, default="example_case/case-6/audio.wav", help='輸入音訊的路徑')
-    parser.add_argument('--seed', type=int, default=42, help='隨機種子，-1 表示隨機')
-    parser.add_argument('--negative_prompt', type=str, default="vivid color, static, blur details, text, style, painting, picture, still, gray, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, malformed, deformed, bad anatomy, fused fingers, still image, messy background, many people in the background, walking backwards", help='負面提示詞')
-    parser.add_argument('--width', type=int, default=512, help='影片寬度')
-    parser.add_argument('--height', type=int, default=512, help='影片高度')
-    parser.add_argument('--num_inference_steps', type=int, default=50, help='推理步數')
-    parser.add_argument('--fps', type=int, default=25, help='影片幀率')
-    parser.add_argument('--gpu_memory_mode', type=str, default="model_cpu_offload", choices=["Normal", "model_cpu_offload"], help='GPU 記憶體優化模式')
-    parser.add_argument('--model_version', type=str, default="square", choices=["square", "rec_vec"], help='StableAvatar 模型版本')
-    args = parser.parse_args()
+    args = parse_args()
+    
+    # 处理输入文件（支持URL或本地路径）
+    image_path = args.input_image
+    audio_path = args.input_audio
+    
+    # 如果是URL，下载到临时文件
+    if image_path.startswith('http'):
+        os.makedirs("temp", exist_ok=True)
+        local_image_path = f"temp/temp_image_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
+        image_path = download_file(image_path, local_image_path)
+    
+    if audio_path.startswith('http'):
+        os.makedirs("temp", exist_ok=True)
+        audio_ext = os.path.splitext(audio_path)[1] or '.wav'
+        local_audio_path = f"temp/temp_audio_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}{audio_ext}"
+        audio_path = download_file(audio_path, local_audio_path)
+    
+    # 检查文件是否存在
+    if not os.path.exists(image_path):
+        print(f"错误：图片文件不存在: {image_path}")
+        return
+    
+    if not os.path.exists(audio_path):
+        print(f"错误：音频文件不存在: {audio_path}")
+        return
+    
+    # 加载指定版本的模型
+    load_transformer_model(args.model_version)
+    
+    # 生成视频
+    output_path, used_seed = generate_video(
+        GPU_memory_mode=args.GPU_memory_mode,
+        teacache_threshold=args.teacache_threshold,
+        num_skip_start_steps=args.num_skip_start_steps,
+        image_path=image_path,
+        audio_path=audio_path,
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        width=args.width,
+        height=args.height,
+        guidance_scale=args.guidance_scale,
+        num_inference_steps=args.num_inference_steps,
+        text_guide_scale=args.text_guide_scale,
+        audio_guide_scale=args.audio_guide_scale,
+        motion_frame=args.motion_frame,
+        fps=args.fps,
+        overlap_window_length=args.overlap_window_length,
+        seed_param=args.seed,
+        overlapping_weight_scheme=args.overlapping_weight_scheme,
+    )
+    
+    print(f"\n=== 生成完成 ===")
+    print(f"输出文件: {output_path}")
+    print(f"使用种子: {used_seed}")
 
-    print("--- 步驟 1: 正在檢查並下載模型與設定檔 ---")
+if __name__ == "__main__":
+    # 初始化模型和配置
+    REPO_ID = "FrancisRing/StableAvatar"
     repo_root = snapshot_download(
-        repo_id="FrancisRing/StableAvatar",
+        repo_id=REPO_ID,
         allow_patterns=[
             "StableAvatar-1.3B/*",
             "Wan2.1-Fun-V1.1-1.3B-InP/*",
             "wav2vec2-base-960h/*",
-            "deepspeed_config/**",
-            "example_case/**"
+            "assets/**",
+            "Kim_Vocal_2.onnx",
         ],
     )
-    print("模型檔案已準備就緒。")
-
-    print("\n--- 步驟 2: 正在解析輸入檔案路徑 ---")
-    final_image_path = resolve_path(args.input_image, repo_root)
-    if not final_image_path:
-        print(f"錯誤：無法找到圖片檔案 {args.input_image}")
-        return
+    pretrained_model_name_or_path = os.path.join(repo_root, "Wan2.1-Fun-V1.1-1.3B-InP")
+    pretrained_wav2vec_path = os.path.join(repo_root, "wav2vec2-base-960h")
 
-    final_audio_path = resolve_path(args.input_audio, repo_root)
-    if not final_audio_path:
-        print(f"錯誤：無法找到音訊檔案 {args.input_audio}")
-        return
+    # 人声分离 onnx
+    audio_separator_model_file = os.path.join(repo_root, "Kim_Vocal_2.onnx")
 
-    print("\n--- 步驟 3: 正在載入模型 ---")
-    pipeline, transformer3d, vae = setup_models(repo_root, args.model_version)
-    print("模型載入完成。")
+    config = OmegaConf.load("deepspeed_config/wan2.1/wan_civitai.yaml")
+    sampler_name = "Flow"
+    clip_sample_n_frames = 81
     
-    print("\n--- 步驟 4: 開始執行推理 ---")
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    run_inference(
-        pipeline=pipeline, transformer3d=transformer3d, vae=vae,
-        image_path=final_image_path, audio_path=final_audio_path,
-        prompt=args.prompt, negative_prompt=args.negative_prompt,
-        seed=args.seed, output_filename=f"output_{timestamp}",
-        gpu_memory_mode=args.gpu_memory_mode, width=args.width,
-        height=args.height, num_inference_steps=args.num_inference_steps,
-        fps=args.fps
+    print("正在初始化模型...")
+    tokenizer = AutoTokenizer.from_pretrained(os.path.join(pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')), )
+    text_encoder = WanT5EncoderModel.from_pretrained(
+        os.path.join(pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
+        additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
+        low_cpu_mem_usage=True,
+        torch_dtype=dtype,
+    )
+    text_encoder = text_encoder.eval()
+    vae = AutoencoderKLWan.from_pretrained(
+        os.path.join(pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')),
+        additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
+    )
+    wav2vec_processor = Wav2Vec2Processor.from_pretrained(pretrained_wav2vec_path)
+    wav2vec = Wav2Vec2Model.from_pretrained(pretrained_wav2vec_path).to("cpu")
+    clip_image_encoder = CLIPModel.from_pretrained(os.path.join(pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')), )
+    clip_image_encoder = clip_image_encoder.eval()
+    transformer3d = WanTransformer3DFantasyModel.from_pretrained(
+        os.path.join(pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+        transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+        low_cpu_mem_usage=False,
+        torch_dtype=dtype,
     )
 
-if __name__ == "__main__":
-    main()
-
+    Choosen_Scheduler = scheduler_dict = {
+        "Flow": FlowMatchEulerDiscreteScheduler,
+    }[sampler_name]
+    scheduler = Choosen_Scheduler(
+        **filter_kwargs(Choosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
+    )
+    pipeline = WanI2VTalkingInferenceLongPipeline(
+        tokenizer=tokenizer,
+        text_encoder=text_encoder,
+        vae=vae,
+        transformer=transformer3d,
+        clip_image_encoder=clip_image_encoder,
+        scheduler=scheduler,
+        wav2vec_processor=wav2vec_processor,
+        wav2vec=wav2vec,
+    )
+    
+    print("模型初始化完成！")
+    main()