nodes/smthemex-ComfyUI_Streamv2v_Plus
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update streamv2v_node.py

Browse Source
This commit is contained in:
smthemex
2024-07-28 12:13:20 +08:00
committed by GitHub
parent 1ccf08f882
commit ba535c91dd
1 changed files with 99 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
streamv2v_node.py
+99 -82
View File
@@ -13,26 +13,30 @@ import folder_paths
from comfy.utils import common_upscale
import cv2
import diffusers
dif_version = str(diffusers.__version__)
dif_version_int = int(dif_version.split(".")[1])
from diffusers import (DDIMScheduler,
KDPM2AncestralDiscreteScheduler, LMSDiscreteScheduler, DPMSolverMultistepScheduler,
DPMSolverSinglestepScheduler,AutoencoderTiny,
DPMSolverSinglestepScheduler, AutoencoderTiny,
EulerDiscreteScheduler, HeunDiscreteScheduler, KDPM2DiscreteScheduler,
EulerAncestralDiscreteScheduler, UniPCMultistepScheduler,DiffusionPipeline,
EulerAncestralDiscreteScheduler, UniPCMultistepScheduler, DiffusionPipeline,
DDPMScheduler, LCMScheduler, StableDiffusionPipeline, StableDiffusionXLPipeline, AutoencoderKL, )
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
path_dir = os.path.dirname(CURRENT_DIR)
file_path = os.path.dirname(path_dir)
BIGMAX = (2 ** 53 - 1)
input_path = folder_paths.get_input_directory()
def tensor_to_image(tensor):
# tensor = tensor.cpu()
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
image = Image.fromarray(image_np, mode='RGB')
return image
def get_video_img(tensor):
outputs = []
for x in tensor:
@@ -40,10 +44,12 @@ def get_video_img(tensor):
outputs.append(x)
return outputs
def phi2narry(img):
img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
return img
def narry_list(list_in):
for i in range(len(list_in)):
value = list_in[i]
@@ -51,6 +57,7 @@ def narry_list(list_in):
list_in[i] = modified_value
return list_in
def nomarl_upscale(img_tensor, width, height):
samples = img_tensor.movedim(-1, 1)
img = common_upscale(samples, width, height, "nearest-exact", "center")
@@ -58,12 +65,14 @@ def nomarl_upscale(img_tensor, width, height):
img_pil = tensor_to_image(samples)
return img_pil
def tensor_upscale(img_tensor, width, height):
samples = img_tensor.movedim(-1, 1)
img = common_upscale(samples, width, height, "nearest-exact", "center")
samples = img.movedim(1, -1)
return samples
def get_local_path(file_path, model_path):
path = os.path.join(file_path, "models", "diffusers", model_path)
model_path = os.path.normpath(path)
@@ -94,25 +103,27 @@ def instance_path(path, repo):
repo = get_instance_path(model_path)
return repo
scheduler_list = ["LCM","DDIM",
"Euler",
"Euler a",
"DDPM",
"DPM++ 2M",
"DPM++ 2M Karras",
"DPM++ 2M SDE",
"DPM++ 2M SDE Karras",
"DPM++ SDE",
"DPM++ SDE Karras",
"DPM2",
"DPM2 Karras",
"DPM2 a",
"DPM2 a Karras",
"Heun",
"LMS",
"LMS Karras",
"UniPC",
]
scheduler_list = ["LCM", "DDIM",
"Euler",
"Euler a",
"DDPM",
"DPM++ 2M",
"DPM++ 2M Karras",
"DPM++ 2M SDE",
"DPM++ 2M SDE Karras",
"DPM++ SDE",
"DPM++ SDE Karras",
"DPM2",
"DPM2 Karras",
"DPM2 a",
"DPM2 a Karras",
"Heun",
"LMS",
"LMS Karras",
"UniPC",
]
def get_sheduler(name):
scheduler = False
@@ -156,10 +167,11 @@ def get_sheduler(name):
scheduler = UniPCMultistepScheduler()
return scheduler
class Stream_Model_Loader:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
@@ -172,12 +184,12 @@ class Stream_Model_Loader:
}
}
RETURN_TYPES = ("MODEL","STRING")
RETURN_NAMES = ("pipe","info")
RETURN_TYPES = ("MODEL", "STRING")
RETURN_NAMES = ("pipe", "info")
FUNCTION = "main"
CATEGORY = "StreamV2V_Plus"
def main(self,ckpt_name,scheduler,use_sdxl,vae_id,lcm_lora,):
def main(self, ckpt_name, scheduler, use_sdxl, vae_id, lcm_lora, ):
# 加载pipe
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
scheduler_used = get_sheduler(scheduler)
@@ -200,49 +212,50 @@ class Stream_Model_Loader:
pipe = StableDiffusionPipeline.from_single_file(
ckpt_path, original_config_file=original_config_file, torch_dtype=torch.float16).to("cuda")
info = "sd15"
if vae_id!="":
if vae_id != "":
pipe.vae = AutoencoderKL.from_pretrained(vae_id, torch_dtype=torch.float16).to("cuda")
pipe.scheduler = scheduler_used.from_config(pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()
pipe.enable_vae_tiling()
lcm_lora_path = get_lora(lcm_lora)
pipe.load_lora_weights(lcm_lora_path, adapter_name="lcm", )
return (pipe,info,)
info = ";".join([info, ckpt_name])
return (pipe, info,)
class Stream_Lora_Loader:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"pipe":("MODEL",),
"pipe": ("MODEL",),
"info": ("STRING", {"forceInput": True, "default": ""}),
"lora": (["none"] + folder_paths.get_filename_list("loras"),),
"lora_scale": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.0, "step": 0.1}),
"trigger_word": ("STRING", {"default": "best quality"}),
}
}
RETURN_TYPES = ("MODEL","STRING",)
RETURN_NAMES = ("pipe","info",)
RETURN_TYPES = ("MODEL", "STRING",)
RETURN_NAMES = ("pipe", "info",)
FUNCTION = "main"
CATEGORY = "StreamV2V_Plus"
def main(self, pipe,info,lora,lora_scale,trigger_word):
def main(self, pipe, info, lora, lora_scale, trigger_word):
info, ckpt_name = info.split(";")
lora_path = get_lora(lora)
if lora == "none" :
if lora == "none":
raise "you need a style lora"
pipe.load_lora_weights(lora_path, adapter_name=trigger_word, )
pipe.set_adapters(adapter_names=["lcm", trigger_word], adapter_weights=[1.0, lora_scale])
pipe.fuse_lora()
print(f"Use style lora: {trigger_word} in {lora}")
info=";".join([info,trigger_word])
return (pipe,info,)
info = ";".join([info, trigger_word, ckpt_name])
return (pipe, info,)
class Stream_Sampler:
@@ -255,12 +268,12 @@ class Stream_Sampler:
os.path.isfile(os.path.join(input_path, f)) and f.split('.')[-1] in ["mp4", "webm", "mkv",
"avi"]]
return {
"required": {
"pipe":("MODEL",),
"info":("STRING",{"forceInput": True, "default": ""}),
"prompt": ("STRING", {"multiline": True, "default": "Claymation, a man is giving a talk"}),
"video": (["none"]+video_files,),
"sampler_type": (["txt2img", "vdieo2vdieo","WebCam2Video"],),
"required": {
"pipe": ("MODEL",),
"info": ("STRING", {"forceInput": True, "default": ""}),
"prompt": ("STRING", {"multiline": True, "default": "Claymation, a man is giving a talk"}),
"video": (["none"] + video_files,),
"sampler_type": (["txt2img", "vdieo2vdieo", "WebCam2Video"],),
"guidance_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.1, "round": 0.01}),
"diffusion_steps": ("INT", {"default": 4, "min": 1, "max": 1000}),
"num_inference_steps": ("INT", {"default": 50, "min": 1, "max": 1000}),
@@ -268,38 +281,40 @@ class Stream_Sampler:
"seed": ("INT", {"default": 2, "min": 0, "max": 0xffffffffffffffff}),
"width": ("INT", {"default": 512, "min": 256, "max": 4096, "step": 64, }),
"height": ("INT", {"default": 512, "min": 256, "max": 4096, "step": 64}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 1024, "step": 1}),
},
"optional":{"image": ("IMAGE",),
}
"acceleration": (["xformers", "tensorrt", ],),
},
"optional": {"image": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE","AUDIO","FLOAT",)
RETURN_NAMES = ("image","audio","fps",)
RETURN_TYPES = ("IMAGE", "AUDIO", "FLOAT",)
RETURN_NAMES = ("image", "audio", "fps",)
FUNCTION = "main_stream"
CATEGORY = "StreamV2V_Plus"
def main_stream(self,pipe,info,prompt,video,sampler_type,guidance_scale,diffusion_steps,num_inference_steps,
noise_strength,seed,width,height,batch_size,**kwargs):
model_type,trigger_word=info.split(";")
if model_type=="sdxl":
using_sdxl=True
def main_stream(self, pipe, info, prompt, video, sampler_type, guidance_scale, diffusion_steps, num_inference_steps,
noise_strength, seed, width, height, acceleration, **kwargs):
model_type, trigger_word, ckpt_name = info.split(";")
if model_type == "sdxl":
using_sdxl = True
else:
using_sdxl=False
image=kwargs.get("image")
if sampler_type=="vdieo2vdieo":
using_sdxl = False
image = kwargs.get("image")
if sampler_type == "vdieo2vdieo":
if video == "none":
raise "need video input"
input_video = get_instance_path(os.path.join(input_path, video))
video_info = read_video(input_video) # input path to the file name "THWC") -> tuple[Tensor, Tensor, dict[str, Any]]
video = video_info[0] / 255 #Tensor THWC
fps = video_info[2]["video_fps"] #dict[str, Any]
video_info = read_video(
input_video) # input path to the file name "THWC") -> tuple[Tensor, Tensor, dict[str, Any]]
video = video_info[0] / 255 # Tensor THWC
fps = video_info[2]["video_fps"] # dict[str, Any]
height_v = int(video.shape[1])
width_v = int(video.shape[2])
waveform=video_info[1]
if waveform.size(1)==0:
audio=None
waveform = video_info[1]
if waveform.size(1) == 0:
audio = None
print("no audio in input video!")
else:
sample_rate = video_info[2]["audio_fps"]
@@ -313,7 +328,7 @@ class Stream_Sampler:
pipe,
mode="img2img",
t_index_list=t_index_list,
frame_buffer_size=batch_size,
frame_buffer_size=1,
width=width_v,
height=height_v,
warmup=10,
@@ -332,7 +347,7 @@ class Stream_Sampler:
seed=seed,
using_sdxl=using_sdxl
)
prompt = prompt + " " + trigger_word + "style"
stream.stream.prepare(
prompt=prompt,
@@ -354,11 +369,11 @@ class Stream_Sampler:
inference_time.append(iteration_end_time - iteration_start_time)
print(f'Avg time: {sum(inference_time[20:]) / len(inference_time[20:])}')
gen = get_video_img(video_result)
gen = narry_list(gen)
gen = narry_list(gen)
images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height_v, width_v, 3)))))
return (images,audio, fps)
return (images, audio, fps)
elif sampler_type=="txt2img":
elif sampler_type == "txt2img":
stream = StreamV2VWrapper(
pipe,
t_index_list=[0, 16, 32, 45],
@@ -366,7 +381,7 @@ class Stream_Sampler:
mode="txt2img",
device="cuda",
dtype=torch.float16,
frame_buffer_size=batch_size,
frame_buffer_size=1,
warmup=10,
width=width,
height=height,
@@ -381,7 +396,7 @@ class Stream_Sampler:
use_tome_cache=True,
seed=seed,
use_safety_checker=False,
using_sdxl = using_sdxl
using_sdxl=using_sdxl
)
prompt = prompt + " " + trigger_word + "style"
stream.stream.prepare(
@@ -389,16 +404,16 @@ class Stream_Sampler:
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
)
output_img=stream(prompt=prompt)
output_img = stream(prompt=prompt)
images = phi2narry(output_img)
return (images,)
else:
stream = StreamV2VWrapper(
pipe,
mode="img2img",
t_index_list=[32, 45],
frame_buffer_size=batch_size,
frame_buffer_size=1,
width=width,
height=height,
warmup=10,
@@ -407,7 +422,7 @@ class Stream_Sampler:
do_add_noise=True,
output_type="pil",
use_denoising_batch=True,
use_cached_attn=True,
use_cached_attn=False,
use_feature_injection=True,
feature_injection_strength=0.8,
feature_similarity_threshold=0.98,
@@ -415,16 +430,18 @@ class Stream_Sampler:
cache_maxframes=1,
use_tome_cache=True,
seed=seed,
using_sdxl=using_sdxl
acceleration=acceleration,
using_sdxl=using_sdxl,
model_id_or_path=ckpt_name
)
prompt = prompt + " " + trigger_word + "style"
stream.stream.prepare(
prompt=prompt,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
) #预处理prompt
) # 预处理prompt
image = nomarl_upscale(image, width, height)
for _ in range(2): #预处理prompt+img
for _ in range(2): # 预处理prompt+img
stream(image)
images = stream(image)
images = phi2narry(images)
@@ -432,13 +449,13 @@ class Stream_Sampler:
NODE_CLASS_MAPPINGS = {
"Stream_Model_Loader":Stream_Model_Loader,
"Stream_Lora_Loader":Stream_Lora_Loader,
"Stream_Model_Loader": Stream_Model_Loader,
"Stream_Lora_Loader": Stream_Lora_Loader,
"Stream_Sampler": Stream_Sampler,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Load_Stream":"Load_Stream",
"Stream_Lora_Loader":"Stream_Lora_Loader",
"Load_Stream": "Load_Stream",
"Stream_Lora_Loader": "Stream_Lora_Loader",
"Stream_Sampler": "Stream_Sampler",
}