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. diff --git a/README.md b/README.md index 12805c2..230d43b 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. @@ -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/__init__.py b/__init__.py index 7a6e35b..1383880 100644 --- a/__init__.py +++ b/__init__.py @@ -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 ") 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..1286e2c 100644 --- a/nodes.py +++ b/nodes.py @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 4efc7d4..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) 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..fc6dfe6 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 @@ -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 diff --git a/workflows/STAR-example-workflow.json b/workflows/STAR-example-workflow.json index 1958ee5..cbaf845 100644 --- a/workflows/STAR-example-workflow.json +++ b/workflows/STAR-example-workflow.json @@ -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,