Compare commits
15
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3cf283eb69 | ||
|
|
695db5e06c | ||
|
|
c78bbccb10 | ||
|
|
01aacb73e0 | ||
|
|
011da8e1f7 | ||
|
|
7e20a4a6eb | ||
|
|
7ed6c75652 | ||
|
|
c3e196822d | ||
|
|
fce8d032f8 | ||
|
|
499096bdb0 | ||
|
|
b1ebd4aa8f | ||
|
|
59035c455a | ||
|
|
da5a043501 | ||
|
|
eedf5d1d1b | ||
|
|
845892959a |
@@ -0,0 +1,28 @@
|
||||
name: Publish to Comfy registry
|
||||
on:
|
||||
workflow_dispatch:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
- master
|
||||
paths:
|
||||
- "pyproject.toml"
|
||||
|
||||
permissions:
|
||||
issues: write
|
||||
|
||||
jobs:
|
||||
publish-node:
|
||||
name: Publish Custom Node to registry
|
||||
runs-on: ubuntu-latest
|
||||
if: ${{ github.repository_owner == 'livepeer' }}
|
||||
steps:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
with:
|
||||
submodules: true
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@v1
|
||||
with:
|
||||
## Add your own personal access token to your Github Repository secrets and reference it here.
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
+17
-83
@@ -1,88 +1,22 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
# Misc
|
||||
__pycache__
|
||||
*.pyc
|
||||
|
||||
# Python packaging
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
.pytest_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
*.egg-info
|
||||
build/
|
||||
dist/
|
||||
|
||||
# Virtual Environment
|
||||
venv/
|
||||
env/
|
||||
ENV/
|
||||
.env/
|
||||
# IDEs
|
||||
.vscode
|
||||
*.code-workspace
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# ML/AI specific
|
||||
*.pt
|
||||
# Model weights
|
||||
*.pb
|
||||
*.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*
|
||||
*.pt
|
||||
*.pth.tar
|
||||
*.pth.tar.gz
|
||||
*.pth.tar.xz
|
||||
|
||||
+39
-14
@@ -1,19 +1,44 @@
|
||||
"""ComfyUI Stream Pack Nodes."""
|
||||
"""Expose all nodes in the Stream Pack."""
|
||||
|
||||
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 logging
|
||||
import pathlib
|
||||
import importlib
|
||||
import traceback
|
||||
|
||||
# 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)
|
||||
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']
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
|
||||
@@ -215,4 +215,4 @@
|
||||
"title": "Load CLIP"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -231,4 +231,4 @@
|
||||
"title": "Feature Bank Attention Processor"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+38
-62
@@ -48,6 +48,36 @@
|
||||
"title": "CLIP Text Encode (Prompt)"
|
||||
}
|
||||
},
|
||||
"7": {
|
||||
"inputs": {
|
||||
"seed": 951969040799634,
|
||||
"steps": 2,
|
||||
"cfg": 1,
|
||||
"sampler_name": "lcm",
|
||||
"scheduler": "normal",
|
||||
"denoise": 1,
|
||||
"model": [
|
||||
"41",
|
||||
0
|
||||
],
|
||||
"positive": [
|
||||
"38",
|
||||
0
|
||||
],
|
||||
"negative": [
|
||||
"37",
|
||||
0
|
||||
],
|
||||
"latent_image": [
|
||||
"16",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "KSampler",
|
||||
"_meta": {
|
||||
"title": "KSampler"
|
||||
}
|
||||
},
|
||||
"8": {
|
||||
"inputs": {
|
||||
"control_net_name": "control_v11f1p_sd15_depth_fp16.safetensors"
|
||||
@@ -128,7 +158,7 @@
|
||||
"14": {
|
||||
"inputs": {
|
||||
"samples": [
|
||||
"42",
|
||||
"7",
|
||||
0
|
||||
],
|
||||
"vae": [
|
||||
@@ -304,73 +334,19 @@
|
||||
},
|
||||
"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,
|
||||
"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": "StreamScheduler",
|
||||
"class_type": "FeatureBankAttentionProcessor",
|
||||
"_meta": {
|
||||
"title": "Stream Scheduler"
|
||||
"title": "Feature Bank Attention Processor"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,334 @@
|
||||
{
|
||||
"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)"
|
||||
}
|
||||
},
|
||||
"7": {
|
||||
"inputs": {
|
||||
"seed": 951969040799634,
|
||||
"steps": 2,
|
||||
"cfg": 1,
|
||||
"sampler_name": "lcm",
|
||||
"scheduler": "normal",
|
||||
"denoise": 1,
|
||||
"model": [
|
||||
"3",
|
||||
0
|
||||
],
|
||||
"positive": [
|
||||
"38",
|
||||
0
|
||||
],
|
||||
"negative": [
|
||||
"37",
|
||||
0
|
||||
],
|
||||
"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": 0.9000000000000001,
|
||||
"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"
|
||||
}
|
||||
},
|
||||
"18": {
|
||||
"inputs": {
|
||||
"stop_at_clip_layer": -2,
|
||||
"clip": [
|
||||
"36",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "CLIPSetLastLayer",
|
||||
"_meta": {
|
||||
"title": "CLIP Set Last Layer"
|
||||
}
|
||||
},
|
||||
"19": {
|
||||
"inputs": {
|
||||
"strength": 0.9000000000000001,
|
||||
"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"
|
||||
},
|
||||
"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.3500000000000001,
|
||||
"min_tracking_confidence": 0.3500000000000001,
|
||||
"draw_mesh": true,
|
||||
"draw_contours": true,
|
||||
"mesh_color": "green",
|
||||
"mesh_thickness": 2,
|
||||
"image": [
|
||||
"26",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "FaceMeshNode",
|
||||
"_meta": {
|
||||
"title": "Face Mesh"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
{
|
||||
"1": {
|
||||
"inputs": {
|
||||
"model_type": "FSRCNN",
|
||||
"scale_factor": "2"
|
||||
},
|
||||
"class_type": "SuperResolutionModelLoader",
|
||||
"_meta": {
|
||||
"title": "SR Model Loader"
|
||||
}
|
||||
},
|
||||
"2": {
|
||||
"inputs": {
|
||||
"use_cuda": "True",
|
||||
"image": [
|
||||
"3",
|
||||
0
|
||||
],
|
||||
"sr_model": [
|
||||
"1",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "SuperResolutionUpscale",
|
||||
"_meta": {
|
||||
"title": "SR Upscale"
|
||||
}
|
||||
},
|
||||
"3": {
|
||||
"inputs": {
|
||||
"image": "example-512x512.png"
|
||||
},
|
||||
"class_type": "LoadImage",
|
||||
"_meta": {
|
||||
"title": "Load Image"
|
||||
}
|
||||
},
|
||||
"4": {
|
||||
"inputs": {
|
||||
"images": [
|
||||
"2",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "PreviewImage",
|
||||
"_meta": {
|
||||
"title": "Preview Image"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
"""
|
||||
Real-time Audio Transcription Node for ComfyUI.
|
||||
|
||||
This node buffers audio segments and performs transcription using faster-whisper,
|
||||
with controlled output timing to prevent message flooding.
|
||||
"""
|
||||
|
||||
# Import node classes
|
||||
from ..src.audio_transcription.transcription_nodes import AudioTranscriptionNode
|
||||
from ..src.audio_transcription.srt_generator_node import SRTGeneratorNode
|
||||
|
||||
# Register the node
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AudioTranscriptionNode": AudioTranscriptionNode,
|
||||
"SRTGeneratorNode": SRTGeneratorNode
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AudioTranscriptionNode": "Audio Transcription (Real-time)",
|
||||
"SRTGeneratorNode": "SRT Generator"
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -1,15 +0,0 @@
|
||||
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,21 @@
|
||||
"""Provides nodes for facial landmark detection and utilities using MediaPipe FaceMesh.
|
||||
Inspired by the ComfyUI ControlNet Auxiliary MediaPipe Face node.
|
||||
|
||||
References:
|
||||
https://github.com/Fannovel16/comfyui_controlnet_aux/blob/main/src/custom_controlnet_aux/mediapipe_face/mediapipe_face_common.py
|
||||
"""
|
||||
|
||||
from ..src.media_pipe.face_mesh 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",
|
||||
}
|
||||
|
||||
__all__ = ["FaceMeshNode", "FaceMeshDrawNode", "FaceMeshMaskNode"]
|
||||
@@ -1,13 +0,0 @@
|
||||
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"
|
||||
}
|
||||
@@ -13,7 +13,6 @@ from ..src.feature_bank.feature_bank import (
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FeatureBankAttentionProcessor": FeatureBankAttentionProcessorNode
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FeatureBankAttentionProcessor": "Feature Bank Attention Processor"
|
||||
}
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
# Import node classes
|
||||
from ..src.super_resolution.nodes.sr_nodes import SuperResolutionModelLoader, SuperResolutionUpscale
|
||||
|
||||
# Define NODE_CLASS_MAPPINGS for ComfyUI
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SuperResolutionModelLoader": SuperResolutionModelLoader,
|
||||
"SuperResolutionUpscale": SuperResolutionUpscale,
|
||||
}
|
||||
|
||||
# Define NODE_DISPLAY_NAME_MAPPINGS for ComfyUI
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"SuperResolutionModelLoader": "SR Model Loader",
|
||||
"SuperResolutionUpscale": "SR Upscale",
|
||||
}
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,22 @@
|
||||
import logging
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def check_opencv_cuda():
|
||||
try:
|
||||
import cv2
|
||||
# Check if CUDA is already available
|
||||
logger.info("ComfyUI-Stream-Pack: Checking OpenCV CUDA availability...")
|
||||
if hasattr(cv2, 'cuda') and cv2.cuda.getCudaEnabledDeviceCount() > 0:
|
||||
logger.info("OpenCV CUDA is installed and working")
|
||||
return True
|
||||
else:
|
||||
logger.warning("OpenCV CUDA is not available or not working. Super Resolution node may not work")
|
||||
return False
|
||||
except ImportError:
|
||||
logger.warning("OpenCV is not installed")
|
||||
return False
|
||||
|
||||
check_opencv_cuda()
|
||||
@@ -0,0 +1,18 @@
|
||||
[project]
|
||||
name = "stream-pack"
|
||||
description = "A suite of custom ComfyUI nodes for building real-time video and audio workflows using ComfyStream."
|
||||
version = "1.0.0"
|
||||
license = {file = "LICENSE"}
|
||||
dynamic = ["dependencies"]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/livepeer/ComfyUI-Stream-Pack"
|
||||
# Used by Comfy Registry https://comfyregistry.org
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = "livepeer-comfystream"
|
||||
DisplayName = "ComfyUI-Stream-Pack"
|
||||
Icon = "https://raw.githubusercontent.com/livepeer/comfystream-docs/main/logo/icon-light-120px.svg"
|
||||
|
||||
[tool.setuptools.dynamic]
|
||||
dependencies = {file = ["requirements.txt"]}
|
||||
@@ -1,3 +1,8 @@
|
||||
mediapipe
|
||||
torch
|
||||
diffusers
|
||||
numpy
|
||||
requests
|
||||
opencv-contrib-python>=4.8.0
|
||||
faster_whisper
|
||||
scipy
|
||||
|
||||
+1
-1
@@ -1 +1 @@
|
||||
# Dummy file ensuring this package will be recognized.
|
||||
# Dummy file ensuring this sub-package will be recognized.
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Audio transcription package for ComfyUI real-time transcription capabilities."""
|
||||
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
SRT subtitle generation node for ComfyStream.
|
||||
Based on project-transcript SRT generation capabilities.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from datetime import timedelta
|
||||
from typing import List, Dict, Any, Union
|
||||
import json
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SRTGeneratorNode:
|
||||
"""
|
||||
Generate SRT subtitle files from transcription with timing information.
|
||||
Compatible with AudioTranscriptionNode output when word timestamps are enabled.
|
||||
"""
|
||||
|
||||
CATEGORY = "text_utils"
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("srt_content",)
|
||||
|
||||
def __init__(self):
|
||||
self.subtitle_counter = 1
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"transcription_data": ("STRING", {
|
||||
"tooltip": "Transcription text or JSON with timing data from AudioTranscriptionNode"
|
||||
}),
|
||||
"segment_start_time": ("FLOAT", {
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 86400.0, # 24 hours max
|
||||
"step": 0.001,
|
||||
"tooltip": "Start time of this segment in seconds (for absolute timing)"
|
||||
}),
|
||||
"segment_duration": ("FLOAT", {
|
||||
"default": 3.0,
|
||||
"min": 0.1,
|
||||
"max": 60.0,
|
||||
"step": 0.1,
|
||||
"tooltip": "Duration of this segment in seconds"
|
||||
}),
|
||||
"use_absolute_time": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "Use absolute timing (True) or segment-relative timing (False)"
|
||||
}),
|
||||
"minimum_duration": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.1,
|
||||
"max": 10.0,
|
||||
"step": 0.1,
|
||||
"tooltip": "Minimum subtitle duration in seconds"
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
FUNCTION = "generate_srt"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float("nan")
|
||||
|
||||
def generate_srt(self,
|
||||
transcription_data: str,
|
||||
segment_start_time: float = 0.0,
|
||||
segment_duration: float = 3.0,
|
||||
use_absolute_time: bool = False,
|
||||
minimum_duration: float = 1.0) -> tuple:
|
||||
"""
|
||||
Generate SRT subtitle content from transcription data.
|
||||
|
||||
Args:
|
||||
transcription_data: Text transcription or JSON with timing data
|
||||
segment_start_time: Start time of this segment (for absolute timing)
|
||||
segment_duration: Duration of this segment
|
||||
use_absolute_time: Whether to use absolute or segment-relative timing
|
||||
minimum_duration: Minimum subtitle duration
|
||||
|
||||
Returns:
|
||||
Tuple containing SRT formatted content
|
||||
"""
|
||||
try:
|
||||
if not transcription_data or not transcription_data.strip() or len(transcription_data.strip()) < 5:
|
||||
logger.debug(f"Empty or minimal transcription data: '{transcription_data}'")
|
||||
return ("",)
|
||||
|
||||
# Pass through warmup sentinel values (they'll be filtered later in the pipeline)
|
||||
if "__WARMUP_SENTINEL__" in transcription_data:
|
||||
logger.debug("Passing through warmup sentinel for pipeline detection")
|
||||
return (transcription_data,)
|
||||
|
||||
# Try to parse as JSON first (if AudioTranscriptionNode supports structured output)
|
||||
segments = self._parse_transcription_data(transcription_data)
|
||||
|
||||
if not segments:
|
||||
logger.debug("No valid segments found in transcription data")
|
||||
return ("",)
|
||||
|
||||
# Generate SRT content
|
||||
srt_content = self._generate_srt_from_segments(
|
||||
segments,
|
||||
segment_start_time,
|
||||
segment_duration,
|
||||
use_absolute_time,
|
||||
minimum_duration
|
||||
)
|
||||
|
||||
logger.debug(f"Generated SRT with {len(segments)} segments")
|
||||
return (srt_content,)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating SRT: {e}")
|
||||
return ("",)
|
||||
|
||||
def _parse_transcription_data(self, data: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Parse transcription data into segments.
|
||||
Enhanced to handle AudioTranscriptionNode JSON output formats.
|
||||
"""
|
||||
segments = []
|
||||
|
||||
try:
|
||||
# Try parsing as JSON first
|
||||
parsed_data = json.loads(data)
|
||||
|
||||
if isinstance(parsed_data, list):
|
||||
# Handle both segment-level and word-level JSON arrays
|
||||
for item in parsed_data:
|
||||
if isinstance(item, dict):
|
||||
# Check if it's word-level data (has 'word' field) or segment-level (has 'text' field)
|
||||
if 'word' in item:
|
||||
# Word-level JSON from AudioTranscriptionNode (json_words format)
|
||||
segments.append({
|
||||
'start': item.get('start', 0.0),
|
||||
'end': item.get('end', 1.0),
|
||||
'text': item.get('word', '').strip()
|
||||
})
|
||||
elif 'text' in item:
|
||||
# Segment-level JSON from AudioTranscriptionNode (json_segments format)
|
||||
segments.append({
|
||||
'start': item.get('start', 0.0),
|
||||
'end': item.get('end', 1.0),
|
||||
'text': item.get('text', '').strip()
|
||||
})
|
||||
|
||||
elif isinstance(parsed_data, dict):
|
||||
# Single segment or word
|
||||
if 'word' in parsed_data:
|
||||
segments.append({
|
||||
'start': parsed_data.get('start', 0.0),
|
||||
'end': parsed_data.get('end', 1.0),
|
||||
'text': parsed_data.get('word', '').strip()
|
||||
})
|
||||
elif 'text' in parsed_data:
|
||||
segments.append({
|
||||
'start': parsed_data.get('start', 0.0),
|
||||
'end': parsed_data.get('end', 1.0),
|
||||
'text': parsed_data.get('text', '').strip()
|
||||
})
|
||||
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
# Fallback to plain text - create a single segment
|
||||
if data.strip():
|
||||
segments.append({
|
||||
'start': 0.0,
|
||||
'end': 3.0, # Default 3-second duration
|
||||
'text': data.strip()
|
||||
})
|
||||
|
||||
return [seg for seg in segments if seg['text']] # Filter empty text
|
||||
|
||||
def _generate_srt_from_segments(self,
|
||||
segments: List[Dict[str, Any]],
|
||||
segment_start_time: float,
|
||||
segment_duration: float,
|
||||
use_absolute_time: bool,
|
||||
minimum_duration: float) -> str:
|
||||
"""
|
||||
Generate SRT content from parsed segments and wrap it in a JSON format.
|
||||
"""
|
||||
subtitles = []
|
||||
|
||||
for i, segment in enumerate(segments):
|
||||
text = segment['text'].strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
# Calculate timing
|
||||
if use_absolute_time:
|
||||
# Absolute timing: add segment start time
|
||||
start_time = segment_start_time + segment['start']
|
||||
end_time = segment_start_time + segment['end']
|
||||
else:
|
||||
# Segment-relative timing
|
||||
start_time = segment['start']
|
||||
end_time = segment['end']
|
||||
|
||||
# Ensure minimum duration
|
||||
if end_time - start_time < minimum_duration:
|
||||
end_time = start_time + minimum_duration
|
||||
|
||||
# Format timing
|
||||
start_srt = self._seconds_to_srt_time(start_time)
|
||||
end_srt = self._seconds_to_srt_time(end_time)
|
||||
|
||||
# Add subtitle entry to list
|
||||
subtitles.append({
|
||||
"id": self.subtitle_counter,
|
||||
"start": start_srt,
|
||||
"end": end_srt,
|
||||
"text": text
|
||||
})
|
||||
|
||||
self.subtitle_counter += 1
|
||||
|
||||
# Return as a JSON string
|
||||
return json.dumps(subtitles, indent=2)
|
||||
|
||||
def _seconds_to_srt_time(self, seconds: float) -> str:
|
||||
"""
|
||||
Convert seconds to SRT time format (HH:MM:SS,mmm).
|
||||
Compatible with project-transcript format.
|
||||
"""
|
||||
td = timedelta(seconds=seconds)
|
||||
total_seconds = int(td.total_seconds())
|
||||
milliseconds = int((seconds - total_seconds) * 1000)
|
||||
|
||||
hours = total_seconds // 3600
|
||||
minutes = (total_seconds % 3600) // 60
|
||||
secs = total_seconds % 60
|
||||
|
||||
return f"{hours:02d}:{minutes:02d}:{secs:02d},{milliseconds:03d}"
|
||||
@@ -0,0 +1,596 @@
|
||||
"""
|
||||
Real-time Audio Transcription Node for ComfyUI.
|
||||
|
||||
This node buffers audio segments and performs transcription using faster-whisper,
|
||||
with controlled output timing to prevent message flooding.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import tempfile
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
from typing import Optional, List, Deque, Dict, Any
|
||||
from collections import deque
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from queue import Queue, Empty
|
||||
from faster_whisper import WhisperModel
|
||||
from scipy import signal
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CriticalTranscriptionError(Exception):
|
||||
"""Raised for unrecoverable errors where the stream should stop."""
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class TranscriptionSegment:
|
||||
"""Represents a transcribed segment with timing information."""
|
||||
start: float # Start time in seconds (buffer-relative)
|
||||
end: float # End time in seconds (buffer-relative)
|
||||
text: str # Transcribed text
|
||||
confidence: float = 0.0 # Confidence score
|
||||
|
||||
|
||||
class AudioTranscriptionNode:
|
||||
"""
|
||||
Real-time audio transcription node that buffers audio and outputs
|
||||
transcribed text on a controlled schedule to prevent message flooding.
|
||||
"""
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Configurable class attributes (easy future exposure via UI):
|
||||
# TRANSCRIPTION_QUEUE_MAXLEN: Max queued outputs waiting for timing window
|
||||
# WORD_TIMESTAMPS / BEAM_SIZE / BEST_OF : Whisper inference params
|
||||
# DEFAULT_VAD : Default Voice Activity Detection toggle
|
||||
# MAX_BUFFER_DURATION : Safety limit to prevent memory exhaustion
|
||||
# MAX_BUFFER_SAMPLES : Hard limit on buffer size in samples
|
||||
#
|
||||
# RECOMMENDED CONFIGURATIONS:
|
||||
# - Fast response: accumulation_duration=2.0, audio_chunk_size_ms=1000.0
|
||||
# - Balanced: accumulation_duration=4.0, audio_chunk_size_ms=2000.0
|
||||
# - High quality: accumulation_duration=8.0, audio_chunk_size_ms=4000.0
|
||||
# ------------------------------------------------------------------
|
||||
TRANSCRIPTION_QUEUE_MAXLEN = 100
|
||||
WORD_TIMESTAMPS = True
|
||||
BEAM_SIZE = 1
|
||||
BEST_OF = 1
|
||||
DEFAULT_VAD = True
|
||||
|
||||
# Safety limits to prevent unbounded buffer growth
|
||||
MAX_BUFFER_DURATION = 60.0 # Maximum 60 seconds of audio in buffer
|
||||
MAX_BUFFER_SAMPLES = 16000 * 60 # 60 seconds at 16kHz (960k samples)
|
||||
|
||||
CATEGORY = "audio_utils"
|
||||
RETURN_TYPES = ("STRING",)
|
||||
|
||||
def __init__(self):
|
||||
# Smart accumulation buffer for better transcription quality (optimized for real-time)
|
||||
self.sample_rate = None
|
||||
self.accumulation_buffer = np.empty(0, dtype=np.float32)
|
||||
self.accumulation_duration = 0.0 # Current buffer duration in seconds
|
||||
self.target_accumulation_duration = 3.0 # Target duration for optimal real-time Whisper results
|
||||
|
||||
# Whisper model (per-instance lazy load)
|
||||
self._whisper_model = None
|
||||
|
||||
# Transcription queue for controlled output
|
||||
self.transcription_queue: Deque[str] = deque(maxlen=self.TRANSCRIPTION_QUEUE_MAXLEN)
|
||||
self.last_output_time = 0.0
|
||||
|
||||
# Processing state
|
||||
self.total_audio_processed = 0.0 # Total audio time processed
|
||||
self.transcription_count = 0
|
||||
|
||||
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("AUDIO",),
|
||||
"transcription_interval": ("FLOAT", {
|
||||
"default": 2.0,
|
||||
"min": 1.0,
|
||||
"max": 10.0,
|
||||
"step": 0.5,
|
||||
"tooltip": "Minimum seconds between transcription outputs (optimized for real-time)"
|
||||
}),
|
||||
"accumulation_duration": ("FLOAT", {
|
||||
"default": 3.0,
|
||||
"min": 2.0,
|
||||
"max": 10.0,
|
||||
"step": 0.5,
|
||||
"tooltip": "Audio accumulation duration for optimal Whisper transcription (reduced for real-time: shorter = faster output, longer = better quality)"
|
||||
}),
|
||||
"whisper_model": (["tiny", "base", "small", "medium", "large-v2"], {
|
||||
"default": "base",
|
||||
"tooltip": "Whisper model size (larger = more accurate but slower)"
|
||||
}),
|
||||
"language": (["auto", "en", "es", "fr", "de", "it", "pt", "ru", "ja", "ko", "zh"], {
|
||||
"default": "auto",
|
||||
"tooltip": "Language for transcription (auto = auto-detect)"
|
||||
}),
|
||||
"enable_vad": ("BOOLEAN", {
|
||||
"default": cls.DEFAULT_VAD,
|
||||
"tooltip": "Enable Voice Activity Detection to filter silence"
|
||||
})
|
||||
},
|
||||
"optional": {
|
||||
"output_format": (["text", "json_segments", "json_words"], {
|
||||
"default": "json_segments",
|
||||
"tooltip": "Output format: text (simple), json_segments (with timing), json_words (word-level timing)"
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
FUNCTION = "execute"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
return float("nan")
|
||||
|
||||
|
||||
|
||||
def _load_whisper_model_now(self, model_size: str):
|
||||
"""Load the Whisper model using faster-whisper's automatic download."""
|
||||
# Use CPU for compatibility, can be changed to CUDA if available
|
||||
device = "cpu"
|
||||
try:
|
||||
# Try CUDA first if available
|
||||
import torch # type: ignore
|
||||
if torch.cuda.is_available():
|
||||
device = "cuda"
|
||||
compute_type = "float16"
|
||||
else:
|
||||
compute_type = "int8"
|
||||
except Exception:
|
||||
compute_type = "int8"
|
||||
|
||||
logger.info(f"Loading Whisper model via faster-whisper: {model_size}")
|
||||
try:
|
||||
self._whisper_model = WhisperModel(
|
||||
model_size,
|
||||
device=device,
|
||||
compute_type=compute_type
|
||||
)
|
||||
# Attach size name for change detection
|
||||
self._whisper_model.model_size_name = model_size # type: ignore[attr-defined]
|
||||
logger.info(f"Whisper model '{model_size}' loaded successfully on {device}")
|
||||
except (MemoryError, RuntimeError) as e:
|
||||
logger.error(f"Critical: Failed to load Whisper model '{model_size}' due to resource issue: {e}")
|
||||
raise CriticalTranscriptionError(f"Model load failed (resources): {e}") from e
|
||||
except OSError as e:
|
||||
logger.error(f"Critical: Filesystem/model download issue for '{model_size}': {e}")
|
||||
raise CriticalTranscriptionError(f"Filesystem/model download failure: {e}") from e
|
||||
except Exception as e:
|
||||
logger.error(f"Critical: Unexpected error loading model '{model_size}': {e}")
|
||||
raise CriticalTranscriptionError(f"Unexpected model load error: {e}") from e
|
||||
|
||||
def _ensure_model_loaded(self, model_size: str) -> bool:
|
||||
"""Ensure the Whisper model is loaded. Returns True if a fresh load occurred."""
|
||||
if (self._whisper_model is None or
|
||||
getattr(self._whisper_model, 'model_size_name', '') != model_size):
|
||||
logger.info(f"Loading Whisper model: {model_size}")
|
||||
self._load_whisper_model_now(model_size)
|
||||
return True
|
||||
return False
|
||||
|
||||
def _process_audio_chunk(self, audio_input: np.ndarray, sample_rate: int, accumulation_duration: float):
|
||||
"""
|
||||
Process a single audio chunk and accumulate it for optimal Whisper transcription.
|
||||
Returns True when we have enough audio for transcription.
|
||||
"""
|
||||
# Initialize sample rate on first call
|
||||
if self.sample_rate is None:
|
||||
self.sample_rate = sample_rate
|
||||
self.target_accumulation_duration = accumulation_duration
|
||||
logger.info(f"Initialized transcription node with sample rate: {sample_rate}Hz, target accumulation: {accumulation_duration}s")
|
||||
|
||||
# Skip empty frames
|
||||
if audio_input is None or audio_input.size == 0:
|
||||
return False
|
||||
|
||||
# Normalize audio for Whisper
|
||||
normalized_audio = self._normalize_audio_for_whisper(audio_input)
|
||||
|
||||
# Add to accumulation buffer
|
||||
if self.accumulation_buffer.size == 0:
|
||||
self.accumulation_buffer = normalized_audio.copy()
|
||||
else:
|
||||
self.accumulation_buffer = np.concatenate([self.accumulation_buffer, normalized_audio])
|
||||
|
||||
# Update accumulation duration
|
||||
self.accumulation_duration = len(self.accumulation_buffer) / self.sample_rate
|
||||
|
||||
# Safety check: enforce buffer limits to prevent memory exhaustion
|
||||
self._enforce_buffer_limits()
|
||||
|
||||
# Check if we have enough audio for transcription
|
||||
ready = self.accumulation_duration >= self.target_accumulation_duration
|
||||
|
||||
if ready:
|
||||
logger.debug(f"Audio accumulation ready: {self.accumulation_duration:.2f}s >= {self.target_accumulation_duration}s ({len(self.accumulation_buffer)} samples)")
|
||||
|
||||
return ready
|
||||
|
||||
def _normalize_audio_for_whisper(self, audio_input: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Minimal audio normalization to avoid artifacts.
|
||||
Assumes LoadAudioTensorStream has already done most processing.
|
||||
"""
|
||||
if audio_input is None or audio_input.size == 0:
|
||||
return np.array([], dtype=np.int16)
|
||||
|
||||
# Ensure numpy array
|
||||
if not isinstance(audio_input, np.ndarray):
|
||||
audio_input = np.array(audio_input)
|
||||
|
||||
# If it's already int16 and 1D, return as-is
|
||||
if audio_input.dtype == np.int16 and audio_input.ndim == 1:
|
||||
return audio_input
|
||||
|
||||
# Handle multi-channel by taking first channel only
|
||||
if audio_input.ndim > 1:
|
||||
audio_input = audio_input.flatten() if audio_input.size < 1000 else audio_input[:, 0] if audio_input.shape[1] > audio_input.shape[0] else audio_input[0, :]
|
||||
|
||||
# Convert to int16 only if necessary
|
||||
if audio_input.dtype != np.int16:
|
||||
if audio_input.dtype in [np.float32, np.float64]:
|
||||
# Very conservative conversion to avoid clipping artifacts
|
||||
audio_input = np.clip(audio_input, -0.99, 0.99) # Leave headroom
|
||||
audio_input = (audio_input * 32000).astype(np.int16) # Slightly lower scale
|
||||
else:
|
||||
audio_input = audio_input.astype(np.int16)
|
||||
|
||||
return audio_input
|
||||
|
||||
def _enforce_buffer_limits(self):
|
||||
"""
|
||||
Enforce safety limits on the accumulation buffer to prevent memory exhaustion.
|
||||
Trims buffer from the beginning if limits are exceeded.
|
||||
"""
|
||||
if self.sample_rate is None or self.accumulation_buffer.size == 0:
|
||||
return
|
||||
|
||||
# Check duration limit
|
||||
if self.accumulation_duration > self.MAX_BUFFER_DURATION:
|
||||
# Calculate how many samples to keep (keep the most recent audio)
|
||||
samples_to_keep = int(self.MAX_BUFFER_DURATION * self.sample_rate)
|
||||
if samples_to_keep < len(self.accumulation_buffer):
|
||||
# Trim from the beginning, keep the end
|
||||
self.accumulation_buffer = self.accumulation_buffer[-samples_to_keep:]
|
||||
self.accumulation_duration = len(self.accumulation_buffer) / self.sample_rate
|
||||
logger.warning(f"Buffer duration limit exceeded, trimmed to {self.accumulation_duration:.2f}s (kept most recent {self.MAX_BUFFER_DURATION}s)")
|
||||
|
||||
# Check sample count limit (secondary safety check)
|
||||
if len(self.accumulation_buffer) > self.MAX_BUFFER_SAMPLES:
|
||||
self.accumulation_buffer = self.accumulation_buffer[-self.MAX_BUFFER_SAMPLES:]
|
||||
self.accumulation_duration = len(self.accumulation_buffer) / self.sample_rate
|
||||
logger.warning(f"Buffer sample limit exceeded, trimmed to {len(self.accumulation_buffer)} samples ({self.accumulation_duration:.2f}s)")
|
||||
|
||||
|
||||
def _transcribe_accumulated_audio(self, model_size: str, language: str, enable_vad: bool, batch_size: int, output_format: str = "json_segments") -> Optional[str]:
|
||||
"""Transcribe the accumulated audio buffer and return combined text."""
|
||||
if self.accumulation_buffer.size == 0:
|
||||
return None
|
||||
try:
|
||||
# Ensure model is loaded; if it was freshly loaded, emit a single warmup sentinel
|
||||
send_warmup_sentinel = self._ensure_model_loaded(model_size)
|
||||
# Emit a single warmup sentinel immediately after model loading completes
|
||||
if send_warmup_sentinel:
|
||||
if output_format == "text":
|
||||
result = "__WARMUP_SENTINEL__"
|
||||
elif output_format == "json_segments":
|
||||
result = json.dumps([{ "start": 0.0, "end": 1.0, "text": "__WARMUP_SENTINEL__"}])
|
||||
elif output_format == "json_words":
|
||||
result = json.dumps([{ "start": 0.0, "end": 1.0, "word": "__WARMUP_SENTINEL__", "probability": 1.0}])
|
||||
else:
|
||||
result = "__WARMUP_SENTINEL__"
|
||||
logger.debug("Warmup sentinel emitted once after model load")
|
||||
return result
|
||||
|
||||
# Use the accumulated audio buffer
|
||||
audio_chunk = self.accumulation_buffer
|
||||
|
||||
# Validate audio buffer quality before transcription
|
||||
if len(audio_chunk) == 0:
|
||||
logger.warning("Audio buffer is empty")
|
||||
return None
|
||||
|
||||
# Check for audio artifacts (all zeros, extreme values, etc.)
|
||||
if np.all(audio_chunk == 0):
|
||||
return ""
|
||||
|
||||
# Check for clipping or extreme values that might cause hallucinations
|
||||
if np.abs(audio_chunk).max() < 100: # Very quiet audio
|
||||
return ""
|
||||
|
||||
# Save audio to temporary WAV file for whisper
|
||||
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
|
||||
temp_path = temp_file.name
|
||||
|
||||
try:
|
||||
# Write clean WAV file to avoid format conversion artifacts
|
||||
try:
|
||||
from scipy.io.wavfile import write # type: ignore
|
||||
# Ensure audio is properly scaled for WAV format
|
||||
clean_audio = np.copy(audio_chunk)
|
||||
write(temp_path, int(self.sample_rate or 16000), clean_audio)
|
||||
except ImportError:
|
||||
# Fallback to manual WAV writing
|
||||
self._write_wav_file(temp_path, audio_chunk, int(self.sample_rate or 16000))
|
||||
|
||||
language_code = None if language == "auto" else language
|
||||
segments, info = self._whisper_model.transcribe( # type: ignore[attr-defined]
|
||||
temp_path,
|
||||
language=language_code,
|
||||
word_timestamps=self.WORD_TIMESTAMPS,
|
||||
vad_filter=enable_vad,
|
||||
beam_size=self.BEAM_SIZE,
|
||||
best_of=self.BEST_OF
|
||||
)
|
||||
|
||||
# Simple format output without complex deduplication (causes accuracy issues)
|
||||
if output_format == "text":
|
||||
texts = [s.text.strip() for s in segments if s.text and s.text.strip()]
|
||||
result = " ".join(texts)
|
||||
elif output_format == "json_segments":
|
||||
seg_list = [
|
||||
{"start": s.start, "end": s.end, "text": s.text.strip()}
|
||||
for s in segments if s.text and s.text.strip()
|
||||
]
|
||||
result = json.dumps(seg_list, ensure_ascii=False)
|
||||
elif output_format == "json_words":
|
||||
words = []
|
||||
for s in segments:
|
||||
s_words = getattr(s, 'words', None)
|
||||
if s_words:
|
||||
for w in s_words:
|
||||
if w.word.strip():
|
||||
words.append({
|
||||
"start": w.start,
|
||||
"end": w.end,
|
||||
"word": w.word.strip(),
|
||||
"probability": getattr(w, 'probability', 1.0)
|
||||
})
|
||||
elif s.text and s.text.strip():
|
||||
words.append({
|
||||
"start": s.start,
|
||||
"end": s.end,
|
||||
"word": s.text.strip(),
|
||||
"probability": 1.0
|
||||
})
|
||||
result = json.dumps(words, ensure_ascii=False)
|
||||
else:
|
||||
result = None
|
||||
|
||||
if not result or (output_format in ["json_segments", "json_words"] and result == "[]"):
|
||||
result = "" # Return empty string for ComfyUI compatibility
|
||||
logger.debug(f"No transcription produced, returning empty")
|
||||
|
||||
|
||||
# Metrics
|
||||
if self.sample_rate and self.accumulation_buffer.size > 0:
|
||||
self.total_audio_processed += len(self.accumulation_buffer) / self.sample_rate
|
||||
self.transcription_count += 1
|
||||
|
||||
# Clear the accumulation buffer after processing (no overlap to prevent duplication)
|
||||
self.accumulation_buffer = np.empty(0, dtype=np.float32)
|
||||
self.accumulation_duration = 0.0
|
||||
logger.debug("Cleared accumulation buffer for next cycle")
|
||||
|
||||
if result:
|
||||
logger.debug(
|
||||
f"Transcribed chunk {self.transcription_count} ({output_format}): '{str(result)[:50]}...' ({len(str(result))} chars)"
|
||||
)
|
||||
else:
|
||||
logger.debug(f"Transcribed chunk {self.transcription_count} produced no result")
|
||||
return result
|
||||
finally:
|
||||
try:
|
||||
os.unlink(temp_path)
|
||||
except FileNotFoundError:
|
||||
logger.debug(f"Temp transcription file already removed: {temp_path}")
|
||||
except PermissionError as e:
|
||||
logger.warning(f"Permission error removing temp file {temp_path}: {e}")
|
||||
except OSError as e:
|
||||
logger.warning(f"OS error removing temp file {temp_path}: {e}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Unexpected error removing temp file {temp_path}: {e}")
|
||||
except CriticalTranscriptionError:
|
||||
# Propagate critical to stop workflow
|
||||
raise
|
||||
except (MemoryError, OSError) as e:
|
||||
logger.error(f"Critical transcription failure: {e}")
|
||||
raise CriticalTranscriptionError(f"Critical transcription failure: {e}") from e
|
||||
except Exception as e:
|
||||
logger.warning(f"Non-critical transcription error, skipping chunk: {e}")
|
||||
return None
|
||||
|
||||
def _resample_to_16khz(self, audio_data: np.ndarray, original_sample_rate: int) -> np.ndarray:
|
||||
"""
|
||||
Resample audio to 16kHz for optimal Whisper performance.
|
||||
|
||||
Args:
|
||||
audio_data: Input audio data as numpy array
|
||||
original_sample_rate: Original sample rate of the audio
|
||||
|
||||
Returns:
|
||||
Resampled audio data at 16kHz
|
||||
"""
|
||||
target_sample_rate = 16000
|
||||
|
||||
if original_sample_rate == target_sample_rate:
|
||||
return audio_data
|
||||
|
||||
if audio_data.size == 0:
|
||||
return audio_data
|
||||
# Use scipy.signal.resample for high-quality resampling
|
||||
num_samples = int(len(audio_data) * target_sample_rate / original_sample_rate)
|
||||
resampled_audio = signal.resample(audio_data, num_samples)
|
||||
|
||||
logger.debug(f"Resampled audio from {original_sample_rate}Hz to {target_sample_rate}Hz "
|
||||
f"({len(audio_data)} -> {len(resampled_audio)} samples)")
|
||||
|
||||
return resampled_audio.astype(np.float32)
|
||||
|
||||
def _write_wav_file(self, filename: str, audio_data: np.ndarray, sample_rate: int):
|
||||
"""Write WAV file manually if scipy is not available."""
|
||||
import struct
|
||||
import wave
|
||||
|
||||
with wave.open(filename, 'wb') as wav_file:
|
||||
wav_file.setnchannels(1) # Mono
|
||||
wav_file.setsampwidth(2) # 16-bit
|
||||
wav_file.setframerate(sample_rate)
|
||||
wav_file.writeframes(audio_data.tobytes())
|
||||
|
||||
def _should_output_transcription(self, transcription_interval: float) -> bool:
|
||||
"""Check if we should output transcription based on queue or timing."""
|
||||
# For real-time streaming, prioritize queue content
|
||||
if len(self.transcription_queue) > 0:
|
||||
# logger.debug(f"Queue has {len(self.transcription_queue)} items - outputting immediately") # Too frequent
|
||||
return True
|
||||
|
||||
# Fallback to timing interval only if queue is empty
|
||||
current_time = time.time()
|
||||
|
||||
effective_interval = transcription_interval
|
||||
|
||||
return (current_time - self.last_output_time) >= effective_interval
|
||||
|
||||
def _get_queued_transcription(self) -> Optional[str]:
|
||||
"""Get transcription from queue if available."""
|
||||
try:
|
||||
return self.transcription_queue.popleft()
|
||||
except IndexError:
|
||||
return None
|
||||
|
||||
def execute(self, audio, transcription_interval, accumulation_duration,
|
||||
whisper_model="base", language="auto", enable_vad=True, output_format="json_segments"):
|
||||
"""
|
||||
Execute transcription on streaming audio input.
|
||||
|
||||
Args:
|
||||
audio: Audio input from LoadAudioTensor - AUDIO dictionary with waveform and sample_rate
|
||||
transcription_interval: Minimum seconds between outputs (prevents message flooding) - optimized for real-time
|
||||
accumulation_duration: Audio accumulation duration for optimal Whisper results - reduced for real-time
|
||||
whisper_model: Whisper model size (tiny/base/small/medium/large-v2)
|
||||
language: Language for transcription (auto for auto-detection)
|
||||
enable_vad: Enable voice activity detection to filter silence
|
||||
output_format: Output format (text/json_segments/json_words)
|
||||
|
||||
Returns:
|
||||
Tuple containing transcribed text (empty string if not ready to output)
|
||||
"""
|
||||
try:
|
||||
# Extract audio data and sample rate from AUDIO dictionary
|
||||
if not isinstance(audio, dict) or 'waveform' not in audio or 'sample_rate' not in audio:
|
||||
logger.warning(f"Invalid audio format: expected AUDIO dict with waveform and sample_rate, got {type(audio)}")
|
||||
return ("",)
|
||||
|
||||
waveform = audio['waveform']
|
||||
sample_rate = audio['sample_rate']
|
||||
|
||||
# Convert torch tensor to numpy if needed
|
||||
if isinstance(waveform, torch.Tensor):
|
||||
audio_data = waveform.detach().cpu().numpy()
|
||||
else:
|
||||
audio_data = waveform
|
||||
|
||||
# Handle tensor dimensions (batch, channels, samples) -> samples
|
||||
if audio_data.ndim == 3:
|
||||
# (batch, channels, samples) -> take first batch, average channels
|
||||
audio_data = audio_data[0] # Remove batch dimension
|
||||
if audio_data.shape[0] > 1: # Multiple channels
|
||||
audio_data = audio_data.mean(axis=0) # Average channels to mono
|
||||
else:
|
||||
audio_data = audio_data[0] # Single channel
|
||||
elif audio_data.ndim == 2:
|
||||
# (channels, samples) -> average channels to mono if needed
|
||||
if audio_data.shape[0] > 1:
|
||||
audio_data = audio_data.mean(axis=0)
|
||||
else:
|
||||
audio_data = audio_data[0]
|
||||
# If already 1D, use as-is
|
||||
|
||||
# Validate inputs
|
||||
if audio_data is None or audio_data.size == 0:
|
||||
logger.warning("Received empty audio data, returning empty")
|
||||
return ("",)
|
||||
|
||||
# Resample to 16kHz for optimal Whisper performance
|
||||
try:
|
||||
audio_data = self._resample_to_16khz(audio_data, sample_rate)
|
||||
except Exception as e:
|
||||
logger.error(f"Audio resampling to 16Khz failed, returning empty: {e}")
|
||||
return ("",)
|
||||
|
||||
sample_rate = 16000 # Update sample rate after resampling
|
||||
|
||||
# Ensure audio_data is a numpy array with proper dtype
|
||||
if not isinstance(audio_data, np.ndarray):
|
||||
logger.debug(f"Converting audio_data from {type(audio_data)} to numpy array.")
|
||||
audio_data = np.array(audio_data, dtype=np.float32)
|
||||
|
||||
|
||||
# Process the audio chunk and accumulate for optimal transcription
|
||||
ready_for_transcription = self._process_audio_chunk(audio_data, sample_rate, accumulation_duration)
|
||||
|
||||
# Transcribe if we have accumulated enough audio
|
||||
if ready_for_transcription:
|
||||
transcription = self._transcribe_accumulated_audio(whisper_model, language, enable_vad, 16, output_format)
|
||||
if transcription and transcription.strip() and len(transcription.strip()) > 5:
|
||||
# Add to queue for controlled output timing (only substantial content or warmup sentinels)
|
||||
self._queue_transcription(transcription)
|
||||
is_sentinel = "__WARMUP_SENTINEL__" in transcription
|
||||
if is_sentinel:
|
||||
logger.debug(f"Queued warmup sentinel ({output_format})")
|
||||
else:
|
||||
logger.debug(f"Queued transcription ({output_format}): '{transcription[:50]}...' (length: {len(transcription)})")
|
||||
elif transcription and transcription != "":
|
||||
logger.debug(f"Skipping minimal transcription: '{transcription}' (too short)")
|
||||
elif transcription == "":
|
||||
logger.debug(f"Empty transcription result (post-warmup silence)")
|
||||
else:
|
||||
logger.debug(f"No transcription result returned")
|
||||
|
||||
# Output transcriptions immediately when available (real-time streaming)
|
||||
if self._should_output_transcription(transcription_interval):
|
||||
output_text = self._get_queued_transcription()
|
||||
if output_text and output_text.strip():
|
||||
self.last_output_time = time.time()
|
||||
|
||||
# Log differently for sentinel vs regular transcription
|
||||
if "__WARMUP_SENTINEL__" in output_text:
|
||||
logger.info(f"Outputting warmup sentinel immediately ({output_format}, {len(output_text)} chars): '{output_text[:100]}...'")
|
||||
else:
|
||||
logger.info(f"Outputting transcription immediately ({output_format}, {len(output_text)} chars): '{output_text[:100]}...' (queue: {len(self.transcription_queue)})")
|
||||
|
||||
return (output_text,)
|
||||
|
||||
# Return empty string if not ready to output (SaveTextTensor will filter empty content)
|
||||
return ("",)
|
||||
|
||||
except CriticalTranscriptionError as e:
|
||||
logger.error(f"Critical error in transcription node: {e}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.warning(f"Recoverable execute error: {e}")
|
||||
return ("",)
|
||||
|
||||
def _queue_transcription(self, transcription: str):
|
||||
"""Safely queue transcription for controlled output."""
|
||||
if len(self.transcription_queue) == self.transcription_queue.maxlen:
|
||||
logger.debug("Transcription queue is full, oldest entry will be replaced.")
|
||||
self.transcription_queue.append(transcription)
|
||||
|
||||
|
||||
|
||||
@@ -1,497 +0,0 @@
|
||||
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,)
|
||||
@@ -1 +1 @@
|
||||
# Dummy file ensuring this package will be recognized.
|
||||
# Dummy file ensuring this sub-package will be recognized.
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Source code for the Feature Bank of the StreamV2V paper.
|
||||
|
||||
References:
|
||||
StreamV2V: A Deep Generative Model for High-fidelity View Synthesis - Liang et al.,
|
||||
StreamV2V: Looking Backward: Streaming Video-to-Video Translation with Feature Banks - Liang et al.,
|
||||
2025, URL: https://arxiv.org/abs/2405.15757
|
||||
"""
|
||||
|
||||
@@ -14,7 +14,7 @@ from torch import nn
|
||||
from diffusers.utils import USE_PEFT_BACKEND
|
||||
|
||||
|
||||
def get_nn_feats(
|
||||
def extract_nn_features(
|
||||
x: torch.Tensor, y: torch.Tensor, threshold: float = 0.9
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
@@ -182,7 +182,6 @@ class FeatureBankAttentionProcessor:
|
||||
"""
|
||||
Forward pass of feature bank attention with feature injection.
|
||||
"""
|
||||
residual = hidden_states
|
||||
args = () if USE_PEFT_BACKEND else (scale,)
|
||||
input_ndim = hidden_states.ndim
|
||||
|
||||
@@ -192,7 +191,7 @@ class FeatureBankAttentionProcessor:
|
||||
batch_size, channel, height * width
|
||||
).transpose(1, 2)
|
||||
|
||||
batch_size, key_tokens, _ = (
|
||||
batch_size, _, _ = (
|
||||
hidden_states.shape
|
||||
if encoder_hidden_states is None
|
||||
else encoder_hidden_states.shape
|
||||
@@ -391,14 +390,14 @@ class FeatureBankAttentionProcessor:
|
||||
"Output",
|
||||
)
|
||||
|
||||
self._debug_print("\n--- DEBUG: Before get_nn_feats ---")
|
||||
self._debug_print("\n--- DEBUG: Before extract_nn_features ---")
|
||||
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(
|
||||
nn_hidden_states = extract_nn_features(
|
||||
hidden_states,
|
||||
cached_output_tensor_from_list,
|
||||
threshold=self.threshold,
|
||||
@@ -460,7 +459,7 @@ class FeatureBankAttentionProcessorNode:
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "update_node"
|
||||
CATEGORY = "feature-bank"
|
||||
CATEGORY = "StreamPack/model_patches/unet"
|
||||
|
||||
def update_node(
|
||||
self,
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
# Dummy file ensuring this sub-package will be recognized.
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,88 @@
|
||||
# ComfyUI Super Resolution
|
||||
|
||||
A collection of high-performance neural network-based Super Resolution models for ComfyUI.
|
||||
|
||||
## Features
|
||||
|
||||
- **Multiple SR Models**: Choose from FSRCNN, ESPCN, LapSRN, and EDSR with different quality/speed tradeoffs
|
||||
- **Modular Design**: Load models once and reuse them across multiple upscaling operations
|
||||
- **CUDA Acceleration**: Automatic GPU acceleration when available
|
||||
- **Multiple Scale Factors**: Support for 2x, 3x, and 4x upscaling
|
||||
|
||||
## Installation
|
||||
|
||||
1. Clone this repository to your ComfyUI custom_nodes directory:
|
||||
```
|
||||
cd ComfyUI/custom_nodes
|
||||
git clone https://github.com/yourusername/ComfyUI_SuperResolution
|
||||
```
|
||||
|
||||
2. Restart ComfyUI - the necessary models will be automatically downloaded on first use
|
||||
|
||||
## Usage
|
||||
|
||||
1. Add a **SR Model Loader** node to your workflow
|
||||
2. Select your desired model type and scale factor
|
||||
3. Connect the model output to a **SR Upscale** node
|
||||
4. Connect an image to the upscale node
|
||||
5. Run your workflow to get high-quality upscaled images!
|
||||
|
||||
## Model Comparison
|
||||
|
||||
| Model | Architecture | Features | Best For | Speed | Quality |
|
||||
|-------|-------------|----------|----------|-------|---------|
|
||||
| **FSRCNN-small** | Lightweight CNN | Fast, minimal memory use | Real-time processing, mobile | ★★★★★ | ★★ |
|
||||
| **FSRCNN** | CNN with larger features | Good balance of speed/quality | General purpose | ★★★★ | ★★★ |
|
||||
| **ESPCN** | Sub-pixel convolutions | Efficient upscaling at end | Text/line drawings | ★★★★ | ★★★ |
|
||||
| **VDSR** | Very deep CNN | Better edge reconstruction | Detailed images with edges | ★★★ | ★★★★ |
|
||||
| **LapSRN** | Laplacian pyramid | Progressive upscaling | Sharp edges, details | ★★★ | ★★★★ |
|
||||
| **EDSR** | Deep residual network | Most parameters, best quality | Maximum detail | ★★ | ★★★★★ |
|
||||
|
||||
## Technical Details
|
||||
|
||||
Each model has unique architectural characteristics:
|
||||
|
||||
1. **FSRCNN / FSRCNN-small**: Direct mapping from low to high resolution, lightweight with fewer parameters.
|
||||
|
||||
2. **ESPCN**: Uses the efficient sub-pixel convolution technique that processes at low resolution and only expands dimensions at the final layer.
|
||||
|
||||
3. **VDSR**: Very Deep Super Resolution network with 20 convolutional layers, uses global residual learning.
|
||||
|
||||
4. **LapSRN**: Uses a Laplacian pyramid structure to progressively upscale images, which preserves edges better than the others.
|
||||
|
||||
5. **EDSR**: Enhanced Deep Super-Resolution, a residual network with significantly more parameters, offering the highest quality but slower processing.
|
||||
|
||||
## Performance Expectations (RTX 4090)
|
||||
|
||||
| Model | 2x Upscaling | 4x Upscaling | Memory Usage |
|
||||
|-------|--------------|--------------|--------------|
|
||||
| FSRCNN-small | 200+ FPS | 150+ FPS | < 100MB |
|
||||
| FSRCNN | 100+ FPS | 80+ FPS | < 150MB |
|
||||
| ESPCN | 90+ FPS | 70+ FPS | < 150MB |
|
||||
| LapSRN | 60+ FPS | 40+ FPS | < 200MB |
|
||||
| EDSR | 40+ FPS | 25+ FPS | < 500MB |
|
||||
|
||||
## When to Use Each Model
|
||||
|
||||
- **FSRCNN-small**: When maximum speed is required
|
||||
- **FSRCNN**: For a good balance of quality and performance
|
||||
- **ESPCN**: For text and line art upscaling
|
||||
- **LapSRN**: When you need better edge preservation than FSRCNN
|
||||
- **EDSR**: When maximum quality is desired and performance is secondary
|
||||
|
||||
## Comparison to Traditional Methods
|
||||
|
||||
These neural network-based upscaling methods offer significantly better quality than traditional algorithms like Lanczos, Bicubic, or Bilinear interpolation, while often being just as fast thanks to CUDA acceleration.
|
||||
|
||||
## Credits
|
||||
|
||||
This nodepack implements models originally created by:
|
||||
|
||||
- FSRCNN: Dong et al. (https://github.com/ryanontheinside/FSRCNN_Tensorflow)
|
||||
- EDSR: Lim et al. (https://github.com/ryanontheinside/EDSR_Tensorflow)
|
||||
- ESPCN: Shi et al. (https://github.com/ryanontheinside/TF-ESPCN)
|
||||
- LapSRN: Lai et al. (https://github.com/ryanontheinside/TF-LapSRN)
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1 @@
|
||||
# Dummy file ensuring this sub-package will be recognized.
|
||||
@@ -0,0 +1,5 @@
|
||||
# Makes models directory a proper Python package
|
||||
|
||||
from .sr_model import SuperResolutionModel
|
||||
|
||||
__all__ = ["SuperResolutionModel"]
|
||||
@@ -0,0 +1,251 @@
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
class SuperResolutionModel:
|
||||
"""Unified Super Resolution Model for all algorithm types"""
|
||||
|
||||
# Configuration registry for all model types
|
||||
MODEL_CONFIGS = {
|
||||
"fsrcnn": {
|
||||
"algorithm": "fsrcnn",
|
||||
"repo_subdir": "FSRCNN_Tensorflow/models",
|
||||
"filename_pattern": "FSRCNN_x{scale}.pb",
|
||||
"supported_scales": [2, 3, 4],
|
||||
},
|
||||
"fsrcnn-small": {
|
||||
"algorithm": "fsrcnn",
|
||||
"repo_subdir": "FSRCNN_Tensorflow/models",
|
||||
"filename_pattern": "FSRCNN-small_x{scale}.pb",
|
||||
"supported_scales": [2, 3, 4],
|
||||
},
|
||||
"edsr": {
|
||||
"algorithm": "edsr",
|
||||
"repo_subdir": "EDSR_Tensorflow/models",
|
||||
"filename_pattern": "EDSR_x{scale}.pb",
|
||||
"supported_scales": [2, 3, 4],
|
||||
},
|
||||
"espcn": {
|
||||
"algorithm": "espcn",
|
||||
"repo_subdir": "TF-ESPCN/export",
|
||||
"filename_pattern": "ESPCN_x{scale}.pb",
|
||||
"supported_scales": [2, 3, 4],
|
||||
},
|
||||
"lapsrn": {
|
||||
"algorithm": "lapsrn",
|
||||
"repo_subdir": "TF-LapSRN/export",
|
||||
"filename_pattern": "LapSRN_x{scale}.pb",
|
||||
"supported_scales": [2, 3, 4],
|
||||
},
|
||||
"vdsr": {
|
||||
"algorithm": "vdsr",
|
||||
"repo_subdir": "TF-VDSR/export",
|
||||
"filename_pattern": "VDSR_x{scale}.pb",
|
||||
"supported_scales": [2, 3, 4],
|
||||
}
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def is_cuda_available(cls):
|
||||
"""Check if CUDA is available in OpenCV"""
|
||||
# Check if CUDA module exists in OpenCV
|
||||
if not hasattr(cv2, "cuda") or cv2.cuda.getCudaEnabledDeviceCount() == 0:
|
||||
return False
|
||||
|
||||
# Check if DNN module has CUDA support
|
||||
try:
|
||||
# Try to get available backends if the method exists
|
||||
if hasattr(cv2.dnn, "getAvailableBackends"):
|
||||
backends = cv2.dnn.getAvailableBackends()
|
||||
if cv2.dnn.DNN_BACKEND_CUDA not in backends:
|
||||
return False
|
||||
|
||||
# Check for CUDA targets if method exists
|
||||
if hasattr(cv2.dnn, "getAvailableTargets"):
|
||||
targets = cv2.dnn.getAvailableTargets(cv2.dnn.DNN_BACKEND_CUDA)
|
||||
if len(targets) == 0:
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def create(cls, model_type, scale_factor, **kwargs):
|
||||
"""Factory method to create a super resolution model by type"""
|
||||
model_type = model_type.lower()
|
||||
|
||||
if model_type not in cls.MODEL_CONFIGS:
|
||||
raise ValueError(f"Unknown model type: {model_type}")
|
||||
|
||||
return cls(model_type, scale_factor, **kwargs)
|
||||
|
||||
def __init__(self, model_name, scale_factor, device="cuda" if torch.cuda.is_available() else "cpu"):
|
||||
self.model_name = model_name.lower()
|
||||
self.scale_factor = int(scale_factor)
|
||||
self.device = device
|
||||
self.model = None
|
||||
self.loaded = False
|
||||
self.using_cuda = False
|
||||
|
||||
# Get configuration for this model
|
||||
self.config = self.MODEL_CONFIGS.get(self.model_name)
|
||||
if not self.config:
|
||||
raise ValueError(f"No configuration found for model: {model_name}")
|
||||
|
||||
# Validate scale factor
|
||||
if self.scale_factor not in self.config["supported_scales"]:
|
||||
raise ValueError(f"Scale factor {scale_factor} not supported for {model_name}. "
|
||||
f"Supported scales: {self.config['supported_scales']}")
|
||||
|
||||
def load_model(self, model_path=None):
|
||||
"""Load the super resolution model"""
|
||||
if model_path is None:
|
||||
# Try to find model in standard locations
|
||||
model_dir = self.model_name.upper()
|
||||
|
||||
# Search paths from highest to lowest priority
|
||||
search_paths = [
|
||||
os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../downloaded_models", model_dir),
|
||||
os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../models"),
|
||||
os.path.expanduser("~/.cache/comfyui/custom_nodes/ComfyUI_SuperResolution/models"),
|
||||
]
|
||||
|
||||
# Add repo-specific path if available
|
||||
if self.config["repo_subdir"]:
|
||||
search_paths.append(
|
||||
os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../../", self.config["repo_subdir"])
|
||||
)
|
||||
|
||||
for path in search_paths:
|
||||
filename = self.config["filename_pattern"].format(scale=self.scale_factor)
|
||||
candidate = os.path.join(path, filename)
|
||||
if os.path.exists(candidate):
|
||||
model_path = candidate
|
||||
break
|
||||
|
||||
if not model_path or not os.path.exists(model_path):
|
||||
raise FileNotFoundError(f"Could not find {self.model_name} x{self.scale_factor} model file")
|
||||
|
||||
# Load model using OpenCV DNN Super Resolution with the correct method
|
||||
try:
|
||||
# Try different methods to create the DNN SuperRes object for cross-platform compatibility
|
||||
# This handles differences between Windows and Linux OpenCV builds
|
||||
try:
|
||||
# Method 1: Windows-style API (most common)
|
||||
self.model = cv2.dnn_superres.DnnSuperResImpl_create()
|
||||
except AttributeError:
|
||||
try:
|
||||
# Method 2: More generic constructor if available
|
||||
self.model = cv2.dnn_superres.SuperResolution_create()
|
||||
except AttributeError:
|
||||
# Method 3: Direct class instantiation if available
|
||||
self.model = cv2.dnn_superres.DnnSuperResImpl()
|
||||
|
||||
self.model.readModel(model_path)
|
||||
self.model.setModel(self.config["algorithm"], self.scale_factor)
|
||||
|
||||
# Check CUDA availability during model loading
|
||||
# Note: We don't raise errors here, as that will be handled during upscale if use_cuda=True
|
||||
system_has_cuda = torch.cuda.is_available()
|
||||
opencv_has_cuda = self.is_cuda_available()
|
||||
|
||||
if system_has_cuda and opencv_has_cuda:
|
||||
self.model.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
|
||||
self.model.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
|
||||
print(f"Model loaded with CUDA backend for {self.model_name}")
|
||||
self.using_cuda = True
|
||||
else:
|
||||
self.model.setPreferableBackend(cv2.dnn.DNN_BACKEND_DEFAULT)
|
||||
self.model.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
|
||||
print(f"Model loaded with CPU backend for {self.model_name}")
|
||||
self.using_cuda = False
|
||||
|
||||
self.loaded = True
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Error with OpenCV DNN Super Resolution: {e}. Check that opencv-contrib-python is installed.")
|
||||
|
||||
return self
|
||||
|
||||
def upscale(self, image, use_cuda=True):
|
||||
"""Upscale an image using the model"""
|
||||
if not self.loaded:
|
||||
self.load_model()
|
||||
|
||||
# Check CUDA requirements if use_cuda is enabled
|
||||
if use_cuda:
|
||||
system_has_cuda = torch.cuda.is_available()
|
||||
opencv_has_cuda = self.is_cuda_available()
|
||||
|
||||
if not system_has_cuda:
|
||||
raise RuntimeError("CUDA acceleration requested but CUDA is not available on your system. "
|
||||
"Either install CUDA or disable CUDA acceleration in the node settings.")
|
||||
|
||||
if not opencv_has_cuda:
|
||||
raise RuntimeError("CUDA acceleration requested but your OpenCV build doesn't support CUDA. "
|
||||
"Install a CUDA-enabled OpenCV build (via conda: 'conda install -c conda-forge opencv cudatoolkit') "
|
||||
"or disable CUDA acceleration in the node settings.")
|
||||
|
||||
# Convert to cv2 format if needed
|
||||
cv2_image = self.convert_comfy_to_cv2(image)
|
||||
|
||||
# Process the image
|
||||
upscaled = self.model.upsample(cv2_image)
|
||||
|
||||
# Convert back to original format
|
||||
return self.convert_cv2_to_comfy(upscaled, image)
|
||||
|
||||
def convert_comfy_to_cv2(self, tensor):
|
||||
"""Convert ComfyUI tensor (BHWC) to cv2 image (HWC)"""
|
||||
# Handle both tensor and numpy array inputs
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
# Ensure tensor is on CPU and convert to numpy
|
||||
image = tensor.detach().cpu().numpy()
|
||||
else:
|
||||
image = tensor
|
||||
|
||||
# Take first image if batch dimension exists
|
||||
if len(image.shape) == 4:
|
||||
image = image[0]
|
||||
|
||||
# Convert float tensor [0,1] to uint8 [0,255]
|
||||
if image.dtype == np.float32 or image.dtype == np.float64:
|
||||
image = (image * 255).astype(np.uint8)
|
||||
|
||||
# Convert RGB to BGR for OpenCV
|
||||
if image.shape[2] == 3: # RGB
|
||||
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
||||
|
||||
return image
|
||||
|
||||
def convert_cv2_to_comfy(self, image, original_tensor=None):
|
||||
"""Convert cv2 image (HWC) back to ComfyUI tensor format (BHWC)"""
|
||||
# Convert BGR to RGB
|
||||
if image.shape[2] == 3: # BGR
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# Convert to float [0,1] if original was float
|
||||
if original_tensor is not None and (
|
||||
isinstance(original_tensor, torch.Tensor) and original_tensor.dtype == torch.float32 or
|
||||
isinstance(original_tensor, np.ndarray) and original_tensor.dtype == np.float32
|
||||
):
|
||||
image = image.astype(np.float32) / 255.0
|
||||
|
||||
# Add batch dimension if needed
|
||||
if original_tensor is not None and len(original_tensor.shape) == 4:
|
||||
image = np.expand_dims(image, axis=0)
|
||||
|
||||
# Convert to torch tensor if original was tensor
|
||||
if isinstance(original_tensor, torch.Tensor):
|
||||
image = torch.from_numpy(image).to(original_tensor.device)
|
||||
|
||||
return image
|
||||
|
||||
def release(self):
|
||||
"""Release resources"""
|
||||
self.model = None
|
||||
self.loaded = False
|
||||
|
||||
def __del__(self):
|
||||
self.release()
|
||||
@@ -0,0 +1,110 @@
|
||||
import cv2
|
||||
import sys
|
||||
import torch
|
||||
|
||||
def test_opencv_version():
|
||||
print(f"OpenCV Version: {cv2.__version__}")
|
||||
print(f"Python Version: {sys.version}")
|
||||
print(f"CUDA available (torch): {torch.cuda.is_available()}")
|
||||
if torch.cuda.is_available():
|
||||
print(f"CUDA version (torch): {torch.version.cuda}")
|
||||
|
||||
def test_cuda_availability():
|
||||
print("\nTesting OpenCV CUDA Support:")
|
||||
|
||||
# Check if CUDA module exists
|
||||
print(f"cv2.cuda exists: {hasattr(cv2, 'cuda')}")
|
||||
if hasattr(cv2, 'cuda'):
|
||||
print(f"CUDA devices: {cv2.cuda.getCudaEnabledDeviceCount()}")
|
||||
|
||||
# Check DNN CUDA backend availability
|
||||
print("\nTesting DNN CUDA Backend:")
|
||||
if hasattr(cv2.dnn, "getAvailableBackends"):
|
||||
backends = cv2.dnn.getAvailableBackends()
|
||||
print(f"Available backends: {backends}")
|
||||
print(f"Has CUDA backend: {cv2.dnn.DNN_BACKEND_CUDA in backends}")
|
||||
else:
|
||||
print("cv2.dnn.getAvailableBackends not available")
|
||||
|
||||
if hasattr(cv2.dnn, "getAvailableTargets"):
|
||||
try:
|
||||
targets = cv2.dnn.getAvailableTargets(cv2.dnn.DNN_BACKEND_CUDA)
|
||||
print(f"Available CUDA targets: {targets}")
|
||||
except Exception as e:
|
||||
print(f"Error getting CUDA targets: {e}")
|
||||
else:
|
||||
print("cv2.dnn.getAvailableTargets not available")
|
||||
|
||||
def test_dnn_superres_availability():
|
||||
print("\nTesting DNN SuperRes availability:")
|
||||
print(f"cv2.dnn_superres exists: {hasattr(cv2, 'dnn_superres')}")
|
||||
|
||||
if hasattr(cv2, 'dnn_superres'):
|
||||
print("\nAvailable attributes in cv2.dnn_superres:")
|
||||
for attr in dir(cv2.dnn_superres):
|
||||
if not attr.startswith('__'):
|
||||
print(f"- {attr}")
|
||||
|
||||
def test_create_methods():
|
||||
print("\nTesting creation methods:")
|
||||
|
||||
methods = [
|
||||
("DnnSuperResImpl()", lambda: cv2.dnn_superres.DnnSuperResImpl()),
|
||||
("DnnSuperResImpl_create()", lambda: cv2.dnn_superres.DnnSuperResImpl_create())
|
||||
]
|
||||
|
||||
working_methods = []
|
||||
|
||||
for method_name, method_func in methods:
|
||||
try:
|
||||
model = method_func()
|
||||
print(f"✓ {method_name} - SUCCESS")
|
||||
working_methods.append((method_name, model))
|
||||
except Exception as e:
|
||||
print(f"✗ {method_name} - FAILED: {str(e)}")
|
||||
|
||||
return working_methods
|
||||
|
||||
def test_model_methods(working_methods):
|
||||
if not working_methods:
|
||||
print("\nNo working creation methods found!")
|
||||
return
|
||||
|
||||
print("\nTesting model methods and CUDA setup on working instances:")
|
||||
for method_name, model in working_methods:
|
||||
print(f"\nTesting {method_name}:")
|
||||
try:
|
||||
print("Basic methods:")
|
||||
print("- hasattr readModel:", hasattr(model, 'readModel'))
|
||||
print("- hasattr setModel:", hasattr(model, 'setModel'))
|
||||
print("- hasattr upsample:", hasattr(model, 'upsample'))
|
||||
|
||||
print("\nCUDA-related methods:")
|
||||
print("- hasattr setPreferableBackend:", hasattr(model, 'setPreferableBackend'))
|
||||
if hasattr(model, 'setPreferableBackend'):
|
||||
try:
|
||||
model.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
|
||||
print(" ✓ Successfully set CUDA backend")
|
||||
except Exception as e:
|
||||
print(f" ✗ Failed to set CUDA backend: {e}")
|
||||
|
||||
print("- hasattr setPreferableTarget:", hasattr(model, 'setPreferableTarget'))
|
||||
if hasattr(model, 'setPreferableTarget'):
|
||||
try:
|
||||
model.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
|
||||
print(" ✓ Successfully set CUDA target")
|
||||
except Exception as e:
|
||||
print(f" ✗ Failed to set CUDA target: {e}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error testing methods: {str(e)}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("OpenCV Super Resolution CUDA Test Script")
|
||||
print("=" * 40)
|
||||
|
||||
test_opencv_version()
|
||||
test_cuda_availability()
|
||||
test_dnn_superres_availability()
|
||||
working_methods = test_create_methods()
|
||||
test_model_methods(working_methods)
|
||||
@@ -0,0 +1 @@
|
||||
# Makes nodes directory a proper Python package
|
||||
@@ -0,0 +1,95 @@
|
||||
class SuperResolutionModelLoader:
|
||||
"""Node that loads a super resolution model and outputs a reference to it"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model_type": (["FSRCNN", "FSRCNN-small", "EDSR", "ESPCN", "LapSRN", "VDSR"], {
|
||||
"default": "FSRCNN",
|
||||
"tooltip": "Select super resolution model type:\n"
|
||||
"• FSRCNN: Good balance of quality and speed\n"
|
||||
"• FSRCNN-small: Fastest, lower quality option\n"
|
||||
"• ESPCN: Efficient for text and line art\n"
|
||||
"• LapSRN: Better edge preservation\n"
|
||||
"• EDSR: Highest quality, slower processing\n"
|
||||
"• VDSR: Very Deep SR with sharp edge reconstruction\n"
|
||||
}),
|
||||
"scale_factor": (["2", "3", "4"], {
|
||||
"default": "2",
|
||||
"tooltip": "Upscaling factor:\n"
|
||||
"• 2: Double resolution (e.g., 512x512 → 1024x1024)\n"
|
||||
"• 3: Triple resolution (e.g., 512x512 → 1536x1536)\n"
|
||||
"• 4: Quadruple resolution (e.g., 512x512 → 2048x2048)"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("SR_MODEL",)
|
||||
RETURN_NAMES = ("sr_model",)
|
||||
FUNCTION = "load_model"
|
||||
CATEGORY = "StreamPack/SuperResolution"
|
||||
|
||||
def load_model(self, model_type, scale_factor):
|
||||
"""Load the selected super resolution model and return a reference"""
|
||||
|
||||
# Import unified model class
|
||||
from ..models import SuperResolutionModel
|
||||
|
||||
# Import model download utilities
|
||||
from ..utils.model_utils import download_model
|
||||
|
||||
# On-demand download the model file if it doesn't exist
|
||||
scale_key = f"x{scale_factor}"
|
||||
model_path = None
|
||||
|
||||
# Download the model if needed
|
||||
model_path = download_model(model_type, scale_key)
|
||||
if not model_path:
|
||||
print(f"Warning: Could not download model {model_type} {scale_key}")
|
||||
|
||||
# Create the model using our factory method
|
||||
model = SuperResolutionModel.create(model_type, int(scale_factor))
|
||||
|
||||
# Load the model with the downloaded path if available
|
||||
if model_path:
|
||||
model.load_model(model_path)
|
||||
else:
|
||||
# Fall back to default paths
|
||||
model.load_model()
|
||||
|
||||
return (model,)
|
||||
|
||||
|
||||
class SuperResolutionUpscale:
|
||||
"""Node that applies a super resolution model to an image"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"sr_model": ("SR_MODEL",),
|
||||
"use_cuda": (["True", "False"], {
|
||||
"default": "True",
|
||||
"tooltip": "Whether to use CUDA for acceleration (requires CUDA-enabled OpenCV build)"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("upscaled_image",)
|
||||
FUNCTION = "upscale_image"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "StreamPack/SuperResolution"
|
||||
|
||||
def upscale_image(self, image, sr_model, use_cuda):
|
||||
"""Upscale image using the provided super resolution model"""
|
||||
|
||||
# Convert use_cuda string to boolean
|
||||
use_cuda = use_cuda == "True"
|
||||
|
||||
# Upscale the image with the use_cuda parameter
|
||||
result = sr_model.upscale(image, use_cuda=use_cuda)
|
||||
|
||||
return (result,)
|
||||
@@ -0,0 +1,2 @@
|
||||
# Utils module for ComfyUI SuperResolution
|
||||
from .model_utils import download_model, get_model_path
|
||||
@@ -0,0 +1,149 @@
|
||||
import os
|
||||
import sys
|
||||
import requests
|
||||
import json
|
||||
from pathlib import Path
|
||||
import itertools
|
||||
import shutil
|
||||
import time
|
||||
|
||||
# Fix the import path to use relative import
|
||||
from ..models.sr_model import SuperResolutionModel
|
||||
|
||||
def get_models_root_dir():
|
||||
"""Get the root directory for downloaded SR models"""
|
||||
# Find the ComfyUI_SuperResolution directory
|
||||
base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
# Create downloaded_models directory separate from code
|
||||
return os.path.join(base_dir, "downloaded_models")
|
||||
|
||||
def get_model_path(model_family, scale):
|
||||
"""Get the path to a model file, ensuring directories exist"""
|
||||
# Normalize model family name to lowercase for consistency
|
||||
model_family = model_family.lower()
|
||||
|
||||
# Check if model family exists in our configurations
|
||||
if model_family not in SuperResolutionModel.MODEL_CONFIGS:
|
||||
print(f"Error: Unknown model family '{model_family}'")
|
||||
return None
|
||||
|
||||
# Get configuration for this model
|
||||
config = SuperResolutionModel.MODEL_CONFIGS[model_family]
|
||||
|
||||
# Convert scale to integer for comparison
|
||||
scale_int = int(scale.replace('x', ''))
|
||||
|
||||
# Check if scale is supported for this model
|
||||
if scale_int not in config["supported_scales"]:
|
||||
print(f"Error: Scale {scale} not supported for {model_family}")
|
||||
return None
|
||||
|
||||
# Get model filename from pattern
|
||||
model_filename = config["filename_pattern"].format(scale=scale_int)
|
||||
|
||||
# Create family-specific subdirectory
|
||||
models_dir = os.path.join(get_models_root_dir(), model_family.upper())
|
||||
os.makedirs(models_dir, exist_ok=True)
|
||||
|
||||
# Return the full path to the model file
|
||||
return os.path.join(models_dir, model_filename)
|
||||
|
||||
def download_model_file(url, destination):
|
||||
"""Download a model file with progress indication"""
|
||||
# Skip if model already exists
|
||||
if os.path.exists(destination):
|
||||
print(f"Model already exists: {os.path.basename(destination)}")
|
||||
return True
|
||||
|
||||
print(f"Downloading {os.path.basename(destination)}...")
|
||||
try:
|
||||
# Download with progress indication for large files
|
||||
with requests.get(url, stream=True) as r:
|
||||
r.raise_for_status()
|
||||
total_size = int(r.headers.get('content-length', 0))
|
||||
with open(destination, 'wb') as f:
|
||||
if total_size > 1024*1024: # Show progress for files > 1MB
|
||||
downloaded = 0
|
||||
start_time = time.time()
|
||||
for chunk in r.iter_content(chunk_size=8192):
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
downloaded += len(chunk)
|
||||
# Show progress every 1MB
|
||||
if downloaded % (1024*1024) < 8192:
|
||||
percent = 100 * downloaded / total_size
|
||||
elapsed = time.time() - start_time
|
||||
speed = downloaded / (1024 * 1024 * elapsed) if elapsed > 0 else 0
|
||||
print(f" {percent:.1f}% ({downloaded/(1024*1024):.1f}MB) - {speed:.1f}MB/s")
|
||||
else:
|
||||
# For small files, just download without progress
|
||||
shutil.copyfileobj(r.raw, f)
|
||||
print(f"Downloaded {os.path.basename(destination)} successfully!")
|
||||
return True
|
||||
except Exception as e:
|
||||
print(f"Error downloading {url}: {e}")
|
||||
# Remove partial file if download failed
|
||||
if os.path.exists(destination):
|
||||
os.remove(destination)
|
||||
return False
|
||||
|
||||
def download_model(model_family, scale):
|
||||
"""Download a specific SR model on-demand if it doesn't exist"""
|
||||
# Normalize model family name to lowercase for consistency
|
||||
model_family_norm = model_family.lower()
|
||||
|
||||
# Check if valid model family
|
||||
if model_family_norm not in SuperResolutionModel.MODEL_CONFIGS:
|
||||
print(f"Error: Unknown model family '{model_family}'")
|
||||
return None
|
||||
|
||||
config = SuperResolutionModel.MODEL_CONFIGS[model_family_norm]
|
||||
|
||||
# Build model download URL map if we don't have a static one
|
||||
model_urls = {
|
||||
"fsrcnn": {
|
||||
"base_url": "https://raw.githubusercontent.com/ryanontheinside/FSRCNN_Tensorflow/master/models/",
|
||||
},
|
||||
"fsrcnn-small": {
|
||||
"base_url": "https://raw.githubusercontent.com/ryanontheinside/FSRCNN_Tensorflow/master/models/",
|
||||
},
|
||||
"edsr": {
|
||||
"base_url": "https://raw.githubusercontent.com/ryanontheinside/EDSR_Tensorflow/master/models/",
|
||||
},
|
||||
"espcn": {
|
||||
"base_url": "https://raw.githubusercontent.com/ryanontheinside/TF-ESPCN/master/export/",
|
||||
},
|
||||
"lapsrn": {
|
||||
"base_url": "https://raw.githubusercontent.com/ryanontheinside/TF-LapSRN/master/export/",
|
||||
},
|
||||
"vdsr": {
|
||||
"base_url": "https://raw.githubusercontent.com/ryanontheinside/TF-VDSR/master/export/",
|
||||
}
|
||||
}
|
||||
|
||||
# Check if we have a download URL for this model
|
||||
if model_family_norm not in model_urls:
|
||||
print(f"Error: No download URL for model family '{model_family}'")
|
||||
return None
|
||||
|
||||
# Get the model path
|
||||
model_path = get_model_path(model_family, scale)
|
||||
if not model_path:
|
||||
return None
|
||||
|
||||
# If model already exists, return its path
|
||||
if os.path.exists(model_path):
|
||||
return model_path
|
||||
|
||||
# Extract the scale number from the format
|
||||
scale_int = int(scale.replace('x', ''))
|
||||
|
||||
# Get the base URL and create the download URL
|
||||
base_url = model_urls[model_family_norm]["base_url"]
|
||||
filename = config["filename_pattern"].format(scale=scale_int)
|
||||
url = base_url + filename
|
||||
|
||||
if download_model_file(url, model_path):
|
||||
return model_path
|
||||
else:
|
||||
return None
|
||||
Reference in New Issue
Block a user