179 lines
9.5 KiB
Python
179 lines
9.5 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
import time
|
|
import torch
|
|
import os
|
|
import gc
|
|
import folder_paths
|
|
from .inference.model.turbo_vaed import TurboVAED, get_turbo_vaed
|
|
from .inference.model.vae2_2 import Wan2_2_VAE, get_vae2_2
|
|
from .inference.common import CPUOffloadWrapper, get_arch_memory
|
|
from .inference.model.sa_audio import SAAudioFeatureExtractor
|
|
from .model_loader_utils import nomarl_upscale
|
|
from .inference.pipeline.entry import load_magihuman
|
|
|
|
from .inference.pipeline.video_process import load_audio_and_encode, resample_audio_sinc, resizecrop
|
|
from .inference.pipeline.video_generate import encode_image_w,decode_video_w
|
|
from .inference.pipeline.prompt_process import pad_or_trim
|
|
from .inference.model.t5_gemma.t5_gemma_model import get_t5_gemma_embedding,get_t5_gemma_encoder
|
|
node_cr_path_ = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
|
def load_model( dit,sr_dit,gguf,sr_gguf):
|
|
dit_path=folder_paths.get_full_path("diffusion_models", dit) if dit != "none" else None
|
|
gguf_path=folder_paths.get_full_path("gguf", gguf) if gguf != "none" else None
|
|
sr_dit_path=folder_paths.get_full_path("diffusion_models", sr_dit) if sr_dit != "none" else None
|
|
sr_gguf_path=folder_paths.get_full_path("gguf", sr_gguf) if sr_gguf != "none" else None
|
|
model=load_magihuman(dit_path,gguf_path,sr_dit_path,sr_gguf_path)
|
|
model.infer_mode="sr" if sr_gguf_path is not None or sr_dit_path is not None else "base"
|
|
return model
|
|
|
|
|
|
def load_clip(clip,gguf,device):
|
|
clip_path=folder_paths.get_full_path("clip", clip) if clip != "none" else None
|
|
gguf_path=folder_paths.get_full_path("gguf", gguf) if gguf != "none" else None
|
|
repo=os.path.join(node_cr_path_,"t5gemma-9b-9b-ul2")
|
|
#assert clip_path is not None,"Please provide a clip_path"
|
|
#clip_path="D:/Downloads/t5gemma-9b-9b-ul2"
|
|
text_encoder=get_t5_gemma_encoder(clip_path,gguf_path,device,torch.bfloat16,repo)
|
|
return text_encoder
|
|
|
|
def encoder_text(text_encoder,prompt,negative_prompt,save_emb,target_length=640):
|
|
|
|
with torch.no_grad():
|
|
txt_feat=get_t5_gemma_embedding(prompt, text_encoder)
|
|
txt_feat, original_len=pad_or_trim(txt_feat, target_size=target_length, dim=1)
|
|
txt_feat=txt_feat.to(torch.float32)
|
|
txt_feat_null=get_t5_gemma_embedding(negative_prompt, text_encoder)
|
|
txt_feat_null, original_len_null=pad_or_trim(txt_feat_null, target_size=target_length, dim=1)
|
|
txt_feat_null=txt_feat_null.to(torch.float32)
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
#print(txt_feat.shape,txt_feat_null.shape) #torch.Size([1, 640, 3584]) torch.Size([1, 640, 3584])
|
|
positive=[[txt_feat,{"pooled_output": original_len}]]
|
|
negative=[[txt_feat_null,{"pooled_output": original_len_null}]]
|
|
if save_emb:
|
|
save_lat_emb("embeds",positive,negative)
|
|
return positive, negative
|
|
|
|
|
|
def load_vae(vae,turbo_vae,device,weight_dtype):
|
|
vae_model_path=folder_paths.get_full_path("vae", vae) if vae != "none" else None
|
|
student_ckpt_path=folder_paths.get_full_path("vae", turbo_vae) if turbo_vae != "none" else None
|
|
|
|
if vae_model_path is not None:
|
|
vae = CPUOffloadWrapper(
|
|
get_vae2_2(vae_model_path, device, weight_dtype=weight_dtype), is_cpu_offload=get_arch_memory() <= 48
|
|
)
|
|
vae.model.use_turbo_vae = False
|
|
elif student_ckpt_path is not None:
|
|
student_config_path=os.path.join(node_cr_path_,"example/TurboV3-Wan22-TinyShallow_7_7.json")
|
|
vae = CPUOffloadWrapper(
|
|
get_turbo_vaed(student_config_path, student_ckpt_path, device, weight_dtype=weight_dtype),
|
|
is_cpu_offload=get_arch_memory() <= 48,
|
|
)
|
|
vae.model.use_turbo_vae = True
|
|
else:
|
|
raise ValueError("Please provide a vae_model_path or student_ckpt_path")
|
|
return vae
|
|
|
|
def load_audio_vae(audio_vae,device):
|
|
#vocoder_path=folder_paths.get_full_path("vae", vocoder) if vocoder != "none" else None
|
|
vae_path=folder_paths.get_full_path("vae", audio_vae) if audio_vae != "none" else None
|
|
repo=os.path.join(node_cr_path_,"stable-audio-open")
|
|
assert vae_path is not None,"Please provide a vae_path"
|
|
audio_vae = SAAudioFeatureExtractor(device=device, model_path=vae_path,repo=repo)
|
|
return audio_vae
|
|
|
|
|
|
|
|
def en_decoder_video(vae,latent):
|
|
lat=latent["samples"]
|
|
with torch.no_grad():
|
|
videos=decode_video_w(vae,lat,torch.bfloat16)
|
|
return videos
|
|
|
|
|
|
def get_latents(vae,image,audio_vae,audio,width,height,sr_width,sr_height,device,seconds,):
|
|
if image is not None and vae is not None:
|
|
br_image=nomarl_upscale(image, width, height)
|
|
br_image = encode_image_w(vae,br_image, height, width,device,torch.bfloat16)
|
|
if sr_width>0 and sr_height>0:
|
|
sr_image = nomarl_upscale(image, sr_width, sr_height)
|
|
sr_image = encode_image_w(vae,sr_image, sr_height, sr_width, device,torch.bfloat16)
|
|
else:
|
|
br_image = None
|
|
sr_image = None
|
|
|
|
if audio is not None and audio_vae is not None:
|
|
latent_audio = load_audio_and_encode(audio_vae, audio, seconds)
|
|
else:
|
|
latent_audio=None
|
|
|
|
output={"latent_audio":latent_audio,"br_image":br_image,"sr_image":sr_image,"br_width":width,"br_height":height,"sr_width":sr_width,"sr_height":sr_height,"seconds":seconds}
|
|
return output
|
|
|
|
|
|
def decoder_audio(audio_vae,audio_latents,device):
|
|
latent_audio=audio_latents["samples"].to(torch.bfloat16)
|
|
latent_audio = latent_audio.squeeze(0)
|
|
audio_output = audio_vae.decode(latent_audio.T)
|
|
audio_output_np = audio_output.squeeze(0).T.cpu().float().numpy()
|
|
audio_output_np = resample_audio_sinc(audio_output_np, 441 / 512)
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
audio_output_np=torch.from_numpy(audio_output_np).to(device)
|
|
print(audio_output_np.shape) #torch.Size([442764, 2])
|
|
|
|
return {"waveform": audio_output_np.permute(1,0).contiguous().reshape(1, -1).cpu().float().unsqueeze(0), "sample_rate": audio_vae.sample_rate}
|
|
|
|
|
|
|
|
def read_lat_emb(prefix, positive, negative,device):
|
|
if prefix =="embeds":
|
|
if not os.path.exists(os.path.join(folder_paths.get_output_directory(),"raw_embeds_MagiHuman_sm.pt")):
|
|
raise Exception("No backup prompt embeddings found. Please run MagiHuman_SM_ENCODER node first.")
|
|
else:
|
|
prompt_embeds=torch.load(os.path.join(folder_paths.get_output_directory(),"raw_embeds_MagiHuman_sm.pt"),weights_only=False)
|
|
if os.path.exists(os.path.join(folder_paths.get_output_directory(),"n_raw_embeds_MagiHuman_sm.pt")):
|
|
negative_prompt_embeds=torch.load(os.path.join(folder_paths.get_output_directory(),"n_raw_embeds_MagiHuman_sm.pt"),weights_only=False)
|
|
else:
|
|
negative_prompt_embeds=[[torch.zeros_like(prompt_embeds[0][0]),prompt_embeds[0][1]]]
|
|
#print("Loaded backup prompt embeddings",prompt_embeds[0][0].shape) # Loaded backup prompt embeddings torch.Size([1, 640, 3584])
|
|
positive=[[prompt_embeds[0][0].to(device,torch.bfloat16),prompt_embeds[0][1]]]
|
|
negative=[[negative_prompt_embeds[0][0].to(device,torch.bfloat16),negative_prompt_embeds[0][1]]]
|
|
#print(positive[0][0].shape,negative[0][0].shape) #torch.Size([1, 640, 3584]) torch.Size([1, 640, 3584])
|
|
#print(negative[0][1]["pooled_output"],positive[0][1]["pooled_output"]) #392 114
|
|
return positive,negative
|
|
|
|
elif prefix =="latents":
|
|
if not os.path.exists(os.path.join(folder_paths.get_output_directory(),"raw_latents_MagiHuman_sm.pt")) or not os.path.exists(os.path.join(folder_paths.get_output_directory(),"raw_audio_latents_MagiHuman_sm.pt")):
|
|
raise Exception("No backup latents found. Please run MagiHuman_SM_KSampler node first.")
|
|
else:
|
|
video_latents=torch.load(os.path.join(folder_paths.get_output_directory(),"raw_latents_MagiHuman_sm.pt"),weights_only=False)
|
|
audio_latents=torch.load(os.path.join(folder_paths.get_output_directory(),"raw_audio_latents_MagiHuman_sm.pt"),weights_only=False)
|
|
#print("Loaded backup latents",video_latents.shape,audio_latents.shape) #([1, 128, 11, 16, 24]) [1, 8, 84, 16] torch.Size([1, 84, 128])
|
|
|
|
video_latents["samples"]=video_latents["samples"].to(device,torch.bfloat16)
|
|
audio_lat=audio_latents["samples"].to(device,torch.bfloat16) # [1, 8, 84, 16]
|
|
print(f"audio shape: {audio_lat.shape}") #audio shape: torch.Size([1, 84, 8, 16])
|
|
print(f"video shape: {video_latents['samples'].shape}")
|
|
# batch, frames, combined_dim = audio_lat.shape
|
|
# reshaped = audio_lat.view(batch, frames, 8, 16)
|
|
#audio_lat = audio_lat.permute(0, 2, 1, 3)
|
|
|
|
audio_latents["samples"]=audio_lat
|
|
return video_latents, audio_latents
|
|
|
|
def save_lat_emb(save_prefix,data1,data2,mode=""):
|
|
data1_prefix, data2_prefix = ("raw_embeds_MagiHuman", "n_raw_embeds_MagiHuman") if save_prefix == "embeds" else ("raw_latents_MagiHuman", "raw_audio_latents_MagiHuman")
|
|
default_data1_path = os.path.join(folder_paths.get_output_directory(),f"{data1_prefix}_sm.pt")
|
|
default_data2_path = os.path.join(folder_paths.get_output_directory(),f"{data2_prefix}_sm.pt")
|
|
prefix = mode+str(int(time.time()))
|
|
if os.path.exists(default_data1_path): # use a different path if the file already exists
|
|
default_data1_path=os.path.join(folder_paths.get_output_directory(),f"{data1_prefix}_sm_{prefix}.pt")
|
|
torch.save(data1,default_data1_path)
|
|
if data2 is not None:
|
|
if os.path.exists(default_data2_path):
|
|
default_data2_path=os.path.join(folder_paths.get_output_directory(),f"{data2_prefix}_sm_{prefix}.pt")
|
|
torch.save(data2,default_data2_path) |