Update app.py

app.py

@@ -23,14 +23,12 @@ import subprocess
 from moviepy.editor import VideoFileClip
 from huggingface_hub import snapshot_download
 import shutil
+import requests
+import uuid
 
-try:
-    from audio_separator.separator import Separator
-except:
-    print("Unable to use vocal separation feature. Please install audio-separator[gpu].")
-
+# Device and dtype setup
 if torch.cuda.is_available():
-    device = "cuda" 
+    device = "cuda"
     if torch.cuda.get_device_capability()[0] >= 8:
         dtype = torch.bfloat16
     else:
@@ -47,37 +45,112 @@ def filter_kwargs(cls, kwargs):
     return filtered_kwargs
 
 def load_transformer_model(model_version, repo_root):
-    """
-    根据选择的模型版本加载对应的transformer模型
-    
-    Args:
-        model_version (str): 模型版本，"square" 或 "rec_vec"
-        repo_root (str): 模型根目录
-    
-    Returns:
-        WanTransformer3DFantasyModel: 加载的transformer模型
-    """
-    if model_version == "square":
-        transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", "transformer3d-square.pt")
-    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}")
-    
+    transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", f"transformer3d-{model_version}.pt")
+    print(f"Loading model: {transformer_path}")
     if os.path.exists(transformer_path):
         state_dict = torch.load(transformer_path, map_location="cpu")
         state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
-        return transformer_path, state_dict
+        m, u = transformer3d.load_state_dict(state_dict, strict=False)
+        print(f"Model loaded successfully: {transformer_path}")
+        print(f"Missing keys: {len(m)}; Unexpected keys: {len(u)}")
+        return transformer3d
     else:
-        raise FileNotFoundError(f"模型文件不存在: {transformer_path}")
+        print(f"Error: Model file does not exist: {transformer_path}")
+        return None
+
+def download_file(url, local_path):
+    """Download file from URL to local path"""
+    try:
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        with open(local_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        return local_path
+    except Exception as e:
+        print(f"Error downloading file from {url}: {e}")
+        return None
+
+def prepare_input_file(input_path, file_type="image"):
+    """Handle local or remote file inputs"""
+    if input_path.startswith("http://") or input_path.startswith("https://"):
+        ext = ".png" if file_type == "image" else ".wav"
+        local_path = os.path.join("temp", f"{uuid.uuid4()}{ext}")
+        os.makedirs("temp", exist_ok=True)
+        return download_file(input_path, local_path)
+    elif os.path.exists(input_path):
+        return input_path
+    else:
+        print(f"Error: {file_type.capitalize()} file {input_path} does not exist")
+        return None
+
+# Initialize model paths
+REPO_ID = "FrancisRing/StableAvatar"
+repo_root = snapshot_download(
+    repo_id=REPO_ID,
+    allow_patterns=[
+        "StableAvatar-1.3B/*",
+        "Wan2.1-Fun-V1.1-1.3B-InP/*",
+        "wav2vec2-base-960h/*",
+        "assets/**",
+        "Kim_Vocal_2.onnx",
+    ],
+)
+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")
+audio_separator_model_file = os.path.join(repo_root, "Kim_Vocal_2.onnx")
+
+# Load configuration and models
+config = OmegaConf.load("deepspeed_config/wan2.1/wan_civitai.yaml")
+sampler_name = "Flow"
+clip_sample_n_frames = 81
 
-def generate_video(
-    transformer3d,
-    pipeline,
-    repo_root,
+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,
+).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'))
+).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,
+)
+
+# Load default transformer model
+load_transformer_model("square", repo_root)
+
+# Initialize scheduler and pipeline
+scheduler_dict = {"Flow": FlowMatchEulerDiscreteScheduler}
+Choosen_Scheduler = scheduler_dict[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,
+)
+
+def generate(
     GPU_memory_mode="model_cpu_offload",
     teacache_threshold=0,
     num_skip_start_steps=5,
@@ -95,32 +168,37 @@ def generate_video(
     fps=25,
     overlap_window_length=10,
     seed_param=42,
-    overlapping_weight_scheme="uniform",
+    overlapping_weight_scheme="uniform"
 ):
+    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)
     else:
         seed = seed_param
-    
-    print(f"使用种子: {seed}")
 
