250 lines
9.5 KiB
Python
250 lines
9.5 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
import os
|
|
import torch
|
|
import gc
|
|
import numpy as np
|
|
from omegaconf import OmegaConf
|
|
from diffusers import AutoencoderKLTemporalDecoder
|
|
from diffusers.schedulers import EulerDiscreteScheduler
|
|
from transformers import WhisperModel, AutoFeatureExtractor
|
|
import random
|
|
import io
|
|
import torchaudio
|
|
from .src.models.base.unet_spatio_temporal_condition import UNetSpatioTemporalConditionModel
|
|
from .sonic import Sonic,sonic_predata,preprocess_face,crop_face_image
|
|
from .src.dataset.test_preprocess import image_audio_to_tensor
|
|
from .src.models.audio_adapter.audio_proj import AudioProjModel
|
|
from .src.models.audio_adapter.audio_to_bucket import Audio2bucketModel
|
|
from .node_utils import tensor2cv,cv2pil
|
|
from .src.dataset.face_align.align import AlignImage
|
|
|
|
import folder_paths
|
|
|
|
|
|
MAX_SEED = np.iinfo(np.int32).max
|
|
current_node_path = os.path.dirname(os.path.abspath(__file__))
|
|
device = torch.device(
|
|
"cuda") if torch.cuda.is_available() else torch.device("cpu")
|
|
|
|
# add checkpoints dir
|
|
SONIC_weigths_path = os.path.join(folder_paths.models_dir, "sonic")
|
|
if not os.path.exists(SONIC_weigths_path):
|
|
os.makedirs(SONIC_weigths_path)
|
|
folder_paths.add_model_folder_path("sonic", SONIC_weigths_path)
|
|
|
|
|
|
|
|
class SONICLoader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"svd_repo": ("STRING", {"default": "F:/test/ComfyUI/models/diffusers/stabilityai/stable-video-diffusion-img2vid-xt-1-1"}),
|
|
"sonic_unet": (["none"] + folder_paths.get_filename_list("sonic"),),
|
|
"ip_audio_scale": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 2.0,"step": 0.1}),
|
|
"use_interframe":("BOOLEAN", {"default": True},),
|
|
"dtype":(["fp16", "fp32","bf16"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL_SONIC",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "loader_main"
|
|
CATEGORY = "SONIC"
|
|
|
|
def loader_main(self,svd_repo,sonic_unet,ip_audio_scale,use_interframe,dtype):
|
|
|
|
if dtype == "fp16":
|
|
weight_dtype = torch.float16
|
|
elif dtype == "fp32":
|
|
weight_dtype = torch.float32
|
|
elif dtype == "bf16":
|
|
weight_dtype = torch.bfloat16
|
|
else:
|
|
raise ValueError(
|
|
f"Do not support weight dtype: {dtype} during training"
|
|
)
|
|
|
|
# check model is exits or not,if not auto downlaod
|
|
flownet_ckpt = os.path.join(SONIC_weigths_path,"RIFE")
|
|
|
|
if sonic_unet!="none":
|
|
sonic_unet=folder_paths.get_full_path("sonic", sonic_unet)
|
|
|
|
# load model
|
|
print("***********Load model ***********")
|
|
vae = AutoencoderKLTemporalDecoder.from_pretrained(
|
|
svd_repo,
|
|
subfolder="vae",
|
|
variant="fp16")
|
|
|
|
val_noise_scheduler = EulerDiscreteScheduler.from_pretrained(
|
|
svd_repo,
|
|
subfolder="scheduler")
|
|
|
|
unet = UNetSpatioTemporalConditionModel.from_pretrained(
|
|
svd_repo,
|
|
subfolder="unet",
|
|
variant="fp16")
|
|
pipe = Sonic(device,weight_dtype,vae,val_noise_scheduler,unet,flownet_ckpt,sonic_unet,
|
|
use_interframe,ip_audio_scale)
|
|
|
|
print("***********Load model done ***********")
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return (pipe,)
|
|
|
|
class SONIC_PreData:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"clip_vision": ("CLIP_VISION",),
|
|
"audio": ("AUDIO",),
|
|
"image": ("IMAGE",),
|
|
"min_resolution": ("INT", {"default": 512, "min": 128, "max": 2048, "step": 64, "display": "number"}),
|
|
"expand_ratio": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0,"step": 0.1}),
|
|
}}
|
|
|
|
|
|
|
|
RETURN_TYPES = ("SONIC_PREDATA",)
|
|
RETURN_NAMES = ("data_dict", )
|
|
FUNCTION = "sampler_main"
|
|
CATEGORY = "SONIC"
|
|
|
|
def sampler_main(self, clip_vision,audio,image,min_resolution,expand_ratio):
|
|
config_file = os.path.join(current_node_path, 'config/inference/sonic.yaml')
|
|
config = OmegaConf.load(config_file)
|
|
|
|
audio2token_ckpt = os.path.join(SONIC_weigths_path,"audio2token.pth")
|
|
audio2bucket_ckpt = os.path.join(SONIC_weigths_path,"audio2bucket.pth")
|
|
yolo_ckpt = os.path.join(SONIC_weigths_path,"yoloface_v5m.pt")
|
|
|
|
if not os.path.exists(audio2bucket_ckpt) or not os.path.exists(audio2token_ckpt) or not os.path.exists(yolo_ckpt):
|
|
raise Exception("Please download the model first")
|
|
#init model
|
|
whisper_repo=os.path.join(SONIC_weigths_path,"whisper-tiny")
|
|
|
|
whisper = WhisperModel.from_pretrained(whisper_repo).to(device).eval()
|
|
whisper.requires_grad_(False)
|
|
|
|
feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_repo)
|
|
|
|
audio2token = AudioProjModel(seq_len=10, blocks=5, channels=384, intermediate_dim=1024, output_dim=1024, context_tokens=32).to(device)
|
|
audio2bucket = Audio2bucketModel(seq_len=50, blocks=1, channels=384, clip_channels=1024, intermediate_dim=1024, output_dim=1, context_tokens=2).to(device)
|
|
|
|
audio2token_dict = torch.load(audio2token_ckpt, map_location="cpu")
|
|
audio2bucket_dict = torch.load(audio2bucket_ckpt, map_location="cpu")
|
|
audio2token.load_state_dict(
|
|
audio2token_dict,
|
|
strict=True,
|
|
)
|
|
|
|
audio2bucket.load_state_dict(
|
|
audio2bucket_dict,
|
|
strict=True,
|
|
)
|
|
del audio2token_dict,audio2bucket_dict
|
|
|
|
audio_file_prefix = ''.join(random.choice("0123456789") for _ in range(6))
|
|
audio_path = os.path.join(folder_paths.get_input_directory(), f"audio_{audio_file_prefix}_temp.wav")
|
|
|
|
# 减少音频数据传递导致的不必要文件存储
|
|
buff = io.BytesIO()
|
|
torchaudio.save(buff, audio["waveform"].squeeze(0), audio["sample_rate"], format="FLAC")
|
|
with open(audio_path, 'wb') as f:
|
|
f.write(buff.getbuffer())
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
face_det = AlignImage(device, det_path=yolo_ckpt)
|
|
|
|
# 先面部裁切处理
|
|
cv_image=tensor2cv(image)
|
|
face_info =preprocess_face(cv_image,face_det, expand_ratio=expand_ratio)
|
|
if face_info['face_num'] > 0:
|
|
crop_image_pil=cv2pil(crop_face_image(cv_image, face_info['crop_bbox']))
|
|
|
|
test_data = image_audio_to_tensor(face_det, feature_extractor, crop_image_pil, audio_path, limit=config.frame_num, image_size=min_resolution, area=config.area)
|
|
|
|
step=2
|
|
for k, v in test_data.items():
|
|
if isinstance(v, torch.Tensor):
|
|
test_data[k] = v.unsqueeze(0).to(device).float()
|
|
ref_img = test_data['ref_img']
|
|
audio_feature = test_data['audio_feature']
|
|
audio_len = test_data['audio_len']
|
|
|
|
ref_tensor_list,audio_tensor_list,uncond_audio_tensor_list,motion_buckets,image_embeddings=sonic_predata(whisper,audio_feature,audio_len,step,audio2bucket,clip_vision,audio2token,ref_img,image,device)
|
|
del clip_vision,face_det,whisper
|
|
audio2bucket.to("cpu")
|
|
audio2token.to("cpu")
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return ({"test_data":test_data,"ref_tensor_list":ref_tensor_list,"config":config,"image_embeddings":image_embeddings,
|
|
"audio_tensor_list":audio_tensor_list,"uncond_audio_tensor_list":uncond_audio_tensor_list,"motion_buckets":motion_buckets},)
|
|
|
|
|
|
|
|
|
|
class SONICSampler:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL_SONIC",),
|
|
"data_dict": ("SONIC_PREDATA",), #{}
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED}),
|
|
"inference_steps": ("INT", {"default": 25, "min": 1, "max": 1024, "step": 1, "display": "number"}),
|
|
"dynamic_scale": ("FLOAT", {"default": 1.0, "min": 0.5, "max": 2.0,"step": 0.1}),
|
|
"fps": ("FLOAT", {"default": 24.0, "min": 5.0, "max": 120.0,"step": 0.5}),
|
|
}}
|
|
|
|
|
|
|
|
RETURN_TYPES = ("IMAGE","FLOAT")
|
|
RETURN_NAMES = ("image","fps" )
|
|
FUNCTION = "sampler_main"
|
|
CATEGORY = "SONIC"
|
|
|
|
def sampler_main(self, model,data_dict,seed,inference_steps,dynamic_scale,fps):
|
|
print("***********Start infer ***********")
|
|
iamge= model.process(data_dict["audio_tensor_list"],
|
|
data_dict["uncond_audio_tensor_list"],
|
|
data_dict["motion_buckets"],
|
|
data_dict["test_data"],
|
|
data_dict["config"],
|
|
image_embeds=data_dict["image_embeddings"],
|
|
fps=fps,
|
|
inference_steps=inference_steps,
|
|
dynamic_scale=dynamic_scale,
|
|
seed=seed)
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return (iamge.permute(0, 2, 3, 4, 1).squeeze(0),fps)
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SONICTLoader":SONICLoader,
|
|
"SONIC_PreData":SONIC_PreData,
|
|
"SONICSampler":SONICSampler,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SONICTLoader":"SONICTLoader",
|
|
"SONIC_PreData":"SONIC_PreData",
|
|
"SONICSampler":"SONICSampler",
|
|
}
|