248 lines
11 KiB
Python
248 lines
11 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
import io
|
|
import os
|
|
import random
|
|
import numpy as np
|
|
import torch
|
|
import torchaudio
|
|
import gc
|
|
import platform
|
|
import subprocess
|
|
from omegaconf import OmegaConf
|
|
|
|
from .node_utils import nomarl_upscale,is_directory_with_files
|
|
from .inference import load_StableAvatar_model,pre_data_process,infer_StableAvatar
|
|
|
|
import folder_paths
|
|
|
|
|
|
MAX_SEED = np.iinfo(np.int32).max
|
|
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
|
|
|
# pre dir
|
|
current_path = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
weigths_current_path = os.path.join(folder_paths.models_dir, "StableAvatar")
|
|
if not os.path.exists(weigths_current_path):
|
|
os.makedirs(weigths_current_path)
|
|
|
|
weigths_DIT_current_path = os.path.join(weigths_current_path, "transformer")
|
|
if not os.path.exists(weigths_DIT_current_path):
|
|
os.makedirs(weigths_DIT_current_path)
|
|
|
|
weigths_audio_current_path = os.path.join(weigths_current_path, "wav2vec2-base-960h")
|
|
if not os.path.exists(weigths_audio_current_path):
|
|
os.makedirs(weigths_audio_current_path)
|
|
|
|
|
|
# ffmpeg
|
|
ffmpeg_path = os.getenv('FFMPEG_PATH')
|
|
if ffmpeg_path is None and platform.system() in ['Linux', 'Darwin']:
|
|
try:
|
|
result = subprocess.run(['which', 'ffmpeg'], capture_output=True, text=True)
|
|
if result.returncode == 0:
|
|
ffmpeg_path = result.stdout.strip()
|
|
print(f"FFmpeg is installed at: {ffmpeg_path}")
|
|
else:
|
|
print("FFmpeg is not installed. Please download ffmpeg-static and export to FFMPEG_PATH.")
|
|
print("For example: export FFMPEG_PATH=/musetalk/ffmpeg-4.4-amd64-static")
|
|
except Exception as e:
|
|
pass
|
|
|
|
if ffmpeg_path is not None and ffmpeg_path not in os.getenv('PATH'):
|
|
print("Adding FFMPEG_PATH to PATH")
|
|
os.environ["PATH"] = f"{ffmpeg_path}:{os.environ['PATH']}"
|
|
|
|
|
|
# *****************main***************
|
|
|
|
class StableAvatar_LoadModel:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"transformer": (folder_paths.get_filename_list("diffusion_models"),),
|
|
"vae": (folder_paths.get_filename_list("vae"),),
|
|
"lora": (["None"]+folder_paths.get_filename_list("loras"),),
|
|
"enable_teacache": ("BOOLEAN", {"default": False},),
|
|
"use_mmgp": (["LowRAM_LowVRAM","None", "VerylowRAM_LowVRAM","LowRAM_HighVRAM","HighRAM_LowVRAM","HighRAM_HighVRAM" ],),
|
|
"GPU_memory_mode": (["None","model_cpu_offload_and_qfloat8", "model_cpu_offload","sequential_cpu_offload", ],),
|
|
"weight_dtype": (["bfloat16", "float16", "float32"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL_PIPE_SA", "MODEL_INFO_SA")
|
|
RETURN_NAMES = ("model", "info")
|
|
FUNCTION = "main_loader"
|
|
CATEGORY = "StableAvatar"
|
|
|
|
def main_loader(self, transformer,vae,lora,enable_teacache,use_mmgp,GPU_memory_mode,weight_dtype):
|
|
|
|
vae_path=folder_paths.get_full_path( "vae",vae)
|
|
transformer_path=folder_paths.get_full_path("diffusion_models",transformer)
|
|
|
|
# 复用模型
|
|
echo_audio_files=os.path.join(folder_paths.models_dir,"echo_mimic/wav2vec2-base-960h")
|
|
pretrained_wav2vec_path= echo_audio_files if is_directory_with_files(echo_audio_files) else weigths_audio_current_path
|
|
|
|
echo_dit_files=os.path.join(folder_paths.models_dir,"echo_mimic/transformer")
|
|
pretrained_dit_path= echo_dit_files if is_directory_with_files(echo_dit_files) else weigths_DIT_current_path
|
|
|
|
config = OmegaConf.load(os.path.join(current_path, "StableAvatar/deepspeed_config/wan2.1/wan_civitai.yaml"))
|
|
args = {"transformer_path": transformer_path,"pretrained_dit_path":pretrained_dit_path,"teacache_offload":True,"sample_steps":25,
|
|
"enable_teacache":enable_teacache,"ulysses_degree":1,"weight_dtype":weight_dtype,"t5_fsdp":True, "t5_cpu":False,"fsdp_dit":True,"ring_degree":1,"num_skip_start_steps":5,
|
|
"pretrained_model_name_or_path":os.path.join(current_path, "Wan2.1-Fun-V1.1-1.3B-InP"),"GPU_memory_mode":GPU_memory_mode,"teacache_threshold":0.10,"local_rank":1,
|
|
"pretrained_wav2vec_path":pretrained_wav2vec_path, "temporal_compression_ratio":4,"input_perturbation":0,
|
|
}
|
|
args=OmegaConf.create(args)
|
|
|
|
if weight_dtype == "bfloat16":
|
|
weight_dtype_ = torch.bfloat16
|
|
elif weight_dtype == "float16":
|
|
weight_dtype_ = torch.float16
|
|
elif weight_dtype == "float32":
|
|
weight_dtype_ = torch.float32
|
|
|
|
lora_path=folder_paths.get_full_path( "loras",lora) if "None"!=lora else None
|
|
|
|
model, tokenizer,temporal_compression_ratio,sampler_name= load_StableAvatar_model(args,vae_path, config, device,weight_dtype_,use_mmgp,lora_path)
|
|
args.temporal_compression_ratio=temporal_compression_ratio
|
|
|
|
info={"args":args,"tokenizer":tokenizer,"weight_dtype":weight_dtype_,"sampler_name":sampler_name}
|
|
|
|
print("##### model loaded #####")
|
|
|
|
return (model,info)
|
|
|
|
|
|
class StableAvatar_Predata:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"info": ("MODEL_INFO_SA",),
|
|
"clip": ("CLIP",),
|
|
"clip_vision": ("CLIP_VISION",),
|
|
"image": ("IMAGE",), # [B,H,W,C], C=3
|
|
"audio": ("AUDIO",),
|
|
"prompt": ("STRING", {"multiline": True,"default":"A middle-aged woman with short light brown hair, wearing pearl earrings and a blue blazer,"
|
|
" is speaking passionately in front of a blurred background resembling a government building. Her mouth is open mid-phrase, her expression is"
|
|
" engaged and energetic, and the lighting is bright and even, suggesting a television interview or live broadcast. The scene gives the impression she is singing with conviction and purpose."}),
|
|
"negative_prompt" :("STRING", {"multiline": True,"default":"色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" }),
|
|
"width": ("INT", {"default": 480, "min": 128, "max": 2048, "step": 16, "display": "number"}),
|
|
"height": ("INT", {"default": 720, "min": 128, "max": 2048, "step": 16, "display": "number"}),
|
|
"fps": ("FLOAT", {"default": 25.0, "min": 5.0, "max": 120.0}),
|
|
"duration": ("INT", {"default": 5, "min": 0, "max": 3600, "step": 1, "display": "number"}),
|
|
"audio_separator": ("BOOLEAN", {"default": True},),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL_EMB_E", )
|
|
RETURN_NAMES = ("emb",)
|
|
FUNCTION = "main_loader"
|
|
CATEGORY = "StableAvatar"
|
|
|
|
def main_loader(self,info,clip,clip_vision, image,audio,prompt,negative_prompt,width,height,fps,duration,audio_separator):
|
|
|
|
args=info.get("args")
|
|
weight_dtype=info.get("weight_dtype")
|
|
args.clip_sample_n_frames=81
|
|
|
|
# pre img
|
|
infer_img = nomarl_upscale(image, width, height)
|
|
|
|
#pre audio
|
|
audio_file_prefix = ''.join(random.choice("0123456789") for _ in range(6))
|
|
audio_file = os.path.join(folder_paths.get_input_directory(), f"audio_{audio_file_prefix}_temp.wav")
|
|
|
|
waveform = audio["waveform"].squeeze(0)
|
|
sample_rate = audio["sample_rate"]
|
|
if duration > 0:
|
|
max_samples = int(duration * sample_rate)
|
|
if waveform.shape[1] > max_samples:
|
|
waveform = waveform[:, :max_samples]
|
|
|
|
buff = io.BytesIO()
|
|
torchaudio.save(buff, waveform, sample_rate, format="FLAC")
|
|
with open(audio_file, 'wb') as f:
|
|
f.write(buff.getbuffer())
|
|
|
|
# get vocal if infer a song
|
|
if audio_separator:
|
|
from .vocal_seperator import separator_audio
|
|
audio_separator_model_file=os.path.join(weigths_current_path,"Kim_Vocal_2.onnx")
|
|
validation_driven_audio_path=separator_audio(audio_file, audio_separator_model_file,folder_paths.get_temp_directory(),audio_file_prefix)
|
|
else:
|
|
validation_driven_audio_path=audio_file
|
|
args.validation_driven_audio_path=validation_driven_audio_path
|
|
|
|
# get emb
|
|
emb = pre_data_process(clip,clip_vision,info.get("tokenizer"),prompt,negative_prompt,infer_img,device,width,height,args,weight_dtype)
|
|
emb.update({"audio_file_prefix":audio_file_prefix,"fps":fps,"height":height,"width":width,"enable_teacache": args.enable_teacache,
|
|
"num_skip_start_steps":args.num_skip_start_steps,
|
|
"teacache_offload":args.teacache_offload,
|
|
"teacache_threshold":args.teacache_threshold,
|
|
"sample_steps":args.sample_steps,
|
|
"pretrained_model_name_or_path":args.pretrained_model_name_or_path,
|
|
"weight_dtype":weight_dtype,
|
|
"sampler_name":info.get("sampler_name"),
|
|
})
|
|
|
|
return (emb,)
|
|
|
|
|
|
class StableAvatar_Sampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL_PIPE_SA",),
|
|
"emb": ("MODEL_EMB_E",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}),
|
|
"cfg": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 10.0, "step": 0.1, "round": 0.01}),
|
|
"sample_text_guide_scale": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 10.0, "step": 0.1, "round": 0.01}),
|
|
"sample_audio_guide_scale": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 10.0, "step": 0.1, "round": 0.01}),
|
|
"steps": ("INT", {"default": 25, "min": 1, "max": 100}),
|
|
"overlap_window_length": ("INT", {"default": 5, "min": 5, "max": 15,"step": 1,}),
|
|
"overlapping_weight_scheme": (["uniform", "log",],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "FLOAT",)
|
|
RETURN_NAMES = ("latent", "frame_rate",)
|
|
FUNCTION = "em_main"
|
|
CATEGORY = "StableAvatar"
|
|
|
|
def em_main(self,model, emb, seed, cfg, sample_text_guide_scale,sample_audio_guide_scale,steps, overlap_window_length,overlapping_weight_scheme):
|
|
|
|
emb["motion_frame"]=25
|
|
emb["sample_steps"]=steps
|
|
frame_rate = float(emb.get("fps"))
|
|
samples = infer_StableAvatar(model, emb,seed, cfg,device, steps,frame_rate, sample_text_guide_scale,sample_audio_guide_scale, overlap_window_length,emb.get("weight_dtype",torch.bfloat16),overlapping_weight_scheme)
|
|
#print("samples.shape:",samples.shape)
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return ({"samples":samples}, frame_rate,)
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"StableAvatar_LoadModel": StableAvatar_LoadModel,
|
|
"StableAvatar_Predata":StableAvatar_Predata,
|
|
"StableAvatar_Sampler": StableAvatar_Sampler,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"StableAvatar_LoadModel": "StableAvatar_LoadModel",
|
|
"StableAvatar_Predata": "StableAvatar_Predata",
|
|
"StableAvatar_Sampler": "StableAvatar_Sampler",
|
|
}
|