diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/_model_index_compat.json b/_model_index_compat.json new file mode 100644 index 0000000..79ce13e --- /dev/null +++ b/_model_index_compat.json @@ -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" + ] +} diff --git a/cptorh01.sh b/cptorh01.sh new file mode 100644 index 0000000..159e409 --- /dev/null +++ b/cptorh01.sh @@ -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 diff --git a/helios/diffusers_version/pipeline_helios_diffusers.py b/helios/diffusers_version/pipeline_helios_diffusers.py index a72f5f0..cb64665 100644 --- a/helios/diffusers_version/pipeline_helios_diffusers.py +++ b/helios/diffusers_version/pipeline_helios_diffusers.py @@ -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 diff --git a/helios/diffusers_version/transformer_helios_diffusers.py b/helios/diffusers_version/transformer_helios_diffusers.py index aec1d75..c23b80c 100644 --- a/helios/diffusers_version/transformer_helios_diffusers.py +++ b/helios/diffusers_version/transformer_helios_diffusers.py @@ -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 diff --git a/helios/modules/helios_kernels/attention_dispatch.py b/helios/modules/helios_kernels/attention_dispatch.py index fb683e5..7096466 100644 --- a/helios/modules/helios_kernels/attention_dispatch.py +++ b/helios/modules/helios_kernels/attention_dispatch.py @@ -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 diff --git a/helios/modules/transformer_helios.py b/helios/modules/transformer_helios.py index 7e1c654..7479a01 100644 --- a/helios/modules/transformer_helios.py +++ b/helios/modules/transformer_helios.py @@ -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 diff --git a/infer_helios.py b/infer_helios.py index 8494f93..614b3cd 100644 --- a/infer_helios.py +++ b/infer_helios.py @@ -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) diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..a26f41a --- /dev/null +++ b/nodes/__init__.py @@ -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", +} diff --git a/nodes/i2v.py b/nodes/i2v.py new file mode 100644 index 0000000..47825d0 --- /dev/null +++ b/nodes/i2v.py @@ -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,) diff --git a/nodes/model_loader.py b/nodes/model_loader.py new file mode 100644 index 0000000..e5b3399 --- /dev/null +++ b/nodes/model_loader.py @@ -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,) diff --git a/nodes/t2v.py b/nodes/t2v.py new file mode 100644 index 0000000..d923fdc --- /dev/null +++ b/nodes/t2v.py @@ -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,) diff --git a/nodes/utils.py b/nodes/utils.py new file mode 100644 index 0000000..f37953b --- /dev/null +++ b/nodes/utils.py @@ -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 diff --git a/nodes/v2v.py b/nodes/v2v.py new file mode 100644 index 0000000..5f7d5cf --- /dev/null +++ b/nodes/v2v.py @@ -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,) diff --git a/rh_config.json b/rh_config.json new file mode 100644 index 0000000..2e53f34 --- /dev/null +++ b/rh_config.json @@ -0,0 +1 @@ +{"enable": true} diff --git a/sync.sh b/sync.sh new file mode 100644 index 0000000..e324552 --- /dev/null +++ b/sync.sh @@ -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