Author SHA1 Message Date
John | Elite Encoder 3cf283eb69 chore: update .gitignore to include Python packaging artifacts (#46) 2025-09-25 12:22:38 -04:00
695db5e06c Add pyproject.toml for Custom Node Registry (#22)
* chore(pyproject): Add pyproject.toml for Custom Node Registry
* use dynamic dependencies

---------

Co-authored-by: snomiao <snomiao+comfy-pr@gmail.com>
Co-authored-by: John | Elite Encoder <john@eliteencoder.net>
2025-09-24 20:11:13 -04:00
Robin Huangandsnomiao c78bbccb10 chore(publish): Add Github Action for Publishing to Comfy Registry (#23)
Co-authored-by: snomiao <snomiao+comfy-pr@gmail.com>
2025-09-24 20:10:47 -04:00
John | Elite Encoder 01aacb73e0 resample to 16Khz in AudioTranscriptionNode, fix missing kwargs on execute (#45)
* resample to 16Khz in AudioTranscriptionNode, fix missing kwargs on execute
* update to float32 for consistency
* fail fast on resampling errors
2025-09-19 18:47:41 -04:00
John | Elite Encoder 011da8e1f7 nodes(transcription): allow up to 192000 sample rate, default to 16000khz for whisper (#42) 2025-09-12 18:20:16 -04:00
John | Elite Encoder 7e20a4a6eb feat: add sample rate as input for transcription node (#41)
* add sample rate as input for transcription node
2025-09-10 20:49:36 -04:00
JasonandJohn | Elite Encoder 7ed6c75652 feat: Add faster whisper node as custom node (#40)
* feat: add real-time audio transcription and SRT generation nodes

* feat: add RETURN_NAMES to AudioTranscriptionNode class

* feat: implement RETURN_NAMES as a class method in AudioTranscriptionNode

* refactor: add real-time audio transcription node and SRT generation capabilities

* refactor: clean up imports and update class attributes in transcription nodes

* refactor: streamline audio transcription logic and improve buffer handling

* refactor: reorganize imports and streamline node registration in audio transcription

* refactor: improve code readability and structure in audio transcription node

* refactor: update SRT generation to return subtitles in JSON format

* fix warmup sentienl in transcription node

* adjust buffer logic for comfystream load_audio_tensor_stream node

* fix return_name in srt_generator node (was breaking lightgraph)

---------

Co-authored-by: John | Elite Encoder <john@eliteencoder.net>
2025-09-04 15:30:09 -07:00
c3e196822d chore: add example workflow for super resolution (#26)
* Added Super Resolution to the streampack

* Super Resolution created by ryan, updated for missing <>

Co-authored-by: ryanontheinside <7623207+ryanontheinside@users.noreply.github.com>

* add prestartup script to check for opencv-cuda

* update log msg

* Add FSRCNN 2x super-resolution GPU workflow example

---------

Co-authored-by: ryanontheinside <7623207+ryanontheinside@users.noreply.github.com>
Co-authored-by: John | Elite Encoder <john@eliteencoder.net>
2025-04-23 12:09:03 +03:00
Prakarsh Kaushik fce8d032f8 docs: correct paper title (#15)
This commit fixes an incorrect reference to the StreamV2V paper in the Python doc comments.
2025-04-22 16:39:59 +02:00
499096bdb0 Super resolution (#18)
* Added Super Resolution to the streampack
* Super Resolution created by ryan, updated for missing 
* add prestartup script to check for opencv-cuda

---------

Co-authored-by: ryanontheinside <7623207+ryanontheinside@users.noreply.github.com>
Co-authored-by: John | Elite Encoder <john@eliteencoder.net>
2025-04-04 17:21:47 -04:00
John | Elite EncoderandRick Staa b1ebd4aa8f Revert "Super resolution (#16)" (#17)
This reverts commit 59035c455a.

Co-authored-by: Rick Staa <rick.staa@outlook.com>
2025-04-02 17:49:43 -04:00
Jason 59035c455a Super resolution (#16)
* Added Super Resolution to the streampack
2025-04-02 14:39:41 -04:00
Rick Staaandryanontheinside da5a043501 feat: add MediaPipe FaceMesh node and utilities (#13)
This commit introduces a ControlNet Auxiliary Preprocessor for facial landmark
detection using MediaPipe FaceMesh. It also adds utility nodes to work with the
detected landmarks:
- Node to draw the facial mesh on an image using landmark data.
- Node to generate masks based on landmark data.

---------

Co-authored-by: ryanontheinside <ryanfosdick87@gmail.com>
2025-03-28 20:05:49 +01:00
Rick StaaandRUFFY-369 eedf5d1d1b feat: expose feature bank to comfyui (#12)
This commit establishes the required file structure and logic to integrate the Feature Bank node into ComfyUI, enabling its use.

---------

Co-authored-by: RUFFY-369 <prakarshkaushik369@gmail.com>
2025-03-28 18:01:53 +01:00
Rick StaaandRUFFY-369 845892959a feat(node): migrate Feature Bank node to new streampack (#3)
This commit migrates the Feature Bank node, originally implemented by @RUFFY-369
on [RUFFY-369/ComfyUI-FeatureBank](https://github.com/RUFFY-369/ComfyUI-FeatureBank),
to the new streampack.

The Feature Bank node is a ComfyUI-native implementation of the method proposed in
[Looking Backward: Streaming Video-to-Video Translation with Feature Banks](https://arxiv.org/abs/2405.15757)
by Liang et al. It enhances the stream diffusion model by providing information from
previous frames, thereby improving temporal consistency and reducing flickering.

---------

Co-authored-by: RUFFY-369 <prakarshkaushik369@gmail.com>
2025-03-17 23:01:32 +01:00
36 changed files with 2680 additions and 937 deletions
+28
View File
@@ -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
View File
@@ -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
View File
@@ -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"]
+1 -1
View File
@@ -215,4 +215,4 @@
"title": "Load CLIP"
}
}
}
}
@@ -231,4 +231,4 @@
"title": "Feature Bank Attention Processor"
}
}
}
}
@@ -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"
}
}
}
+23
View File
@@ -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"]
-15
View File
@@ -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"]
+21
View File
@@ -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"]
-13
View File
@@ -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"
}
-1
View File
@@ -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"
}
+16
View File
@@ -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"]
+22
View File
@@ -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()
+18
View File
@@ -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"]}
+5
View File
@@ -1,3 +1,8 @@
mediapipe
torch
diffusers
numpy
requests
opencv-contrib-python>=4.8.0
faster_whisper
scipy
+1 -1
View File
@@ -1 +1 @@
# Dummy file ensuring this package will be recognized.
# Dummy file ensuring this sub-package will be recognized.
+1
View File
@@ -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)
View File
-497
View File
@@ -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
View File
@@ -1 +1 @@
# Dummy file ensuring this package will be recognized.
# Dummy file ensuring this sub-package will be recognized.
+6 -7
View File
@@ -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,
+1
View File
@@ -0,0 +1 @@
# Dummy file ensuring this sub-package will be recognized.
File diff suppressed because it is too large Load Diff
+88
View File
@@ -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
+1
View File
@@ -0,0 +1 @@
# Dummy file ensuring this sub-package will be recognized.
+5
View File
@@ -0,0 +1,5 @@
# Makes models directory a proper Python package
from .sr_model import SuperResolutionModel
__all__ = ["SuperResolutionModel"]
+251
View File
@@ -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)
+1
View File
@@ -0,0 +1 @@
# Makes nodes directory a proper Python package
+95
View File
@@ -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,)
+2
View File
@@ -0,0 +1,2 @@
# Utils module for ComfyUI SuperResolution
from .model_utils import download_model, get_model_path
+149
View File
@@ -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