CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Project Overview

STARVSRWrapper is a ComfyUI custom node that wraps the STAR (Spatio-Temporal Adaptive Restoration) video super-resolution model. It enables video upscaling with light or heavy degradation models through a ComfyUI node interface.

Architecture

Core Components

nodes.py - ComfyUI V3 node implementations
- Uses io.ComfyNode, define_schema, classmethod execute, and io.NodeOutput
- Provides modular model-loader, video-preparation, text-encoding, sampling, VAE-decoding, and color-fix nodes
- Retains STARVSRNode as a deprecated V3 all-in-one compatibility node for existing workflows
- Handles model loading from Hugging Face Hub or ComfyUI's models/STAR/ cache
- Converts between ComfyUI image format [B, H, W, C] (0-1 range) and STAR's internal format
- Two model variants: "Light Degradation" and "Heavy Degradation" downloaded from HF repo SherryX/STAR
inference_utils.py - Preprocessing and data utilities
- preprocess(): Converts input frames to normalized tensors (mean=0.5, std=0.5)
- tensor2vid(): Converts output tensors back to video frames
- collate_fn(): Recursive device placement for nested data structures
- load_video() / save_video(): File I/O utilities (uses ffmpeg for encoding)
color_fix.py - Post-processing color correction
- adain_color_fix(): Adaptive Instance Normalization to match source video colors
- wavelet_color_fix(): Alternative wavelet-based color correction
- Currently uses AdaIN by default in the node pipeline
video_to_video/ - STAR model implementation
- video_to_video_model.py: Main VideoToVideo_sr class that orchestrates inference
- diffusion/: Diffusion sampling and noise schedules (based on SDEdit)
- modules/: UNet architecture (ControlledV2VUNet) and text embedder (OpenCLIP)
- utils/: Configuration, logging, and seed utilities

Data Flow

ComfyUI Images [B,H,W,C] (0-1)
    ↓ STAR Prepare Video (CPU, normalized BCHW)
Prepared STAR video
    ↓ STAR Sample (VAE encode + diffusion)
Sampled STAR latent (CPU)
    ↓ STAR VAE Decode
Raw ComfyUI Images [B,H,W,C] (0-1)
    ↓ STAR Color Fix + original images
Color-corrected ComfyUI Images [B,H,W,C] (0-1)

Model Components

Text Encoder: FrozenOpenCLIPEmbedder (laion2b_s32b_b79k) for caption conditioning
UNet: ControlledV2VUNet - temporal-aware UNet with ControlNet for LR guidance
VAE: AutoencoderKLTemporalDecoder from Stable Video Diffusion
Diffusion: GaussianDiffusion with logsnr_cosine_interp schedule, supports heun and dpmpp_2m_sde samplers
Chunking: Long videos are processed in overlapping chunks (default max_chunk_len=32)

Development Commands

Environment Setup

# Create and activate virtual environment
python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

Code Formatting

# Format code using ruff
ruff format .

# Check formatting without making changes
ruff format --check .

Important Notes

CUDA Requirement: Model runs on GPU (cuda:0). Hardcoded in several places.
Model Caching: Models are downloaded to ./pretrained_weight/ by Hugging Face Hub and expected at ./models/STAR/ for local loading
Resolution Constraints: Model pads inputs to fit minimum 720x1280 or multiples of 64 pixels
Memory Management: Uses chunked processing and explicit torch.cuda.empty_cache() calls
FP16: Model runs in half precision for efficiency
Color Space: Internal processing uses BGR (cv2 format), but ComfyUI expects RGB

Integration with ComfyUI

This is a ComfyUI custom node that should be placed in ComfyUI's custom_nodes/ directory. The node:

Category: "video/upscaling"
Display name: "STAR Video Super Resolution"
Inputs: IMAGE tensor, model type, prompt, upscale factor, sampling parameters
Outputs: Upscaled IMAGE tensor

The ComfyExtension and async comfy_entrypoint in __init__.py register all V3 nodes with ComfyUI.

4.2 KiB Raw Permalink Blame History