Add ComfyUI node integration and fix compatibility issues
- Add ComfyUI custom nodes: model loader, t2v, i2v, v2v - Fix kernels import: move to try block for graceful fallback when not installed - Fix diffusers compatibility: add fallback imports for HeliosLoraLoaderMixin, HeliosTransformer3DModel, HeliosScheduler, and apply_lora_scale - Fix infer_helios: cast num_blocks_per_group to int, add low_cpu_mem_usage - Add rh_config.json, model index compat, and deployment scripts Made-with: Cursor
This commit is contained in:
wenjian
@@ -0,0 +1,3 @@
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,25 @@
|
||||
{
|
||||
"_class_name": "WanPipeline",
|
||||
"_diffusers_version": "0.36.0",
|
||||
"is_distilled": true,
|
||||
"scheduler": [
|
||||
"diffusers",
|
||||
"FlowMatchEulerDiscreteScheduler"
|
||||
],
|
||||
"text_encoder": [
|
||||
"transformers",
|
||||
"UMT5EncoderModel"
|
||||
],
|
||||
"tokenizer": [
|
||||
"transformers",
|
||||
"T5TokenizerFast"
|
||||
],
|
||||
"transformer": [
|
||||
"diffusers",
|
||||
"WanTransformer3DModel"
|
||||
],
|
||||
"vae": [
|
||||
"diffusers",
|
||||
"AutoencoderKLWan"
|
||||
]
|
||||
}
|
||||
+37
@@ -0,0 +1,37 @@
|
||||
#!/bin/bash
|
||||
|
||||
# 创建目标目录
|
||||
mkdir -p /root/custom_nodes
|
||||
|
||||
# 简化的mount命令
|
||||
mount -t nfs4 -o rw rh-nfs.runninghub.cn:/data/rh_storage/global/custom_nodes_rel /root/custom_nodes
|
||||
|
||||
# 获取当前目录名作为NODE_NAME
|
||||
NODE_NAME=$(basename "$PWD")
|
||||
|
||||
# 创建目标目录
|
||||
mkdir -p /root/custom_nodes/${NODE_NAME}
|
||||
|
||||
# 显示将要执行的rsync命令
|
||||
echo "准备执行以下rsync命令："
|
||||
echo "rsync -av --include=\"*/\" --include=\"*.py\" --exclude=\"*\" ./ /root/custom_nodes/${NODE_NAME}/"
|
||||
echo ""
|
||||
echo "此命令将同步当前目录及子目录中的所有 .py 文件到 /root/custom_nodes/${NODE_NAME}/"
|
||||
echo ""
|
||||
read -p "是否继续执行？(输入 Y 确认): " -n 1 -r
|
||||
echo ""
|
||||
|
||||
if [[ $REPLY =~ ^[Yy]$ ]]; then
|
||||
echo "开始同步 Python 文件..."
|
||||
rsync -av --include="*/" --include="*.py" --exclude="*" ./ /root/custom_nodes/${NODE_NAME}/
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "Python 文件同步完成！"
|
||||
else
|
||||
echo "Python 文件同步失败！"
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
echo "取消同步操作。"
|
||||
exit 1
|
||||
fi
|
||||
@@ -25,10 +25,21 @@ from transformers import AutoTokenizer, UMT5EncoderModel
|
||||
|
||||
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
|
||||
from diffusers.image_processor import PipelineImageInput
|
||||
from diffusers.loaders import HeliosLoraLoaderMixin
|
||||
from diffusers.models import AutoencoderKLWan, HeliosTransformer3DModel
|
||||
try:
|
||||
from diffusers.loaders import HeliosLoraLoaderMixin
|
||||
except ImportError:
|
||||
class HeliosLoraLoaderMixin:
|
||||
pass
|
||||
from diffusers.models import AutoencoderKLWan
|
||||
try:
|
||||
from diffusers.models import HeliosTransformer3DModel
|
||||
except ImportError:
|
||||
from .transformer_helios_diffusers import HeliosTransformer3DModel
|
||||
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
|
||||
from diffusers.schedulers import HeliosScheduler
|
||||
try:
|
||||
from diffusers.schedulers import HeliosScheduler
|
||||
except ImportError:
|
||||
from .scheduling_helios_diffusers import HeliosScheduler
|
||||
from diffusers.utils import is_ftfy_available, is_torch_xla_available, logging, replace_example_docstring
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from diffusers.video_processor import VideoProcessor
|
||||
|
||||
@@ -30,7 +30,14 @@ from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbed
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.models.normalization import FP32LayerNorm
|
||||
from diffusers.utils import apply_lora_scale, logging
|
||||
from diffusers.utils import logging
|
||||
try:
|
||||
from diffusers.utils import apply_lora_scale
|
||||
except ImportError:
|
||||
def apply_lora_scale(key):
|
||||
def decorator(fn):
|
||||
return fn
|
||||
return decorator
|
||||
from diffusers.utils.torch_utils import maybe_allow_in_graph
|
||||
|
||||
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
import torch
|
||||
from kernels import get_kernel
|
||||
|
||||
|
||||
try:
|
||||
# raise NotImplementedError
|
||||
from kernels import get_kernel
|
||||
|
||||
try:
|
||||
flash_attn3 = get_kernel("kernels-community/flash-attn3")
|
||||
flash_attn_func = flash_attn3.flash_attn_func
|
||||
@@ -16,7 +17,7 @@ try:
|
||||
flash_attn_varlen_func = flash_attn2.flash_attn_varlen_func
|
||||
|
||||
print("Flash Attn 2 is installed!")
|
||||
except ImportError:
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
print("Flash Attn 2 / 3 is not installed!")
|
||||
flash_attn_varlen_func = None
|
||||
flash_attn_func = None
|
||||
|
||||
@@ -38,7 +38,14 @@ from diffusers.models.embeddings import (
|
||||
from diffusers.models.modeling_outputs import Transformer2DModelOutput
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.models.normalization import FP32LayerNorm
|
||||
from diffusers.utils import apply_lora_scale, deprecate, logging
|
||||
from diffusers.utils import deprecate, logging
|
||||
try:
|
||||
from diffusers.utils import apply_lora_scale
|
||||
except ImportError:
|
||||
def apply_lora_scale(key):
|
||||
def decorator(fn):
|
||||
return fn
|
||||
return decorator
|
||||
from diffusers.utils.torch_utils import maybe_allow_in_graph
|
||||
|
||||
from .helios_kernels import attn_varlen_func, create_navit_attention_masks
|
||||
|
||||
+2
-1
@@ -277,9 +277,10 @@ def main():
|
||||
onload_device=torch.device("cuda"),
|
||||
offload_device=torch.device("cpu"),
|
||||
offload_type=args.group_offloading_type,
|
||||
num_blocks_per_group=args.num_blocks_per_group if args.group_offloading_type == "block_level" else None,
|
||||
num_blocks_per_group=int(args.num_blocks_per_group) if args.group_offloading_type == "block_level" else None,
|
||||
use_stream=True,
|
||||
record_stream=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
else:
|
||||
pipe = pipe.to(device)
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
from .model_loader import RunningHubHeliosModelLoader
|
||||
from .t2v import RunningHubHeliosT2V
|
||||
from .i2v import RunningHubHeliosI2V
|
||||
from .v2v import RunningHubHeliosV2V
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"RunningHubHeliosModelLoader": RunningHubHeliosModelLoader,
|
||||
"RunningHubHeliosT2V": RunningHubHeliosT2V,
|
||||
"RunningHubHeliosI2V": RunningHubHeliosI2V,
|
||||
"RunningHubHeliosV2V": RunningHubHeliosV2V,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"RunningHubHeliosModelLoader": "RunningHub HeliosModelLoader",
|
||||
"RunningHubHeliosT2V": "RunningHub HeliosT2V",
|
||||
"RunningHubHeliosI2V": "RunningHub HeliosI2V",
|
||||
"RunningHubHeliosV2V": "RunningHub HeliosV2V",
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
import torch
|
||||
|
||||
from .utils import (
|
||||
comfyui_image_to_pil,
|
||||
estimate_total_steps,
|
||||
helios_output_to_video,
|
||||
make_progress_callback,
|
||||
parse_pyramid_steps,
|
||||
)
|
||||
|
||||
|
||||
class RunningHubHeliosI2V:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"helios_pipe": ("HELIOS_PIPE",),
|
||||
"image": ("IMAGE",),
|
||||
"prompt": ("STRING", {"default": "", "multiline": True}),
|
||||
"width": ("INT", {"default": 640, "min": 128, "max": 1920, "step": 16}),
|
||||
"height": ("INT", {"default": 384, "min": 128, "max": 1088, "step": 16}),
|
||||
"num_frames": ("INT", {"default": 99, "min": 1, "max": 480, "step": 1}),
|
||||
"num_inference_steps": ("INT", {"default": 50, "min": 1, "max": 200, "step": 1}),
|
||||
"guidance_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 20.0, "step": 0.1}),
|
||||
"seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFF}),
|
||||
"is_enable_stage2": ("BOOLEAN", {"default": True}),
|
||||
"pyramid_steps": ("STRING", {"default": "2,2,2"}),
|
||||
"is_amplify_first_chunk": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
"optional": {
|
||||
"negative_prompt": ("STRING", {"default": "", "multiline": True}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("video",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "RunningHub/Helios"
|
||||
|
||||
def generate(self, helios_pipe, image, prompt, width, height, num_frames,
|
||||
num_inference_steps, guidance_scale, seed,
|
||||
is_enable_stage2, pyramid_steps, is_amplify_first_chunk,
|
||||
negative_prompt=""):
|
||||
pipe = helios_pipe
|
||||
pyramid_list = parse_pyramid_steps(pyramid_steps)
|
||||
|
||||
total_steps = estimate_total_steps(
|
||||
num_frames, num_inference_steps, is_enable_stage2,
|
||||
pyramid_list, is_amplify_first_chunk,
|
||||
)
|
||||
progress_callback = make_progress_callback(total_steps)
|
||||
|
||||
pil_image = comfyui_image_to_pil(image, index=0)
|
||||
pil_image = pil_image.resize((width, height))
|
||||
|
||||
with torch.no_grad():
|
||||
output = pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt if negative_prompt else None,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
generator=torch.Generator(device="cuda").manual_seed(seed),
|
||||
history_sizes=[16, 2, 1],
|
||||
num_latent_frames_per_chunk=9,
|
||||
keep_first_frame=True,
|
||||
is_enable_stage2=is_enable_stage2,
|
||||
pyramid_num_inference_steps_list=pyramid_list,
|
||||
is_skip_first_chunk=False,
|
||||
is_amplify_first_chunk=is_amplify_first_chunk,
|
||||
use_zero_init=False,
|
||||
zero_steps=1,
|
||||
image=pil_image,
|
||||
video=None,
|
||||
callback_on_step_end=progress_callback,
|
||||
)
|
||||
|
||||
video = helios_output_to_video(output)
|
||||
return (video,)
|
||||
@@ -0,0 +1,140 @@
|
||||
import os
|
||||
import sys
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import folder_paths
|
||||
|
||||
logger = logging.getLogger("RunningHub.Helios")
|
||||
|
||||
HELIOS_MODEL_DIR = os.path.join(folder_paths.models_dir, "Helios-Distilled")
|
||||
|
||||
_cached_pipe = None
|
||||
_cached_config_hash = None
|
||||
|
||||
|
||||
def _get_helios_root():
|
||||
return os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
|
||||
def _ensure_helios_in_path():
|
||||
helios_root = _get_helios_root()
|
||||
if helios_root not in sys.path:
|
||||
sys.path.insert(0, helios_root)
|
||||
|
||||
|
||||
def load_helios_pipeline(weight_dtype_str, enable_low_vram, offloading_type):
|
||||
global _cached_pipe, _cached_config_hash
|
||||
|
||||
config_hash = hash((weight_dtype_str, enable_low_vram, offloading_type))
|
||||
if _cached_pipe is not None and config_hash == _cached_config_hash:
|
||||
logger.info("Using cached Helios pipeline")
|
||||
return _cached_pipe
|
||||
|
||||
if _cached_pipe is not None:
|
||||
logger.info("Config changed, releasing old pipeline")
|
||||
del _cached_pipe
|
||||
_cached_pipe = None
|
||||
_cached_config_hash = None
|
||||
import gc
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
_ensure_helios_in_path()
|
||||
|
||||
model_path = HELIOS_MODEL_DIR
|
||||
if not os.path.isdir(model_path):
|
||||
raise FileNotFoundError(f"Model directory not found: {model_path}")
|
||||
|
||||
dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16}
|
||||
weight_dtype = dtype_map.get(weight_dtype_str, torch.bfloat16)
|
||||
|
||||
logger.info(f"Loading Helios pipeline from {model_path} (dtype={weight_dtype_str})")
|
||||
|
||||
from helios.diffusers_version.pipeline_helios_diffusers import HeliosPipeline
|
||||
from helios.diffusers_version.scheduling_helios_diffusers import HeliosScheduler
|
||||
from helios.diffusers_version.transformer_helios_diffusers import HeliosTransformer3DModel
|
||||
from helios.modules.helios_kernels import (
|
||||
replace_all_norms_with_flash_norms,
|
||||
replace_rmsnorm_with_fp32,
|
||||
replace_rope_with_flash_rope,
|
||||
)
|
||||
from diffusers.models import AutoencoderKLWan
|
||||
|
||||
logger.info("Loading transformer...")
|
||||
transformer = HeliosTransformer3DModel.from_pretrained(
|
||||
model_path,
|
||||
subfolder="transformer",
|
||||
torch_dtype=weight_dtype,
|
||||
)
|
||||
|
||||
transformer = replace_rmsnorm_with_fp32(transformer)
|
||||
transformer = replace_all_norms_with_flash_norms(transformer)
|
||||
replace_rope_with_flash_rope()
|
||||
|
||||
try:
|
||||
transformer.set_attention_backend("_flash_3_hub")
|
||||
except Exception:
|
||||
transformer.set_attention_backend("flash_hub")
|
||||
|
||||
logger.info("Loading VAE...")
|
||||
vae = AutoencoderKLWan.from_pretrained(
|
||||
model_path,
|
||||
subfolder="vae",
|
||||
torch_dtype=torch.float32,
|
||||
)
|
||||
|
||||
logger.info("Loading scheduler...")
|
||||
scheduler = HeliosScheduler.from_pretrained(
|
||||
model_path,
|
||||
subfolder="scheduler",
|
||||
)
|
||||
|
||||
logger.info("Assembling pipeline...")
|
||||
pipe = HeliosPipeline.from_pretrained(
|
||||
model_path,
|
||||
transformer=transformer,
|
||||
vae=vae,
|
||||
scheduler=scheduler,
|
||||
torch_dtype=weight_dtype,
|
||||
)
|
||||
|
||||
if enable_low_vram:
|
||||
logger.info(f"Enabling group offload ({offloading_type})...")
|
||||
pipe.enable_group_offload(
|
||||
onload_device=torch.device("cuda"),
|
||||
offload_device=torch.device("cpu"),
|
||||
offload_type=offloading_type,
|
||||
num_blocks_per_group=1 if offloading_type == "block_level" else None,
|
||||
use_stream=True,
|
||||
record_stream=True,
|
||||
low_cpu_mem_usage=True,
|
||||
)
|
||||
else:
|
||||
pipe = pipe.to("cuda")
|
||||
|
||||
_cached_pipe = pipe
|
||||
_cached_config_hash = config_hash
|
||||
logger.info("Helios pipeline loaded successfully")
|
||||
return pipe
|
||||
|
||||
|
||||
class RunningHubHeliosModelLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"weight_dtype": (["bf16", "fp16"], {"default": "bf16"}),
|
||||
"enable_low_vram_mode": ("BOOLEAN", {"default": True}),
|
||||
"group_offloading_type": (["block_level", "leaf_level"], {"default": "block_level"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("HELIOS_PIPE",)
|
||||
RETURN_NAMES = ("helios_pipe",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "RunningHub/Helios"
|
||||
|
||||
def load_model(self, weight_dtype, enable_low_vram_mode, group_offloading_type):
|
||||
pipe = load_helios_pipeline(weight_dtype, enable_low_vram_mode, group_offloading_type)
|
||||
return (pipe,)
|
||||
@@ -0,0 +1,76 @@
|
||||
import torch
|
||||
|
||||
from .utils import (
|
||||
estimate_total_steps,
|
||||
helios_output_to_video,
|
||||
make_progress_callback,
|
||||
parse_pyramid_steps,
|
||||
)
|
||||
|
||||
|
||||
class RunningHubHeliosT2V:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"helios_pipe": ("HELIOS_PIPE",),
|
||||
"prompt": ("STRING", {"default": "", "multiline": True}),
|
||||
"width": ("INT", {"default": 640, "min": 128, "max": 1920, "step": 16}),
|
||||
"height": ("INT", {"default": 384, "min": 128, "max": 1088, "step": 16}),
|
||||
"num_frames": ("INT", {"default": 99, "min": 1, "max": 480, "step": 1}),
|
||||
"num_inference_steps": ("INT", {"default": 50, "min": 1, "max": 200, "step": 1}),
|
||||
"guidance_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 20.0, "step": 0.1}),
|
||||
"seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFF}),
|
||||
"is_enable_stage2": ("BOOLEAN", {"default": True}),
|
||||
"pyramid_steps": ("STRING", {"default": "2,2,2"}),
|
||||
"is_amplify_first_chunk": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
"optional": {
|
||||
"negative_prompt": ("STRING", {"default": "", "multiline": True}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("video",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "RunningHub/Helios"
|
||||
|
||||
def generate(self, helios_pipe, prompt, width, height, num_frames,
|
||||
num_inference_steps, guidance_scale, seed,
|
||||
is_enable_stage2, pyramid_steps, is_amplify_first_chunk,
|
||||
negative_prompt=""):
|
||||
pipe = helios_pipe
|
||||
pyramid_list = parse_pyramid_steps(pyramid_steps)
|
||||
|
||||
total_steps = estimate_total_steps(
|
||||
num_frames, num_inference_steps, is_enable_stage2,
|
||||
pyramid_list, is_amplify_first_chunk,
|
||||
)
|
||||
progress_callback = make_progress_callback(total_steps)
|
||||
|
||||
with torch.no_grad():
|
||||
output = pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt if negative_prompt else None,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
generator=torch.Generator(device="cuda").manual_seed(seed),
|
||||
history_sizes=[16, 2, 1],
|
||||
num_latent_frames_per_chunk=9,
|
||||
keep_first_frame=True,
|
||||
is_enable_stage2=is_enable_stage2,
|
||||
pyramid_num_inference_steps_list=pyramid_list,
|
||||
is_skip_first_chunk=False,
|
||||
is_amplify_first_chunk=is_amplify_first_chunk,
|
||||
use_zero_init=False,
|
||||
zero_steps=1,
|
||||
image=None,
|
||||
video=None,
|
||||
callback_on_step_end=progress_callback,
|
||||
)
|
||||
|
||||
video = helios_output_to_video(output)
|
||||
return (video,)
|
||||
+136
@@ -0,0 +1,136 @@
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
|
||||
try:
|
||||
from comfy_api.input_impl.video_types import VideoFromFile
|
||||
except ImportError:
|
||||
VideoFromFile = None
|
||||
|
||||
logger = logging.getLogger("RunningHub.Helios")
|
||||
|
||||
|
||||
def helios_output_to_comfyui_image(output):
|
||||
"""Convert Helios pipeline output to ComfyUI IMAGE tensor [N, H, W, C] float32 0~1."""
|
||||
frames = output.frames[0]
|
||||
|
||||
if isinstance(frames, torch.Tensor):
|
||||
# (N, C, H, W) -> (N, H, W, C)
|
||||
if frames.ndim == 4 and frames.shape[1] in (1, 3):
|
||||
frames = frames.permute(0, 2, 3, 1)
|
||||
frames = frames.float()
|
||||
if frames.max() > 1.0:
|
||||
frames = frames / 255.0
|
||||
return frames.cpu()
|
||||
|
||||
if isinstance(frames, np.ndarray):
|
||||
if frames.ndim == 4 and frames.shape[1] in (1, 3):
|
||||
frames = np.transpose(frames, (0, 2, 3, 1))
|
||||
tensor = torch.from_numpy(frames).float()
|
||||
if tensor.max() > 1.0:
|
||||
tensor = tensor / 255.0
|
||||
return tensor
|
||||
|
||||
if isinstance(frames, list):
|
||||
np_frames = []
|
||||
for f in frames:
|
||||
if isinstance(f, Image.Image):
|
||||
np_frames.append(np.array(f))
|
||||
elif isinstance(f, np.ndarray):
|
||||
np_frames.append(f)
|
||||
else:
|
||||
np_frames.append(np.array(f))
|
||||
stacked = np.stack(np_frames)
|
||||
tensor = torch.from_numpy(stacked).float()
|
||||
if tensor.max() > 1.0:
|
||||
tensor = tensor / 255.0
|
||||
return tensor
|
||||
|
||||
raise TypeError(f"Unexpected frames type: {type(frames)}")
|
||||
|
||||
|
||||
def comfyui_image_to_pil(image_tensor, index=0):
|
||||
"""Convert ComfyUI IMAGE tensor [B, H, W, C] to a single PIL Image."""
|
||||
img = image_tensor[index].cpu().numpy()
|
||||
img = (img * 255).clip(0, 255).astype(np.uint8)
|
||||
return Image.fromarray(img)
|
||||
|
||||
|
||||
def comfyui_images_to_pil_list(image_tensor):
|
||||
"""Convert ComfyUI IMAGE tensor [N, H, W, C] to a list of PIL Images."""
|
||||
frames = []
|
||||
for i in range(image_tensor.shape[0]):
|
||||
img = image_tensor[i].cpu().numpy()
|
||||
img = (img * 255).clip(0, 255).astype(np.uint8)
|
||||
frames.append(Image.fromarray(img))
|
||||
return frames
|
||||
|
||||
|
||||
def parse_pyramid_steps(steps_str):
|
||||
"""Parse comma-separated pyramid steps string to list of ints."""
|
||||
parts = [s.strip() for s in steps_str.split(",")]
|
||||
return [int(p) for p in parts if p]
|
||||
|
||||
|
||||
def estimate_total_steps(num_frames, num_inference_steps, is_enable_stage2,
|
||||
pyramid_steps, is_amplify_first_chunk,
|
||||
num_latent_frames_per_chunk=9, vae_temporal_factor=4):
|
||||
"""Estimate total denoising steps for ComfyUI progress bar."""
|
||||
window_num_frames = (num_latent_frames_per_chunk - 1) * vae_temporal_factor + 1
|
||||
num_chunks = max(1, math.ceil(num_frames / window_num_frames))
|
||||
|
||||
if is_enable_stage2:
|
||||
base_steps = sum(pyramid_steps)
|
||||
first_chunk_steps = base_steps * 2 if is_amplify_first_chunk else base_steps
|
||||
other_chunk_steps = base_steps
|
||||
total = first_chunk_steps + max(0, num_chunks - 1) * other_chunk_steps
|
||||
else:
|
||||
total = num_chunks * num_inference_steps
|
||||
|
||||
return total
|
||||
|
||||
|
||||
def make_progress_callback(total_steps):
|
||||
"""Create a ComfyUI-compatible progress callback for Helios pipeline."""
|
||||
pbar = comfy.utils.ProgressBar(total_steps)
|
||||
step_counter = [0]
|
||||
|
||||
def callback(pipe, step, timestep, callback_kwargs):
|
||||
step_counter[0] += 1
|
||||
pbar.update_absolute(step_counter[0], total_steps)
|
||||
return callback_kwargs
|
||||
|
||||
return callback
|
||||
|
||||
|
||||
def helios_output_to_video(output, fps=24):
|
||||
"""Convert Helios pipeline output to ComfyUI VIDEO object.
|
||||
|
||||
Saves frames to an mp4 file and wraps it with VideoFromFile.
|
||||
"""
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
frames = output.frames[0]
|
||||
output_dir = folder_paths.get_output_directory()
|
||||
filename = f"helios_{uuid.uuid4().hex[:12]}.mp4"
|
||||
video_path = os.path.join(output_dir, filename)
|
||||
|
||||
export_to_video(frames, video_path, fps=fps)
|
||||
logger.info(f"Video saved to {video_path}")
|
||||
|
||||
return create_video_object(video_path)
|
||||
|
||||
|
||||
def create_video_object(video_path):
|
||||
"""Create ComfyUI VIDEO object from a file path."""
|
||||
if VideoFromFile is not None:
|
||||
return VideoFromFile(video_path)
|
||||
return video_path
|
||||
@@ -0,0 +1,80 @@
|
||||
import torch
|
||||
|
||||
from .utils import (
|
||||
comfyui_images_to_pil_list,
|
||||
estimate_total_steps,
|
||||
helios_output_to_video,
|
||||
make_progress_callback,
|
||||
parse_pyramid_steps,
|
||||
)
|
||||
|
||||
|
||||
class RunningHubHeliosV2V:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"helios_pipe": ("HELIOS_PIPE",),
|
||||
"video_frames": ("IMAGE",),
|
||||
"prompt": ("STRING", {"default": "", "multiline": True}),
|
||||
"width": ("INT", {"default": 640, "min": 128, "max": 1920, "step": 16}),
|
||||
"height": ("INT", {"default": 384, "min": 128, "max": 1088, "step": 16}),
|
||||
"num_frames": ("INT", {"default": 99, "min": 1, "max": 480, "step": 1}),
|
||||
"num_inference_steps": ("INT", {"default": 50, "min": 1, "max": 200, "step": 1}),
|
||||
"guidance_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 20.0, "step": 0.1}),
|
||||
"seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFF}),
|
||||
"is_enable_stage2": ("BOOLEAN", {"default": True}),
|
||||
"pyramid_steps": ("STRING", {"default": "2,2,2"}),
|
||||
"is_amplify_first_chunk": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
"optional": {
|
||||
"negative_prompt": ("STRING", {"default": "", "multiline": True}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("video",)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "RunningHub/Helios"
|
||||
|
||||
def generate(self, helios_pipe, video_frames, prompt, width, height, num_frames,
|
||||
num_inference_steps, guidance_scale, seed,
|
||||
is_enable_stage2, pyramid_steps, is_amplify_first_chunk,
|
||||
negative_prompt=""):
|
||||
pipe = helios_pipe
|
||||
pyramid_list = parse_pyramid_steps(pyramid_steps)
|
||||
|
||||
total_steps = estimate_total_steps(
|
||||
num_frames, num_inference_steps, is_enable_stage2,
|
||||
pyramid_list, is_amplify_first_chunk,
|
||||
)
|
||||
progress_callback = make_progress_callback(total_steps)
|
||||
|
||||
pil_frames = comfyui_images_to_pil_list(video_frames)
|
||||
|
||||
with torch.no_grad():
|
||||
output = pipe(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt if negative_prompt else None,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
num_inference_steps=num_inference_steps,
|
||||
guidance_scale=guidance_scale,
|
||||
generator=torch.Generator(device="cuda").manual_seed(seed),
|
||||
history_sizes=[16, 2, 1],
|
||||
num_latent_frames_per_chunk=9,
|
||||
keep_first_frame=True,
|
||||
is_enable_stage2=is_enable_stage2,
|
||||
pyramid_num_inference_steps_list=pyramid_list,
|
||||
is_skip_first_chunk=False,
|
||||
is_amplify_first_chunk=is_amplify_first_chunk,
|
||||
use_zero_init=False,
|
||||
zero_steps=1,
|
||||
image=None,
|
||||
video=pil_frames,
|
||||
callback_on_step_end=progress_callback,
|
||||
)
|
||||
|
||||
video = helios_output_to_video(output)
|
||||
return (video,)
|
||||
@@ -0,0 +1 @@
|
||||
{"enable": true}
|
||||
@@ -0,0 +1,154 @@
|
||||
#!/bin/bash
|
||||
|
||||
# 配置变量
|
||||
NODE_NAME=$(basename "$PWD")
|
||||
REMOTE_HOST="root@10.132.208.7"
|
||||
DOCKER_CONTAINER="RunningHub"
|
||||
DOCKER_PATH="/workspace/ComfyUI/custom_nodes/${NODE_NAME}/"
|
||||
# 使用固定的远程同步目录，保持 rsync 状态
|
||||
REMOTE_SYNC_DIR="/root/${NODE_NAME}"
|
||||
|
||||
echo "开始增量同步到Docker容器..."
|
||||
|
||||
# 自动生成 cptorh01.sh 脚本
|
||||
echo "正在生成 cptorh01.sh 脚本..."
|
||||
cat > cptorh01.sh << 'EOF'
|
||||
#!/bin/bash
|
||||
|
||||
# 创建目标目录
|
||||
mkdir -p /root/custom_nodes
|
||||
|
||||
# 简化的mount命令
|
||||
mount -t nfs4 -o rw rh-nfs.runninghub.cn:/data/rh_storage/global/custom_nodes_rel /root/custom_nodes
|
||||
|
||||
# 获取当前目录名作为NODE_NAME
|
||||
NODE_NAME=$(basename "$PWD")
|
||||
|
||||
# 创建目标目录
|
||||
mkdir -p /root/custom_nodes/${NODE_NAME}
|
||||
|
||||
# 显示将要执行的rsync命令
|
||||
echo "准备执行以下rsync命令："
|
||||
echo "rsync -av --include=\"*/\" --include=\"*.py\" --exclude=\"*\" ./ /root/custom_nodes/${NODE_NAME}/"
|
||||
echo ""
|
||||
echo "此命令将同步当前目录及子目录中的所有 .py 文件到 /root/custom_nodes/${NODE_NAME}/"
|
||||
echo ""
|
||||
read -p "是否继续执行？(输入 Y 确认): " -n 1 -r
|
||||
echo ""
|
||||
|
||||
if [[ $REPLY =~ ^[Yy]$ ]]; then
|
||||
echo "开始同步 Python 文件..."
|
||||
rsync -av --include="*/" --include="*.py" --exclude="*" ./ /root/custom_nodes/${NODE_NAME}/
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "Python 文件同步完成！"
|
||||
else
|
||||
echo "Python 文件同步失败！"
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
echo "取消同步操作。"
|
||||
exit 1
|
||||
fi
|
||||
EOF
|
||||
|
||||
# 给 cptorh01.sh 添加执行权限
|
||||
chmod +x cptorh01.sh
|
||||
echo "cptorh01.sh 脚本生成完成！"
|
||||
|
||||
# 检查本地是否安装了rsync
|
||||
if ! command -v rsync &> /dev/null; then
|
||||
echo "错误：本地未安装 rsync，请先安装: apt-get install rsync 或 yum install rsync"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# 检查远程连接和Docker容器
|
||||
echo "检查远程连接和Docker容器..."
|
||||
ssh ${REMOTE_HOST} "docker ps | grep ${DOCKER_CONTAINER}" > /dev/null
|
||||
if [ $? -ne 0 ]; then
|
||||
echo "Docker容器 ${DOCKER_CONTAINER} 未运行！"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# 检查远程是否安装了rsync
|
||||
echo "检查远程rsync..."
|
||||
ssh ${REMOTE_HOST} "which rsync" > /dev/null
|
||||
if [ $? -ne 0 ]; then
|
||||
echo "错误：远程主机未安装 rsync，请在远程主机安装: apt-get install rsync 或 yum install rsync"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# 确保Docker容器内目标目录存在
|
||||
echo "确保Docker容器内目标目录存在..."
|
||||
ssh ${REMOTE_HOST} "docker exec ${DOCKER_CONTAINER} mkdir -p ${DOCKER_PATH}"
|
||||
if [ $? -ne 0 ]; then
|
||||
echo "无法在Docker容器内创建目录！"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# 确保远程同步目录存在
|
||||
echo "确保远程同步目录存在..."
|
||||
ssh ${REMOTE_HOST} "mkdir -p ${REMOTE_SYNC_DIR}"
|
||||
|
||||
# 定义要排除的文件和目录
|
||||
EXCLUDE_PATTERNS=(
|
||||
"--exclude=.git"
|
||||
"--exclude=__pycache__"
|
||||
"--exclude=*.pyc"
|
||||
"--exclude=*.pyo"
|
||||
"--exclude=.DS_Store"
|
||||
"--exclude=Thumbs.db"
|
||||
"--exclude=*.tmp"
|
||||
"--exclude=*.swp"
|
||||
"--exclude=.vscode"
|
||||
"--exclude=.idea"
|
||||
"--exclude=*.log"
|
||||
"--exclude=weights/"
|
||||
"--exclude=save_audio/"
|
||||
)
|
||||
|
||||
# 使用rsync进行增量同步到远程固定目录
|
||||
echo "开始增量同步文件..."
|
||||
echo "正在比较文件差异，只传输变化的文件..."
|
||||
|
||||
rsync -avz --progress --delete \
|
||||
"${EXCLUDE_PATTERNS[@]}" \
|
||||
--rsync-path="rsync" \
|
||||
./ ${REMOTE_HOST}:${REMOTE_SYNC_DIR}/
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "增量同步到远程主机完成，正在同步到Docker容器..."
|
||||
|
||||
# 直接从远程同步目录同步到Docker容器，保持文件权限
|
||||
ssh ${REMOTE_HOST} "cd ${REMOTE_SYNC_DIR} && tar -cpf - . | docker exec -i ${DOCKER_CONTAINER} tar -xpf - -C ${DOCKER_PATH}"
|
||||
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "增量同步完成！只传输了变化的文件。"
|
||||
echo "提示：下次同步将更快，因为只会传输修改过的文件。"
|
||||
echo "远程同步目录：${REMOTE_SYNC_DIR}"
|
||||
|
||||
# 修改文件权限为644
|
||||
echo "正在修改文件权限为644..."
|
||||
ssh ${REMOTE_HOST} "docker exec ${DOCKER_CONTAINER} find ${DOCKER_PATH} -type f -exec chmod 644 {} \;"
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "文件权限修改完成！"
|
||||
else
|
||||
echo "文件权限修改失败，请手动检查！"
|
||||
fi
|
||||
|
||||
# 重启ComfyUI服务
|
||||
echo "正在重启ComfyUI服务..."
|
||||
ssh ${REMOTE_HOST} "pkill -9 -f 'python main.py'"
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "ComfyUI服务已重启！"
|
||||
else
|
||||
echo "重启ComfyUI服务失败，请手动检查！"
|
||||
fi
|
||||
else
|
||||
echo "同步到Docker容器失败！"
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
echo "增量同步失败！"
|
||||
exit 1
|
||||
fi
|
||||
Reference in New Issue
Block a user