From 6b73e1d50f16357c0f669910120b990097ac7b9e Mon Sep 17 00:00:00 2001 From: VjumpKunG Date: Wed, 24 Dec 2025 21:25:19 +0700 Subject: [PATCH 1/8] test quantization --- README.md | 2 +- __init__.py | 2 + install.py | 4 +- nodes.py | 25 +++++-- pyproject.toml | 2 +- requirements.txt | 3 +- video_to_video/modules/unet_v2v.py | 7 +- video_to_video/video_to_video_model.py | 93 +++++++++++++++++++++++++- 8 files changed, 121 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index 12805c2..bf10249 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# STAR Video Super Resolution - ComfyUI Node +# ComfyUI Wrapper nodes for STAR Video Super-Resolution A ComfyUI custom node wrapper for [STAR (Spatial-Temporal Augmentation with Text-to-Video Models for Real-World Video Super-Resolution)](https://github.com/NJU-PCALab/STAR), enabling high-quality video upscaling with AI-powered super-resolution. diff --git a/__init__.py b/__init__.py index 7a6e35b..4718a2b 100644 --- a/__init__.py +++ b/__init__.py @@ -22,3 +22,5 @@ except ImportError: from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] + +print("🎉 Loading STAR Video Super-Resolution Completed ") diff --git a/install.py b/install.py index 75b1f84..327832a 100644 --- a/install.py +++ b/install.py @@ -4,9 +4,9 @@ Installation script for STARVSRWrapper. Checks for uv/pip, matches xformers with PyTorch version, and installs requirements. """ -import sys -import subprocess import re +import subprocess +import sys from pathlib import Path diff --git a/nodes.py b/nodes.py index 651ee24..0e33519 100644 --- a/nodes.py +++ b/nodes.py @@ -21,6 +21,10 @@ class STARVSRNode: "required": { "images": ("IMAGE",), "model": (["Light Degradation", "Heavy Degradation"],), + "precision": ( + ["fp16", "fp8", "nf4"], + {"default": "fp16"}, + ), "prompt": ("STRING", {"default": "a good video", "multiline": True}), "upscale": ("INT", {"default": 4, "min": 2, "max": 4, "step": 1}), "max_chunk_len": ( @@ -51,10 +55,15 @@ class STARVSRNode: def __init__(self): self.model = None self.current_model_type = None + self.current_precision = None - def load_model(self, model_type): - """Load the appropriate model based on degradation type""" - if self.model is not None and self.current_model_type == model_type: + def load_model(self, model_type, precision="fp16"): + """Load the appropriate model based on degradation type and precision""" + if ( + self.model is not None + and self.current_model_type == model_type + and self.current_precision == precision + ): return self.model # Map model type to Hugging Face repo files @@ -98,8 +107,12 @@ class STARVSRNode: model_cfg = EasyDict(__name__="model_cfg") model_cfg.model_path = model_path - self.model = VideoToVideo_sr(model_cfg) + + # Pass precision parameter to VideoToVideo_sr + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + self.model = VideoToVideo_sr(model_cfg, device=device, precision=precision) self.current_model_type = model_type + self.current_precision = precision return self.model @@ -107,6 +120,7 @@ class STARVSRNode: self, images, model, + precision, prompt, upscale, max_chunk_len, @@ -122,6 +136,7 @@ class STARVSRNode: Args: images: Tensor of shape [B, H, W, C] in ComfyUI format (0-1 range) model: Model type ("Light Degradation" or "Heavy Degradation") + precision: Model precision ("fp16", "fp8", or "nf4") prompt: Text prompt for enhancement upscale: Upscale factor (2, 3, or 4) max_chunk_len: Maximum chunk length for processing @@ -134,7 +149,7 @@ class STARVSRNode: Enhanced images tensor in ComfyUI format [B, H, W, C] """ # Load the appropriate model - model_instance = self.load_model(model) + model_instance = self.load_model(model, precision) # Convert ComfyUI format [B, H, W, C] to our format # ComfyUI images are in range [0, 1], convert to [0, 255] for preprocessing diff --git a/pyproject.toml b/pyproject.toml index 4efc7d4..e769311 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ classifiers = [ "Environment :: GPU :: NVIDIA CUDA", # NVIDIA CUDA support ] -dependencies = ["accelerate", "av", "opencv-python", "easydict", "einops", "open-clip-torch", "fairscale", "torchsde", "pytorch-lightning", "diffusers[torch]", "huggingface_hub", "numpy<2", "ruff", "transformers"] +dependencies = ["accelerate", "av", "opencv-python", "easydict", "einops", "open-clip-torch", "fairscale", "torchsde", "pytorch-lightning", "diffusers[torch]", "huggingface_hub", "numpy<2", "ruff", "transformers", "bitsandbytes"] [project.urls] Repository = "https://github.com/vjumpkung/ComfyUI-STARWrapper" diff --git a/requirements.txt b/requirements.txt index ffcfb02..91acb3b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,4 +11,5 @@ diffusers[torch] huggingface_hub numpy<2 ruff -transformers \ No newline at end of file +transformers +bitsandbytes \ No newline at end of file diff --git a/video_to_video/modules/unet_v2v.py b/video_to_video/modules/unet_v2v.py index 9058891..0061484 100644 --- a/video_to_video/modules/unet_v2v.py +++ b/video_to_video/modules/unet_v2v.py @@ -7,7 +7,6 @@ from abc import abstractmethod import torch import torch.nn as nn import torch.nn.functional as F - import xformers import xformers.ops from einops import rearrange, repeat @@ -1708,9 +1707,9 @@ class Vid2VidSDUNet(nn.Module): elif isinstance(module, CrossAttention): module = checkpoint_wrapper(module) if self.use_checkpoint else module x = module(x, context) - # elif isinstance(module, MemoryEfficientCrossAttention): - # module = checkpoint_wrapper(module) if self.use_checkpoint else module - # x = module(x, context) + elif isinstance(module, MemoryEfficientCrossAttention): + module = checkpoint_wrapper(module) if self.use_checkpoint else module + x = module(x, context) elif isinstance(module, BasicTransformerBlock): module = checkpoint_wrapper(module) if self.use_checkpoint else module x = module(x, context) diff --git a/video_to_video/video_to_video_model.py b/video_to_video/video_to_video_model.py index 5e34cf5..f33e046 100644 --- a/video_to_video/video_to_video_model.py +++ b/video_to_video/video_to_video_model.py @@ -1,7 +1,7 @@ import os import os.path as osp import random -from typing import Any, Dict +from typing import Any, Dict, Literal import torch import torch.nn.functional as F @@ -15,11 +15,28 @@ from ..video_to_video.utils.logger import get_logger logger = get_logger() +# Try to import bitsandbytes for quantization +try: + import bitsandbytes as bnb + + BITSANDBYTES_AVAILABLE = True +except ImportError: + BITSANDBYTES_AVAILABLE = False + logger.warning( + "bitsandbytes not available. NF4 quantization will not be supported." + ) + class VideoToVideo_sr: - def __init__(self, opt, device=torch.device(f"cuda:0")): + def __init__( + self, + opt, + device=torch.device(f"cuda:0"), + precision: Literal["fp16", "fp8", "nf4"] = "fp16", + ): self.opt = opt self.device = device # torch.device(f'cuda:0') + self.precision = precision # text_encoder text_encoder = FrozenOpenCLIPEmbedder( @@ -40,7 +57,21 @@ class VideoToVideo_sr: load_dict = load_dict["state_dict"] ret = generator.load_state_dict(load_dict, strict=False) - self.generator = generator.half() + # Apply quantization based on precision parameter + if self.precision == "nf4": + if not BITSANDBYTES_AVAILABLE: + logger.warning("bitsandbytes not available. Falling back to FP16.") + self.generator = generator.half() + else: + logger.info("Quantizing model to NF4 (4-bit)") + self.generator = self._quantize_nf4(generator) + elif self.precision == "fp8": + logger.info("Quantizing model to FP8") + self.generator = self._quantize_fp8(generator) + else: + logger.info("Using FP16 precision") + self.generator = generator.half() + logger.info( "Load model path {}, with local status {}".format(cfg.model_path, ret) ) @@ -179,6 +210,62 @@ class VideoToVideo_sr: z = rearrange(z, "(b f) c h w -> b c f h w", f=num_f) return z * self.vae.config.scaling_factor + def _quantize_nf4(self, model): + """ + Recursively quantize Linear layers to NF4 (4-bit NormalFloat) + Requires bitsandbytes library + """ + if not BITSANDBYTES_AVAILABLE: + logger.error("bitsandbytes is not available. Cannot quantize to NF4.") + return model.half() + + for name, module in model.named_children(): + if isinstance(module, torch.nn.Linear): + # Replace with 4-bit quantized version + quant_layer = bnb.nn.Linear4bit( + module.in_features, + module.out_features, + bias=module.bias is not None, + compute_dtype=torch.bfloat16, + compress_statistics=True, + quant_type="nf4", + ) + + # Copy and quantize weights + quant_layer.weight = bnb.nn.Params4bit( + module.weight.data.to("cpu"), requires_grad=False, quant_type="nf4" + ) + if module.bias is not None: + quant_layer.bias = module.bias + + setattr(model, name, quant_layer) + elif len(list(module.children())) > 0: + # Recursively quantize child modules + self._quantize_nf4(module) + + return model + + def _quantize_fp8(self, model): + """ + Quantize model to FP8 (requires PyTorch 2.1+ and compatible GPU) + Falls back to FP16 if FP8 is not supported + """ + if hasattr(torch, "float8_e4m3fn"): + try: + logger.info("Converting model to FP8 (E4M3)") + for param in model.parameters(): + if param.dtype in [torch.float32, torch.float16]: + param.data = param.data.to(torch.float8_e4m3fn) + return model + except Exception as e: + logger.warning(f"FP8 conversion failed: {e}. Falling back to FP16.") + return model.half() + else: + logger.warning( + "FP8 not supported on this PyTorch version. Falling back to FP16." + ) + return model.half() + def pad_to_fit(h, w): BEST_H, BEST_W = 720, 1280 From 1ef42ec3b30c027c56c031bce8cbedfab361b606 Mon Sep 17 00:00:00 2001 From: VjumpKunG Date: Wed, 24 Dec 2025 21:50:07 +0700 Subject: [PATCH 2/8] fix nf4 err --- video_to_video/video_to_video_model.py | 80 +++++++++++++++++++------- 1 file changed, 60 insertions(+), 20 deletions(-) diff --git a/video_to_video/video_to_video_model.py b/video_to_video/video_to_video_model.py index f33e046..965a661 100644 --- a/video_to_video/video_to_video_model.py +++ b/video_to_video/video_to_video_model.py @@ -214,34 +214,74 @@ class VideoToVideo_sr: """ Recursively quantize Linear layers to NF4 (4-bit NormalFloat) Requires bitsandbytes library + + Note: Only quantizes larger Linear layers (>256 input features) to avoid + compatibility issues with small layers and time embeddings. """ if not BITSANDBYTES_AVAILABLE: logger.error("bitsandbytes is not available. Cannot quantize to NF4.") return model.half() - for name, module in model.named_children(): + quantized_count = 0 + skipped_count = 0 + + def quantize_layer(parent_module, name, module): + nonlocal quantized_count, skipped_count + if isinstance(module, torch.nn.Linear): - # Replace with 4-bit quantized version - quant_layer = bnb.nn.Linear4bit( - module.in_features, - module.out_features, - bias=module.bias is not None, - compute_dtype=torch.bfloat16, - compress_statistics=True, - quant_type="nf4", - ) + # Skip small layers (like time embeddings) that may cause issues + # Only quantize larger layers where memory savings matter + if module.in_features < 256 or module.out_features < 256: + skipped_count += 1 + return - # Copy and quantize weights - quant_layer.weight = bnb.nn.Params4bit( - module.weight.data.to("cpu"), requires_grad=False, quant_type="nf4" - ) - if module.bias is not None: - quant_layer.bias = module.bias + try: + # Create 4-bit quantized layer + quant_layer = bnb.nn.Linear4bit( + module.in_features, + module.out_features, + bias=module.bias is not None, + compute_dtype=torch.bfloat16, + compress_statistics=True, + quant_type="nf4", + ) - setattr(model, name, quant_layer) - elif len(list(module.children())) > 0: - # Recursively quantize child modules - self._quantize_nf4(module) + # Move original weights to CPU and quantize + with torch.no_grad(): + weight_data = module.weight.data.to("cpu") + quant_layer.weight = bnb.nn.Params4bit( + weight_data, requires_grad=False, quant_type="nf4" + ) + + if module.bias is not None: + quant_layer.bias = torch.nn.Parameter( + module.bias.data.clone() + ) + + # Replace the layer + setattr(parent_module, name, quant_layer) + quantized_count += 1 + + except Exception as e: + logger.warning( + f"Failed to quantize layer {name}: {e}. Keeping as FP16." + ) + skipped_count += 1 + + # Recursively process child modules + for child_name, child_module in list(module.named_children()): + quantize_layer(module, child_name, child_module) + + # Start quantization from the root + for name, module in list(model.named_children()): + quantize_layer(model, name, module) + + logger.info( + f"NF4 Quantization: {quantized_count} layers quantized, {skipped_count} layers kept in FP16" + ) + + # Convert non-quantized layers to FP16 + model = model.half() return model From 952f267d47a63767a2800df809660bfa4dd871c4 Mon Sep 17 00:00:00 2001 From: VjumpKunG Date: Wed, 24 Dec 2025 21:52:42 +0700 Subject: [PATCH 3/8] fix nf4 move to device --- video_to_video/video_to_video_model.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/video_to_video/video_to_video_model.py b/video_to_video/video_to_video_model.py index 965a661..51c4de5 100644 --- a/video_to_video/video_to_video_model.py +++ b/video_to_video/video_to_video_model.py @@ -258,6 +258,10 @@ class VideoToVideo_sr: module.bias.data.clone() ) + # Move quantized layer to device BEFORE replacing + # This initializes the quantization state properly + quant_layer = quant_layer.to(self.device) + # Replace the layer setattr(parent_module, name, quant_layer) quantized_count += 1 From 4704fa263aa8083b9b65bef0ab0e720354cb10b2 Mon Sep 17 00:00:00 2001 From: VjumpKunG Date: Wed, 24 Dec 2025 21:57:37 +0700 Subject: [PATCH 4/8] change local path --- nodes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nodes.py b/nodes.py index 0e33519..03d53ad 100644 --- a/nodes.py +++ b/nodes.py @@ -71,12 +71,12 @@ class STARVSRNode: "Light Degradation": { "repo_id": "SherryX/STAR", "filename": "I2VGen-XL-based/light_deg.pt", - "local_path": "./models/STAR/light_degradation.pt", + "local_path": "./models/STAR/light_deg.pt", }, "Heavy Degradation": { "repo_id": "SherryX/STAR", "filename": "I2VGen-XL-based/heavy_deg.pt", - "local_path": "./models/STAR/heavy_degradation.pt", + "local_path": "./models/STAR/heavy_deg.pt", }, } From 9f67ff578f83b0938133f1681c1e8d4e00dd3b76 Mon Sep 17 00:00:00 2001 From: VjumpKunG Date: Wed, 24 Dec 2025 22:04:12 +0700 Subject: [PATCH 5/8] remove nf4 due to OOM --- __init__.py | 8 - nodes.py | 8 +- pyproject.toml | 4 +- requirements.txt | 3 +- video_to_video/video_to_video_model.py | 99 +-------- workflows/STAR-example-workflow.json | 291 +++++++++++++------------ 6 files changed, 154 insertions(+), 259 deletions(-) diff --git a/__init__.py b/__init__.py index 4718a2b..1383880 100644 --- a/__init__.py +++ b/__init__.py @@ -11,14 +11,6 @@ print(r""" _____ _______ _____ |_____/ |_/_/ \_\_| \_\ """) -try: - import xformers -except ImportError: - raise ImportError( - "Please install xformers manually version before use this custom nodes see instruction at https://github.com/vjumpkung/ComfyUI-STARWrapper" - ) - - from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py index 03d53ad..c30fe02 100644 --- a/nodes.py +++ b/nodes.py @@ -22,7 +22,7 @@ class STARVSRNode: "images": ("IMAGE",), "model": (["Light Degradation", "Heavy Degradation"],), "precision": ( - ["fp16", "fp8", "nf4"], + ["fp16", "fp8"], {"default": "fp16"}, ), "prompt": ("STRING", {"default": "a good video", "multiline": True}), @@ -71,12 +71,12 @@ class STARVSRNode: "Light Degradation": { "repo_id": "SherryX/STAR", "filename": "I2VGen-XL-based/light_deg.pt", - "local_path": "./models/STAR/light_deg.pt", + "local_path": "./models/STAR/I2VGen-XL-based/light_deg.pt", }, "Heavy Degradation": { "repo_id": "SherryX/STAR", "filename": "I2VGen-XL-based/heavy_deg.pt", - "local_path": "./models/STAR/heavy_deg.pt", + "local_path": "./models/STAR/I2VGen-XL-based/heavy_deg.pt", }, } @@ -136,7 +136,7 @@ class STARVSRNode: Args: images: Tensor of shape [B, H, W, C] in ComfyUI format (0-1 range) model: Model type ("Light Degradation" or "Heavy Degradation") - precision: Model precision ("fp16", "fp8", or "nf4") + precision: Model precision ("fp16", "fp8") prompt: Text prompt for enhancement upscale: Upscale factor (2, 3, or 4) max_chunk_len: Maximum chunk length for processing diff --git a/pyproject.toml b/pyproject.toml index e769311..dddeca2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-starwrapper" description = "A ComfyUI custom node wrapper for [STAR (Spatial-Temporal Augmentation with Text-to-Video Models for Real-World Video Super-Resolution)](https://github.com/NJU-PCALab/STAR), enabling high-quality video upscaling with AI-powered super-resolution." -version = "0.1.4" +version = "0.2.0" license = {file = "LICENSE"} classifiers = [ # For OS-independent nodes (works on all operating systems) @@ -10,7 +10,7 @@ classifiers = [ "Environment :: GPU :: NVIDIA CUDA", # NVIDIA CUDA support ] -dependencies = ["accelerate", "av", "opencv-python", "easydict", "einops", "open-clip-torch", "fairscale", "torchsde", "pytorch-lightning", "diffusers[torch]", "huggingface_hub", "numpy<2", "ruff", "transformers", "bitsandbytes"] +dependencies = ["accelerate", "av", "opencv-python", "easydict", "einops", "open-clip-torch", "fairscale", "torchsde", "pytorch-lightning", "diffusers[torch]", "huggingface_hub", "numpy<2", "ruff", "transformers"] [project.urls] Repository = "https://github.com/vjumpkung/ComfyUI-STARWrapper" diff --git a/requirements.txt b/requirements.txt index 91acb3b..ffcfb02 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,5 +11,4 @@ diffusers[torch] huggingface_hub numpy<2 ruff -transformers -bitsandbytes \ No newline at end of file +transformers \ No newline at end of file diff --git a/video_to_video/video_to_video_model.py b/video_to_video/video_to_video_model.py index 51c4de5..fc6dfe6 100644 --- a/video_to_video/video_to_video_model.py +++ b/video_to_video/video_to_video_model.py @@ -15,17 +15,6 @@ from ..video_to_video.utils.logger import get_logger logger = get_logger() -# Try to import bitsandbytes for quantization -try: - import bitsandbytes as bnb - - BITSANDBYTES_AVAILABLE = True -except ImportError: - BITSANDBYTES_AVAILABLE = False - logger.warning( - "bitsandbytes not available. NF4 quantization will not be supported." - ) - class VideoToVideo_sr: def __init__( @@ -58,14 +47,7 @@ class VideoToVideo_sr: ret = generator.load_state_dict(load_dict, strict=False) # Apply quantization based on precision parameter - if self.precision == "nf4": - if not BITSANDBYTES_AVAILABLE: - logger.warning("bitsandbytes not available. Falling back to FP16.") - self.generator = generator.half() - else: - logger.info("Quantizing model to NF4 (4-bit)") - self.generator = self._quantize_nf4(generator) - elif self.precision == "fp8": + if self.precision == "fp8": logger.info("Quantizing model to FP8") self.generator = self._quantize_fp8(generator) else: @@ -210,85 +192,6 @@ class VideoToVideo_sr: z = rearrange(z, "(b f) c h w -> b c f h w", f=num_f) return z * self.vae.config.scaling_factor - def _quantize_nf4(self, model): - """ - Recursively quantize Linear layers to NF4 (4-bit NormalFloat) - Requires bitsandbytes library - - Note: Only quantizes larger Linear layers (>256 input features) to avoid - compatibility issues with small layers and time embeddings. - """ - if not BITSANDBYTES_AVAILABLE: - logger.error("bitsandbytes is not available. Cannot quantize to NF4.") - return model.half() - - quantized_count = 0 - skipped_count = 0 - - def quantize_layer(parent_module, name, module): - nonlocal quantized_count, skipped_count - - if isinstance(module, torch.nn.Linear): - # Skip small layers (like time embeddings) that may cause issues - # Only quantize larger layers where memory savings matter - if module.in_features < 256 or module.out_features < 256: - skipped_count += 1 - return - - try: - # Create 4-bit quantized layer - quant_layer = bnb.nn.Linear4bit( - module.in_features, - module.out_features, - bias=module.bias is not None, - compute_dtype=torch.bfloat16, - compress_statistics=True, - quant_type="nf4", - ) - - # Move original weights to CPU and quantize - with torch.no_grad(): - weight_data = module.weight.data.to("cpu") - quant_layer.weight = bnb.nn.Params4bit( - weight_data, requires_grad=False, quant_type="nf4" - ) - - if module.bias is not None: - quant_layer.bias = torch.nn.Parameter( - module.bias.data.clone() - ) - - # Move quantized layer to device BEFORE replacing - # This initializes the quantization state properly - quant_layer = quant_layer.to(self.device) - - # Replace the layer - setattr(parent_module, name, quant_layer) - quantized_count += 1 - - except Exception as e: - logger.warning( - f"Failed to quantize layer {name}: {e}. Keeping as FP16." - ) - skipped_count += 1 - - # Recursively process child modules - for child_name, child_module in list(module.named_children()): - quantize_layer(module, child_name, child_module) - - # Start quantization from the root - for name, module in list(model.named_children()): - quantize_layer(model, name, module) - - logger.info( - f"NF4 Quantization: {quantized_count} layers quantized, {skipped_count} layers kept in FP16" - ) - - # Convert non-quantized layers to FP16 - model = model.half() - - return model - def _quantize_fp8(self, model): """ Quantize model to FP8 (requires PyTorch 2.1+ and compatible GPU) diff --git a/workflows/STAR-example-workflow.json b/workflows/STAR-example-workflow.json index 1958ee5..beb8610 100644 --- a/workflows/STAR-example-workflow.json +++ b/workflows/STAR-example-workflow.json @@ -1,150 +1,19 @@ { "id": "f0bce436-ca7c-453d-8f96-d287c25cf130", "revision": 0, - "last_node_id": 4, - "last_link_id": 2, + "last_node_id": 6, + "last_link_id": 5, "nodes": [ - { - "id": 4, - "type": "STARVSRNode", - "pos": [ - 756.214876033058, - 290.9504132231405 - ], - "size": [ - 413.2231404958677, - 485.81818181818176 - ], - "flags": {}, - "order": 1, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 1 - } - ], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 2 - ] - } - ], - "properties": { - "aux_id": "vjumpkung/ComfyUI-STARWrapper", - "ver": "584e291d82f140dc12916a91e33a37f83acd2399", - "Node name for S&R": "STARVSRNode" - }, - "widgets_values": [ - "Light Degradation", - "A girl walking and wearing sunglasses.", - 2, - 32, - 7.5, - "dpmpp_2m_sde", - "fast", - 15, - 42, - "fixed" - ] - }, - { - "id": 3, - "type": "VHS_LoadVideo", - "pos": [ - 432.74380165289244, - 291.5206611570246 - ], - "size": [ - 255.8291015625, - 457.1858620289278 - ], - "flags": {}, - "order": 0, - "mode": 0, - "inputs": [ - { - "name": "meta_batch", - "shape": 7, - "type": "VHS_BatchManager", - "link": null - }, - { - "name": "vae", - "shape": 7, - "type": "VAE", - "link": null - } - ], - "outputs": [ - { - "name": "IMAGE", - "type": "IMAGE", - "links": [ - 1 - ] - }, - { - "name": "frame_count", - "type": "INT", - "links": null - }, - { - "name": "audio", - "type": "AUDIO", - "links": null - }, - { - "name": "video_info", - "type": "VHS_VIDEOINFO", - "links": null - } - ], - "properties": { - "cnr_id": "comfyui-videohelpersuite", - "ver": "1.7.9", - "Node name for S&R": "VHS_LoadVideo" - }, - "widgets_values": { - "video": "video_00001.mp4", - "force_rate": 0, - "custom_width": 0, - "custom_height": 0, - "frame_load_cap": 0, - "skip_first_frames": 0, - "select_every_nth": 1, - "format": "AnimateDiff", - "videopreview": { - "hidden": false, - "paused": false, - "params": { - "force_rate": 0, - "custom_width": 0, - "custom_height": 0, - "frame_load_cap": 0, - "skip_first_frames": 0, - "select_every_nth": 1, - "filename": "video_00001.mp4", - "type": "input", - "format": "video/mp4" - } - } - } - }, { "id": 2, "type": "VHS_VideoCombine", "pos": [ - 1217.1239669421477, - 290.23140495867796 + 1154.958677685951, + 306.6942148760332 ], "size": [ 334.1807744573441, - 525.8381742211382 + 334 ], "flags": {}, "order": 2, @@ -153,7 +22,7 @@ { "name": "images", "type": "IMAGE", - "link": 2 + "link": 4 }, { "name": "audio", @@ -194,7 +63,7 @@ "pix_fmt": "yuv420p", "crf": 19, "save_metadata": true, - "trim_to_audio": false, + "trim_to_audio": true, "pingpong": false, "save_output": true, "videopreview": { @@ -211,20 +80,152 @@ } } } + }, + { + "id": 3, + "type": "VHS_LoadVideo", + "pos": [ + 400, + 310 + ], + "size": [ + 255.8291015625, + 740.5813428874735 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 3 + ] + }, + { + "name": "frame_count", + "type": "INT", + "links": null + }, + { + "name": "audio", + "type": "AUDIO", + "links": null + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": null + } + ], + "properties": { + "cnr_id": "comfyui-videohelpersuite", + "ver": "1.7.9", + "Node name for S&R": "VHS_LoadVideo" + }, + "widgets_values": { + "video": "video_00010.mp4", + "force_rate": 0, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "format": "AnimateDiff", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "force_rate": 0, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "filename": "video_00010.mp4", + "type": "input", + "format": "video/mp4" + } + } + } + }, + { + "id": 5, + "type": "STARVSRNode", + "pos": [ + 702.5652892561982, + 306.85702479338863 + ], + "size": [ + 403.3057851239671, + 447.8347107438017 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 3 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 4 + ] + } + ], + "properties": { + "cnr_id": "comfyui-starwrapper", + "ver": "6b73e1d50f16357c0f669910120b990097ac7b9e", + "Node name for S&R": "STARVSRNode" + }, + "widgets_values": [ + "Light Degradation", + "fp8", + "An alpaca turning it's head", + 2, + 32, + 5.5, + "dpmpp_2m_sde", + "fast", + 15, + 42, + "fixed" + ] } ], "links": [ [ - 1, + 3, 3, 0, - 4, + 5, 0, "IMAGE" ], [ - 2, 4, + 5, 0, 2, 0, @@ -235,14 +236,14 @@ "config": {}, "extra": { "ds": { - "scale": 1.1000000000000003, + "scale": 0.9090909090909091, "offset": [ - -149.28766272203228, - -141.35987067640917 + 93.97723448237635, + -134.0369738564675 ] }, + "frontendVersion": "1.35.9", "workflowRendererVersion": "LG", - "frontendVersion": "1.34.9", "VHS_latentpreview": false, "VHS_latentpreviewrate": 0, "VHS_MetadataImage": true, From f9e14a6506028423b32ce0a1f277d94c9525e7da Mon Sep 17 00:00:00 2001 From: VjumpKunG Date: Wed, 24 Dec 2025 22:11:00 +0700 Subject: [PATCH 6/8] change args to resolution --- nodes.py | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/nodes.py b/nodes.py index c30fe02..1286e2c 100644 --- a/nodes.py +++ b/nodes.py @@ -26,7 +26,10 @@ class STARVSRNode: {"default": "fp16"}, ), "prompt": ("STRING", {"default": "a good video", "multiline": True}), - "upscale": ("INT", {"default": 4, "min": 2, "max": 4, "step": 1}), + "resolution": ( + "INT", + {"default": 720, "min": 16, "max": 16384, "step": 2}, + ), "max_chunk_len": ( "INT", {"default": 32, "min": 1, "max": 128, "step": 1}, @@ -122,7 +125,7 @@ class STARVSRNode: model, precision, prompt, - upscale, + resolution, max_chunk_len, cfg, sampler, @@ -138,9 +141,10 @@ class STARVSRNode: model: Model type ("Light Degradation" or "Heavy Degradation") precision: Model precision ("fp16", "fp8") prompt: Text prompt for enhancement - upscale: Upscale factor (2, 3, or 4) + resolution: Target resolution for the shorter side (maintains aspect ratio) max_chunk_len: Maximum chunk length for processing cfg: Guidance scale + sampler: Sampling method ("heun" or "dpmpp_2m_sde") solver_mode: Solver mode ("fast" or "normal") steps: Number of denoising steps seed: Random seed for reproducibility @@ -165,7 +169,21 @@ class STARVSRNode: _, _, h, w = video_data.shape logger.info(f"Input resolution: {(h, w)}") - target_h, target_w = h * upscale, w * upscale + # Calculate target resolution maintaining aspect ratio + # Use resolution parameter as the target for the shorter side + aspect_ratio = w / h + + if h < w: # Height is shorter + target_h = resolution + target_w = int(resolution * aspect_ratio) + else: # Width is shorter or equal + target_w = resolution + target_h = int(resolution / aspect_ratio) + + # Ensure dimensions are even (required by model) + target_h = target_h + (target_h % 2) + target_w = target_w + (target_w % 2) + logger.info(f"Target resolution: {(target_h, target_w)}") # Prepare caption From 4e963520e55ed93ba0d71f0a1b74c81978e8d4f9 Mon Sep 17 00:00:00 2001 From: VjumpKunG Date: Wed, 24 Dec 2025 22:24:51 +0700 Subject: [PATCH 7/8] update workflow --- README.md | 4 +- workflows/STAR-example-workflow.json | 187 +++++++++++++++------------ 2 files changed, 107 insertions(+), 84 deletions(-) diff --git a/README.md b/README.md index bf10249..230d43b 100644 --- a/README.md +++ b/README.md @@ -49,7 +49,7 @@ pip install -r requirements.txt 3. Configure the parameters: - **Model**: Choose "Light Degradation" or "Heavy Degradation" - **Prompt**: Describe the desired output (e.g., "a high quality video") - - **Upscale**: Select upscaling factor (2-4x) + - **Resolution**: Target Resolution (720p, 1080p, 2160p, ...) - **Steps**: Number of denoising steps (15-50 recommended) - **CFG**: Guidance scale (7.5 default) - **Sampler**: Choose sampling method @@ -62,7 +62,7 @@ pip install -r requirements.txt | ------------- | ----------------------------- | ----------------- | ----------------- | | model | Degradation type | Light Degradation | Light/Heavy | | prompt | Text guidance for enhancement | "a good video" | - | -| upscale | Upscaling factor | 4 | 2-4 | +| resolution | Target Resolution | 720 | 16-16384 | | max_chunk_len | Maximum frames per chunk | 32 | 1-128 | | cfg | Guidance scale | 7.5 | 0.0-20.0 | | sampler | Sampling method | dpmpp_2m_sde | heun/dpmpp_2m_sde | diff --git a/workflows/STAR-example-workflow.json b/workflows/STAR-example-workflow.json index beb8610..cbaf845 100644 --- a/workflows/STAR-example-workflow.json +++ b/workflows/STAR-example-workflow.json @@ -1,86 +1,9 @@ { "id": "f0bce436-ca7c-453d-8f96-d287c25cf130", "revision": 0, - "last_node_id": 6, + "last_node_id": 7, "last_link_id": 5, "nodes": [ - { - "id": 2, - "type": "VHS_VideoCombine", - "pos": [ - 1154.958677685951, - 306.6942148760332 - ], - "size": [ - 334.1807744573441, - 334 - ], - "flags": {}, - "order": 2, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 4 - }, - { - "name": "audio", - "shape": 7, - "type": "AUDIO", - "link": null - }, - { - "name": "meta_batch", - "shape": 7, - "type": "VHS_BatchManager", - "link": null - }, - { - "name": "vae", - "shape": 7, - "type": "VAE", - "link": null - } - ], - "outputs": [ - { - "name": "Filenames", - "type": "VHS_FILENAMES", - "links": null - } - ], - "properties": { - "cnr_id": "comfyui-videohelpersuite", - "ver": "1.7.9", - "Node name for S&R": "VHS_VideoCombine" - }, - "widgets_values": { - "frame_rate": 16, - "loop_count": 0, - "filename_prefix": "STAR-Upscaled", - "format": "video/h264-mp4", - "pix_fmt": "yuv420p", - "crf": 19, - "save_metadata": true, - "trim_to_audio": true, - "pingpong": false, - "save_output": true, - "videopreview": { - "hidden": false, - "paused": false, - "params": { - "filename": "AnimateDiff_00002.mp4", - "subfolder": "", - "type": "output", - "format": "video/h264-mp4", - "frame_rate": 15, - "workflow": "AnimateDiff_00002.png", - "fullpath": "/notebooks/output_images/AnimateDiff_00002.mp4" - } - } - } - }, { "id": 3, "type": "VHS_LoadVideo", @@ -164,6 +87,83 @@ } } }, + { + "id": 2, + "type": "VHS_VideoCombine", + "pos": [ + 1154.958677685951, + 306.6942148760332 + ], + "size": [ + 334.1807744573441, + 334 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 4 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": null + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "cnr_id": "comfyui-videohelpersuite", + "ver": "1.7.9", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 16, + "loop_count": 0, + "filename_prefix": "STAR-Upscaled", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00002.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 15, + "workflow": "AnimateDiff_00002.png", + "fullpath": "/notebooks/output_images/AnimateDiff_00002.mp4" + } + } + } + }, { "id": 5, "type": "STARVSRNode", @@ -176,7 +176,7 @@ 447.8347107438017 ], "flags": {}, - "order": 1, + "order": 2, "mode": 0, "inputs": [ { @@ -203,7 +203,7 @@ "Light Degradation", "fp8", "An alpaca turning it's head", - 2, + 720, 32, 5.5, "dpmpp_2m_sde", @@ -212,6 +212,29 @@ 42, "fixed" ] + }, + { + "id": 7, + "type": "MarkdownNote", + "pos": [ + 700, + 810 + ], + "size": [ + 407.29999999999995, + 232.70000000000005 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "# Recommended for Video Upscale Only Single Frame Output is weird.\n\n# VRAM Comsumption\n\n## 960p 45 GB\n\n## 720p 33 GB" + ], + "color": "#432", + "bgcolor": "#653" } ], "links": [ @@ -238,8 +261,8 @@ "ds": { "scale": 0.9090909090909091, "offset": [ - 93.97723448237635, - -134.0369738564675 + 188.19998248026926, + -115.73694602483096 ] }, "frontendVersion": "1.35.9", From 6c525cb7270846a160f652e646e257aaa566fa75 Mon Sep 17 00:00:00 2001 From: VjumpKunG Date: Wed, 24 Dec 2025 22:33:09 +0700 Subject: [PATCH 8/8] add github actions --- .github/workflows/push_to_comfy_registry.yaml | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 .github/workflows/push_to_comfy_registry.yaml diff --git a/.github/workflows/push_to_comfy_registry.yaml b/.github/workflows/push_to_comfy_registry.yaml new file mode 100644 index 0000000..e3c5c78 --- /dev/null +++ b/.github/workflows/push_to_comfy_registry.yaml @@ -0,0 +1,20 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + paths: + - "pyproject.toml" + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} ## Add your own personal access token to your Github Repository secrets and reference it here.