Spaces:
Runtime error
Runtime error
| import torch | |
| import psutil | |
| import argparse | |
| import gradio as gr | |
| import os | |
| from diffusers import FlowMatchEulerDiscreteScheduler | |
| from diffusers.utils import load_image | |
| from transformers import AutoTokenizer, Wav2Vec2Model, Wav2Vec2Processor | |
| from omegaconf import OmegaConf | |
| from wan.models.cache_utils import get_teacache_coefficients | |
| from wan.models.wan_fantasy_transformer3d_1B import WanTransformer3DFantasyModel | |
| from wan.models.wan_text_encoder import WanT5EncoderModel | |
| from wan.models.wan_vae import AutoencoderKLWan | |
| from wan.models.wan_image_encoder import CLIPModel | |
| from wan.pipeline.wan_inference_long_pipeline import WanI2VTalkingInferenceLongPipeline | |
| from wan.utils.fp8_optimization import replace_parameters_by_name, convert_weight_dtype_wrapper, convert_model_weight_to_float8 | |
| from wan.utils.utils import get_image_to_video_latent, save_videos_grid | |
| import numpy as np | |
| import librosa | |
| import datetime | |
| import random | |
| import math | |
| import subprocess | |
| from moviepy.editor import VideoFileClip | |
| from huggingface_hub import snapshot_download | |
| import shutil | |
| import spaces | |
| 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" | |
| if torch.cuda.get_device_capability()[0] >= 8: | |
| dtype = torch.bfloat16 | |
| else: | |
| dtype = torch.float16 | |
| else: | |
| device = "cpu" | |
| 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 | |
| def load_transformer_model(model_version): | |
| """ | |
| 根据选择的模型版本加载对应的transformer模型 | |
| Args: | |
| model_version (str): 模型版本，"square" 或 "rec_vec" | |
| Returns: | |
| WanTransformer3DFantasyModel: 加载的transformer模型 | |
| """ | |
| global transformer3d | |
| 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}") | |
| 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 | |
| m, u = transformer3d.load_state_dict(state_dict, strict=False) | |
| print(f"模型加载成功: {transformer_path}") | |
| print(f"Missing keys: {len(m)}; Unexpected keys: {len(u)}") | |
| return transformer3d | |
| else: | |
| print(f"错误：模型文件不存在: {transformer_path}") | |
| return None | |
| 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") | |
| # 人声分离 onnx | |
| audio_separator_model_file = os.path.join(repo_root, "Kim_Vocal_2.onnx") | |
| # model_path = "/datadrive/stableavatar/checkpoints" | |
| # pretrained_model_name_or_path = f"{model_path}/Wan2.1-Fun-V1.1-1.3B-InP" | |
| # pretrained_wav2vec_path = f"{model_path}/wav2vec2-base-960h" | |
| # transformer_path = f"{model_path}/StableAvatar-1.3B/transformer3d-square.pt" | |
| config = OmegaConf.load("deepspeed_config/wan2.1/wan_civitai.yaml") | |
| sampler_name = "Flow" | |
| clip_sample_n_frames = 81 | |
| 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, | |
| ) | |
| # 默认加载square版本模型 | |
| load_transformer_model("square") | |
| 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, | |
| ) | |
| def generate( | |
| GPU_memory_mode, | |
| teacache_threshold, | |
| num_skip_start_steps, | |
| image_path, | |
| audio_path, | |
| prompt, | |
| negative_prompt, | |
| width, | |
| height, | |
| guidance_scale, | |
| num_inference_steps, | |
| text_guide_scale, | |
| audio_guide_scale, | |
| motion_frame, | |
| fps, | |
| overlap_window_length, | |
| seed_param, | |
| overlapping_weight_scheme, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| 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 | |
| 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_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(): | |
| 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) | |
| sample = pipeline( | |
| 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 | |
| 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") | |
| subprocess.run([ | |
| "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, f"Generated outputs/{timestamp}.mp4 / 已生成outputs/{timestamp}.mp4" | |
| def exchange_width_height(width, height): | |
| return height, width, "✅ Width and Height Swapped / 宽高交换完毕" | |
| def adjust_width_height(image): | |
| image = load_image(image) | |
| width, height = image.size | |
| original_area = width * height | |
| default_area = 512*512 | |
| ratio = math.sqrt(original_area / default_area) | |
| width = width / ratio // 16 * 16 | |
| height = height / ratio // 16 * 16 | |
| return int(width), int(height), "✅ Adjusted Size Based on Image / 根据图片调整宽高" | |
| def audio_extractor(video_path): | |
| timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") | |
| os.makedirs("outputs", exist_ok=True) # 确保目录存在 | |
| out_wav = os.path.abspath(os.path.join("outputs", f"{timestamp}.wav")) | |
| video = VideoFileClip(video_path) | |
| audio = video.audio | |
| audio.write_audiofile(out_wav, codec="pcm_s16le") | |
| return out_wav, f"Generated {out_wav} / 已生成 {out_wav}", out_wav # ← 第3个返回给 gr.File | |
| def vocal_separation(audio_path): | |
| timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") | |
| os.makedirs("outputs", exist_ok=True) | |
| # audio_separator_model_file = "checkpoints/Kim_Vocal_2.onnx" | |
| audio_separator = Separator( | |
| output_dir=os.path.abspath(os.path.join("outputs", 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("outputs", f"{timestamp}.wav")) | |
| shutil.copy(vocal_audio_file, destination_file) | |
| os.remove(vocal_audio_file) | |
| return destination_file, f"Generated {destination_file} / 已生成 {destination_file}", destination_file | |
| def update_language(language): | |
| if language == "English": | |
| return { | |
| GPU_memory_mode: gr.Dropdown(label="GPU Memory Mode", info="Normal uses 25G VRAM, model_cpu_offload uses 13G VRAM"), | |
| teacache_threshold: gr.Slider(label="TeaCache Threshold", info="Recommended 0.1, 0 disables TeaCache acceleration"), | |
| num_skip_start_steps: gr.Slider(label="Skip Start Steps", info="Recommended 5"), | |
| model_version: gr.Dropdown(label="Model Version", choices=["square", "rec_vec"], value="square"), | |
| image_path: gr.Image(label="Upload Image"), | |
| audio_path: gr.Audio(label="Upload Audio"), | |
| prompt: gr.Textbox(label="Prompt"), | |
| negative_prompt: gr.Textbox(label="Negative Prompt"), | |
| generate_button: gr.Button("🎬 Start Generation"), | |
| width: gr.Slider(label="Width"), | |
| height: gr.Slider(label="Height"), | |
| exchange_button: gr.Button("🔄 Swap Width/Height"), | |
| adjust_button: gr.Button("Adjust Size Based on Image"), | |
| guidance_scale: gr.Slider(label="Guidance Scale"), | |
| num_inference_steps: gr.Slider(label="Sampling Steps (Recommended 50)"), | |
| text_guide_scale: gr.Slider(label="Text Guidance Scale"), | |
| audio_guide_scale: gr.Slider(label="Audio Guidance Scale"), | |
| motion_frame: gr.Slider(label="Motion Frame"), | |
| fps: gr.Slider(label="FPS"), | |
| overlap_window_length: gr.Slider(label="Overlap Window Length"), | |
| seed_param: gr.Number(label="Seed (positive integer, -1 for random)"), | |
| overlapping_weight_scheme: gr.Dropdown(label="Overlapping Weight Scheme", choices=["uniform", "log"], value="uniform"), | |
| info: gr.Textbox(label="Status"), | |
| video_output: gr.Video(label="Generated Result"), | |
| seed_output: gr.Textbox(label="Seed"), | |
| video_path: gr.Video(label="Upload Video"), | |
| extractor_button: gr.Button("🎬 Start Extraction"), | |
| info2: gr.Textbox(label="Status"), | |
| audio_output: gr.Audio(label="Generated Result"), | |
| audio_path3: gr.Audio(label="Upload Audio"), | |
| separation_button: gr.Button("🎬 Start Separation"), | |
| info3: gr.Textbox(label="Status"), | |
| audio_output3: gr.Audio(label="Generated Result"), | |
| example_title: gr.Markdown(value="### Select the following example cases for testing:"), | |
| example1_label: gr.Markdown(value="**Example 1**"), | |
| example2_label: gr.Markdown(value="**Example 2**"), | |
| example3_label: gr.Markdown(value="**Example 3**"), | |
| example4_label: gr.Markdown(value="**Example 4**"), | |
| example5_label: gr.Markdown(value="**Example 5**"), | |
| example1_btn: gr.Button("🚀 Use Example 1", variant="secondary"), | |
| example2_btn: gr.Button("🚀 Use Example 2", variant="secondary"), | |
| example3_btn: gr.Button("🚀 Use Example 3", variant="secondary"), | |
| example4_btn: gr.Button("🚀 Use Example 4", variant="secondary"), | |
| example5_btn: gr.Button("🚀 Use Example 5", variant="secondary"), | |
| parameter_settings_title: gr.Accordion(label="Parameter Settings", open=True), | |
| example_cases_title: gr.Accordion(label="Example Cases", open=True), | |
| stableavatar_title: gr.TabItem(label="StableAvatar"), | |
| audio_extraction_title: gr.TabItem(label="Audio Extraction"), | |
| vocal_separation_title: gr.TabItem(label="Vocal Separation") | |
| } | |
| else: | |
| return { | |
| GPU_memory_mode: gr.Dropdown(label="显存模式", info="Normal占用25G显存，model_cpu_offload占用13G显存"), | |
| teacache_threshold: gr.Slider(label="teacache threshold", info="推荐参数0.1，0为禁用teacache加速"), | |
| num_skip_start_steps: gr.Slider(label="跳过开始步数", info="推荐参数5"), | |
| model_version: gr.Dropdown(label="模型版本", choices=["square", "rec_vec"], value="square"), | |
| image_path: gr.Image(label="上传图片"), | |
| audio_path: gr.Audio(label="上传音频"), | |
| prompt: gr.Textbox(label="提示词"), | |
| negative_prompt: gr.Textbox(label="负面提示词"), | |
| generate_button: gr.Button("🎬 开始生成"), | |
| width: gr.Slider(label="宽度"), | |
| height: gr.Slider(label="高度"), | |
| exchange_button: gr.Button("🔄 交换宽高"), | |
| adjust_button: gr.Button("根据图片调整宽高"), | |
| guidance_scale: gr.Slider(label="guidance scale"), | |
| num_inference_steps: gr.Slider(label="采样步数（推荐50步）", minimum=1, maximum=100, step=1, value=50), | |
| text_guide_scale: gr.Slider(label="text guidance scale"), | |
| audio_guide_scale: gr.Slider(label="audio guidance scale"), | |
| motion_frame: gr.Slider(label="motion frame"), | |
| fps: gr.Slider(label="帧率"), | |
| overlap_window_length: gr.Slider(label="overlap window length"), | |
| seed_param: gr.Number(label="种子，请输入正整数，-1为随机"), | |
| overlapping_weight_scheme: gr.Dropdown(label="Overlapping Weight Scheme", choices=["uniform", "log"], value="uniform"), | |
| info: gr.Textbox(label="提示信息"), | |
| video_output: gr.Video(label="生成结果"), | |
| seed_output: gr.Textbox(label="种子"), | |
| video_path: gr.Video(label="上传视频"), | |
| extractor_button: gr.Button("🎬 开始提取"), | |
| info2: gr.Textbox(label="提示信息"), | |
| audio_output: gr.Audio(label="生成结果"), | |
| audio_path3: gr.Audio(label="上传音频"), | |
| separation_button: gr.Button("🎬 开始分离"), | |
| info3: gr.Textbox(label="提示信息"), | |
| audio_output3: gr.Audio(label="生成结果"), | |
| example_title: gr.Markdown(value="### 选择以下示例案例进行测试："), | |
| example1_label: gr.Markdown(value="**示例 1**"), | |
| example2_label: gr.Markdown(value="**示例 2**"), | |
| example3_label: gr.Markdown(value="**示例 3**"), | |
| example4_label: gr.Markdown(value="**示例 4**"), | |
| example5_label: gr.Markdown(value="**示例 5**"), | |
| example1_btn: gr.Button("🚀 使用示例 1", variant="secondary"), | |
| example2_btn: gr.Button("🚀 使用示例 2", variant="secondary"), | |
| example3_btn: gr.Button("🚀 使用示例 3", variant="secondary"), | |
| example4_btn: gr.Button("🚀 使用示例 4", variant="secondary"), | |
| example5_btn: gr.Button("🚀 使用示例 5", variant="secondary"), | |
| parameter_settings_title: gr.Accordion(label="参数设置", open=True), | |
| example_cases_title: gr.Accordion(label="示例案例", open=True), | |
| stableavatar_title: gr.TabItem(label="StableAvatar"), | |
| audio_extraction_title: gr.TabItem(label="音频提取"), | |
| vocal_separation_title: gr.TabItem(label="人声分离") | |
| } | |
| BANNER_HTML = """ | |
| <div class="hero"> | |
| <div class="brand"> | |
| <!-- 如有项目 logo，可放到仓库并换成你的地址；没有就删这一行 --> | |
| <!-- <img src="https://raw.githubusercontent.com/Francis-Rings/StableAvatar/main/assets/logo.png" alt="StableAvatar Logo"> --> | |
| <span class="brand-text">STABLEAVATAR</span> | |
| </div> | |
| <div class="titles"> | |
| <h1>StableAvatar</h1> | |
| <div class="badges"> | |
| <a class="badge" href="https://arxiv.org/abs/2508.08248" target="_blank" rel="noopener"> | |
| <img src="https://img.shields.io/badge/arXiv-2508.08248-b31b1b"> | |
| </a> | |
| <a class="badge" href="https://francis-rings.github.io/StableAvatar/" target="_blank" rel="noopener"> | |
| <img src="https://img.shields.io/badge/Webpage-Visit-2266ee"> | |
| </a> | |
| <a class="badge" href="https://github.com/Francis-Rings/StableAvatar" target="_blank" rel="noopener"> | |
| <img src="https://img.shields.io/badge/GitHub-Repo-181717?logo=github&logoColor=white"> | |
| </a> | |
| <a class="badge" href="https://www.youtube.com/watch?v=6lhvmbzvv3Y" target="_blank" rel="noopener"> | |
| <img src="https://img.shields.io/badge/YouTube-Demo-ff0000?logo=youtube&logoColor=white"> | |
| </a> | |
| </div> | |
| </div> | |
| </div> | |
| <hr class="divider"> | |
| """ | |
| BANNER_CSS = """ | |
| .hero{display:flex;align-items:center;gap:18px;padding:18px;border-radius:14px; | |
| color:inherit;margin-bottom:12px} | |
| .brand-text{font-weight:800;letter-spacing:2px} | |
| .brand img{height:46px} | |
| .titles h1{font-size:28px;margin:0 0 6px 0} | |
| .badges{display:flex;gap:10px;flex-wrap:wrap} | |
| .badge img{height:22px} | |
| .divider{border:0;border-top:1px solid rgba(0,0,0,0.12);margin:6px 0 18px} | |
| """ | |
| # with gr.Blocks(theme=gr.themes.Base()) as demo: | |
| # gr.Markdown(""" | |
| # <div> | |
| # <h2 style="font-size: 30px;text-align: center;">StableAvatar</h2> | |
| # </div> | |
| # """) | |
| with gr.Blocks(theme=gr.themes.Base(), css=BANNER_CSS) as demo: | |
| gr.HTML(BANNER_HTML) | |
| language_radio = gr.Radio( | |
| choices=["English", "中文"], | |
| value="English", | |
| label="Language / 语言" | |
| ) | |
| with gr.Accordion("Model Settings / 模型设置", open=False): | |
| with gr.Row(): | |
| GPU_memory_mode = gr.Dropdown( | |
| label = "显存模式", | |
| info = "Normal占用25G显存，model_cpu_offload占用13G显存", | |
| choices = ["Normal", "model_cpu_offload", "model_cpu_offloadand_qfloat8", "sequential_cpu_offload"], | |
| value = "model_cpu_offload" | |
| ) | |
| teacache_threshold = gr.Slider(label="teacache threshold", info = "推荐参数0.1，0为禁用teacache加速", minimum=0, maximum=1, step=0.01, value=0) | |
| num_skip_start_steps = gr.Slider(label="跳过开始步数", info = "推荐参数5", minimum=0, maximum=100, step=1, value=5) | |
| with gr.Row(): | |
| model_version = gr.Dropdown( | |
| label = "模型版本", | |
| choices = ["square","rec_vec"], | |
| value = "square" | |
| ) | |
| stableavatar_title = gr.TabItem(label="StableAvatar") | |
| with stableavatar_title: | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Row(): | |
| image_path = gr.Image(label="上传图片", type="filepath", height=280) | |
| audio_path = gr.Audio(label="上传音频", type="filepath") | |
| prompt = gr.Textbox(label="提示词", value="") | |
| negative_prompt = gr.Textbox(label="负面提示词", value="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走") | |
| generate_button = gr.Button("🎬 开始生成", variant='primary') | |
| parameter_settings_title = gr.Accordion(label="参数设置", open=True) | |
| with parameter_settings_title: | |
| with gr.Row(): | |
| width = gr.Slider(label="宽度", minimum=256, maximum=2048, step=16, value=512) | |
| height = gr.Slider(label="高度", minimum=256, maximum=2048, step=16, value=512) | |
| with gr.Row(): | |
| exchange_button = gr.Button("🔄 交换宽高") | |
| adjust_button = gr.Button("根据图片调整宽高") | |
| with gr.Row(): | |
| guidance_scale = gr.Slider(label="guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=6.0) | |
| num_inference_steps = gr.Slider(label="采样步数（推荐50步）", minimum=1, maximum=100, step=1, value=50) | |
| with gr.Row(): | |
| text_guide_scale = gr.Slider(label="text guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=3.0) | |
| audio_guide_scale = gr.Slider(label="audio guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=5.0) | |
| with gr.Row(): | |
| motion_frame = gr.Slider(label="motion frame", minimum=1, maximum=50, step=1, value=25) | |
| fps = gr.Slider(label="帧率", minimum=1, maximum=60, step=1, value=25) | |
| with gr.Row(): | |
| overlap_window_length = gr.Slider(label="overlap window length", minimum=1, maximum=20, step=1, value=10) | |
| seed_param = gr.Number(label="种子，请输入正整数，-1为随机", value=42) | |
| with gr.Row(): | |
| overlapping_weight_scheme = gr.Dropdown(label="Overlapping Weight Scheme", choices=["uniform", "log"], value="uniform") | |
| with gr.Column(): | |
| info = gr.Textbox(label="提示信息", interactive=False) | |
| video_output = gr.Video(label="生成结果", interactive=False) | |
| seed_output = gr.Textbox(label="种子") | |
| # 示例案例部分移到StableAvatar标签页内部 | |
| example_cases_title = gr.Accordion(label="示例案例", open=True) | |
| with example_cases_title: | |
| example_title = gr.Markdown(value="### 选择以下示例案例进行测试：") | |
| with gr.Row(): | |
| with gr.Column(): | |
| example1_label = gr.Markdown(value="**示例 1**") | |
| example1_image = gr.Image(value="example_case/case-1/reference.png", label="", interactive=False, height=120, show_label=False) | |
| example1_audio = gr.Audio(value="example_case/case-1/audio.wav", label="", interactive=False, show_label=False) | |
| example1_btn = gr.Button("🚀 使用示例 1", variant="secondary", size="sm") | |
| with gr.Column(): | |
| example2_label = gr.Markdown(value="**示例 2**") | |
| example2_image = gr.Image(value="example_case/case-2/reference.png", label="", interactive=False, height=120, show_label=False) | |
| example2_audio = gr.Audio(value="example_case/case-2/audio.wav", label="", interactive=False, show_label=False) | |
| example2_btn = gr.Button("🚀 使用示例 2", variant="secondary", size="sm") | |
| with gr.Column(): | |
| example3_label = gr.Markdown(value="**示例 3**") | |
| example3_image = gr.Image(value="example_case/case-6/reference.png", label="", interactive=False, height=120, show_label=False) | |
| example3_audio = gr.Audio(value="example_case/case-6/audio.wav", label="", interactive=False, show_label=False) | |
| example3_btn = gr.Button("🚀 使用示例 3", variant="secondary", size="sm") | |
| with gr.Column(): | |
| example4_label = gr.Markdown(value="**示例 4**") | |
| example4_image = gr.Image(value="example_case/case-45/reference.png", label="", interactive=False, height=120, show_label=False) | |
| example4_audio = gr.Audio(value="example_case/case-45/audio.wav", label="", interactive=False, show_label=False) | |
| example4_btn = gr.Button("🚀 使用示例 4", variant="secondary", size="sm") | |
| with gr.Column(): | |
| example5_label = gr.Markdown(value="**示例 5**") | |
| example5_image = gr.Image(value="example_case/case-3/reference.jpg", label="", interactive=False, height=120, show_label=False) | |
| example5_audio = gr.Audio(value="example_case/case-3/audio.wav", label="", interactive=False, show_label=False) | |
| example5_btn = gr.Button("🚀 使用示例 5", variant="secondary", size="sm") | |
| audio_extraction_title = gr.TabItem(label="音频提取") | |
| with audio_extraction_title: | |
| with gr.Row(): | |
| with gr.Column(): | |
| video_path = gr.Video(label="上传视频", height=500) | |
| extractor_button = gr.Button("🎬 开始提取", variant='primary') | |
| with gr.Column(): | |
| info2 = gr.Textbox(label="提示信息", interactive=False) | |
| audio_output = gr.Audio(label="生成结果", interactive=False) | |
| audio_file = gr.File(label="download audio file") | |
| vocal_separation_title = gr.TabItem(label="人声分离") | |
| with vocal_separation_title: | |
| with gr.Row(): | |
| with gr.Column(): | |
| audio_path3 = gr.Audio(label="上传音频", type="filepath") | |
| separation_button = gr.Button("🎬 开始分离", variant='primary') | |
| with gr.Column(): | |
| info3 = gr.Textbox(label="提示信息", interactive=False) | |
| audio_output3 = gr.Audio(label="生成结果", interactive=False) | |
| audio_file3 = gr.File(label="download audio file") | |
| # 示例案例部分移到末尾 | |
| # example_cases_title = gr.Accordion(label="示例案例", open=True) | |
| # with example_cases_title: | |
| # example_title = gr.Markdown(value="### 选择以下示例案例进行测试：") | |
| # with gr.Row(): | |
| # with gr.Column(): | |
| # example1_label = gr.Markdown(value="**示例 1**") | |
| # example1_image = gr.Image(value="example_case/case-1/reference.png", label="", interactive=False, height=120, show_label=False) | |
| # example1_audio = gr.Audio(value="example_case/case-1/audio.wav", label="", interactive=False, show_label=False) | |
| # example1_btn = gr.Button("🚀 使用示例 1", variant="secondary", size="sm") | |
| # with gr.Column(): | |
| # example2_label = gr.Markdown(value="**示例 2**") | |
| # example2_image = gr.Image(value="example_case/case-2/reference.png", label="", interactive=False, height=120, show_label=False) | |
| # example2_audio = gr.Audio(value="example_case/case-2/audio.wav", label="", interactive=False, show_label=False) | |
| # example2_btn = gr.Button("🚀 使用示例 2", variant="secondary", size="sm") | |
| # with gr.Column(): | |
| # example3_label = gr.Markdown(value="**示例 3**") | |
| # example3_image = gr.Image(value="example_case/case-6/reference.png", label="", interactive=False, height=120, show_label=False) | |
| # example3_audio = gr.Audio(value="example_case/case-6/audio.wav", label="", interactive=False, show_label=False) | |
| # example3_btn = gr.Button("🚀 使用示例 3", variant="secondary", size="sm") | |
| # with gr.Column(): | |
| # example4_label = gr.Markdown(value="**示例 4**") | |
| # example4_image = gr.Image(value="example_case/case-45/reference.png", label="", interactive=False, height=120, show_label=False) | |
| # example4_audio = gr.Audio(value="example_case/case-45/audio.wav", label="", interactive=False, show_label=False) | |
| # example4_btn = gr.Button("🚀 使用示例 4", variant="secondary", size="sm") | |
| # with gr.Column(): | |
| # example5_label = gr.Markdown(value="**示例 5**") | |
| # example5_image = gr.Image(value="example_case/case-3/reference.jpg", label="", interactive=False, height=120, show_label=False) | |
| # example5_audio = gr.Audio(value="example_case/case-3/audio.wav", label="", interactive=False, show_label=False) | |
| # example5_btn = gr.Button("🚀 使用示例 5", variant="secondary", size="sm") | |
| all_components = [GPU_memory_mode, teacache_threshold, num_skip_start_steps, model_version, image_path, audio_path, prompt, negative_prompt, generate_button, width, height, exchange_button, adjust_button, guidance_scale, num_inference_steps, text_guide_scale, audio_guide_scale, motion_frame, fps, overlap_window_length, seed_param, overlapping_weight_scheme, info, video_output, seed_output, video_path, extractor_button, info2, audio_output, audio_path3, separation_button, info3, audio_output3, example_title, example1_label, example2_label, example3_label, example4_label, example1_btn, example2_btn, example3_btn, example4_btn, example5_label, example5_btn, parameter_settings_title, example_cases_title, stableavatar_title, audio_extraction_title, vocal_separation_title] | |
| language_radio.change( | |
| fn=update_language, | |
| inputs=[language_radio], | |
| outputs=all_components | |
| ) | |
| # 添加模型版本选择的事件处理 | |
| def on_model_version_change(model_version): | |
| """当模型版本改变时，重新加载对应的模型""" | |
| result = load_transformer_model(model_version) | |
| if result is not None: | |
| return f"✅ 模型已切换到 {model_version} 版本" | |
| else: | |
| return f"❌ 模型切换失败，请检查文件是否存在" | |
| model_version.change( | |
| fn=on_model_version_change, | |
| inputs=[model_version], | |
| outputs=[info] | |
| ) | |
| demo.load(fn=update_language, inputs=[language_radio], outputs=all_components) | |
| # 添加示例案例按钮的事件处理 | |
| def load_example1(): | |
| try: | |
| with open("example_case/case-1/prompt.txt", "r", encoding="utf-8") as f: | |
| prompt_text = f.read().strip() | |
| except: | |
| prompt_text = "" | |
| return "example_case/case-1/reference.png", "example_case/case-1/audio.wav", prompt_text | |
| def load_example2(): | |
| try: | |
| with open("example_case/case-2/prompt.txt", "r", encoding="utf-8") as f: | |
| prompt_text = f.read().strip() | |
| except: | |
| prompt_text = "" | |
| return "example_case/case-2/reference.png", "example_case/case-2/audio.wav", prompt_text | |
| def load_example3(): | |
| try: | |
| with open("example_case/case-6/prompt.txt", "r", encoding="utf-8") as f: | |
| prompt_text = f.read().strip() | |
| except: | |
| prompt_text = "" | |
| return "example_case/case-6/reference.png", "example_case/case-6/audio.wav", prompt_text | |
| def load_example4(): | |
| try: | |
| with open("example_case/case-45/prompt.txt", "r", encoding="utf-8") as f: | |
| prompt_text = f.read().strip() | |
| except: | |
| prompt_text = "" | |
| return "example_case/case-45/reference.png", "example_case/case-45/audio.wav", prompt_text | |
| def load_example5(): | |
| try: | |
| with open("example_case/case-3/prompt.txt", "r", encoding="utf-8") as f: | |
| prompt_text = f.read().strip() | |
| except: | |
| prompt_text = "" | |
| return "example_case/case-3/reference.jpg", "example_case/case-3/audio.wav", prompt_text | |
| example1_btn.click(fn=load_example1, outputs=[image_path, audio_path, prompt]) | |
| example2_btn.click(fn=load_example2, outputs=[image_path, audio_path, prompt]) | |
| example3_btn.click(fn=load_example3, outputs=[image_path, audio_path, prompt]) | |
| example4_btn.click(fn=load_example4, outputs=[image_path, audio_path, prompt]) | |
| example5_btn.click(fn=load_example5, outputs=[image_path, audio_path, prompt]) | |
| gr.on( | |
| triggers=[generate_button.click, prompt.submit, negative_prompt.submit], | |
| fn = generate, | |
| inputs = [ | |
| GPU_memory_mode, | |
| teacache_threshold, | |
| num_skip_start_steps, | |
| image_path, | |
| audio_path, | |
| prompt, | |
| negative_prompt, | |
| width, | |
| height, | |
| guidance_scale, | |
| num_inference_steps, | |
| text_guide_scale, | |
| audio_guide_scale, | |
| motion_frame, | |
| fps, | |
| overlap_window_length, | |
| seed_param, | |
| overlapping_weight_scheme, | |
| ], | |
| outputs = [video_output, seed_output, info] | |
| ) | |
| exchange_button.click( | |
| fn=exchange_width_height, | |
| inputs=[width, height], | |
| outputs=[width, height, info] | |
| ) | |
| adjust_button.click( | |
| fn=adjust_width_height, | |
| inputs=[image_path], | |
| outputs=[width, height, info] | |
| ) | |
| extractor_button.click( | |
| fn=audio_extractor, | |
| inputs=[video_path], | |
| outputs=[audio_output, info2, audio_file] | |
| ) | |
| separation_button.click( | |
| fn=vocal_separation, | |
| inputs=[audio_path3], | |
| outputs=[audio_output3, info3, audio_file3] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=int(os.getenv("PORT", 7860)), | |
| share=False, | |
| inbrowser=False, | |
| ) | |