239 lines
8.0 KiB
Python
239 lines
8.0 KiB
Python
import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
from easydict import EasyDict
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
from .color_fix import adain_color_fix
|
|
from .inference_utils import collate_fn, preprocess, 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
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class STARVSRNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"model": (["Light Degradation", "Heavy Degradation"],),
|
|
"precision": (
|
|
["fp16", "fp8"],
|
|
{"default": "fp16"},
|
|
),
|
|
"prompt": ("STRING", {"default": "a good video", "multiline": True}),
|
|
"resolution": (
|
|
"INT",
|
|
{"default": 720, "min": 16, "max": 16384, "step": 2},
|
|
),
|
|
"max_chunk_len": (
|
|
"INT",
|
|
{"default": 32, "min": 1, "max": 128, "step": 1},
|
|
),
|
|
"cfg": (
|
|
"FLOAT",
|
|
{"default": 7.5, "min": 0.0, "max": 20.0, "step": 0.1},
|
|
),
|
|
"sampler": (
|
|
["heun", "dpmpp_2m_sde"],
|
|
{"default": "dpmpp_2m_sde"},
|
|
),
|
|
"solver_mode": (
|
|
["fast", "normal"],
|
|
{"default": "fast"},
|
|
),
|
|
"steps": ("INT", {"default": 15, "min": 1, "max": 100, "step": 1}),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "enhance_video"
|
|
CATEGORY = "video/upscaling"
|
|
|
|
def __init__(self):
|
|
self.model = None
|
|
self.current_model_type = None
|
|
self.current_precision = None
|
|
|
|
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
|
|
hf_model_config = {
|
|
"Light Degradation": {
|
|
"repo_id": "SherryX/STAR",
|
|
"filename": "I2VGen-XL-based/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/I2VGen-XL-based/heavy_deg.pt",
|
|
},
|
|
}
|
|
|
|
config = hf_model_config.get(model_type)
|
|
if config is None:
|
|
raise ValueError(f"Unknown model type: {model_type}")
|
|
|
|
model_path = config["local_path"]
|
|
|
|
# Download from Hugging Face if not exists locally
|
|
if not os.path.exists(model_path):
|
|
logger.info(
|
|
f"Model not found locally. Downloading from Hugging Face: {config['repo_id']}/{config['filename']}"
|
|
)
|
|
os.makedirs(os.path.dirname(model_path), exist_ok=True)
|
|
|
|
# Download the model from Hugging Face Hub
|
|
downloaded_path = hf_hub_download(
|
|
repo_id=config["repo_id"],
|
|
filename=config["filename"],
|
|
local_dir="./models/STAR/",
|
|
local_dir_use_symlinks=False,
|
|
)
|
|
logger.info(f"Model downloaded to: {downloaded_path}")
|
|
model_path = downloaded_path
|
|
|
|
logger.info(f"Loading model: {model_path}")
|
|
|
|
model_cfg = EasyDict(__name__="model_cfg")
|
|
model_cfg.model_path = model_path
|
|
|
|
# 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
|
|
|
|
def enhance_video(
|
|
self,
|
|
images,
|
|
model,
|
|
precision,
|
|
prompt,
|
|
resolution,
|
|
max_chunk_len,
|
|
cfg,
|
|
sampler,
|
|
solver_mode,
|
|
steps,
|
|
seed,
|
|
):
|
|
"""
|
|
Enhance video frames using STAR model
|
|
|
|
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
|
|
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
|
|
|
|
Returns:
|
|
Enhanced images tensor in ComfyUI format [B, H, W, C]
|
|
"""
|
|
# Load the appropriate 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
|
|
input_frames = []
|
|
for i in range(images.shape[0]):
|
|
frame = (images[i].cpu().numpy() * 255).astype(np.uint8)
|
|
# Convert RGB to BGR for processing (matching cv2 format)
|
|
frame_bgr = frame[:, :, ::-1]
|
|
input_frames.append(frame_bgr)
|
|
|
|
# Preprocess frames
|
|
video_data = preprocess(input_frames)
|
|
_, _, h, w = video_data.shape
|
|
logger.info(f"Input resolution: {(h, w)}")
|
|
|
|
# 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
|
|
caption = prompt or model_instance.positive_prompt
|
|
logger.info(f"Caption: {caption}")
|
|
|
|
# Prepare data
|
|
pre_data = {
|
|
"video_data": video_data,
|
|
"y": caption,
|
|
"target_res": (target_h, target_w),
|
|
}
|
|
|
|
total_noise_levels = 900
|
|
setup_seed(seed)
|
|
|
|
# Run inference
|
|
with torch.no_grad():
|
|
data_tensor = collate_fn(
|
|
pre_data, "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
)
|
|
output = model_instance.test(
|
|
data_tensor,
|
|
total_noise_levels,
|
|
steps=steps,
|
|
solver=sampler,
|
|
solver_mode=solver_mode,
|
|
guide_scale=cfg,
|
|
max_chunk_len=max_chunk_len,
|
|
)
|
|
|
|
# Convert output to video frames
|
|
output = tensor2vid(output)
|
|
|
|
# Apply color fix
|
|
output = adain_color_fix(output, video_data)
|
|
|
|
# Convert back to ComfyUI format [B, H, W, C] in range [0, 1]
|
|
# output is currently [T, H, W, C] in range [0, 255]
|
|
output_normalized = output / 255.0
|
|
|
|
# Ensure it's a torch tensor and move to CPU
|
|
if not isinstance(output_normalized, torch.Tensor):
|
|
output_normalized = torch.tensor(output_normalized)
|
|
output_normalized = output_normalized.cpu().float()
|
|
|
|
return (output_normalized,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"STARVSRNode": STARVSRNode}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"STARVSRNode": "STAR Video Super Resolution"}
|