245 lines
11 KiB
Python
245 lines
11 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
|
|
import numpy as np
|
|
import torch
|
|
import os
|
|
from omegaconf import OmegaConf
|
|
from diffusers.hooks import apply_group_offloading
|
|
from .SteadyDancer.generate_dancer import load_models,pre_data_process,inference
|
|
import comfy.utils
|
|
import folder_paths
|
|
from typing_extensions import override
|
|
from comfy_api.latest import ComfyExtension, io
|
|
import nodes
|
|
import comfy.model_management as mm
|
|
from .SteadyDancer.wan import WanI2VDancer
|
|
from .SteadyDancer.wan.modules.vae import WanVAE
|
|
from .SteadyDancer.wan.pipeline_wan_i2v import I2V14B_VAE_CONFIG
|
|
|
|
|
|
MAX_SEED = np.iinfo(np.int32).max
|
|
|
|
device = torch.device(
|
|
"cuda:0") if torch.cuda.is_available() else torch.device(
|
|
"mps") if torch.backends.mps.is_available() else torch.device("cpu")
|
|
|
|
node_cr_path = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
weigths_gguf_current_path = os.path.join(folder_paths.models_dir, "gguf")
|
|
if not os.path.exists(weigths_gguf_current_path):
|
|
os.makedirs(weigths_gguf_current_path)
|
|
folder_paths.add_model_folder_path("gguf", weigths_gguf_current_path) # gguf dir
|
|
|
|
class SteadyDancer_SM_Model(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="SteadyDancer_SM_Model",
|
|
display_name="SteadyDancer_SM_Model",
|
|
category="SteadyDancer_SM",
|
|
inputs=[
|
|
io.Combo.Input("dit",options= ["none"] +folder_paths.get_filename_list("diffusion_models") ),
|
|
io.Combo.Input("gguf",options= ["none"] +folder_paths.get_filename_list("gguf") ),
|
|
],
|
|
outputs=[
|
|
io.Custom("SteadyDancer_SM_Model").Output(display_name="model"),
|
|
],
|
|
)
|
|
@classmethod
|
|
def execute(cls, dit,gguf) -> io.NodeOutput:
|
|
origin_dict={}
|
|
args=OmegaConf.create(origin_dict)
|
|
args.repo=os.path.join(node_cr_path, "SteadyDancer/wan_repo")
|
|
args.config_path=os.path.join(node_cr_path, "SteadyDancer/wan_repo/transformer")
|
|
args.dit_path=folder_paths.get_full_path("diffusion_models", dit) if dit != "none" else None
|
|
args.ckpt_dir=folder_paths.get_full_path("diffusion_models", dit) if dit != "none" else None
|
|
args.gguf_path=folder_paths.get_full_path("gguf", gguf) if gguf != "none" else None
|
|
model=load_models(args)
|
|
return io.NodeOutput(model)
|
|
|
|
|
|
class SteadyDancer_SM_VAE(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
|
|
return io.Schema(
|
|
node_id="SteadyDancer_SM_VAE",
|
|
display_name="SteadyDancer_SM_VAE",
|
|
category="SteadyDancer_SM",
|
|
inputs=[
|
|
io.Combo.Input("vae",options= ["none"] +folder_paths.get_filename_list("vae") ),
|
|
io.Latent.Input("latents",optional=True),
|
|
],
|
|
outputs=[
|
|
io.Vae.Output(),
|
|
io.Image.Output(),
|
|
],
|
|
)
|
|
@classmethod
|
|
def execute(cls, vae, latents=None) -> io.NodeOutput:
|
|
vae_path=folder_paths.get_full_path("vae", vae) if vae != "none" else None
|
|
model = WanVAE(vae_pth=vae_path,device=device)
|
|
if latents is not None:
|
|
import torchvision
|
|
model_device = next(model.model.parameters()).device
|
|
if model_device != device:
|
|
model.model.to(device)
|
|
tensor=model.decode( [i.to(model_device) for i in latents["samples"]])[0][None]
|
|
try:
|
|
print(tensor.shape)
|
|
except:
|
|
pass
|
|
# preprocess
|
|
value_range=(-1, 1)
|
|
tensor = tensor.clamp(min(value_range), max(value_range))
|
|
tensor = torch.stack([
|
|
torchvision.utils.make_grid(
|
|
u, nrow=8, normalize=True, value_range=value_range)
|
|
for u in tensor.unbind(2)
|
|
],dim=1).permute(1, 2, 3, 0)
|
|
tensor = (tensor * 255).type(torch.uint8).cpu()
|
|
try:
|
|
print(tensor.shape)
|
|
except:
|
|
pass
|
|
model.model.to("cpu")
|
|
else:
|
|
tensor=None
|
|
return io.NodeOutput(model,tensor)
|
|
|
|
class SteadyDancer_SM_Cond(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="SteadyDancer_SM_Cond",
|
|
display_name="SteadyDancer_SM_Cond",
|
|
category="SteadyDancer_SM",
|
|
inputs=[
|
|
io.Vae.Input("vae"),
|
|
io.ClipVision.Input("clip_vision"),
|
|
io.Conditioning.Input("positive"),
|
|
io.Conditioning.Input("negative"),
|
|
io.Image.Input("image",),
|
|
io.Image.Input("pose_video",),
|
|
io.Image.Input("pose_neg_lvideo",),
|
|
io.Combo.Input("wh_size",options= ['720*1280', '1280*720', '480*832', '832*480', '576*1024', '1024*576', '1024*800']),
|
|
io.Int.Input("num_frames", default=81, min=25, max=10240,step=4,display_mode=io.NumberDisplay.number),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output("cond"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, vae,clip_vision,positive,negative,image,pose_video,pose_neg_lvideo,wh_size,num_frames) -> io.NodeOutput:
|
|
|
|
width,height = map(int, wh_size.split('*'))
|
|
if image.shape[1]!=height or image.shape[2]!=width:
|
|
image=comfy.utils.common_upscale(image.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
|
|
if pose_video.shape[1]!=height or pose_video.shape[2]!=width:
|
|
pose_video=comfy.utils.common_upscale(pose_video.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
|
|
if pose_neg_lvideo.shape[1]!=height or pose_neg_lvideo.shape[2]!=width:
|
|
pose_neg_lvideo=comfy.utils.common_upscale(pose_neg_lvideo.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
|
|
if pose_video.shape[0]!=num_frames:
|
|
pose_video=pose_video[:num_frames]
|
|
if pose_neg_lvideo.shape[0]!=num_frames:
|
|
pose_neg_lvideo=pose_neg_lvideo[:num_frames]
|
|
num_frames=min(num_frames,pose_video.shape[0],pose_neg_lvideo.shape[0])
|
|
print(f"########## Final num_frames is : {num_frames} #########")
|
|
# pre-process
|
|
image_embeds=clip_vision.encode_image(image)["penultimate_hidden_states"]
|
|
image_embeds_c=clip_vision.encode_image(pose_video[:1,:,:,:3])["penultimate_hidden_states"]
|
|
if isinstance(vae,WanVAE):
|
|
# Offload model to CPU before custom vae encoding
|
|
cf_models=mm.loaded_models()
|
|
for model in cf_models:
|
|
model.unpatch_model(device_to=torch.device("cpu"))
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
cond=pre_data_process(vae,image,pose_video,pose_neg_lvideo,num_frames,width,height,device)
|
|
cond["image_embeds"]=image_embeds
|
|
cond["image_embeds_c"]=image_embeds_c
|
|
cond["prompt_embeds"]=positive[0][0]
|
|
cond["negative_prompt_embeds"]=negative[0][0]
|
|
cond["size"]=(height,width)
|
|
cond["num_frames"]=num_frames
|
|
|
|
# Offload model to CPU after encoding
|
|
cf_models=mm.loaded_models()
|
|
for model in cf_models:
|
|
model.unpatch_model(device_to=torch.device("cpu"))
|
|
mm.soft_empty_cache()
|
|
torch.cuda.empty_cache()
|
|
|
|
return io.NodeOutput (cond)
|
|
|
|
|
|
class SteadyDancer_SM_KSampler(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="SteadyDancer_SM_KSampler",
|
|
display_name="SteadyDancer_SM_KSampler",
|
|
category="SteadyDancer_SM",
|
|
inputs=[
|
|
io.Custom("SteadyDancer_SM_Model").Input("model"),
|
|
io.Conditioning.Input("cond"),
|
|
io.Combo.Input("lora",options= ["none"] + folder_paths.get_filename_list("loras") ),
|
|
io.Int.Input("steps", default=1, min=1, max=10000,display_mode=io.NumberDisplay.number),
|
|
io.Float.Input("sample_shift", default=5.0, min=0, max=10.0,step=0.01,display_mode=io.NumberDisplay.number),
|
|
io.Float.Input("sample_guide_scale", default=1.0, min=0.0, max=100.0,step=0.01,display_mode=io.NumberDisplay.number),
|
|
io.Float.Input("condition_guide_scale", default=1.5, min=0.0, max=10.0,step=0.01,display_mode=io.NumberDisplay.number),
|
|
io.Int.Input("seed", default=0, min=0, max=MAX_SEED),
|
|
io.Boolean.Input("offload_model", default=True),
|
|
io.Int.Input("block_num", default=1, min=1, max=100,display_mode=io.NumberDisplay.number),
|
|
],
|
|
outputs=[
|
|
io.Latent.Output(display_name="latents"),
|
|
],
|
|
)
|
|
@classmethod
|
|
def execute(cls, model,cond,lora,steps,sample_shift,sample_guide_scale,condition_guide_scale,seed,offload_model,block_num ) -> io.NodeOutput:
|
|
lora_path=folder_paths.get_full_path("loras", lora) if lora != "none" else None
|
|
if lora_path is not None and not isinstance(model, WanI2VDancer):
|
|
model.load_lora_weights_(lora_path)
|
|
model.set_lora_adapter(adapter_name="default", weight=1.0)
|
|
max_area=cond["size"]
|
|
# Upscale
|
|
if max_area[0]==832 or max_area[0]==480:
|
|
sample_shift = 3.0
|
|
frame_num=cond["num_frames"]
|
|
if not isinstance(model, WanI2VDancer):
|
|
apply_group_offloading(model.transformer, onload_device=torch.device("cuda"), offload_type="block_level", num_blocks_per_group=block_num)
|
|
videos=inference( model,cond,steps,frame_num,sample_shift,sample_guide_scale,condition_guide_scale,seed,max_area,offload_model)
|
|
videos=videos.unsqueeze(0)
|
|
#print(f"########## Final video shape is : {videos.shape} #########")
|
|
latents_mean = (
|
|
torch.tensor(I2V14B_VAE_CONFIG["latents_mean"]).view(1, 16, 1, 1, 1).to(videos.device, videos.dtype))
|
|
latents_std = 1.0 / torch.tensor(I2V14B_VAE_CONFIG['latents_std']).view(1, 16, 1, 1, 1).to(videos.device, videos.dtype)
|
|
videos = videos / latents_std + latents_mean
|
|
output={}
|
|
output["samples"]=videos
|
|
return io.NodeOutput(output)
|
|
|
|
|
|
from aiohttp import web
|
|
from server import PromptServer
|
|
@PromptServer.instance.routes.get("/SteadyDancer_SM_Extension")
|
|
async def get_hello(request):
|
|
return web.json_response("SteadyDancer_SM_Extension")
|
|
|
|
class SteadyDancer_SM_Extension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
SteadyDancer_SM_Model,
|
|
SteadyDancer_SM_VAE,
|
|
SteadyDancer_SM_Cond,
|
|
SteadyDancer_SM_KSampler,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> SteadyDancer_SM_Extension: # ComfyUI calls this to load your extension and its nodes.
|
|
return SteadyDancer_SM_Extension()
|