feat: expose feature bank to comfyui (#12)
This commit establishes the required file structure and logic to integrate the Feature Bank node into ComfyUI, enabling its use. --------- Co-authored-by: RUFFY-369 <prakarshkaushik369@gmail.com>
This commit is contained in:
Rick Staa
@@ -0,0 +1,7 @@
|
||||
# Misc
|
||||
__pycache__
|
||||
*.pyc
|
||||
|
||||
# IDEs
|
||||
.vscode
|
||||
*.code-workspace
|
||||
+44
@@ -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"]
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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"]
|
||||
@@ -0,0 +1,2 @@
|
||||
torch
|
||||
diffusers
|
||||
@@ -0,0 +1 @@
|
||||
# Dummy file ensuring this sub-package will be recognized.
|
||||
@@ -0,0 +1 @@
|
||||
# Dummy file ensuring this sub-package will be recognized.
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user