diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..27425f4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +# Misc +__pycache__ +*.pyc + +# IDEs +.vscode +*.code-workspace diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..6cc7cc2 --- /dev/null +++ b/__init__.py @@ -0,0 +1,44 @@ +"""Expose all nodes in the Stream Pack.""" + +import logging +import pathlib +import importlib +import traceback + +logging.basicConfig(level=logging.INFO) +log = logging.getLogger(__name__) + + +def load_nodes() -> tuple[dict[str, type], dict[str, str]]: + """Load all nodes in the Stream Pack.""" + nodes_dir = pathlib.Path(__file__).parent / "node_wrappers" + node_class_mappings, node_display_name_mappings = {}, {} + + # Dynamically import all Python modules in the node_wrappers directory. + for module_path in nodes_dir.iterdir(): + if ( + module_path.is_file() + and module_path.suffix == ".py" + and module_path.stem != "__init__" + ): + try: + module_name = f"{__package__}.node_wrappers.{module_path.stem}" + module = importlib.import_module(module_name) + + # Update mappings if defined in the module. + if hasattr(module, "NODE_CLASS_MAPPINGS"): + node_class_mappings.update(module.NODE_CLASS_MAPPINGS) + if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"): + node_display_name_mappings.update(module.NODE_DISPLAY_NAME_MAPPINGS) + except Exception: + format_exc = traceback.format_exc() + log.error(f"Failed to load module {module_name}:\n{format_exc}") + + return node_class_mappings, node_display_name_mappings + + +# Dynamically collect mappings from submodules. +NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = load_nodes() +NODE_DISPLAY_NAME_MAPPINGS["StreamPack"] = "Stream Pack Nodes" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/examples/workflows/sd15-tensorrt-api.json b/examples/workflows/sd15-tensorrt-api.json new file mode 100644 index 0000000..4cfb707 --- /dev/null +++ b/examples/workflows/sd15-tensorrt-api.json @@ -0,0 +1,218 @@ +{ + "1": { + "inputs": { + "image": "example.png", + "upload": "image" + }, + "class_type": "LoadImage", + "_meta": { + "title": "Load Image" + } + }, + "2": { + "inputs": { + "engine": "depth_anything_vitl14-fp16.engine", + "images": [ + "1", + 0 + ] + }, + "class_type": "DepthAnythingTensorrt", + "_meta": { + "title": "Depth Anything Tensorrt" + } + }, + "3": { + "inputs": { + "unet_name": "static-dreamshaper8_SD15_$stat-b-1-h-512-w-512_00001_.engine", + "model_type": "SD15" + }, + "class_type": "TensorRTLoader", + "_meta": { + "title": "TensorRT Loader" + } + }, + "5": { + "inputs": { + "text": "the hulk", + "clip": [ + "23", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "6": { + "inputs": { + "text": "", + "clip": [ + "23", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "7": { + "inputs": { + "seed": 905056445574169, + "steps": 1, + "cfg": 1, + "sampler_name": "lcm", + "scheduler": "normal", + "denoise": 1, + "model": [ + "3", + 0 + ], + "positive": [ + "9", + 0 + ], + "negative": [ + "9", + 1 + ], + "latent_image": [ + "16", + 0 + ] + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "8": { + "inputs": { + "control_net_name": "control_v11f1p_sd15_depth_fp16.safetensors" + }, + "class_type": "ControlNetLoader", + "_meta": { + "title": "Load ControlNet Model" + } + }, + "9": { + "inputs": { + "strength": 1, + "start_percent": 0, + "end_percent": 1, + "positive": [ + "5", + 0 + ], + "negative": [ + "6", + 0 + ], + "control_net": [ + "10", + 0 + ], + "image": [ + "2", + 0 + ] + }, + "class_type": "ControlNetApplyAdvanced", + "_meta": { + "title": "Apply ControlNet" + } + }, + "10": { + "inputs": { + "backend": "inductor", + "fullgraph": false, + "mode": "reduce-overhead", + "controlnet": [ + "8", + 0 + ] + }, + "class_type": "TorchCompileLoadControlNet", + "_meta": { + "title": "TorchCompileLoadControlNet" + } + }, + "11": { + "inputs": { + "vae_name": "taesd" + }, + "class_type": "VAELoader", + "_meta": { + "title": "Load VAE" + } + }, + "13": { + "inputs": { + "backend": "inductor", + "fullgraph": true, + "mode": "reduce-overhead", + "compile_encoder": true, + "compile_decoder": true, + "vae": [ + "11", + 0 + ] + }, + "class_type": "TorchCompileLoadVAE", + "_meta": { + "title": "TorchCompileLoadVAE" + } + }, + "14": { + "inputs": { + "samples": [ + "7", + 0 + ], + "vae": [ + "13", + 0 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "15": { + "inputs": { + "images": [ + "14", + 0 + ] + }, + "class_type": "PreviewImage", + "_meta": { + "title": "Preview Image" + } + }, + "16": { + "inputs": { + "width": 512, + "height": 512, + "batch_size": 1 + }, + "class_type": "EmptyLatentImage", + "_meta": { + "title": "Empty Latent Image" + } + }, + "23": { + "inputs": { + "clip_name": "CLIPText/model.fp16.safetensors", + "type": "stable_diffusion", + "device": "default" + }, + "class_type": "CLIPLoader", + "_meta": { + "title": "Load CLIP" + } + } +} diff --git a/examples/workflows/sd15-tensorrt-feature-bank-api.json b/examples/workflows/sd15-tensorrt-feature-bank-api.json new file mode 100644 index 0000000..a4ba405 --- /dev/null +++ b/examples/workflows/sd15-tensorrt-feature-bank-api.json @@ -0,0 +1,234 @@ +{ + "1": { + "inputs": { + "image": "example.png" + }, + "class_type": "LoadImage", + "_meta": { + "title": "Load Image" + } + }, + "2": { + "inputs": { + "engine": "depth_anything_vitl14-fp16.engine", + "images": [ + "1", + 0 + ] + }, + "class_type": "DepthAnythingTensorrt", + "_meta": { + "title": "Depth Anything Tensorrt" + } + }, + "3": { + "inputs": { + "unet_name": "static-dreamshaper8_SD15_$stat-b-1-h-512-w-512_00001_.engine", + "model_type": "SD15" + }, + "class_type": "TensorRTLoader", + "_meta": { + "title": "TensorRT Loader" + } + }, + "5": { + "inputs": { + "text": "the hulk", + "clip": [ + "23", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "6": { + "inputs": { + "text": "", + "clip": [ + "23", + 0 + ] + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "CLIP Text Encode (Prompt)" + } + }, + "7": { + "inputs": { + "seed": 785664736216738, + "steps": 1, + "cfg": 1, + "sampler_name": "lcm", + "scheduler": "normal", + "denoise": 1, + "model": [ + "24", + 0 + ], + "positive": [ + "9", + 0 + ], + "negative": [ + "9", + 1 + ], + "latent_image": [ + "16", + 0 + ] + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "8": { + "inputs": { + "control_net_name": "control_v11f1p_sd15_depth_fp16.safetensors" + }, + "class_type": "ControlNetLoader", + "_meta": { + "title": "Load ControlNet Model" + } + }, + "9": { + "inputs": { + "strength": 1, + "start_percent": 0, + "end_percent": 1, + "positive": [ + "5", + 0 + ], + "negative": [ + "6", + 0 + ], + "control_net": [ + "10", + 0 + ], + "image": [ + "2", + 0 + ] + }, + "class_type": "ControlNetApplyAdvanced", + "_meta": { + "title": "Apply ControlNet" + } + }, + "10": { + "inputs": { + "backend": "inductor", + "fullgraph": false, + "mode": "reduce-overhead", + "controlnet": [ + "8", + 0 + ] + }, + "class_type": "TorchCompileLoadControlNet", + "_meta": { + "title": "TorchCompileLoadControlNet" + } + }, + "11": { + "inputs": { + "vae_name": "taesd" + }, + "class_type": "VAELoader", + "_meta": { + "title": "Load VAE" + } + }, + "13": { + "inputs": { + "backend": "inductor", + "fullgraph": true, + "mode": "reduce-overhead", + "compile_encoder": true, + "compile_decoder": true, + "vae": [ + "11", + 0 + ] + }, + "class_type": "TorchCompileLoadVAE", + "_meta": { + "title": "TorchCompileLoadVAE" + } + }, + "14": { + "inputs": { + "samples": [ + "7", + 0 + ], + "vae": [ + "13", + 0 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "15": { + "inputs": { + "images": [ + "14", + 0 + ] + }, + "class_type": "PreviewImage", + "_meta": { + "title": "Preview Image" + } + }, + "16": { + "inputs": { + "width": 512, + "height": 512, + "batch_size": 1 + }, + "class_type": "EmptyLatentImage", + "_meta": { + "title": "Empty Latent Image" + } + }, + "23": { + "inputs": { + "clip_name": "CLIPText/model.fp16.safetensors", + "type": "stable_diffusion", + "device": "default" + }, + "class_type": "CLIPLoader", + "_meta": { + "title": "Load CLIP" + } + }, + "24": { + "inputs": { + "use_feature_injection": false, + "feature_injection_strength": 0.8, + "feature_similarity_threshold": 0.98, + "feature_cache_interval": 4, + "feature_bank_max_frames": 4, + "model": [ + "3", + 0 + ] + }, + "class_type": "FeatureBankAttentionProcessor", + "_meta": { + "title": "Feature Bank Attention Processor" + } + } +} diff --git a/node_wrappers/feature_bank.py b/node_wrappers/feature_bank.py new file mode 100644 index 0000000..d771f29 --- /dev/null +++ b/node_wrappers/feature_bank.py @@ -0,0 +1,20 @@ +"""A ComfyUI native node implementation of the Feature Bank of the StreamV2V paper. + + +References: + StreamV2V: A Deep Generative Model for High-fidelity View Synthesis - Liang et al., + 2025, URL: https://arxiv.org/abs/2405.15757 +""" + +from ..src.feature_bank.feature_bank import ( + FeatureBankAttentionProcessorNode, +) + +NODE_CLASS_MAPPINGS = { + "FeatureBankAttentionProcessor": FeatureBankAttentionProcessorNode +} +NODE_DISPLAY_NAME_MAPPINGS = { + "FeatureBankAttentionProcessor": "Feature Bank Attention Processor" +} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d8e21cf --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +torch +diffusers diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e6119e5 --- /dev/null +++ b/src/__init__.py @@ -0,0 +1 @@ +# Dummy file ensuring this sub-package will be recognized. diff --git a/src/feature_bank/__init__.py b/src/feature_bank/__init__.py new file mode 100644 index 0000000..e6119e5 --- /dev/null +++ b/src/feature_bank/__init__.py @@ -0,0 +1 @@ +# Dummy file ensuring this sub-package will be recognized. diff --git a/src/feature_bank/feature_bank.py b/src/feature_bank/feature_bank.py index 19040f1..bd09c0a 100644 --- a/src/feature_bank/feature_bank.py +++ b/src/feature_bank/feature_bank.py @@ -191,7 +191,7 @@ class FeatureBankAttentionProcessor: batch_size, channel, height * width ).transpose(1, 2) - batch_size, _ , _ = ( + batch_size, _, _ = ( hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape @@ -459,7 +459,7 @@ class FeatureBankAttentionProcessorNode: RETURN_TYPES = ("MODEL",) FUNCTION = "update_node" - CATEGORY = "feature-bank" + CATEGORY = "StreamPack/model_patches/unet" def update_node( self,