nodes/vjumpkung-ComfyUI-STARWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #1 from vjumpkung/develop

Browse Source 
Develop -> Main
This commit is contained in:
Chanrich Pisitjing
2025-12-24 22:33:30 +07:00
committed by GitHub
parent 91a58167cc 6c525cb727
commit 2da37076a3
9 changed files with 205 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/push_to_comfy_registry.yaml
+20
View File
@@ -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.
README.md
+3 -3
View File
@@ -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.
@@ -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 |
__init__.py
+2 -8
View File
@@ -11,14 +11,8 @@ 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"]
print("🎉 Loading STAR Video Super-Resolution Completed ")
install.py
+2 -2
View File
@@ -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
nodes.py
+44 -11
View File
@@ -21,8 +21,15 @@ class STARVSRNode:
"required": {
"images": ("IMAGE",),
"model": (["Light Degradation", "Heavy Degradation"],),
"precision": (
["fp16", "fp8"],
{"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},
@@ -51,10 +58,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
@@ -62,12 +74,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/I2VGen-XL-based/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/I2VGen-XL-based/heavy_deg.pt",
},
}
@@ -98,8 +110,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,8 +123,9 @@ class STARVSRNode:
self,
images,
model,
precision,
prompt,
upscale,
resolution,
max_chunk_len,
cfg,
sampler,
@@ -122,10 +139,12 @@ 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")
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
@@ -134,7 +153,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
@@ -150,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
pyproject.toml
+1 -1
View File
@@ -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)
video_to_video/modules/unet_v2v.py
+3 -4
View File
@@ -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)
video_to_video/video_to_video_model.py
+37 -3
View File
@@ -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
@@ -17,9 +17,15 @@ logger = get_logger()
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 +46,14 @@ 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 == "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 +192,27 @@ 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_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
workflows/STAR-example-workflow.json
+93 -69
View File
@@ -1,67 +1,19 @@
{
"id": "f0bce436-ca7c-453d-8f96-d287c25cf130",
"revision": 0,
"last_node_id": 4,
"last_link_id": 2,
"last_node_id": 7,
"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
400,
310
],
"size": [
255.8291015625,
457.1858620289278
740.5813428874735
],
"flags": {},
"order": 0,
@@ -85,7 +37,7 @@
"name": "IMAGE",
"type": "IMAGE",
"links": [
1
3
]
},
{
@@ -110,7 +62,7 @@
"Node name for S&R": "VHS_LoadVideo"
},
"widgets_values": {
"video": "video_00001.mp4",
"video": "video_00010.mp4",
"force_rate": 0,
"custom_width": 0,
"custom_height": 0,
@@ -128,7 +80,7 @@
"frame_load_cap": 0,
"skip_first_frames": 0,
"select_every_nth": 1,
"filename": "video_00001.mp4",
"filename": "video_00010.mp4",
"type": "input",
"format": "video/mp4"
}
@@ -139,21 +91,21 @@
"id": 2,
"type": "VHS_VideoCombine",
"pos": [
1217.1239669421477,
290.23140495867796
1154.958677685951,
306.6942148760332
],
"size": [
334.1807744573441,
525.8381742211382
334
],
"flags": {},
"order": 2,
"order": 3,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 2
"link": 4
},
{
"name": "audio",
@@ -194,7 +146,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 +163,92 @@
}
}
}
},
{
"id": 5,
"type": "STARVSRNode",
"pos": [
702.5652892561982,
306.85702479338863
],
"size": [
403.3057851239671,
447.8347107438017
],
"flags": {},
"order": 2,
"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",
720,
32,
5.5,
"dpmpp_2m_sde",
"fast",
15,
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": [
[
1,
3,
3,
0,
4,
5,
0,
"IMAGE"
],
[
2,
4,
5,
0,
2,
0,
@@ -235,14 +259,14 @@
"config": {},
"extra": {
"ds": {
"scale": 1.1000000000000003,
"scale": 0.9090909090909091,
"offset": [
-149.28766272203228,
-141.35987067640917
188.19998248026926,
-115.73694602483096
]
},
"frontendVersion": "1.35.9",
"workflowRendererVersion": "LG",
"frontendVersion": "1.34.9",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,