644 lines
20 KiB
Python
644 lines
20 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from threading import Lock
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import comfy.model_management as model_management
|
|
from comfy_api.latest import io
|
|
from easydict import EasyDict
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
from .color_fix import adain_color_fix, wavelet_color_fix
|
|
from .inference_utils import tensor2vid
|
|
from .video_to_video.utils.logger import get_logger
|
|
from .video_to_video.utils.seed import setup_seed
|
|
from .video_to_video.video_to_video_model import VideoToVideo_sr
|
|
|
|
try:
|
|
import folder_paths
|
|
except ImportError: # Allows lightweight tooling outside a ComfyUI checkout.
|
|
folder_paths = None
|
|
|
|
|
|
logger = get_logger()
|
|
|
|
CATEGORY = "video/upscaling/STAR"
|
|
TOTAL_NOISE_LEVELS = 900
|
|
|
|
MODEL_SPECS = {
|
|
"Light Degradation": {
|
|
"repo_id": "SherryX/STAR",
|
|
"filename": "I2VGen-XL-based/light_deg.pt",
|
|
},
|
|
"Heavy Degradation": {
|
|
"repo_id": "SherryX/STAR",
|
|
"filename": "I2VGen-XL-based/heavy_deg.pt",
|
|
},
|
|
}
|
|
MODEL_OPTIONS = list(MODEL_SPECS)
|
|
PRECISION_OPTIONS = ["fp16", "fp8"]
|
|
SAMPLER_OPTIONS = ["heun", "dpmpp_2m_sde"]
|
|
SOLVER_MODE_OPTIONS = ["fast", "normal"]
|
|
COLOR_FIX_OPTIONS = ["AdaIN", "wavelet", "none"]
|
|
|
|
STAR_MODEL = io.Custom("STARVSR_MODEL")
|
|
STAR_VIDEO = io.Custom("STARVSR_VIDEO")
|
|
STAR_CONDITIONING = io.Custom("STARVSR_CONDITIONING")
|
|
STAR_LATENT = io.Custom("STARVSR_LATENT")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class STARModelHandle:
|
|
model: VideoToVideo_sr
|
|
model_type: str
|
|
precision: str
|
|
checkpoint_path: str
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class STARPreparedVideo:
|
|
video_data: torch.Tensor
|
|
target_res: tuple[int, int]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class STARConditioning:
|
|
embeddings: torch.Tensor
|
|
prompt: str
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class STARSampledLatent:
|
|
samples: torch.Tensor
|
|
padding: tuple[int, int, int, int]
|
|
output_size: tuple[int, int]
|
|
batch_size: int
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _CachedModel:
|
|
identity: tuple[str, int, int]
|
|
handle: STARModelHandle
|
|
|
|
|
|
_MODEL_CACHE: dict[tuple[str, str], _CachedModel] = {}
|
|
_MODEL_CACHE_LOCK = Lock()
|
|
|
|
|
|
def _models_root() -> Path:
|
|
if folder_paths is not None:
|
|
return Path(folder_paths.models_dir)
|
|
return Path.cwd() / "models"
|
|
|
|
|
|
def _checkpoint_path(model_type: str) -> Path:
|
|
try:
|
|
filename = MODEL_SPECS[model_type]["filename"]
|
|
except KeyError as exc:
|
|
raise ValueError(f"Unknown STAR model type: {model_type}") from exc
|
|
return _models_root() / "STAR" / filename
|
|
|
|
|
|
def _checkpoint_identity(path: Path) -> tuple[str, int, int]:
|
|
stat = path.stat()
|
|
return str(path.resolve()), stat.st_mtime_ns, stat.st_size
|
|
|
|
|
|
def _checkpoint_fingerprint(model_type: str) -> tuple[str, int, int] | tuple[str, str]:
|
|
path = _checkpoint_path(model_type)
|
|
if path.is_file():
|
|
return _checkpoint_identity(path)
|
|
spec = MODEL_SPECS[model_type]
|
|
return "missing", f"{spec['repo_id']}:{spec['filename']}"
|
|
|
|
|
|
def _ensure_checkpoint(model_type: str) -> Path:
|
|
path = _checkpoint_path(model_type)
|
|
if path.is_file():
|
|
return path
|
|
|
|
spec = MODEL_SPECS[model_type]
|
|
local_dir = _models_root() / "STAR"
|
|
local_dir.mkdir(parents=True, exist_ok=True)
|
|
logger.info(
|
|
"Model not found locally. Downloading from Hugging Face: "
|
|
f"{spec['repo_id']}/{spec['filename']}"
|
|
)
|
|
downloaded_path = Path(
|
|
hf_hub_download(
|
|
repo_id=spec["repo_id"],
|
|
filename=spec["filename"],
|
|
local_dir=str(local_dir),
|
|
)
|
|
)
|
|
logger.info(f"Model downloaded to: {downloaded_path}")
|
|
return downloaded_path
|
|
|
|
|
|
def _load_model(model_type: str, precision: str) -> STARModelHandle:
|
|
cache_key = model_type, precision
|
|
with _MODEL_CACHE_LOCK:
|
|
path = _ensure_checkpoint(model_type)
|
|
identity = _checkpoint_identity(path)
|
|
cached = _MODEL_CACHE.get(cache_key)
|
|
if cached is not None and cached.identity == identity:
|
|
return cached.handle
|
|
|
|
logger.info(f"Loading model: {path}")
|
|
model_cfg = EasyDict(__name__="model_cfg", model_path=str(path))
|
|
device = model_management.get_torch_device()
|
|
handle = STARModelHandle(
|
|
model=VideoToVideo_sr(model_cfg, device=device, precision=precision),
|
|
model_type=model_type,
|
|
precision=precision,
|
|
checkpoint_path=str(path),
|
|
)
|
|
|
|
# Match the old node's one-model cache. References held by ComfyUI's
|
|
# output cache remain valid when a different model is loaded.
|
|
_MODEL_CACHE.clear()
|
|
_MODEL_CACHE[cache_key] = _CachedModel(identity=identity, handle=handle)
|
|
return handle
|
|
|
|
|
|
def _require_type(value: Any, expected_type: type, input_name: str):
|
|
if not isinstance(value, expected_type):
|
|
raise TypeError(
|
|
f"{input_name} must come from the matching STAR node; "
|
|
f"received {type(value).__name__}."
|
|
)
|
|
return value
|
|
|
|
|
|
def _prepare_video(images: torch.Tensor, resolution: int) -> STARPreparedVideo:
|
|
if not isinstance(images, torch.Tensor):
|
|
raise TypeError("images must be a ComfyUI IMAGE tensor.")
|
|
if images.ndim != 4 or images.shape[-1] != 3:
|
|
raise ValueError(
|
|
"STAR expects RGB IMAGE data in [frames, height, width, 3] format; "
|
|
f"received {tuple(images.shape)}."
|
|
)
|
|
if images.shape[0] < 1 or images.shape[1] < 1 or images.shape[2] < 1:
|
|
raise ValueError("STAR requires at least one non-empty video frame.")
|
|
|
|
# Preserve the legacy wrapper's uint8 preprocessing while avoiding the
|
|
# RGB -> BGR -> RGB and PIL round trips. Cache this tensor on the CPU so it
|
|
# does not reserve VRAM between executions.
|
|
video_data = (
|
|
images.detach()
|
|
.clamp(0.0, 1.0)
|
|
.mul(255)
|
|
.to(device="cpu", dtype=torch.uint8)
|
|
.to(dtype=torch.float32)
|
|
.div(255)
|
|
.permute(0, 3, 1, 2)
|
|
.contiguous()
|
|
.mul(2)
|
|
.sub(1)
|
|
)
|
|
|
|
_, _, height, width = video_data.shape
|
|
aspect_ratio = width / height
|
|
if height < width:
|
|
target_h = resolution
|
|
target_w = int(resolution * aspect_ratio)
|
|
else:
|
|
target_w = resolution
|
|
target_h = int(resolution / aspect_ratio)
|
|
|
|
target_h += target_h % 2
|
|
target_w += target_w % 2
|
|
logger.info(f"Input resolution: {(height, width)}")
|
|
logger.info(f"Target resolution: {(target_h, target_w)}")
|
|
return STARPreparedVideo(video_data=video_data, target_res=(target_h, target_w))
|
|
|
|
|
|
def _decode_video(
|
|
model: STARModelHandle,
|
|
latent: STARSampledLatent,
|
|
vae_decode_chunk: int,
|
|
) -> torch.Tensor:
|
|
video = model.model.decode_latent(
|
|
latent.samples,
|
|
latent.padding,
|
|
latent.output_size,
|
|
batch_size=latent.batch_size,
|
|
chunk_size=vae_decode_chunk,
|
|
)
|
|
return tensor2vid(video).div(255).to(dtype=torch.float32, device="cpu")
|
|
|
|
|
|
def _color_fix(
|
|
images: torch.Tensor,
|
|
reference_images: torch.Tensor,
|
|
method: str,
|
|
) -> torch.Tensor:
|
|
if method == "none":
|
|
return images
|
|
if images.ndim != 4 or images.shape[-1] != 3:
|
|
raise ValueError("images must use ComfyUI's [frames, height, width, 3] format.")
|
|
if reference_images.ndim != 4 or reference_images.shape[-1] != 3:
|
|
raise ValueError(
|
|
"reference_images must use ComfyUI's [frames, height, width, 3] format."
|
|
)
|
|
if images.shape[0] != reference_images.shape[0]:
|
|
raise ValueError(
|
|
"images and reference_images must contain the same number of frames."
|
|
)
|
|
|
|
original_device = images.device
|
|
original_dtype = images.dtype
|
|
target = images.detach().to(dtype=torch.float32).clamp(0.0, 1.0)
|
|
source = (
|
|
reference_images.detach()
|
|
.to(device=target.device)
|
|
.clamp(0.0, 1.0)
|
|
.mul(255)
|
|
.to(dtype=torch.uint8)
|
|
.to(dtype=torch.float32)
|
|
.div(255)
|
|
.permute(0, 3, 1, 2)
|
|
.contiguous()
|
|
.mul(2)
|
|
.sub(1)
|
|
)
|
|
|
|
if method == "AdaIN":
|
|
result = adain_color_fix(target.mul(255), source).div(255)
|
|
elif method == "wavelet":
|
|
if source.shape[-2:] != target.shape[1:3]:
|
|
source = F.interpolate(
|
|
source, size=target.shape[1:3], mode="bilinear", align_corners=False
|
|
)
|
|
result = wavelet_color_fix(target.mul(255), source).div(255)
|
|
else:
|
|
raise ValueError(f"Unknown color-fix method: {method}")
|
|
|
|
return result.to(device=original_device, dtype=original_dtype)
|
|
|
|
|
|
def _model_inputs() -> list[io.Input]:
|
|
return [
|
|
io.Combo.Input(
|
|
"model",
|
|
options=MODEL_OPTIONS,
|
|
default="Light Degradation",
|
|
tooltip="STAR checkpoint trained for the source degradation level.",
|
|
),
|
|
io.Combo.Input(
|
|
"precision",
|
|
options=PRECISION_OPTIONS,
|
|
default="fp16",
|
|
tooltip="UNet weight precision. FP8 requires compatible PyTorch and hardware.",
|
|
),
|
|
]
|
|
|
|
|
|
def _sampling_inputs() -> list[io.Input]:
|
|
return [
|
|
io.Int.Input(
|
|
"max_chunk_len",
|
|
default=32,
|
|
min=1,
|
|
max=128,
|
|
step=1,
|
|
tooltip="Maximum temporal window processed by the diffusion model.",
|
|
),
|
|
io.Float.Input(
|
|
"cfg",
|
|
default=7.5,
|
|
min=0.0,
|
|
max=20.0,
|
|
step=0.1,
|
|
tooltip="Classifier-free guidance strength.",
|
|
),
|
|
io.Combo.Input(
|
|
"sampler",
|
|
options=SAMPLER_OPTIONS,
|
|
default="dpmpp_2m_sde",
|
|
),
|
|
io.Combo.Input(
|
|
"solver_mode",
|
|
options=SOLVER_MODE_OPTIONS,
|
|
default="fast",
|
|
),
|
|
io.Int.Input("steps", default=15, min=1, max=100, step=1),
|
|
io.Int.Input(
|
|
"seed",
|
|
default=42,
|
|
min=0,
|
|
max=0xFFFFFFFFFFFFFFFF,
|
|
control_after_generate=True,
|
|
),
|
|
]
|
|
|
|
|
|
def _vae_decode_chunk_input(*, optional: bool = False) -> io.Int.Input:
|
|
return io.Int.Input(
|
|
"vae_decode_chunk",
|
|
default=1,
|
|
min=1,
|
|
max=8,
|
|
step=1,
|
|
optional=optional,
|
|
tooltip=(
|
|
"Frames decoded per VAE pass. One minimizes peak VRAM; increase only "
|
|
"when enough VRAM is available. This value no longer invalidates "
|
|
"diffusion sampling in the modular workflow."
|
|
),
|
|
)
|
|
|
|
|
|
class STARModelLoader(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="STARVSR_ModelLoader",
|
|
display_name="STAR Model Loader",
|
|
category=CATEGORY,
|
|
description=(
|
|
"Loads a STAR degradation checkpoint. Keep this separate so prompt and "
|
|
"sampling changes reuse the loaded model through ComfyUI's cache."
|
|
),
|
|
inputs=_model_inputs(),
|
|
outputs=[STAR_MODEL.Output("star_model", display_name="STAR model")],
|
|
)
|
|
|
|
@classmethod
|
|
def fingerprint_inputs(cls, model: str, precision: str, **_kwargs):
|
|
return model, precision, _checkpoint_fingerprint(model)
|
|
|
|
@classmethod
|
|
def execute(cls, model: str, precision: str) -> io.NodeOutput:
|
|
return io.NodeOutput(_load_model(model, precision))
|
|
|
|
|
|
class STARPrepareVideo(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="STARVSR_PrepareVideo",
|
|
display_name="STAR Prepare Video",
|
|
category=CATEGORY,
|
|
description=(
|
|
"Quantizes and normalizes an IMAGE frame batch and calculates the target "
|
|
"resolution. The cached output is CPU-resident."
|
|
),
|
|
inputs=[
|
|
io.Image.Input(
|
|
"images", tooltip="Video frames as a ComfyUI IMAGE batch."
|
|
),
|
|
io.Int.Input(
|
|
"resolution",
|
|
default=720,
|
|
min=16,
|
|
max=16384,
|
|
step=2,
|
|
tooltip="Target size of the shorter side; aspect ratio is preserved.",
|
|
),
|
|
],
|
|
outputs=[STAR_VIDEO.Output("video", display_name="prepared video")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, images: torch.Tensor, resolution: int) -> io.NodeOutput:
|
|
return io.NodeOutput(_prepare_video(images, resolution))
|
|
|
|
|
|
class STARTextEncode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="STARVSR_TextEncode",
|
|
display_name="STAR Text Encode",
|
|
category=CATEGORY,
|
|
description=(
|
|
"Encodes the prompt once so sampler-setting changes can reuse the "
|
|
"conditioning from ComfyUI's cache."
|
|
),
|
|
inputs=[
|
|
STAR_MODEL.Input("star_model"),
|
|
io.String.Input("prompt", default="a good video", multiline=True),
|
|
],
|
|
outputs=[
|
|
STAR_CONDITIONING.Output(
|
|
"conditioning", display_name="STAR conditioning"
|
|
)
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, star_model: STARModelHandle, prompt: str) -> io.NodeOutput:
|
|
star_model = _require_type(star_model, STARModelHandle, "star_model")
|
|
caption = prompt or star_model.model.positive_prompt
|
|
logger.info(f"Caption: {caption}")
|
|
embeddings = star_model.model.encode_prompt(caption)
|
|
return io.NodeOutput(STARConditioning(embeddings=embeddings, prompt=caption))
|
|
|
|
|
|
class STARSample(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="STARVSR_Sample",
|
|
display_name="STAR Sample",
|
|
category=CATEGORY,
|
|
description=(
|
|
"VAE-encodes the prepared video and runs STAR diffusion. Returns a "
|
|
"small CPU latent so decoding can be cached separately."
|
|
),
|
|
inputs=[
|
|
STAR_MODEL.Input("star_model"),
|
|
STAR_VIDEO.Input("video"),
|
|
STAR_CONDITIONING.Input("conditioning"),
|
|
*_sampling_inputs(),
|
|
],
|
|
outputs=[STAR_LATENT.Output("latent", display_name="STAR latent")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
star_model: STARModelHandle,
|
|
video: STARPreparedVideo,
|
|
conditioning: STARConditioning,
|
|
max_chunk_len: int,
|
|
cfg: float,
|
|
sampler: str,
|
|
solver_mode: str,
|
|
steps: int,
|
|
seed: int,
|
|
) -> io.NodeOutput:
|
|
star_model = _require_type(star_model, STARModelHandle, "star_model")
|
|
video = _require_type(video, STARPreparedVideo, "video")
|
|
conditioning = _require_type(conditioning, STARConditioning, "conditioning")
|
|
setup_seed(seed)
|
|
samples, padding, output_size, batch_size = star_model.model.sample_latent(
|
|
{"video_data": video.video_data, "target_res": video.target_res},
|
|
conditioning.embeddings,
|
|
total_noise_levels=TOTAL_NOISE_LEVELS,
|
|
steps=steps,
|
|
solver=sampler,
|
|
solver_mode=solver_mode,
|
|
guide_scale=cfg,
|
|
max_chunk_len=max_chunk_len,
|
|
)
|
|
return io.NodeOutput(
|
|
STARSampledLatent(
|
|
samples=samples,
|
|
padding=padding,
|
|
output_size=output_size,
|
|
batch_size=batch_size,
|
|
)
|
|
)
|
|
|
|
|
|
class STARDecode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="STARVSR_Decode",
|
|
display_name="STAR VAE Decode",
|
|
category=CATEGORY,
|
|
description=(
|
|
"Decodes a STAR latent to raw enhanced frames. Changing decode chunk "
|
|
"size reruns only this node, not diffusion sampling."
|
|
),
|
|
inputs=[
|
|
STAR_MODEL.Input("star_model"),
|
|
STAR_LATENT.Input("latent"),
|
|
_vae_decode_chunk_input(),
|
|
],
|
|
outputs=[io.Image.Output("images", display_name="raw images")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
star_model: STARModelHandle,
|
|
latent: STARSampledLatent,
|
|
vae_decode_chunk: int,
|
|
) -> io.NodeOutput:
|
|
star_model = _require_type(star_model, STARModelHandle, "star_model")
|
|
latent = _require_type(latent, STARSampledLatent, "latent")
|
|
return io.NodeOutput(_decode_video(star_model, latent, vae_decode_chunk))
|
|
|
|
|
|
class STARColorFix(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="STARVSR_ColorFix",
|
|
display_name="STAR Color Fix",
|
|
category=CATEGORY,
|
|
description=(
|
|
"Matches enhanced-frame colors to the source. Keeping post-processing "
|
|
"separate allows method changes without rerunning STAR."
|
|
),
|
|
inputs=[
|
|
io.Image.Input(
|
|
"images", tooltip="Raw enhanced frames from STAR decode."
|
|
),
|
|
io.Image.Input(
|
|
"reference_images",
|
|
tooltip="Original source frames used as color reference.",
|
|
),
|
|
io.Combo.Input("method", options=COLOR_FIX_OPTIONS, default="AdaIN"),
|
|
],
|
|
outputs=[io.Image.Output("images", display_name="images")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
images: torch.Tensor,
|
|
reference_images: torch.Tensor,
|
|
method: str,
|
|
) -> io.NodeOutput:
|
|
return io.NodeOutput(_color_fix(images, reference_images, method))
|
|
|
|
|
|
class STARVSRNode(io.ComfyNode):
|
|
"""V3 all-in-one compatibility node retaining the released workflow ID."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="STARVSRNode",
|
|
display_name="STAR Video Super Resolution (All-in-One)",
|
|
category="video/upscaling",
|
|
description=(
|
|
"V3 compatibility node for existing workflows. New workflows should use "
|
|
"the modular STAR nodes for finer ComfyUI caching."
|
|
),
|
|
search_aliases=["STAR Video Super Resolution"],
|
|
inputs=[
|
|
io.Image.Input("images"),
|
|
*_model_inputs(),
|
|
io.String.Input("prompt", default="a good video", multiline=True),
|
|
io.Int.Input("resolution", default=720, min=16, max=16384, step=2),
|
|
*_sampling_inputs(),
|
|
_vae_decode_chunk_input(optional=True),
|
|
],
|
|
outputs=[io.Image.Output("images", display_name="images")],
|
|
is_deprecated=True,
|
|
)
|
|
|
|
@classmethod
|
|
def fingerprint_inputs(cls, model: str, precision: str, **_kwargs):
|
|
return model, precision, _checkpoint_fingerprint(model)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
images: torch.Tensor,
|
|
model: str,
|
|
precision: str,
|
|
prompt: str,
|
|
resolution: int,
|
|
max_chunk_len: int,
|
|
cfg: float,
|
|
sampler: str,
|
|
solver_mode: str,
|
|
steps: int,
|
|
seed: int,
|
|
vae_decode_chunk: int = 1,
|
|
) -> io.NodeOutput:
|
|
star_model = _load_model(model, precision)
|
|
video = _prepare_video(images, resolution)
|
|
caption = prompt or star_model.model.positive_prompt
|
|
logger.info(f"Caption: {caption}")
|
|
conditioning = STARConditioning(
|
|
embeddings=star_model.model.encode_prompt(caption), prompt=caption
|
|
)
|
|
|
|
setup_seed(seed)
|
|
samples, padding, output_size, batch_size = star_model.model.sample_latent(
|
|
{"video_data": video.video_data, "target_res": video.target_res},
|
|
conditioning.embeddings,
|
|
total_noise_levels=TOTAL_NOISE_LEVELS,
|
|
steps=steps,
|
|
solver=sampler,
|
|
solver_mode=solver_mode,
|
|
guide_scale=cfg,
|
|
max_chunk_len=max_chunk_len,
|
|
)
|
|
latent = STARSampledLatent(samples, padding, output_size, batch_size)
|
|
decoded = _decode_video(star_model, latent, vae_decode_chunk)
|
|
return io.NodeOutput(_color_fix(decoded, images, "AdaIN"))
|
|
|
|
|
|
STAR_NODE_CLASSES = [
|
|
STARModelLoader,
|
|
STARPrepareVideo,
|
|
STARTextEncode,
|
|
STARSample,
|
|
STARDecode,
|
|
STARColorFix,
|
|
STARVSRNode,
|
|
]
|