Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 RyanOnTheInside
|
f6e1990905 | ||
|
ryanontheinside
|
06975653ae | ||
|
ryanontheinside
|
0992414295 | ||
|
ryanontheinside
|
33d26c4808 | ||
|
Rick Staa
|
87e609d13a | ||
|
RUFFY-369
|
6122e62f6e |
+88
@@ -0,0 +1,88 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
.pytest_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
|
||||
# Virtual Environment
|
||||
venv/
|
||||
env/
|
||||
ENV/
|
||||
.env/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# ML/AI specific
|
||||
*.pt
|
||||
*.pth
|
||||
*.onnx
|
||||
*.bin
|
||||
*.safetensors
|
||||
*.ckpt
|
||||
output/
|
||||
outputs/
|
||||
results/
|
||||
runs/
|
||||
logs/
|
||||
checkpoints/
|
||||
# IDE/Editor specific
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# OS specific
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
desktop.ini
|
||||
|
||||
# ComfyUI specific
|
||||
temp/
|
||||
cache/
|
||||
input/
|
||||
output/
|
||||
|
||||
# Media files that might be large
|
||||
*.jpg
|
||||
*.jpeg
|
||||
*.png
|
||||
*.gif
|
||||
*.bmp
|
||||
*.tif
|
||||
*.tiff
|
||||
*.mp4
|
||||
*.avi
|
||||
*.mov
|
||||
!examples/*.jpg
|
||||
!examples/*.png
|
||||
|
||||
# Dependencies
|
||||
node_modules/
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
npm-debug.log*
|
||||
yarn-debug.log*
|
||||
yarn-error.log*
|
||||
+19
@@ -0,0 +1,19 @@
|
||||
"""ComfyUI Stream Pack Nodes."""
|
||||
|
||||
from .node_wrappers import feature_bank
|
||||
from .node_wrappers import facemesh
|
||||
from .node_wrappers import batched_denoising
|
||||
# Collect all NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS from submodules.
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
# Import and update mappings from submodules.
|
||||
for module in [feature_bank, facemesh, batched_denoising]:
|
||||
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)
|
||||
|
||||
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,376 @@
|
||||
{
|
||||
"2": {
|
||||
"inputs": {
|
||||
"engine": "depth_anything_vitl14-fp16.engine",
|
||||
"images": [
|
||||
"26",
|
||||
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": "beautiful pencil sketch, masterpiece ",
|
||||
"clip": [
|
||||
"18",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {
|
||||
"title": "CLIP Text Encode (Prompt)"
|
||||
}
|
||||
},
|
||||
"6": {
|
||||
"inputs": {
|
||||
"text": "",
|
||||
"clip": [
|
||||
"18",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {
|
||||
"title": "CLIP Text Encode (Prompt)"
|
||||
}
|
||||
},
|
||||
"8": {
|
||||
"inputs": {
|
||||
"control_net_name": "control_v11f1p_sd15_depth_fp16.safetensors"
|
||||
},
|
||||
"class_type": "ControlNetLoader",
|
||||
"_meta": {
|
||||
"title": "Load ControlNet Model"
|
||||
}
|
||||
},
|
||||
"9": {
|
||||
"inputs": {
|
||||
"strength": 0.9,
|
||||
"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": [
|
||||
"42",
|
||||
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"
|
||||
}
|
||||
},
|
||||
"18": {
|
||||
"inputs": {
|
||||
"stop_at_clip_layer": -2,
|
||||
"clip": [
|
||||
"36",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "CLIPSetLastLayer",
|
||||
"_meta": {
|
||||
"title": "CLIP Set Last Layer"
|
||||
}
|
||||
},
|
||||
"19": {
|
||||
"inputs": {
|
||||
"strength": 0.9,
|
||||
"start_percent": 0,
|
||||
"end_percent": 1,
|
||||
"positive": [
|
||||
"5",
|
||||
0
|
||||
],
|
||||
"negative": [
|
||||
"6",
|
||||
0
|
||||
],
|
||||
"control_net": [
|
||||
"21",
|
||||
0
|
||||
],
|
||||
"image": [
|
||||
"39",
|
||||
1
|
||||
]
|
||||
},
|
||||
"class_type": "ControlNetApplyAdvanced",
|
||||
"_meta": {
|
||||
"title": "Apply ControlNet"
|
||||
}
|
||||
},
|
||||
"20": {
|
||||
"inputs": {
|
||||
"control_net_name": "control_v2p_sd15_mediapipe_face.safetensors"
|
||||
},
|
||||
"class_type": "ControlNetLoader",
|
||||
"_meta": {
|
||||
"title": "Load ControlNet Model"
|
||||
}
|
||||
},
|
||||
"21": {
|
||||
"inputs": {
|
||||
"backend": "inductor",
|
||||
"fullgraph": false,
|
||||
"mode": "reduce-overhead",
|
||||
"controlnet": [
|
||||
"20",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "TorchCompileLoadControlNet",
|
||||
"_meta": {
|
||||
"title": "TorchCompileLoadControlNet"
|
||||
}
|
||||
},
|
||||
"26": {
|
||||
"inputs": {
|
||||
"image": "example.png",
|
||||
"upload": "image"
|
||||
},
|
||||
"class_type": "LoadImage",
|
||||
"_meta": {
|
||||
"title": "Load Image"
|
||||
}
|
||||
},
|
||||
"36": {
|
||||
"inputs": {
|
||||
"clip_name": "CLIPText/model.fp16.safetensors",
|
||||
"type": "stable_diffusion",
|
||||
"device": "default"
|
||||
},
|
||||
"class_type": "CLIPLoader",
|
||||
"_meta": {
|
||||
"title": "Load CLIP"
|
||||
}
|
||||
},
|
||||
"37": {
|
||||
"inputs": {
|
||||
"conditioning_to": [
|
||||
"9",
|
||||
1
|
||||
],
|
||||
"conditioning_from": [
|
||||
"19",
|
||||
1
|
||||
]
|
||||
},
|
||||
"class_type": "ConditioningConcat",
|
||||
"_meta": {
|
||||
"title": "Conditioning (Concat)"
|
||||
}
|
||||
},
|
||||
"38": {
|
||||
"inputs": {
|
||||
"conditioning_to": [
|
||||
"9",
|
||||
0
|
||||
],
|
||||
"conditioning_from": [
|
||||
"19",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "ConditioningConcat",
|
||||
"_meta": {
|
||||
"title": "Conditioning (Concat)"
|
||||
}
|
||||
},
|
||||
"39": {
|
||||
"inputs": {
|
||||
"static_image_mode": false,
|
||||
"max_num_faces": 1,
|
||||
"refine_landmarks": false,
|
||||
"min_detection_confidence": 0.35000000000000003,
|
||||
"min_tracking_confidence": 0.35000000000000003,
|
||||
"draw_mesh": true,
|
||||
"draw_contours": true,
|
||||
"mesh_color": "green",
|
||||
"mesh_thickness": 2,
|
||||
"image": [
|
||||
"26",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "FaceMeshNode",
|
||||
"_meta": {
|
||||
"title": "Face Mesh"
|
||||
}
|
||||
},
|
||||
"41": {
|
||||
"inputs": {
|
||||
"num_steps": 3
|
||||
},
|
||||
"class_type": "StreamBatchSampler",
|
||||
"_meta": {
|
||||
"title": "Stream Batch Sampler"
|
||||
}
|
||||
},
|
||||
"42": {
|
||||
"inputs": {
|
||||
"add_noise": true,
|
||||
"noise_seed": 977567268563782,
|
||||
"cfg": 1,
|
||||
"model": [
|
||||
"3",
|
||||
0
|
||||
],
|
||||
"positive": [
|
||||
"38",
|
||||
0
|
||||
],
|
||||
"negative": [
|
||||
"37",
|
||||
0
|
||||
],
|
||||
"sampler": [
|
||||
"41",
|
||||
0
|
||||
],
|
||||
"sigmas": [
|
||||
"44",
|
||||
0
|
||||
],
|
||||
"latent_image": [
|
||||
"43",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "SamplerCustom",
|
||||
"_meta": {
|
||||
"title": "SamplerCustom"
|
||||
}
|
||||
},
|
||||
"43": {
|
||||
"inputs": {
|
||||
"buffer_size": 3,
|
||||
"latent": [
|
||||
"16",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "StreamFrameBuffer",
|
||||
"_meta": {
|
||||
"title": "Stream Frame Buffer"
|
||||
}
|
||||
},
|
||||
"44": {
|
||||
"inputs": {
|
||||
"t_index_list": "0,16,32",
|
||||
"num_inference_steps": 50,
|
||||
"model": [
|
||||
"3",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "StreamScheduler",
|
||||
"_meta": {
|
||||
"title": "Stream Scheduler"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,15 @@
|
||||
from ..src.batched_denoising.batched_denoising import StreamBatchSampler, StreamScheduler, StreamFrameBuffer
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"StreamBatchSampler": StreamBatchSampler,
|
||||
"StreamScheduler": StreamScheduler,
|
||||
"StreamFrameBuffer": StreamFrameBuffer,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"StreamBatchSampler": "Stream Batch Sampler",
|
||||
"StreamScheduler": "Stream Scheduler",
|
||||
"StreamFrameBuffer": "Stream Frame Buffer",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,13 @@
|
||||
from ..src.media_pipe.facemesh import *
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FaceMeshNode": FaceMeshNode,
|
||||
"FaceMeshDrawNode": FaceMeshDrawNode,
|
||||
"FaceMeshMaskNode": FaceMeshMaskNode
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FaceMeshNode": "Face Mesh",
|
||||
"FaceMeshDrawNode": "Face Mesh Draw",
|
||||
"FaceMeshMaskNode": "Face Region Mask"
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
"""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,3 @@
|
||||
mediapipe
|
||||
torch
|
||||
diffusers
|
||||
@@ -0,0 +1 @@
|
||||
# Dummy file ensuring this package will be recognized.
|
||||
@@ -0,0 +1,497 @@
|
||||
import torch
|
||||
import comfy.model_management
|
||||
import comfy.samplers
|
||||
import math
|
||||
|
||||
class StreamBatchSampler:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"cfg_type": (["none", "full", "self", "initialize"], {"default": "self", "tooltip": "'self' (RCFG) is fastest, 'full' is standard SD but slower, 'initialize' is a middle ground, 'none' disables guidance"}),
|
||||
"delta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1, "tooltip": "Only used with 'self' CFG. Controls strength of self-guidance. Higher = stronger guidance but more artifacts. Default 1.0", "round": 0.01}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SAMPLER",)
|
||||
FUNCTION = "update"
|
||||
CATEGORY = "StreamPack/sampling"
|
||||
DESCRIPTION = "Implements batched denoising with different CFG options (including RCFG) for faster inference and temporal consistency."
|
||||
|
||||
def __init__(self):
|
||||
self.frame_buffer = []
|
||||
self.x_t_latent_buffer = None
|
||||
self.stock_noise = None
|
||||
self.is_txt2img_mode = False
|
||||
self.delta = 1.0
|
||||
self.cfg_type = "self"
|
||||
|
||||
# Initialize all buffers
|
||||
self.zeros_reference = None
|
||||
self.random_noise_buffer = None
|
||||
self.sigmas_view_buffer = None
|
||||
self.expanded_stock_noise = None
|
||||
self.working_buffer = None
|
||||
self.output_buffer = None
|
||||
|
||||
# Add RCFG state from reference
|
||||
self.cached_sigmas = None
|
||||
self.alpha_prod_t = None
|
||||
self.beta_prod_t = None
|
||||
self.alpha_prod_t_sqrt = None
|
||||
self.beta_prod_t_sqrt = None
|
||||
self.init_noise = None # Will be initialized in sample based on input shape
|
||||
self.step_counter = 0
|
||||
|
||||
def compute_alpha_beta(self, sigmas):
|
||||
"""Pre-compute alpha and beta terms for the given sigmas (LCM-like using cosine schedule)"""
|
||||
self.cached_sigmas = sigmas.clone()
|
||||
print(f"\n[Alpha/Beta] Input sigmas: {sigmas}")
|
||||
|
||||
num_train_timesteps = 1000 # Standard for many diffusion models
|
||||
timesteps = torch.linspace(
|
||||
1, 0, num_train_timesteps, dtype=torch.float32, device=sigmas.device
|
||||
)
|
||||
alphas_cumprod = torch.cos((timesteps + 0.008) / 1.008 * math.pi / 2).pow(2)
|
||||
alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
|
||||
|
||||
# Map sigmas to timesteps
|
||||
scaled_sigmas = (
|
||||
sigmas / sigmas.max() if sigmas.max() > 0 else torch.zeros_like(sigmas)
|
||||
)
|
||||
timesteps_for_sigma = (1 - scaled_sigmas) * (num_train_timesteps - 1)
|
||||
indices = (
|
||||
torch.round(timesteps_for_sigma).long().clamp(0, num_train_timesteps - 1)
|
||||
)
|
||||
|
||||
alpha_prod_t = alphas_cumprod[indices].view(-1, 1, 1, 1)
|
||||
beta_prod_t = 1 - alpha_prod_t
|
||||
|
||||
self.alpha_prod_t = alpha_prod_t
|
||||
self.beta_prod_t = beta_prod_t
|
||||
self.alpha_prod_t_sqrt = torch.sqrt(alpha_prod_t)
|
||||
self.beta_prod_t_sqrt = torch.sqrt(beta_prod_t)
|
||||
|
||||
print(f"[Alpha/Beta] Final stats:")
|
||||
print(
|
||||
f" alpha_prod_t: min={self.alpha_prod_t.min():.4f}, max={self.alpha_prod_t.max():.4f}, mean={self.alpha_prod_t.mean():.4f}"
|
||||
)
|
||||
print(
|
||||
f" beta_prod_t: min={self.beta_prod_t.min():.4f}, max={self.beta_prod_t.max():.4f}, mean={self.beta_prod_t.mean():.4f}"
|
||||
)
|
||||
print(
|
||||
f" alpha_sqrt: min={self.alpha_prod_t_sqrt.min():.4f}, max={self.alpha_prod_t_sqrt.max():.4f}, mean={self.alpha_prod_t_sqrt.mean():.4f}"
|
||||
)
|
||||
print(
|
||||
f" beta_sqrt: min={self.beta_prod_t_sqrt.min():.4f}, max={self.beta_prod_t_sqrt.max():.4f}, mean={self.beta_prod_t_sqrt.mean():.4f}\n"
|
||||
)
|
||||
|
||||
def scheduler_step_batch(self, denoised_batch, scaled_noise):
|
||||
"""Compute scheduler step for the entire batch (LCM-like)"""
|
||||
print("\n[Scheduler] Starting scheduler step (LCM-like):")
|
||||
print(f" Denoised shape: {denoised_batch.shape}")
|
||||
print(f" Scaled noise shape: {scaled_noise.shape}")
|
||||
print(
|
||||
f" Denoised stats: min={denoised_batch.min():.4f}, max={denoised_batch.max():.4f}, mean={denoised_batch.mean():.4f}"
|
||||
)
|
||||
print(
|
||||
f" Scaled noise stats: min={scaled_noise.min():.4f}, max={scaled_noise.max():.4f}, mean={scaled_noise.mean():.4f}"
|
||||
)
|
||||
|
||||
# Assuming self.alpha_prod_t and self.beta_prod_t are for the current timestep
|
||||
alpha_t_sqrt = self.alpha_prod_t_sqrt[: denoised_batch.shape[0]]
|
||||
beta_t_sqrt = self.beta_prod_t_sqrt[: denoised_batch.shape[0]]
|
||||
|
||||
# Calculate the predicted original sample (x_0) from the noisy sample and model output
|
||||
# This is similar to the inverse of the forward diffusion step
|
||||
pred_original_sample = (
|
||||
denoised_batch - beta_t_sqrt.view(-1, 1, 1, 1) * scaled_noise
|
||||
) / (alpha_t_sqrt.view(-1, 1, 1, 1) + 1e-7)
|
||||
print(
|
||||
f" Predicted original sample stats: min={pred_original_sample.min():.4f}, max={pred_original_sample.max():.4f}, mean={pred_original_sample.mean():.4f}"
|
||||
)
|
||||
|
||||
# For the reverse step, we need the alpha and beta for the previous timestep
|
||||
# Since we are operating on a batch of timesteps, we'll use the values corresponding to the first timestep in the batch for updating stock_noise
|
||||
# In StreamDiffusion paper, the update happens based on the result of the last denoising step
|
||||
|
||||
# Get alpha and beta for the next step: (which would be the previous in the denoising sequence)
|
||||
alpha_next_sqrt = torch.cat(
|
||||
[
|
||||
self.alpha_prod_t_sqrt[1 : denoised_batch.shape[0]],
|
||||
torch.ones_like(self.alpha_prod_t_sqrt[0:1]),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
|
||||
# Calculate the coefficient for the predicted original sample in the next step
|
||||
sigma_next = (1 - alpha_next_sqrt.pow(2)) / (alpha_next_sqrt + 1e-7)
|
||||
sigma_t = (1 - alpha_t_sqrt.pow(2)) / (alpha_t_sqrt + 1e-7)
|
||||
|
||||
# Guidance scale (set to 1 for unconditional prediction)
|
||||
guidance_scale = 1
|
||||
|
||||
# Calculate delta_x (the change in the noisy sample)
|
||||
delta_x = (sigma_next - sigma_t) * guidance_scale * pred_original_sample
|
||||
|
||||
print(
|
||||
f" Final delta_x stats: min={delta_x.min():.4f}, max={delta_x.max():.4f}, mean={delta_x.mean():.4f}\n"
|
||||
)
|
||||
|
||||
return delta_x
|
||||
|
||||
def sample(self, model, noise, sigmas, extra_args=None, callback=None, disable=None):
|
||||
"""Sample with staggered batch denoising steps - Optimized version with multiple CFG types"""
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
cond_scale = extra_args.get("model_options", {}).get("cfg_scale", 1.0) # Get cond_scale from model options
|
||||
cfg_type = self.cfg_type
|
||||
|
||||
# Get number of frames in batch and available sigmas
|
||||
batch_size = noise.shape[0]
|
||||
num_sigmas = len(sigmas) - 1 # Subtract 1 because last sigma is the target (0.0)
|
||||
|
||||
# Precompute alpha/beta if needed
|
||||
self.compute_alpha_beta(sigmas) # Sigmas include the final 0.0, compute uses all
|
||||
|
||||
# Reuse zeros buffer for txt2img detection
|
||||
if self.zeros_reference is None or self.zeros_reference.device != noise.device or self.zeros_reference.dtype != noise.dtype:
|
||||
self.zeros_reference = torch.zeros(1, device=noise.device, dtype=noise.dtype)
|
||||
|
||||
# Check if noise tensor is all zeros
|
||||
self.is_txt2img_mode = torch.abs(noise).sum() < 1e-5
|
||||
|
||||
# Initialize init_noise sequence like reference if needed (once or if shape changes)
|
||||
noise_shape_base = noise.shape[1:] # C, H, W
|
||||
if self.init_noise is None or self.init_noise.shape[0] != num_sigmas or self.init_noise.shape[1:] != noise_shape_base:
|
||||
self.init_noise = torch.randn(
|
||||
(num_sigmas, *noise_shape_base),
|
||||
device=noise.device,
|
||||
dtype=noise.dtype,
|
||||
)
|
||||
# Initialize stock_noise to zeros when init_noise is created (like reference)
|
||||
self.stock_noise = torch.zeros_like(self.init_noise[0])
|
||||
self.step_counter = 0 # Reset step counter when noise shape changes
|
||||
|
||||
# Shift noise sequence
|
||||
if self.step_counter > 0:
|
||||
self.init_noise = torch.cat([self.init_noise[1:], self.init_noise[0:1]], dim=0)
|
||||
# Update stock noise based on the *first* element of the *shifted* init_noise
|
||||
self.stock_noise = self.init_noise[0].clone()
|
||||
|
||||
|
||||
if self.is_txt2img_mode:
|
||||
# If txt2img mode, use the init_noise sequence directly as the base noise for the batch
|
||||
# This ensures consistent noise across steps/frames for txt2img
|
||||
if self.random_noise_buffer is None or self.random_noise_buffer.shape != noise.shape:
|
||||
self.random_noise_buffer = torch.empty_like(noise)
|
||||
|
||||
# Copy from the init_noise sequence for the current batch size
|
||||
self.random_noise_buffer.copy_(self.init_noise[:batch_size])
|
||||
x = self.random_noise_buffer # Use pre-allocated buffer
|
||||
else:
|
||||
# If not txt2img, we'll still need to add noise later based on sigmas
|
||||
x = noise # Start with the input noise (img2img)
|
||||
|
||||
# Verify batch size matches number of timesteps
|
||||
if batch_size != num_sigmas:
|
||||
raise ValueError(f"Batch size ({batch_size}) must match number of sigmas ({num_sigmas})")
|
||||
|
||||
# Pre-allocate and reuse view buffer for sigmas
|
||||
# Ensure buffer matches the number of steps *in this batch* (num_sigmas)
|
||||
if self.sigmas_view_buffer is None or self.sigmas_view_buffer.shape[0] != num_sigmas:
|
||||
self.sigmas_view_buffer = torch.empty((num_sigmas, 1, 1, 1),
|
||||
device=sigmas.device,
|
||||
dtype=sigmas.dtype)
|
||||
# In-place copy of sigmas view (use sigmas[:-1] as these are the steps)
|
||||
self.sigmas_view_buffer.copy_(sigmas[:-1].view(-1, 1, 1, 1))
|
||||
|
||||
# Apply noise with pre-allocated buffers - no new memory allocation
|
||||
if not self.is_txt2img_mode:
|
||||
# For img2img, add noise based on sigmas and init_noise sequence
|
||||
# Expand the relevant part of init_noise for the batch
|
||||
current_batch_init_noise = self.init_noise[:batch_size].to(x.device) # Ensure device match
|
||||
|
||||
# If x is the same object as noise, use a working buffer
|
||||
if id(x) == id(noise):
|
||||
if self.working_buffer is None or self.working_buffer.shape != noise.shape:
|
||||
self.working_buffer = torch.empty_like(noise)
|
||||
x = self.working_buffer
|
||||
torch.add(noise, current_batch_init_noise * self.sigmas_view_buffer, out=x)
|
||||
else: # Otherwise, add noise directly to x
|
||||
torch.add(noise, current_batch_init_noise * self.sigmas_view_buffer, out=x)
|
||||
|
||||
# Initialize and manage latent buffer with memory optimization
|
||||
# Only buffer if more than one step
|
||||
if num_sigmas > 1:
|
||||
if self.x_t_latent_buffer is None or self.x_t_latent_buffer.shape != x[0].shape or self.is_txt2img_mode:
|
||||
# Pre-allocate or resize as needed
|
||||
if self.x_t_latent_buffer is None or self.x_t_latent_buffer.shape != x[0].shape:
|
||||
self.x_t_latent_buffer = torch.empty_like(x[0])
|
||||
self.x_t_latent_buffer.copy_(x[0])
|
||||
|
||||
# Use buffer for first frame to maintain temporal consistency
|
||||
x[0].copy_(self.x_t_latent_buffer)
|
||||
else: # Single step, no buffering needed
|
||||
self.x_t_latent_buffer = None
|
||||
|
||||
with torch.no_grad():
|
||||
# Process all frames in parallel
|
||||
sigma_batch = sigmas[:-1] # Use sigmas for the steps, excluding the final 0.0
|
||||
|
||||
# --- CFG Handling ---
|
||||
denoised_batch = None
|
||||
stock_noise_update_value = None # Track value to update stock noise with (used in self-cfg)
|
||||
|
||||
# 1. No CFG
|
||||
if cfg_type == "none" or cond_scale <= 1.0:
|
||||
denoised_batch = model(x, sigma_batch, **extra_args)
|
||||
stock_noise_update_value = denoised_batch[0].clone().detach()
|
||||
|
||||
# 2. Full CFG
|
||||
elif cfg_type == "full":
|
||||
x_double = torch.cat([x, x], dim=0)
|
||||
sigma_double = torch.cat([sigma_batch, sigma_batch], dim=0)
|
||||
model_output = model(x_double, sigma_double, **extra_args)
|
||||
noise_pred_uncond, noise_pred_text = model_output.chunk(2)
|
||||
denoised_batch = noise_pred_uncond + cond_scale * (noise_pred_text - noise_pred_uncond)
|
||||
# Update stock noise with uncond prediction (using first frame)
|
||||
stock_noise_update_value = noise_pred_uncond[0].clone().detach()
|
||||
|
||||
# 3. Initialize CFG
|
||||
elif cfg_type == "initialize":
|
||||
x_plus_uc = torch.cat([x[0:1], x], dim=0) # Add first frame again for uncond
|
||||
sigma_plus_uc = torch.cat([sigma_batch[0:1], sigma_batch], dim=0)
|
||||
model_output = model(x_plus_uc, sigma_plus_uc, **extra_args)
|
||||
noise_pred_uncond = model_output[0:1] # Uncond is the first item
|
||||
noise_pred_cond = model_output[1:] # Cond are the rest
|
||||
# Calculate final denoised batch using CFG scale
|
||||
denoised_batch = noise_pred_uncond + cond_scale * (noise_pred_cond - noise_pred_uncond)
|
||||
# Update stock noise with uncond prediction
|
||||
stock_noise_update_value = noise_pred_uncond[0].clone().detach()
|
||||
|
||||
# 4. Self CFG (RCFG)
|
||||
elif cfg_type == "self":
|
||||
eps_c = model(x, sigma_batch, **extra_args) # Conditional prediction
|
||||
|
||||
# Initialize stock noise if needed (first run after init_noise creation)
|
||||
if self.stock_noise is None: # Should have been initialized with init_noise
|
||||
# Fallback: initialize with first frame prediction if somehow missed
|
||||
self.stock_noise = eps_c[0].clone().detach()
|
||||
|
||||
# RCFG equation
|
||||
# Use current stock_noise (which should be based on init_noise[0])
|
||||
noise_pred_uncond = self.stock_noise.to(eps_c.device) * self.delta
|
||||
denoised_batch = noise_pred_uncond + cond_scale * (eps_c - noise_pred_uncond)
|
||||
|
||||
# Defer stock noise update until after scheduler step below
|
||||
# Store the conditional prediction for the scheduler step
|
||||
stock_noise_update_value = eps_c.clone().detach() # Need the whole batch for scheduler
|
||||
|
||||
|
||||
# --- CFG Handling End ---
|
||||
|
||||
# --- RCFG Stock Noise Update ---
|
||||
if cfg_type == "self" and stock_noise_update_value is not None:
|
||||
# Get the noise sequence for the current batch
|
||||
current_batch_init_noise = self.init_noise[:batch_size].to(denoised_batch.device)
|
||||
|
||||
|
||||
beta_sqrt_batch = self.beta_prod_t_sqrt[:batch_size].to(denoised_batch.device)
|
||||
scaled_noise = beta_sqrt_batch * current_batch_init_noise # Use the init_noise for this step
|
||||
|
||||
delta_x = self.scheduler_step_batch(denoised_batch, scaled_noise)
|
||||
num_steps_in_batch = denoised_batch.shape[0]
|
||||
if self.alpha_prod_t_sqrt.shape[0] >= num_steps_in_batch:
|
||||
one_tensor = torch.ones_like(self.alpha_prod_t_sqrt[0:1])
|
||||
zero_tensor = torch.zeros_like(self.beta_prod_t_sqrt[0:1])
|
||||
alpha_next_sqrt = torch.cat([self.alpha_prod_t_sqrt[1:num_steps_in_batch], one_tensor], dim=0).to(delta_x.device)
|
||||
beta_next_sqrt = torch.cat([self.beta_prod_t_sqrt[1:num_steps_in_batch], zero_tensor], dim=0).to(delta_x.device)
|
||||
else:
|
||||
raise ValueError("Precomputed alpha/beta shapes mismatch for final scaling.")
|
||||
|
||||
# Apply the scaling
|
||||
scaled_delta_x = (alpha_next_sqrt / (beta_next_sqrt + 1e-7)) * delta_x
|
||||
|
||||
# And the *scaled* delta_x calculated for the *first frame* of the batch
|
||||
self.stock_noise += scaled_delta_x[0].to(self.stock_noise.device) # Add scaled delta for first frame
|
||||
elif stock_noise_update_value is not None:
|
||||
# For non-RCFG modes, update stock noise simply with the first frame's value
|
||||
if self.stock_noise is None or self.stock_noise.shape != stock_noise_update_value.shape:
|
||||
self.stock_noise = stock_noise_update_value.clone().detach()
|
||||
else:
|
||||
self.stock_noise.copy_(stock_noise_update_value.detach())
|
||||
|
||||
|
||||
# Update buffer with intermediate results
|
||||
if num_sigmas > 1:
|
||||
# Store result from first frame as buffer for next iteration
|
||||
self.x_t_latent_buffer.copy_(denoised_batch[0]) # In-place update
|
||||
|
||||
# Pre-allocate output buffer
|
||||
output_shape = (1, *denoised_batch[-1].shape)
|
||||
if self.output_buffer is None or self.output_buffer.shape != output_shape:
|
||||
self.output_buffer = torch.empty(output_shape,
|
||||
device=denoised_batch.device,
|
||||
dtype=denoised_batch.dtype)
|
||||
# Copy the result directly to pre-allocated buffer
|
||||
self.output_buffer[0].copy_(denoised_batch[-1])
|
||||
x_0_pred_out = self.output_buffer
|
||||
else: # Single step
|
||||
x_0_pred_out = denoised_batch.unsqueeze(0) # Ensure batch dim [1,C,H,W]
|
||||
self.x_t_latent_buffer = None # No buffer needed
|
||||
|
||||
# Call callback if provided
|
||||
if callback is not None:
|
||||
callback({'x': x_0_pred_out, 'i': self.step_counter, 'sigma': sigmas[0], 'sigma_hat': sigmas[0], 'denoised': denoised_batch[-1:]})
|
||||
|
||||
# Increment step counter at the end of the sample method
|
||||
self.step_counter += 1
|
||||
|
||||
return x_0_pred_out
|
||||
|
||||
def update(self, cfg_type="self", delta=1.0):
|
||||
"""Create sampler with specified settings"""
|
||||
self.cfg_type = cfg_type
|
||||
self.delta = delta
|
||||
sampler = comfy.samplers.KSAMPLER(self.sample)
|
||||
return (sampler,)
|
||||
|
||||
|
||||
class StreamScheduler:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"t_index_list": ("STRING", {
|
||||
"default": "0,16,32,49",
|
||||
"tooltip": "Comma-separated list of timesteps to actually use for denoising. Examples: '32,45' for img2img or '0,16,32,45' for txt2img"
|
||||
}),
|
||||
"num_inference_steps": ("INT", {
|
||||
"default": 50,
|
||||
"min": 1,
|
||||
"max": 1000,
|
||||
"step": 1,
|
||||
"tooltip": "Total number of timesteps in schedule. StreamDiffusion uses 50 by default. Only timesteps specified in t_index_list are actually used."
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SIGMAS",)
|
||||
FUNCTION = "update"
|
||||
CATEGORY = "StreamPack/sampling"
|
||||
DESCRIPTION = "Implements StreamDiffusion's efficient timestep selection. Use in conjunction with StreamBatchSampler."
|
||||
|
||||
def update(self, model, t_index_list="32,45", num_inference_steps=50):
|
||||
# Get model's sampling parameters
|
||||
model_sampling = model.get_model_object("model_sampling")
|
||||
|
||||
# Parse timestep list
|
||||
try:
|
||||
t_index_list = [int(t.strip()) for t in t_index_list.split(",")]
|
||||
except ValueError as e:
|
||||
|
||||
t_index_list = [32, 45]
|
||||
|
||||
# Create full schedule using normal scheduler
|
||||
full_sigmas = comfy.samplers.normal_scheduler(model_sampling, num_inference_steps)
|
||||
|
||||
# Select only the sigmas at our desired indices, but in reverse order
|
||||
# This ensures we go from high noise to low noise
|
||||
selected_sigmas = []
|
||||
for t in sorted(t_index_list, reverse=True): # Sort in reverse to go from high noise to low
|
||||
if t < 0 or t >= num_inference_steps:
|
||||
|
||||
continue
|
||||
selected_sigmas.append(float(full_sigmas[t]))
|
||||
|
||||
# Add final sigma
|
||||
selected_sigmas.append(0.0)
|
||||
|
||||
# Convert to tensor and move to appropriate device
|
||||
selected_sigmas = torch.FloatTensor(selected_sigmas).to(comfy.model_management.get_torch_device())
|
||||
return (selected_sigmas,)
|
||||
|
||||
|
||||
class StreamFrameBuffer:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"latent": ("LATENT",),
|
||||
"buffer_size": ("INT", {
|
||||
"default": 4,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"tooltip": "Number of frames to buffer before starting batch processing. Should match number of denoising steps."
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "update"
|
||||
CATEGORY = "StreamPack/sampling"
|
||||
DESCRIPTION = "Accumulates frames to enable staggered batch denoising like StreamDiffusion. Use in conjunction with StreamBatchSampler"
|
||||
|
||||
|
||||
def __init__(self):
|
||||
self.frame_buffer = None # Tensor of shape [buffer_size, C, H, W]
|
||||
self.buffer_size = None
|
||||
self.buffer_pos = 0 # Current position
|
||||
self.is_initialized = False # Track buffer initialization
|
||||
self.is_txt2img_mode = False
|
||||
|
||||
def update(self, latent, buffer_size=4):
|
||||
"""Add new frame to buffer and return batch when ready"""
|
||||
self.buffer_size = buffer_size
|
||||
|
||||
# Extract latent tensor from input and remove batch dimension if present
|
||||
x = latent["samples"]
|
||||
|
||||
# Check if this is a txt2img (zeros tensor) or img2img mode
|
||||
# In ComfyUI, EmptyLatentImage returns a zeros tensor with shape [batch_size, 4, h//8, w//8]
|
||||
# We consider it txt2img mode if the tensor contains all zeros
|
||||
is_txt2img_mode = torch.sum(torch.abs(x)) < 1e-6
|
||||
self.is_txt2img_mode = is_txt2img_mode
|
||||
|
||||
# If it's a batch with size 1, remove the batch dimension for our buffer
|
||||
if x.dim() == 4 and x.shape[0] == 1: # [1,C,H,W]
|
||||
x = x.squeeze(0) # Remove batch dimension -> [C,H,W]
|
||||
|
||||
# Initialize or resize frame_buffer as a tensor
|
||||
if not self.is_initialized or self.frame_buffer.shape[0] != self.buffer_size or \
|
||||
self.frame_buffer.shape[1:] != x.shape:
|
||||
# Pre-allocate buffer with correct shape
|
||||
self.frame_buffer = torch.zeros(
|
||||
(self.buffer_size, *x.shape),
|
||||
device=x.device,
|
||||
dtype=x.dtype
|
||||
)
|
||||
if self.is_txt2img_mode or not self.is_initialized:
|
||||
# Optimization: Use broadcasting to fill buffer with copies
|
||||
self.frame_buffer[:] = x.unsqueeze(0) # Broadcast x to [buffer_size, C, H, W]
|
||||
|
||||
self.is_initialized = True
|
||||
self.buffer_pos = 0
|
||||
else:
|
||||
# Add new frame to buffer using ring buffer logic
|
||||
self.frame_buffer[self.buffer_pos] = x # In-place update
|
||||
self.buffer_pos = (self.buffer_pos + 1) % self.buffer_size # Circular increment
|
||||
|
||||
# Optimization: frame_buffer is already a tensor batch, no need to stack
|
||||
batch = self.frame_buffer
|
||||
|
||||
|
||||
# Return as latent dict with preserved dimensions
|
||||
result = {"samples": batch}
|
||||
|
||||
# Preserve height and width if present in input
|
||||
if "height" in latent:
|
||||
result["height"] = latent["height"]
|
||||
if "width" in latent:
|
||||
result["width"] = latent["width"]
|
||||
|
||||
return (result,)
|
||||
@@ -0,0 +1 @@
|
||||
# Dummy file ensuring this package will be recognized.
|
||||
@@ -0,0 +1,536 @@
|
||||
"""Source code for 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 typing import Optional, Deque, List, Tuple, Dict, Any
|
||||
from collections import deque
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from diffusers.utils import USE_PEFT_BACKEND
|
||||
|
||||
|
||||
def get_nn_feats(
|
||||
x: torch.Tensor, y: torch.Tensor, threshold: float = 0.9
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Selects features from x or nearest neighbors from y based on cosine similarity.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): Input features tensor.
|
||||
y (torch.Tensor): Neighbor features tensor.
|
||||
threshold (float): Cosine similarity threshold for neighbor selection.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Selected features tensor.
|
||||
"""
|
||||
if isinstance(x, deque):
|
||||
x = torch.cat(list(x), dim=1)
|
||||
if isinstance(y, deque):
|
||||
y = torch.cat(list(y), dim=1)
|
||||
|
||||
x_norm = F.normalize(x, p=2, dim=-1)
|
||||
y_norm = F.normalize(y, p=2, dim=-1)
|
||||
|
||||
cosine_similarity = torch.matmul(x_norm, y_norm.transpose(1, 2))
|
||||
max_cosine_values, nearest_neighbors_indices = torch.max(cosine_similarity, dim=-1)
|
||||
mask = max_cosine_values < threshold
|
||||
|
||||
indices_expanded = nearest_neighbors_indices.unsqueeze(-1).expand(
|
||||
-1, -1, x_norm.size(-1)
|
||||
)
|
||||
nearest_neighbor_tensor = torch.gather(y, 1, indices_expanded)
|
||||
selected_tensor = torch.where(mask.unsqueeze(-1), x, nearest_neighbor_tensor)
|
||||
|
||||
return selected_tensor
|
||||
|
||||
|
||||
class FeatureBank:
|
||||
"""
|
||||
Feature bank to cache and retrieve attention keys, values, and outputs.
|
||||
|
||||
Args:
|
||||
max_frames (int): Maximum number of frames to store in the cache.
|
||||
"""
|
||||
|
||||
def __init__(self, max_frames: int = 4):
|
||||
self.cached_key: Deque[torch.FloatTensor] = deque(maxlen=max_frames)
|
||||
self.cached_value: Deque[torch.FloatTensor] = deque(maxlen=max_frames)
|
||||
self.cached_output: Deque[torch.FloatTensor] = deque(maxlen=max_frames)
|
||||
|
||||
def add(
|
||||
self,
|
||||
key: torch.FloatTensor,
|
||||
value: torch.FloatTensor,
|
||||
output: torch.FloatTensor,
|
||||
):
|
||||
"""Add key, value, and output tensors to the feature bank."""
|
||||
self.cached_key.append(key.clone())
|
||||
self.cached_value.append(value.clone())
|
||||
self.cached_output.append(output.clone())
|
||||
|
||||
def get(
|
||||
self,
|
||||
) -> Tuple[
|
||||
List[torch.FloatTensor], List[torch.FloatTensor], List[torch.FloatTensor]
|
||||
]:
|
||||
"""Retrieve cached keys, values, and outputs as lists."""
|
||||
return list(self.cached_key), list(self.cached_value), list(self.cached_output)
|
||||
|
||||
def clear(self):
|
||||
"""Clear all cached features."""
|
||||
self.cached_key.clear()
|
||||
self.cached_value.clear()
|
||||
self.cached_output.clear()
|
||||
|
||||
|
||||
class FeatureBankAttentionProcessor:
|
||||
"""
|
||||
Attention processor with feature injection.
|
||||
"""
|
||||
|
||||
DEBUG_PRINTS = (
|
||||
False # Class-level debug flag, turn on/off all debug prints from this class
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
use_feature_injection=True,
|
||||
feature_injection_strength=0.8,
|
||||
feature_similarity_threshold=0.98,
|
||||
interval=4,
|
||||
max_frames=4,
|
||||
):
|
||||
self.use_feature_injection = use_feature_injection
|
||||
self.fi_strength = feature_injection_strength
|
||||
self.threshold = feature_similarity_threshold
|
||||
self.frame_id = 0
|
||||
self.interval = interval
|
||||
self.feature_bank = FeatureBank(max_frames=max_frames)
|
||||
self.key_projection: Optional[nn.Linear] = (
|
||||
None # Linear layer for key projection
|
||||
)
|
||||
self.value_projection: Optional[nn.Linear] = (
|
||||
None # Linear layer for value projection
|
||||
)
|
||||
self.output_projection: Optional[nn.Linear] = (
|
||||
None # Linear layer for output projection
|
||||
)
|
||||
|
||||
def _debug_print(self, *args, **kwargs):
|
||||
"""Conditional debug print based on class-level DEBUG_PRINTS flag."""
|
||||
if self.DEBUG_PRINTS:
|
||||
print("[FeatureBankAttentionProcessor DEBUG]", *args, **kwargs)
|
||||
|
||||
def _project_tensor(
|
||||
self,
|
||||
tensor: torch.Tensor,
|
||||
projection_layer: Optional[nn.Linear],
|
||||
expected_dim: int,
|
||||
layer_type: str,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Handles tensor projection if dimension mismatch is detected.
|
||||
Creates projection layer if needed and applies it.
|
||||
"""
|
||||
if (
|
||||
projection_layer is None
|
||||
or projection_layer.in_features != tensor.shape[-1]
|
||||
or projection_layer.out_features != expected_dim
|
||||
):
|
||||
self._debug_print(
|
||||
f"\n--- DEBUG: {layer_type} dimension mismatch! Projecting cached {layer_type.lower()}s. ---"
|
||||
)
|
||||
self._debug_print(
|
||||
f"Original {layer_type.lower()} dimension: {tensor.shape[-1]}, Expected: {expected_dim}"
|
||||
)
|
||||
new_projection_layer = (
|
||||
nn.Linear(tensor.shape[-1], expected_dim)
|
||||
.to(tensor.device)
|
||||
.to(tensor.dtype)
|
||||
)
|
||||
self._debug_print(
|
||||
f"Creating new {layer_type.lower()} projection layer: In Features={tensor.shape[-1]}, Out Features={expected_dim}"
|
||||
)
|
||||
if layer_type == "Key":
|
||||
self.key_projection = new_projection_layer
|
||||
elif layer_type == "Value":
|
||||
self.value_projection = new_projection_layer
|
||||
elif layer_type == "Output":
|
||||
self.output_projection = new_projection_layer
|
||||
projection_layer = (
|
||||
new_projection_layer # Use the newly created or updated layer
|
||||
)
|
||||
|
||||
projected_tensor = projection_layer(tensor)
|
||||
self._debug_print(
|
||||
f"Shape of cached {layer_type.lower()}s after projection: {projected_tensor.shape}"
|
||||
)
|
||||
return projected_tensor
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: nn.Module,
|
||||
hidden_states: torch.FloatTensor,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
scale: float = 1.0,
|
||||
) -> torch.FloatTensor:
|
||||
"""
|
||||
Forward pass of feature bank attention with feature injection.
|
||||
"""
|
||||
residual = hidden_states
|
||||
args = () if USE_PEFT_BACKEND else (scale,)
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
if input_ndim == 4:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(
|
||||
batch_size, channel, height * width
|
||||
).transpose(1, 2)
|
||||
|
||||
batch_size, key_tokens, _ = (
|
||||
hidden_states.shape
|
||||
if encoder_hidden_states is None
|
||||
else encoder_hidden_states.shape
|
||||
)
|
||||
|
||||
query = attn.to_q(hidden_states)
|
||||
|
||||
# Self-attention detection (heuristic based on input feature dimensions)
|
||||
is_self_attn = attn.to_k.in_features == attn.to_q.in_features
|
||||
current_encoder_hidden_states = (
|
||||
hidden_states if is_self_attn else encoder_hidden_states
|
||||
)
|
||||
|
||||
self._debug_print("\n--- Attention Module Info ---")
|
||||
self._debug_print("Module type:", type(attn))
|
||||
self._debug_print("Module:", attn)
|
||||
self._debug_print("Is self-attention (heuristic):", is_self_attn)
|
||||
self._debug_print("Shape of hidden_states:", hidden_states.shape)
|
||||
self._debug_print(
|
||||
"Shape of current_encoder_hidden_states before attn.to_k:",
|
||||
(
|
||||
current_encoder_hidden_states.shape
|
||||
if current_encoder_hidden_states is not None
|
||||
else "None"
|
||||
),
|
||||
)
|
||||
self._debug_print("Shape of attn.to_k.weight:", attn.to_k.weight.shape)
|
||||
|
||||
key = attn.to_k(current_encoder_hidden_states)
|
||||
value = attn.to_v(current_encoder_hidden_states)
|
||||
|
||||
cached_key_tensor = key.clone()
|
||||
cached_value_tensor = value.clone()
|
||||
|
||||
num_heads = attn.heads # Assuming attn.heads exists in CrossAttention
|
||||
head_dim = query.shape[-1] // num_heads
|
||||
query = (
|
||||
query.view(batch_size, -1, num_heads, head_dim).transpose(1, 2).contiguous()
|
||||
) # [batch_size, num_heads, seq_len_q, head_dim]
|
||||
|
||||
if is_self_attn:
|
||||
(
|
||||
cached_keys_list,
|
||||
cached_values_list,
|
||||
cached_outputs,
|
||||
) = self.feature_bank.get()
|
||||
|
||||
self._debug_print(
|
||||
"\n--- Feature Bank Get Output ---"
|
||||
) # Separator for clarity
|
||||
self._debug_print(
|
||||
f"Type of cached_keys_list from feature_bank.get(): {type(cached_keys_list)}"
|
||||
)
|
||||
self._debug_print(
|
||||
f"Content of cached_keys_list from feature_bank.get(): {cached_keys_list}"
|
||||
)
|
||||
|
||||
if (
|
||||
cached_keys_list
|
||||
and cached_keys_list[0] is not None
|
||||
and torch.is_tensor(cached_keys_list[0])
|
||||
):
|
||||
cached_key_tensor_from_list = cached_keys_list[0]
|
||||
cached_key_tensor_from_list = self._project_tensor(
|
||||
cached_key_tensor_from_list,
|
||||
self.key_projection,
|
||||
key.shape[-1],
|
||||
"Key",
|
||||
)
|
||||
|
||||
self._debug_print("\n--- DEBUG: Before torch.cat (Keys) ---")
|
||||
self._debug_print(
|
||||
"Module:", attn
|
||||
) # Print the module itself for identification
|
||||
self._debug_print("Type of key:", type(key))
|
||||
self._debug_print("Shape of key:", key.shape)
|
||||
self._debug_print(
|
||||
"Type of cached_key_tensor_from_list:",
|
||||
type(cached_key_tensor_from_list),
|
||||
)
|
||||
self._debug_print(
|
||||
"Shape of cached_key_tensor_from_list:",
|
||||
cached_key_tensor_from_list.shape,
|
||||
)
|
||||
|
||||
key = torch.cat([key, cached_key_tensor_from_list], dim=1)
|
||||
|
||||
if (
|
||||
cached_values_list
|
||||
and cached_values_list[0] is not None
|
||||
and torch.is_tensor(cached_values_list[0])
|
||||
):
|
||||
cached_value_tensor_from_list = cached_values_list[0]
|
||||
cached_value_tensor_from_list = self._project_tensor(
|
||||
cached_value_tensor_from_list,
|
||||
self.value_projection,
|
||||
value.shape[-1],
|
||||
"Value",
|
||||
)
|
||||
|
||||
self._debug_print("\n--- DEBUG: Before torch.cat (Values) ---")
|
||||
self._debug_print(
|
||||
"Module:", attn
|
||||
) # Print the module itself for identification
|
||||
self._debug_print("Type of value:", type(value))
|
||||
self._debug_print("Shape of value:", value.shape)
|
||||
self._debug_print(
|
||||
"Type of cached_value_tensor_from_list:",
|
||||
type(cached_value_tensor_from_list),
|
||||
)
|
||||
self._debug_print(
|
||||
"Shape of cached_value_tensor_from_list:",
|
||||
cached_value_tensor_from_list.shape,
|
||||
)
|
||||
|
||||
value = torch.cat([value, cached_value_tensor_from_list], dim=1)
|
||||
|
||||
key = (
|
||||
key.view(-1, key.shape[1], num_heads, head_dim)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
) # [batch_size, num_heads, seq_len_kv, head_dim] (seq_len_kv == key_tokens)
|
||||
value = (
|
||||
value.view(-1, value.shape[1], num_heads, head_dim)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
) # [batch_size, num_heads, seq_len_kv, head_dim] (seq_len_kv == key_tokens)
|
||||
else:
|
||||
key = (
|
||||
key.view(batch_size, -1, num_heads, head_dim)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
) # [batch_size, num_heads, seq_len_kv, head_dim] (seq_len_kv == key_tokens)
|
||||
value = (
|
||||
value.view(batch_size, -1, num_heads, head_dim)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
) # [batch_size, num_heads, seq_len_kv, head_dim] (seq_len_kv == key_tokens)
|
||||
|
||||
# --- Manual attention implementation ---
|
||||
scale_factor = 1.0 / query.shape[-1] ** 0.5
|
||||
query = query * scale_factor
|
||||
attn_weights = torch.bmm(
|
||||
query.view(batch_size * num_heads, -1, head_dim),
|
||||
key.view(batch_size * num_heads, -1, head_dim).transpose(-2, -1),
|
||||
) # [batch_size * num_heads, seq_len_q, seq_len_kv]
|
||||
attn_weights = torch.softmax(attn_weights, dim=-1)
|
||||
attn_output = torch.bmm(
|
||||
attn_weights, value.view(batch_size * num_heads, -1, head_dim)
|
||||
) # [batch_size * num_heads, seq_len_q, head_dim]
|
||||
hidden_states = (
|
||||
attn_output.view(batch_size, num_heads, -1, head_dim)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
) # [batch_size, seq_len_q, num_heads * head_dim]
|
||||
# --- End of manual attention implementation ---
|
||||
|
||||
self._debug_print(f"Type of attn before to_out: {type(attn)}") # Debug print
|
||||
|
||||
hidden_states = hidden_states.to(
|
||||
query.dtype
|
||||
) # Keep this line to ensure correct dtype
|
||||
hidden_states = hidden_states.reshape(
|
||||
batch_size, -1, num_heads * head_dim
|
||||
).contiguous() # Modified reshape - removed transpose after reshape
|
||||
|
||||
hidden_states = attn.to_out[0](hidden_states, *args)
|
||||
hidden_states = attn.to_out[1](hidden_states)
|
||||
|
||||
if input_ndim == 4:
|
||||
hidden_states = hidden_states.transpose(-1, -2).reshape(
|
||||
batch_size, channel, height, width
|
||||
)
|
||||
|
||||
rescale_output_factor = 1.0
|
||||
cached_output_tensor = hidden_states.clone()
|
||||
|
||||
if self.use_feature_injection:
|
||||
_, _, cached_outputs = self.feature_bank.get()
|
||||
self._debug_print(
|
||||
f"Type of cached_outputs: {type(cached_outputs)}"
|
||||
) # Debug print
|
||||
|
||||
if (
|
||||
cached_outputs
|
||||
and cached_outputs[0] is not None
|
||||
and cached_outputs[0].numel() > 0
|
||||
): # Corrected condition
|
||||
cached_output_tensor_from_list = cached_outputs[
|
||||
0
|
||||
] # Get the cached output tensor
|
||||
cached_output_tensor_from_list = self._project_tensor(
|
||||
cached_output_tensor_from_list,
|
||||
self.output_projection,
|
||||
hidden_states.shape[-1],
|
||||
"Output",
|
||||
)
|
||||
|
||||
self._debug_print("\n--- DEBUG: Before get_nn_feats ---")
|
||||
self._debug_print("Shape of hidden_states:", hidden_states.shape)
|
||||
self._debug_print(
|
||||
"Shape of cached_output_tensor_from_list:",
|
||||
cached_output_tensor_from_list.shape,
|
||||
)
|
||||
|
||||
nn_hidden_states = get_nn_feats(
|
||||
hidden_states,
|
||||
cached_output_tensor_from_list,
|
||||
threshold=self.threshold,
|
||||
) # Pass PROJECTED tensor
|
||||
hidden_states = (
|
||||
hidden_states * (1 - self.fi_strength)
|
||||
+ self.fi_strength * nn_hidden_states
|
||||
)
|
||||
|
||||
hidden_states = hidden_states / rescale_output_factor
|
||||
|
||||
if self.frame_id % self.interval == 0 and is_self_attn:
|
||||
self.feature_bank.add(
|
||||
cached_key_tensor, cached_value_tensor, cached_output_tensor
|
||||
)
|
||||
|
||||
self.frame_id += 1
|
||||
return hidden_states
|
||||
|
||||
|
||||
class FeatureBankAttentionProcessorNode:
|
||||
"""
|
||||
ComfyUI Node for Feature Bank Attention Processor.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.cross_attention_hook = None
|
||||
self.feature_bank_processor = FeatureBankAttentionProcessor()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s) -> Dict[str, Dict[str, Tuple[str, Dict[str, Any]]]]:
|
||||
"""Input types for FeatureBankAttentionProcessorNode node."""
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"use_feature_injection": (
|
||||
"BOOLEAN",
|
||||
{"default": True},
|
||||
), # Configurable FI
|
||||
"feature_injection_strength": (
|
||||
"FLOAT",
|
||||
{"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.01},
|
||||
), # Configurable FI strength
|
||||
"feature_similarity_threshold": (
|
||||
"FLOAT",
|
||||
{"default": 0.98, "min": 0.0, "max": 1.0, "step": 0.01},
|
||||
), # Configurable threshold
|
||||
"feature_cache_interval": (
|
||||
"INT",
|
||||
{"default": 4, "min": 1, "max": 64},
|
||||
), # Configurable interval
|
||||
"feature_bank_max_frames": (
|
||||
"INT",
|
||||
{"default": 4, "min": 1, "max": 16},
|
||||
), # Configurable max frames in bank
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "update_node"
|
||||
CATEGORY = "feature-bank"
|
||||
|
||||
def update_node(
|
||||
self,
|
||||
model: Any,
|
||||
use_feature_injection: bool,
|
||||
feature_injection_strength: float,
|
||||
feature_similarity_threshold: float,
|
||||
feature_cache_interval: int,
|
||||
feature_bank_max_frames: int,
|
||||
) -> Tuple[Any]:
|
||||
"""
|
||||
Update method for FeatureBankAttentionProcessorNode. Sets up feature injection and attention hooks.
|
||||
"""
|
||||
print(f"[FeatureBankAttentionProcessorNode] Initializing ")
|
||||
|
||||
# Update FeatureBankProcessor configuration from node inputs on each update call
|
||||
self.feature_bank_processor.use_feature_injection = use_feature_injection
|
||||
self.feature_bank_processor.fi_strength = feature_injection_strength
|
||||
self.feature_bank_processor.threshold = feature_similarity_threshold
|
||||
self.feature_bank_processor.interval = feature_cache_interval
|
||||
self.feature_bank_processor.feature_bank = FeatureBank(
|
||||
max_frames=feature_bank_max_frames
|
||||
) # Re-instantiate feature bank with new max_frames if it changes
|
||||
|
||||
def cross_attention_forward(
|
||||
module: nn.Module,
|
||||
x: torch.Tensor,
|
||||
context: Optional[torch.Tensor] = None,
|
||||
mask: Optional[torch.Tensor] = None,
|
||||
value: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
) -> (
|
||||
torch.Tensor
|
||||
): # Added **kwargs to accept potential extra args from original forward
|
||||
"""Optimized cross attention using FeatureBankAttentionProcessor."""
|
||||
context = x if context is None else context
|
||||
encoder_hidden_states = context # Assuming 'context' in cross_attention_forward corresponds to encoder_hidden_states
|
||||
|
||||
return self.feature_bank_processor(
|
||||
attn=module, # Pass the current module as 'attn'
|
||||
hidden_states=x,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
scale=1.0, # Or pass scale if it's relevant, currently hardcoded to 1.0 in processor
|
||||
)
|
||||
|
||||
def hook_cross_attention(
|
||||
module: nn.Module, input: Tuple[torch.Tensor], output: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""Forward hook to replace cross-attention forward method."""
|
||||
if isinstance(module, torch.nn.Module) and hasattr(module, "to_q"):
|
||||
if not hasattr(module, "_original_forward"):
|
||||
module._original_forward = module.forward
|
||||
module.forward = lambda *args, **kwargs: cross_attention_forward(
|
||||
module, *args, **kwargs
|
||||
)
|
||||
return output
|
||||
|
||||
# Remove old hooks if they exist
|
||||
if self.cross_attention_hook is not None:
|
||||
self.cross_attention_hook.remove()
|
||||
|
||||
# Clone model and apply hooks
|
||||
m = model.clone()
|
||||
|
||||
def register_hooks(module: nn.Module):
|
||||
"""Recursively register forward hooks for cross-attention modules."""
|
||||
if isinstance(module, torch.nn.Module) and hasattr(module, "to_q"):
|
||||
self.cross_attention_hook = module.register_forward_hook(
|
||||
hook_cross_attention
|
||||
)
|
||||
|
||||
m.model.apply(register_hooks)
|
||||
return (m,)
|
||||
@@ -0,0 +1,608 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
import mediapipe as mp
|
||||
import time
|
||||
import torch
|
||||
|
||||
class FaceMeshUtils:
|
||||
"""
|
||||
Utility class for shared face mesh functionality.
|
||||
Contains drawing methods and specifications used by multiple nodes.
|
||||
"""
|
||||
def __init__(self):
|
||||
# Initialize MediaPipe drawing utilities
|
||||
self.mp_drawing = mp.solutions.drawing_utils
|
||||
self.mp_drawing_styles = mp.solutions.drawing_styles
|
||||
self.mp_face_mesh = mp.solutions.face_mesh
|
||||
|
||||
# Define drawing specifications for different facial features
|
||||
self.f_thick = 2
|
||||
self.f_rad = 1
|
||||
self.right_iris_draw = self.mp_drawing.DrawingSpec(color=(10, 200, 250), thickness=self.f_thick, circle_radius=self.f_rad)
|
||||
self.right_eye_draw = self.mp_drawing.DrawingSpec(color=(10, 200, 180), thickness=self.f_thick, circle_radius=self.f_rad)
|
||||
self.right_eyebrow_draw = self.mp_drawing.DrawingSpec(color=(10, 220, 180), thickness=self.f_thick, circle_radius=self.f_rad)
|
||||
self.left_iris_draw = self.mp_drawing.DrawingSpec(color=(250, 200, 10), thickness=self.f_thick, circle_radius=self.f_rad)
|
||||
self.left_eye_draw = self.mp_drawing.DrawingSpec(color=(180, 200, 10), thickness=self.f_thick, circle_radius=self.f_rad)
|
||||
self.left_eyebrow_draw = self.mp_drawing.DrawingSpec(color=(180, 220, 10), thickness=self.f_thick, circle_radius=self.f_rad)
|
||||
self.mouth_draw = self.mp_drawing.DrawingSpec(color=(10, 180, 10), thickness=self.f_thick, circle_radius=self.f_rad)
|
||||
self.head_draw = self.mp_drawing.DrawingSpec(color=(10, 200, 10), thickness=self.f_thick, circle_radius=self.f_rad)
|
||||
|
||||
# Create face connection specifications
|
||||
self.face_connection_spec = {}
|
||||
for edge in self.mp_face_mesh.FACEMESH_FACE_OVAL:
|
||||
self.face_connection_spec[edge] = self.head_draw
|
||||
for edge in self.mp_face_mesh.FACEMESH_LEFT_EYE:
|
||||
self.face_connection_spec[edge] = self.left_eye_draw
|
||||
for edge in self.mp_face_mesh.FACEMESH_LEFT_EYEBROW:
|
||||
self.face_connection_spec[edge] = self.left_eyebrow_draw
|
||||
for edge in self.mp_face_mesh.FACEMESH_RIGHT_EYE:
|
||||
self.face_connection_spec[edge] = self.right_eye_draw
|
||||
for edge in self.mp_face_mesh.FACEMESH_RIGHT_EYEBROW:
|
||||
self.face_connection_spec[edge] = self.right_eyebrow_draw
|
||||
for edge in self.mp_face_mesh.FACEMESH_LIPS:
|
||||
self.face_connection_spec[edge] = self.mouth_draw
|
||||
|
||||
# Iris landmark specification
|
||||
self.iris_landmark_spec = {468: self.right_iris_draw, 473: self.left_iris_draw}
|
||||
|
||||
def reverse_channels(self, image):
|
||||
"""Given a numpy array in RGB form, convert to BGR. Will also convert from BGR to RGB."""
|
||||
# im[:,:,::-1] is a neat hack to convert BGR to RGB by reversing the indexing order.
|
||||
return image[:, :, ::-1]
|
||||
|
||||
def draw_pupils(self, image, landmark_list, drawing_spec, halfwidth=2):
|
||||
"""Draw pupils with custom function for better control"""
|
||||
if len(image.shape) != 3:
|
||||
raise ValueError("Input image must be H,W,C.")
|
||||
image_rows, image_cols, image_channels = image.shape
|
||||
if image_channels != 3: # BGR channels
|
||||
raise ValueError('Input image must contain three channel bgr data.')
|
||||
|
||||
# Ensure image is contiguous in memory for OpenCV compatibility
|
||||
image = np.ascontiguousarray(image, dtype=np.uint8)
|
||||
|
||||
for idx, landmark in enumerate(landmark_list.landmark):
|
||||
if ((landmark.HasField('visibility') and landmark.visibility < 0.9) or
|
||||
(landmark.HasField('presence') and landmark.presence < 0.5)):
|
||||
continue
|
||||
if landmark.x >= 1.0 or landmark.x < 0 or landmark.y >= 1.0 or landmark.y < 0:
|
||||
continue
|
||||
image_x = int(image_cols*landmark.x)
|
||||
image_y = int(image_rows*landmark.y)
|
||||
draw_color = None
|
||||
if isinstance(drawing_spec, dict):
|
||||
if drawing_spec.get(idx) is None:
|
||||
continue
|
||||
else:
|
||||
draw_color = drawing_spec[idx].color
|
||||
elif isinstance(drawing_spec, self.mp_drawing.DrawingSpec):
|
||||
draw_color = drawing_spec.color
|
||||
image[image_y-halfwidth:image_y+halfwidth, image_x-halfwidth:image_x+halfwidth, :] = draw_color
|
||||
|
||||
def draw_facemesh(self, img_rgb, results, thickness=1):
|
||||
"""Draw face mesh visualization using the ControlNet Aux methodology on a black background"""
|
||||
# Create an empty image for drawing
|
||||
empty = np.zeros_like(img_rgb)
|
||||
|
||||
# Ensure the image is contiguous in memory and has the correct data type for OpenCV
|
||||
empty = np.ascontiguousarray(empty, dtype=np.uint8)
|
||||
|
||||
# Update thickness based on input parameter
|
||||
for spec in self.face_connection_spec.values():
|
||||
spec.thickness = thickness
|
||||
|
||||
# Draw detected faces
|
||||
if results and results.multi_face_landmarks:
|
||||
for face_landmarks in results.multi_face_landmarks:
|
||||
self.mp_drawing.draw_landmarks(
|
||||
empty,
|
||||
face_landmarks,
|
||||
connections=self.face_connection_spec.keys(),
|
||||
landmark_drawing_spec=None,
|
||||
connection_drawing_spec=self.face_connection_spec
|
||||
)
|
||||
self.draw_pupils(empty, face_landmarks, self.iris_landmark_spec, 2)
|
||||
|
||||
# Flip BGR back to RGB
|
||||
empty = self.reverse_channels(empty).copy()
|
||||
|
||||
return empty
|
||||
|
||||
def draw_facemesh_on_image(self, img_rgb, results, thickness=1):
|
||||
"""Draw face mesh visualization on the provided image"""
|
||||
# Make a copy of the input image to avoid modifying the original
|
||||
img_copy = img_rgb.copy()
|
||||
|
||||
# Convert RGB to BGR for OpenCV and MediaPipe drawing functions
|
||||
img_copy = self.reverse_channels(img_copy)
|
||||
|
||||
# Ensure the image is contiguous in memory and has the correct data type for OpenCV
|
||||
img_copy = np.ascontiguousarray(img_copy, dtype=np.uint8)
|
||||
|
||||
# Update thickness based on input parameter
|
||||
for spec in self.face_connection_spec.values():
|
||||
spec.thickness = thickness
|
||||
|
||||
# Draw detected faces
|
||||
if results and results.multi_face_landmarks:
|
||||
for face_landmarks in results.multi_face_landmarks:
|
||||
self.mp_drawing.draw_landmarks(
|
||||
img_copy,
|
||||
face_landmarks,
|
||||
connections=self.face_connection_spec.keys(),
|
||||
landmark_drawing_spec=None,
|
||||
connection_drawing_spec=self.face_connection_spec
|
||||
)
|
||||
self.draw_pupils(img_copy, face_landmarks, self.iris_landmark_spec, 2)
|
||||
|
||||
# Convert back from BGR to RGB
|
||||
img_copy = self.reverse_channels(img_copy).copy()
|
||||
|
||||
return img_copy
|
||||
|
||||
# Create a global instance of the utils class
|
||||
face_mesh_utils = FaceMeshUtils()
|
||||
|
||||
class FaceMeshDrawNode:
|
||||
"""
|
||||
A node that draws MediaPipe face mesh on images.
|
||||
Takes face mesh data from FaceMeshDetectNode and visualizes it.
|
||||
Optimized for visualization only, requires an input image.
|
||||
Uses the same drawing methodology as ControlNet Aux for optimal performance.
|
||||
"""
|
||||
def __init__(self):
|
||||
# Use the shared utils instance
|
||||
self.utils = face_mesh_utils
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"facemesh_data": ("FACEMESH_DATA",),
|
||||
"draw_contours": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"description": "Draw contour lines connecting mesh points"
|
||||
}),
|
||||
"mesh_thickness": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.1,
|
||||
"max": 5.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
"description": "Thickness of mesh lines"
|
||||
}),
|
||||
"mesh_color": (["green", "red", "blue", "yellow", "white"], {
|
||||
"default": "green",
|
||||
"description": "Color of the face mesh lines"
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "draw_facemesh"
|
||||
CATEGORY = "ComfyUI_FaceMesh"
|
||||
|
||||
def draw_facemesh(self, image, facemesh_data, draw_contours, mesh_thickness, mesh_color):
|
||||
"""Draw face mesh visualization on the input image"""
|
||||
# Get batch size
|
||||
batch_size = image.shape[0]
|
||||
output_images = []
|
||||
|
||||
# Process each image in the batch
|
||||
for b in range(batch_size):
|
||||
# Get current image
|
||||
img = image[b].numpy()
|
||||
|
||||
# Get corresponding facemesh data
|
||||
current_facemesh_data = facemesh_data[b]
|
||||
|
||||
# If no face detected, add original image unchanged
|
||||
if not current_facemesh_data["face_detected"]:
|
||||
output_images.append(img)
|
||||
continue
|
||||
|
||||
# Convert from float [0,1] to uint8 [0,255] for drawing
|
||||
img_rgb = (img * 255).astype(np.uint8)
|
||||
|
||||
# Ensure the image is contiguous in memory for OpenCV compatibility
|
||||
img_rgb = np.ascontiguousarray(img_rgb, dtype=np.uint8)
|
||||
|
||||
# Use original results for drawing
|
||||
results = current_facemesh_data["original_results"]
|
||||
|
||||
# Draw the face mesh on the input image using the shared utility
|
||||
output_img_rgb = self.utils.draw_facemesh_on_image(img_rgb, results, thickness=int(mesh_thickness))
|
||||
|
||||
# Convert back to float [0,1] for ComfyUI
|
||||
output_img = output_img_rgb.astype(np.float32) / 255.0
|
||||
|
||||
# Add to output list
|
||||
output_images.append(output_img)
|
||||
|
||||
# Convert output images list to tensor with batch dimension
|
||||
output_batch = torch.from_numpy(np.stack(output_images, axis=0))
|
||||
|
||||
return (output_batch,)
|
||||
|
||||
class FaceMeshMaskNode:
|
||||
"""
|
||||
Creates binary masks for different regions of the face detected by FaceMeshDetectNode.
|
||||
"""
|
||||
# Define facial region indices
|
||||
FACIAL_REGIONS = {
|
||||
# These are approximate regions based on MediaPipe Face Mesh indices
|
||||
"whole_face": [
|
||||
# Face contour/oval - same as face_oval for proper face filling
|
||||
10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288, 397, 365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136, 172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109, 10
|
||||
],
|
||||
"mouth": [
|
||||
61, 185, 40, 39, 37, 0, 267, 269, 270, 409, 291, 375, 321, 405, 314, 17, 84, 181, 91, 146
|
||||
],
|
||||
"left_eye": [
|
||||
# Left eye outline (from the person's perspective)
|
||||
263, 249, 390, 373, 374, 380, 381, 382, 362, 263, 466, 388, 387, 386, 385, 384, 398, 362
|
||||
],
|
||||
"right_eye": [
|
||||
# Right eye outline
|
||||
33, 7, 163, 144, 145, 153, 154, 155, 133, 33, 246, 161, 160, 159, 158, 157, 173, 133
|
||||
],
|
||||
"nose": [
|
||||
# Nose area
|
||||
168, 6, 197, 195, 5, 4, 1, 19, 94, 2, 164, 0, 11, 12, 13, 14, 15, 16, 17, 18, 200, 199, 175
|
||||
],
|
||||
"left_eyebrow": [
|
||||
# Left eyebrow
|
||||
276, 283, 282, 295, 285, 300, 293, 334, 296, 336
|
||||
],
|
||||
"right_eyebrow": [
|
||||
# Right eyebrow
|
||||
46, 53, 52, 65, 55, 70, 63, 105, 66, 107
|
||||
],
|
||||
"face_oval": [
|
||||
# Face contour/oval
|
||||
10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288, 397, 365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136, 172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109, 10
|
||||
]
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
regions = list(s.FACIAL_REGIONS.keys())
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"facemesh_data": ("FACEMESH_DATA",),
|
||||
"face_index": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 9,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"description": "Which face to mask (if multiple detected)"
|
||||
}),
|
||||
"region": (regions, {
|
||||
"default": "mouth",
|
||||
"description": "Facial region to mask"
|
||||
}),
|
||||
"fill_region": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"description": "Fill the region or just draw contour"
|
||||
}),
|
||||
"mask_value": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01,
|
||||
"display": "slider",
|
||||
"description": "Value/intensity of the mask (0.0-1.0)"
|
||||
}),
|
||||
"dilation_pixels": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 50,
|
||||
"step": 1,
|
||||
"display": "slider",
|
||||
"description": "Dilate mask by this many pixels"
|
||||
}),
|
||||
"blur_radius": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 20,
|
||||
"step": 1,
|
||||
"display": "slider",
|
||||
"description": "Apply gaussian blur with this radius"
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
RETURN_NAMES = ("mask",)
|
||||
FUNCTION = "create_mask"
|
||||
CATEGORY = "ComfyUI_FaceMesh"
|
||||
|
||||
def create_mask(self, facemesh_data, face_index, region, fill_region, mask_value, dilation_pixels, blur_radius):
|
||||
"""Create a mask for the specified facial region with custom intensity value"""
|
||||
# Get number of images in batch
|
||||
batch_size = len(facemesh_data)
|
||||
output_masks = []
|
||||
|
||||
# Process each image in the batch
|
||||
for b in range(batch_size):
|
||||
# Get current facemesh data
|
||||
current_facemesh_data = facemesh_data[b]
|
||||
|
||||
# If no face detected, return blank image
|
||||
if not current_facemesh_data["face_detected"] or face_index >= current_facemesh_data["num_faces_detected"]:
|
||||
# Create a blank mask with original image dimensions
|
||||
h, w = current_facemesh_data["image_shape"]
|
||||
mask_float = np.zeros((h, w), dtype=np.float32)
|
||||
output_masks.append(mask_float)
|
||||
continue
|
||||
|
||||
# Get image shape
|
||||
h, w = current_facemesh_data["image_shape"]
|
||||
|
||||
# Create blank mask
|
||||
mask = np.zeros((h, w), dtype=np.uint8)
|
||||
|
||||
# Get the landmarks for the specified face
|
||||
landmarks = current_facemesh_data["landmarks"][face_index]
|
||||
|
||||
# Get indices for the specified region
|
||||
region_indices = self.FACIAL_REGIONS[region]
|
||||
|
||||
# Extract points for the region - directly use numpy for better performance
|
||||
landmarks_array = np.array(landmarks)
|
||||
if len(region_indices) > 0:
|
||||
# Get subset of landmarks for the region
|
||||
region_landmarks = np.array([landmarks_array[i] for i in region_indices if i < len(landmarks_array)])
|
||||
|
||||
if len(region_landmarks) > 2: # Need at least 3 points to create a contour
|
||||
# Convert normalized coordinates to pixel coordinates
|
||||
points = np.zeros((len(region_landmarks), 2), dtype=np.int32)
|
||||
points[:, 0] = (region_landmarks[:, 0] * w).astype(np.int32) # x coordinates
|
||||
points[:, 1] = (region_landmarks[:, 1] * h).astype(np.int32) # y coordinates
|
||||
|
||||
# Ensure mask is contiguous in memory for OpenCV compatibility
|
||||
mask = np.ascontiguousarray(mask, dtype=np.uint8)
|
||||
|
||||
# Use 255 for drawing, we'll scale to mask_value later
|
||||
fill_value = 255
|
||||
|
||||
# Special handling for mouth region - create a convex hull to ensure proper filling
|
||||
if region == "mouth" and fill_region:
|
||||
# Use convex hull to get a proper fillable shape
|
||||
hull = cv2.convexHull(points)
|
||||
cv2.fillConvexPoly(mask, hull, fill_value)
|
||||
# For whole_face and face_oval, always fill regardless of fill_region setting
|
||||
elif (region == "whole_face" or region == "face_oval") and fill_region:
|
||||
# Use convex hull for better filling of the face
|
||||
hull = cv2.convexHull(points)
|
||||
cv2.fillConvexPoly(mask, hull, fill_value)
|
||||
else:
|
||||
# Create the mask based on fill_region option
|
||||
if fill_region:
|
||||
# Fill the polygon with white (255)
|
||||
cv2.fillPoly(mask, [points], fill_value)
|
||||
else:
|
||||
# Just draw the outline with white (255)
|
||||
cv2.polylines(mask, [points], True, fill_value, 2)
|
||||
|
||||
# Apply dilation if specified - use optimized kernel
|
||||
if dilation_pixels > 0:
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (dilation_pixels, dilation_pixels))
|
||||
mask = cv2.dilate(mask, kernel, iterations=1)
|
||||
|
||||
# Apply gaussian blur if specified
|
||||
if blur_radius > 0:
|
||||
mask = cv2.GaussianBlur(mask, (blur_radius * 2 + 1, blur_radius * 2 + 1), 0)
|
||||
|
||||
# Convert mask to float [0,1] for ComfyUI and apply mask_value
|
||||
mask_float = mask.astype(np.float32) / 255.0 * mask_value
|
||||
|
||||
# Add to output list
|
||||
output_masks.append(mask_float)
|
||||
|
||||
# Convert to batch tensor - BHW format
|
||||
mask_batch = torch.from_numpy(np.stack(output_masks, axis=0))
|
||||
|
||||
return (mask_batch,)
|
||||
|
||||
class FaceMeshNode:
|
||||
"""
|
||||
A node that uses MediaPipe to detect face mesh from images and draws it.
|
||||
Optimized for real-time video processing.
|
||||
Uses the same drawing methodology as ControlNet Aux for optimal performance.
|
||||
"""
|
||||
def __init__(self):
|
||||
# Initialize face mesh detector once during node creation
|
||||
self.mp_face_mesh = mp.solutions.face_mesh
|
||||
self.face_mesh = self.mp_face_mesh.FaceMesh(
|
||||
static_image_mode=False,
|
||||
max_num_faces=1,
|
||||
refine_landmarks=False,
|
||||
min_detection_confidence=0.5,
|
||||
min_tracking_confidence=0.5
|
||||
)
|
||||
self.last_results = None
|
||||
# Track current settings to know when to reinitialize
|
||||
self.current_settings = {
|
||||
"static_image_mode": False,
|
||||
"max_num_faces": 1,
|
||||
"refine_landmarks": False,
|
||||
"min_detection_confidence": 0.5,
|
||||
"min_tracking_confidence": 0.5
|
||||
}
|
||||
# Use the shared utils instance
|
||||
self.utils = face_mesh_utils
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"static_image_mode": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"description": "Set to true for static images, false for video streams (faster)"
|
||||
}),
|
||||
"max_num_faces": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 10,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
"description": "Maximum number of faces to detect (lower is faster)"
|
||||
}),
|
||||
"refine_landmarks": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"description": "Refine landmarks around eyes and lips (slower but more accurate)"
|
||||
}),
|
||||
"min_detection_confidence": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.1,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"display": "slider",
|
||||
"description": "Minimum confidence for face detection"
|
||||
}),
|
||||
"min_tracking_confidence": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.1,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"display": "slider",
|
||||
"description": "Minimum confidence for face tracking"
|
||||
}),
|
||||
"draw_mesh": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"description": "Draw face mesh on output image"
|
||||
}),
|
||||
"draw_contours": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"description": "Draw contour lines connecting mesh points"
|
||||
}),
|
||||
"mesh_color": (["green", "red", "blue", "yellow", "white"], {
|
||||
"default": "green",
|
||||
"description": "Color of the face mesh lines"
|
||||
}),
|
||||
"mesh_thickness": ("FLOAT", {
|
||||
"default": 2.0,
|
||||
"min": 0.1,
|
||||
"max": 5.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
"description": "Thickness of mesh lines"
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FACEMESH_DATA", "IMAGE")
|
||||
RETURN_NAMES = ("facemesh_data", "annotations")
|
||||
FUNCTION = "detect_facemesh"
|
||||
CATEGORY = "ComfyUI_FaceMesh"
|
||||
|
||||
def detect_facemesh(self, image, static_image_mode, max_num_faces, refine_landmarks,
|
||||
min_detection_confidence, min_tracking_confidence,
|
||||
draw_mesh, draw_contours, mesh_color, mesh_thickness):
|
||||
"""Detect face mesh in the input image and optionally draw it using ControlNet Aux methodology"""
|
||||
# Check if any parameters have changed and we need to reinitialize
|
||||
new_settings = {
|
||||
"static_image_mode": static_image_mode,
|
||||
"max_num_faces": max_num_faces,
|
||||
"refine_landmarks": refine_landmarks,
|
||||
"min_detection_confidence": min_detection_confidence,
|
||||
"min_tracking_confidence": min_tracking_confidence
|
||||
}
|
||||
|
||||
# Reinitialize if settings changed
|
||||
if new_settings != self.current_settings:
|
||||
self.face_mesh = self.mp_face_mesh.FaceMesh(
|
||||
static_image_mode=static_image_mode,
|
||||
max_num_faces=max_num_faces,
|
||||
refine_landmarks=refine_landmarks,
|
||||
min_detection_confidence=min_detection_confidence,
|
||||
min_tracking_confidence=min_tracking_confidence
|
||||
)
|
||||
self.current_settings = new_settings
|
||||
|
||||
# Get batch size
|
||||
batch_size = image.shape[0]
|
||||
|
||||
# Initialize empty lists to store batch results
|
||||
facemesh_data_batch = []
|
||||
output_images = []
|
||||
|
||||
# Process each image in the batch
|
||||
for b in range(batch_size):
|
||||
# Get current image from batch
|
||||
img = image[b].numpy()
|
||||
|
||||
# Convert from float [0,1] to uint8 [0,255] for mediapipe
|
||||
img_rgb = (img * 255).astype(np.uint8)
|
||||
|
||||
# Ensure the image is contiguous in memory for OpenCV compatibility
|
||||
img_rgb = np.ascontiguousarray(img_rgb, dtype=np.uint8)
|
||||
|
||||
# Process the image
|
||||
results = self.face_mesh.process(img_rgb)
|
||||
|
||||
# Prepare facemesh data structure for this image
|
||||
facemesh_data = {
|
||||
"landmarks": [],
|
||||
"face_detected": False,
|
||||
"num_faces_detected": 0,
|
||||
"tesselation": self.mp_face_mesh.FACEMESH_TESSELATION,
|
||||
"contours": self.mp_face_mesh.FACEMESH_CONTOURS,
|
||||
"image_shape": img.shape[:2],
|
||||
"original_results": results # Store original results for native drawing
|
||||
}
|
||||
|
||||
# Process detection results
|
||||
if results and results.multi_face_landmarks:
|
||||
facemesh_data["face_detected"] = True
|
||||
facemesh_data["num_faces_detected"] = len(results.multi_face_landmarks)
|
||||
|
||||
all_landmarks = []
|
||||
|
||||
for face_idx, face_landmarks in enumerate(results.multi_face_landmarks):
|
||||
# Extract all landmarks at once for more efficiency
|
||||
face_points = [(lm.x, lm.y, lm.z) for lm in face_landmarks.landmark]
|
||||
all_landmarks.append(face_points)
|
||||
|
||||
facemesh_data["landmarks"] = all_landmarks
|
||||
|
||||
# Draw annotations if requested - on a black background for FaceMeshNode
|
||||
if draw_mesh and facemesh_data["face_detected"]:
|
||||
# Draw the face mesh on a black background using the shared utility
|
||||
empty = self.utils.draw_facemesh(img_rgb, results, thickness=int(mesh_thickness))
|
||||
|
||||
# Convert back to float [0,1] for ComfyUI
|
||||
output_img = empty.astype(np.float32) / 255.0
|
||||
|
||||
# Add image to output list
|
||||
output_images.append(output_img)
|
||||
else:
|
||||
# Use original image if no drawing requested or no face detected
|
||||
output_images.append(img)
|
||||
|
||||
# Add facemesh data to batch results
|
||||
facemesh_data_batch.append(facemesh_data)
|
||||
|
||||
# Convert output images list to tensor with batch dimension
|
||||
output_batch = torch.from_numpy(np.stack(output_images, axis=0))
|
||||
|
||||
return (facemesh_data_batch, output_batch)
|
||||
|
||||
# Node mapping
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FaceMeshNode": FaceMeshNode,
|
||||
"FaceMeshDrawNode": FaceMeshDrawNode,
|
||||
"FaceMeshMaskNode": FaceMeshMaskNode
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FaceMeshNode": "Face Mesh",
|
||||
"FaceMeshDrawNode": "Face Mesh Draw",
|
||||
"FaceMeshMaskNode": "Face Region Mask"
|
||||
}
|
||||
Reference in New Issue
Block a user