105 lines
4.1 KiB
Python
105 lines
4.1 KiB
Python
import folder_paths
|
|
import os
|
|
import comfy.model_management as mm
|
|
import time
|
|
import torchaudio
|
|
import torchvision.utils as vutils
|
|
|
|
from .generate import InferenceAgent
|
|
from .options.base_options import BaseOptionsJson
|
|
|
|
class LoadFloatModels:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": (['float.pth'],)
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT_PIPE",)
|
|
RETURN_NAMES = ("float_pipe",)
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "FLOAT"
|
|
DESCRIPTION = "Models are auto-downloaded to /ComfyUI/models/float"
|
|
|
|
def loadmodel(self, model):
|
|
# download models if not exist
|
|
float_models_dir = os.path.join(folder_paths.models_dir, "float")
|
|
os.makedirs(float_models_dir, exist_ok=True)
|
|
|
|
wav2vec2_base_960h_models_dir = os.path.join(float_models_dir,"wav2vec2-base-960h")
|
|
wav2vec_english_speech_emotion_recognition_models_dir = os.path.join(float_models_dir,"wav2vec-english-speech-emotion-recognition")
|
|
float_model_path = os.path.join(float_models_dir,"float.pth")
|
|
|
|
if not os.path.exists(float_model_path) or not os.path.isdir(wav2vec2_base_960h_models_dir) or not os.path.isdir(wav2vec_english_speech_emotion_recognition_models_dir):
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(repo_id="yuvraj108c/float", local_dir=float_models_dir, local_dir_use_symlinks=False)
|
|
|
|
# use custom dictionary instead of original parser for arguments
|
|
opt = BaseOptionsJson
|
|
opt.rank, opt.ngpus = 0,1
|
|
opt.ckpt_path = float_model_path
|
|
opt.pretrained_dir = float_models_dir
|
|
opt.wav2vec_model_path = wav2vec2_base_960h_models_dir
|
|
opt.audio2emotion_path = wav2vec_english_speech_emotion_recognition_models_dir
|
|
agent = InferenceAgent(opt)
|
|
|
|
return (agent,)
|
|
|
|
class FloatProcess:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"ref_image": ("IMAGE",),
|
|
"ref_audio": ("AUDIO",),
|
|
"float_pipe": ("FLOAT_PIPE",),
|
|
"a_cfg_scale": ("FLOAT", {"default": 2.0,"min": 1.0, "step": 0.1}),
|
|
"r_cfg_scale": ("FLOAT", {"default": 1.0,"min": 1.0, "step": 0.1}),
|
|
"e_cfg_scale": ("FLOAT", {"default": 1.0,"min": 1.0, "step": 0.1}),
|
|
"fps": ("FLOAT", {"default": 25, "step": 1}),
|
|
"emotion": (['none', 'angry', 'disgust', 'fear', 'happy', 'neutral', 'sad', 'surprise'], {"default": "none"}),
|
|
"crop": ("BOOLEAN",{"default":False},),
|
|
"seed": ("INT", {"default": 62064758300528, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "floatprocess"
|
|
CATEGORY = "FLOAT"
|
|
DESCRIPTION = "Float Processing"
|
|
|
|
def floatprocess(self, ref_image, ref_audio, float_pipe, a_cfg_scale, r_cfg_scale, e_cfg_scale, fps, emotion, crop, seed):
|
|
# save audio
|
|
temp_dir = folder_paths.get_temp_directory()
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
audio_save_path = os.path.join(temp_dir, f"{int(time.time())}.wav")
|
|
torchaudio.save(audio_save_path, ref_audio['waveform'].squeeze(0), ref_audio["sample_rate"])
|
|
|
|
# save image
|
|
if ref_image.shape[0] != 1:
|
|
raise Exception("Only a single image is supported.")
|
|
ref_image_bchw = ref_image.permute(0, 3, 1, 2)
|
|
image_save_path = os.path.join(temp_dir, f"{int(time.time())}.png")
|
|
vutils.save_image(ref_image_bchw[0], image_save_path)
|
|
|
|
float_pipe.G.to(float_pipe.rank)
|
|
|
|
float_pipe.opt.fps = fps
|
|
images_bhwc = float_pipe.run_inference(
|
|
None,
|
|
image_save_path,
|
|
audio_save_path,
|
|
a_cfg_scale = a_cfg_scale,
|
|
r_cfg_scale = r_cfg_scale,
|
|
e_cfg_scale = e_cfg_scale,
|
|
emo = None if emotion == "none" else emotion,
|
|
no_crop = not crop,
|
|
seed = seed
|
|
)
|
|
|
|
float_pipe.G.to(mm.unet_offload_device())
|
|
|
|
return (images_bhwc,) |