+    # Handle input files
+    image_path = prepare_input_file(image_path, "image")
+    audio_path = prepare_input_file(audio_path, "audio")
+    if not image_path or not audio_path:
+        return None, None, "Error: Invalid input file paths"
+
+    # Configure pipeline based on GPU memory mode
     if GPU_memory_mode == "sequential_cpu_offload":
-        replace_parameters_by_name(transformer3d, ["modulation", ], device=device)
+        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_and_qfloat8":
-        convert_model_weight_to_float8(transformer3d, exclude_module_name=["modulation", ])
+        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)
-        
+
+    # Enable TeaCache if specified
     if teacache_threshold > 0:
-        pretrained_model_name_or_path = os.path.join(repo_root, "Wan2.1-Fun-V1.1-1.3B-InP")
         coefficients = get_teacache_coefficients(pretrained_model_name_or_path)
         pipeline.transformer.enable_teacache(
             coefficients,
@@ -129,15 +207,12 @@ def generate_video(
             num_skip_start_steps=num_skip_start_steps,
         )
 
+    # Perform inference
     with torch.no_grad():
-        clip_sample_n_frames = 81
-        vae = pipeline.vae
         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, sample_rate = librosa.load(audio_path, sr=sr)
-        
-        print("开始生成视频...")
         sample = pipeline(
             prompt,
             num_frames=video_length,
@@ -161,195 +236,43 @@ def generate_video(
             seed=seed,
             overlapping_weight_scheme=overlapping_weight_scheme,
         ).videos
-        
         os.makedirs("outputs", exist_ok=True)
         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")
-        
-        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)
-        
-    return output_video_with_audio, seed
-
-def audio_extractor(video_path, output_dir="outputs"):
-    """从视频中提取音频"""
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    os.makedirs(output_dir, exist_ok=True)
-    out_wav = os.path.abspath(os.path.join(output_dir, f"{timestamp}.wav"))
-    video = VideoFileClip(video_path)
-    audio = video.audio
-    audio.write_audiofile(out_wav, codec="pcm_s16le")
-    return out_wav
 
-def vocal_separation(audio_path, audio_separator_model_file, output_dir="outputs"):
-    """人声分离"""
-    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-    os.makedirs(output_dir, exist_ok=True)
-    
-    audio_separator = Separator(
-        output_dir=os.path.abspath(os.path.join(output_dir, timestamp)),
-        output_single_stem="vocals",
-        model_file_dir=os.path.dirname(audio_separator_model_file),
-    )
-    audio_separator.load_model(os.path.basename(audio_separator_model_file))
-    assert audio_separator.model_instance is not None, "Fail to load audio separate model."
-    outputs = audio_separator.separate(audio_path)
-    vocal_audio_file = os.path.join(audio_separator.output_dir, outputs[0])
-    destination_file = os.path.abspath(os.path.join(output_dir, f"{timestamp}.wav"))
-    shutil.copy(vocal_audio_file, destination_file)
-    os.remove(vocal_audio_file)
-    return destination_file
+    return output_video_with_audio, seed, f"Generated outputs/{timestamp}.mp4"
 
 def main():
-    parser = argparse.ArgumentParser(description="StableAvatar 命令行推理工具")
-    
-    # 主要参数
-    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("--model_version", type=str, default="square", choices=["square", "rec_vec"], 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, help="TeaCache阈值")
-    parser.add_argument("--num_skip_start_steps", type=int, default=5, help="跳过开始步数")
-    
-    # 生成参数
-    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("--extract_audio", type=str, help="从视频提取音频，提供视频路径")
-    parser.add_argument("--separate_vocal", type=str, help="人声分离，提供音频路径")
-    parser.add_argument("--output_dir", type=str, default="outputs", help="输出目录")
-    
+    parser = argparse.ArgumentParser(description="StableAvatar Inference Script")
+    parser.add_argument("--prompt", type=str, default="", help="Text prompt for generation")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed, -1 for random")
+    parser.add_argument("--input_image", type=str, required=True, help="Path or URL to input image (e.g., ./image.png or https://example.com/image.png)")
+    parser.add_argument("--input_audio", type=str, required=True, help="Path or URL to input audio (e.g., ./audio.wav or https://example.com/audio.wav)")
+    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 memory mode")
+    parser.add_argument("--teacache_threshold", type=float, default=0, help="TeaCache threshold, 0 to disable")
+    parser.add_argument("--num_skip_start_steps", type=int, default=5, help="Number of start steps to skip")
+    parser.add_argument("--negative_prompt", type=str, default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", help="Negative prompt")
+    parser.add_argument("--width", type=int, default=512, help="Output video width")
+    parser.add_argument("--height", type=int, default=512, help="Output video height")
+    parser.add_argument("--guidance_scale", type=float, default=6.0, help="Guidance scale")
+    parser.add_argument("--num_inference_steps", type=int, default=50, help="Number of inference steps")
+    parser.add_argument("--text_guide_scale", type=float, default=3.0, help="Text guidance scale")
+    parser.add_argument("--audio_guide_scale", type=float, default=5.0, help="Audio guidance scale")
+    parser.add_argument("--motion_frame", type=int, default=25, help="Motion frame")
+    parser.add_argument("--fps", type=int, default=25, help="Frames per second")
+    parser.add_argument("--overlap_window_length", type=int, default=10, help="Overlap window length")
+    parser.add_argument("--overlapping_weight_scheme", type=str, default="uniform", choices=["uniform", "log"], help="Overlapping weight scheme")
+
     args = parser.parse_args()
-    
-    # 下载模型
-    print("正在下载模型...")
-    REPO_ID = "FrancisRing/StableAvatar"
-    repo_root = snapshot_download(
-        repo_id=REPO_ID,
-        allow_patterns=[
-            "StableAvatar-1.3B/*",
-            "Wan2.1-Fun-V1.1-1.3B-InP/*", 
-            "wav2vec2-base-960h/*",
-            "assets/**",
-            "Kim_Vocal_2.onnx",
-        ],
-    )
-    
-    # 工具功能
-    if args.extract_audio:
-        print(f"从视频提取音频: {args.extract_audio}")
-        output_audio = audio_extractor(args.extract_audio, args.output_dir)
-        print(f"音频已保存到: {output_audio}")
-        return
-    
-    if args.separate_vocal:
-        print(f"人声分离: {args.separate_vocal}")
-        audio_separator_model_file = os.path.join(repo_root, "Kim_Vocal_2.onnx")
-        output_audio = vocal_separation(args.separate_vocal, audio_separator_model_file, args.output_dir)
-        print(f"分离后的人声已保存到: {output_audio}")
-        return
-    
-    # 检查必要参数
-    if not args.input_image or not args.input_audio:
-        print("错误: 必须提供 --input_image 和 --input_audio 参数")
-        return
-    
-    # 初始化模型
-    print("正在初始化模型...")
-    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 = OmegaConf.load("deepspeed_config/wan2.1/wan_civitai.yaml")
-    sampler_name = "Flow"
-    
-    # 加载各个组件
-    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,
-    )
-    
-    # 加载指定版本的transformer模型
-    transformer_path, state_dict = load_transformer_model(args.model_version, repo_root)
-    m, u = transformer3d.load_state_dict(state_dict, strict=False)
-    print(f"模型加载成功: {transformer_path}")
-    print(f"Missing keys: {len(m)}; Unexpected keys: {len(u)}")
-    
-    Choosen_Scheduler = {
-        "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("开始生成...")
-    output_video, used_seed = generate_video(
-        transformer3d=transformer3d,
-        pipeline=pipeline,
-        repo_root=repo_root,
-        GPU_memory_mode=args.gpu_memory_mode,
+
+    video_path, seed, message = generate(
+        GPU_memory_mode=args.GPU_memory_mode,
         teacache_threshold=args.teacache_threshold,
         num_skip_start_steps=args.num_skip_start_steps,
         image_path=args.input_image,
@@ -366,12 +289,13 @@ def main():
         fps=args.fps,
         overlap_window_length=args.overlap_window_length,
         seed_param=args.seed,
-        overlapping_weight_scheme=args.overlapping_weight_scheme,
+        overlapping_weight_scheme=args.overlapping_weight_scheme
     )
-    
-    print(f"生成完成!")
-    print(f"输出视频: {output_video}")
-    print(f"使用种子: {used_seed}")
+
+    if video_path:
+        print(f"{message}\nSeed: {seed}")
+    else:
+        print("Generation failed.")
 
 if __name__ == "__main__":
     main()