573 lines
19 KiB
Python
573 lines
19 KiB
Python
import math
|
||
import re
|
||
from fractions import Fraction
|
||
from typing import Dict, List, Optional, Tuple
|
||
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image
|
||
|
||
|
||
WAN_PRESETS: Dict[str, List[Tuple[int, int, str]]] = {
|
||
"1:1": [
|
||
(480, 480, "Fast Draft"),
|
||
(640, 640, "Preview"),
|
||
(832, 832, "High Detail"),
|
||
(960, 960, "Wan 2.2 Native"),
|
||
],
|
||
"2:3": [
|
||
(384, 576, "Fast Draft"),
|
||
(512, 768, "Preview"),
|
||
(672, 1008, "High Detail"),
|
||
(768, 1168, "Wan 2.2 Native"),
|
||
],
|
||
"3:2": [
|
||
(576, 384, "Fast Draft"),
|
||
(768, 512, "Preview"),
|
||
(1008, 672, "High Detail"),
|
||
(1168, 768, "Wan 2.2 Native"),
|
||
],
|
||
"3:4": [
|
||
(432, 576, "Fast Draft"),
|
||
(576, 768, "Preview"),
|
||
(720, 960, "High Detail"),
|
||
(816, 1104, "Wan 2.2 Native"),
|
||
],
|
||
"4:3": [
|
||
(576, 432, "Fast Draft"),
|
||
(768, 576, "Preview"),
|
||
(960, 720, "High Detail"),
|
||
(1104, 816, "Wan 2.2 Native"),
|
||
],
|
||
"9:16": [
|
||
(352, 624, "Fast Draft"),
|
||
(480, 848, "Preview"),
|
||
(624, 1104, "High Detail"),
|
||
(720, 1280, "Wan 2.2 Native"),
|
||
],
|
||
"16:9": [
|
||
(624, 352, "Fast Draft"),
|
||
(848, 480, "Preview"),
|
||
(1104, 624, "High Detail"),
|
||
(1280, 720, "Wan 2.2 Native"),
|
||
],
|
||
}
|
||
|
||
|
||
MINIMAX_PRESETS: Dict[str, List[Tuple[int, int, str]]] = {
|
||
"1:1": [
|
||
(512, 512, "Fast Draft"),
|
||
(640, 640, "Preview"),
|
||
(768, 768, "High Detail"),
|
||
(1024, 1024, "Native"),
|
||
(1152, 1152, "1.5K"),
|
||
(1440, 1440, "1080P Class"),
|
||
(1536, 1536, "2K"),
|
||
],
|
||
"3:4": [
|
||
(448, 576, "Fast Draft"),
|
||
(576, 736, "Preview"),
|
||
(672, 896, "High Detail"),
|
||
(864, 1184, "Native"),
|
||
(992, 1344, "1.5K"),
|
||
(1248, 1664, "1080P Class"),
|
||
(1344, 1760, "2K"),
|
||
],
|
||
"4:3": [
|
||
(576, 448, "Fast Draft"),
|
||
(736, 576, "Preview"),
|
||
(896, 672, "High Detail"),
|
||
(1184, 864, "Native"),
|
||
(1344, 992, "1.5K"),
|
||
(1664, 1248, "1080P Class"),
|
||
(1760, 1344, "2K"),
|
||
],
|
||
"9:16": [
|
||
(384, 672, "Fast Draft"),
|
||
(480, 864, "Preview"),
|
||
(576, 1024, "High Detail"),
|
||
(768, 1344, "Native"),
|
||
(864, 1536, "1.5K"),
|
||
(1088, 1920, "1080P Class"),
|
||
(1152, 2048, "2K"),
|
||
],
|
||
"16:9": [
|
||
(672, 384, "Fast Draft"),
|
||
(864, 480, "Preview"),
|
||
(1024, 576, "High Detail"),
|
||
(1344, 768, "Native"),
|
||
(1536, 864, "1.5K"),
|
||
(1920, 1088, "1080P Class"),
|
||
(2048, 1152, "2K"),
|
||
],
|
||
"21:9": [
|
||
(768, 320, "Fast Draft"),
|
||
(992, 416, "Preview"),
|
||
(1184, 512, "High Detail"),
|
||
(1536, 672, "Native"),
|
||
(1760, 768, "1.5K"),
|
||
(2208, 960, "1080P Class"),
|
||
(2336, 992, "2K"),
|
||
],
|
||
}
|
||
|
||
MINIMAX_TARGET_PIXELS = (
|
||
512 * 512,
|
||
int(0.40 * 1024 * 1024),
|
||
1024 * 576,
|
||
1344 * 768,
|
||
1536 * 864,
|
||
int(2.00 * 1024 * 1024),
|
||
2048 * 1152,
|
||
)
|
||
|
||
|
||
LTX_PRESETS: Dict[str, List[Tuple[int, int, str]]] = {
|
||
"1:1": [
|
||
(320, 320, "Fast Draft"),
|
||
(640, 640, "Preview"),
|
||
(768, 768, "High Detail"),
|
||
(960, 960, "Native"),
|
||
(1184, 1184, "HD Output"),
|
||
(1440, 1440, "Full HD Output"),
|
||
],
|
||
"2:3": [
|
||
(256, 384, "Fast Draft"),
|
||
(512, 768, "Preview"),
|
||
(640, 960, "High Detail"),
|
||
(768, 1152, "Native"),
|
||
(960, 1440, "HD Output"),
|
||
(1152, 1728, "Full HD Output"),
|
||
],
|
||
"3:2": [
|
||
(384, 256, "Fast Draft"),
|
||
(768, 512, "Preview"),
|
||
(960, 640, "High Detail"),
|
||
(1152, 768, "Native"),
|
||
(1440, 960, "HD Output"),
|
||
(1728, 1152, "Full HD Output"),
|
||
],
|
||
"3:4": [
|
||
(256, 352, "Fast Draft"),
|
||
(512, 704, "Preview"),
|
||
(640, 864, "High Detail"),
|
||
(864, 1152, "Native"),
|
||
(1056, 1408, "HD Output"),
|
||
(1248, 1664, "Full HD Output"),
|
||
],
|
||
"4:3": [
|
||
(352, 256, "Fast Draft"),
|
||
(704, 512, "Preview"),
|
||
(864, 640, "High Detail"),
|
||
(1152, 864, "Native"),
|
||
(1408, 1056, "HD Output"),
|
||
(1664, 1248, "Full HD Output"),
|
||
],
|
||
"9:16": [
|
||
(288, 512, "Fast Draft"),
|
||
(544, 960, "Preview"),
|
||
(672, 1184, "High Detail"),
|
||
(736, 1312, "Native"),
|
||
(864, 1536, "HD Output"),
|
||
(1088, 1920, "Full HD Output"),
|
||
],
|
||
"16:9": [
|
||
(512, 288, "Fast Draft"),
|
||
(960, 544, "Preview"),
|
||
(1184, 672, "High Detail"),
|
||
(1312, 736, "Native"),
|
||
(1536, 864, "HD Output"),
|
||
(1920, 1088, "Full HD Output"),
|
||
],
|
||
}
|
||
|
||
|
||
MODEL_CONFIG: Dict[str, Dict] = {
|
||
"WAN 2.2": {
|
||
"aspect_order": ("1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9"),
|
||
"fallback_aspect": "1:1",
|
||
"presets": WAN_PRESETS,
|
||
"default_divisible_by": 16,
|
||
},
|
||
"MiniMax H3": {
|
||
"aspect_order": ("1:1", "3:4", "4:3", "9:16", "16:9", "21:9"),
|
||
"fallback_aspect": "16:9",
|
||
"presets": MINIMAX_PRESETS,
|
||
"default_divisible_by": 32,
|
||
},
|
||
"LTX": {
|
||
"aspect_order": ("1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9"),
|
||
"fallback_aspect": "1:1",
|
||
"presets": LTX_PRESETS,
|
||
"default_divisible_by": 32,
|
||
},
|
||
}
|
||
|
||
|
||
def _normalize_text(value: str) -> str:
|
||
return re.sub(r"\s+", " ", re.sub(r"[\(\)]", "", value or "").lower()).strip()
|
||
|
||
|
||
def _parse_size(value: str) -> Optional[Tuple[int, int]]:
|
||
m = re.search(r"(\d+)\s*[x×]\s*(\d+)", value or "", flags=re.IGNORECASE)
|
||
if not m:
|
||
return None
|
||
return int(m.group(1)), int(m.group(2))
|
||
|
||
|
||
def _parse_aspect_ratio_value(aspect_ratio: str) -> Optional[float]:
|
||
m = re.match(r"^\s*(\d+)\s*:\s*(\d+)\s*$", aspect_ratio or "")
|
||
if not m:
|
||
return None
|
||
w = max(1, int(m.group(1)))
|
||
h = max(1, int(m.group(2)))
|
||
return float(w) / float(h)
|
||
|
||
|
||
def _best_aspect_ratio(width: int, height: int, aspect_order: Tuple[str, ...], fallback: str) -> str:
|
||
if width <= 0 or height <= 0:
|
||
return fallback
|
||
target = float(width) / float(height)
|
||
best_ar = fallback
|
||
best_delta = float("inf")
|
||
for ar in aspect_order:
|
||
ratio = _parse_aspect_ratio_value(ar)
|
||
if ratio is None:
|
||
continue
|
||
delta = abs(ratio - target)
|
||
if delta < best_delta:
|
||
best_delta = delta
|
||
best_ar = ar
|
||
return best_ar
|
||
|
||
|
||
def _rows_for(presets: Dict[str, List[Tuple[int, int, str]]], aspect_ratio: str, fallback: str) -> List[Tuple[int, int, str]]:
|
||
return presets.get(aspect_ratio) or presets.get(fallback) or list(next(iter(presets.values())))
|
||
|
||
|
||
def _labels_for(presets: Dict[str, List[Tuple[int, int, str]]], aspect_ratio: str, fallback: str) -> List[str]:
|
||
rows = _rows_for(presets, aspect_ratio, fallback)
|
||
return [f"{note} — {w}×{h}" for w, h, note in rows]
|
||
|
||
|
||
def _tier_index_for_value(presets: Dict, aspect_ratio: str, value: str, fallback: str) -> Optional[int]:
|
||
rows = _rows_for(presets, aspect_ratio, fallback)
|
||
normalized_value = _normalize_text(value)
|
||
|
||
note_matches = sorted(
|
||
enumerate(rows),
|
||
key=lambda indexed_row: len(_normalize_text(indexed_row[1][2])),
|
||
reverse=True,
|
||
)
|
||
for i, (_, _, note) in note_matches:
|
||
if _normalize_text(note) in normalized_value:
|
||
return i
|
||
|
||
m = re.match(r"\s*(\d+)\s*[\.\)]", value or "")
|
||
if m:
|
||
idx = int(m.group(1)) - 1
|
||
return max(0, min(idx, len(rows) - 1))
|
||
|
||
parsed_size = _parse_size(value)
|
||
if parsed_size is not None:
|
||
width, height = parsed_size
|
||
for i, (row_w, row_h, _) in enumerate(rows):
|
||
if row_w == width and row_h == height:
|
||
return i
|
||
|
||
return None
|
||
|
||
|
||
def _parse_resolution(presets: Dict, aspect_ratio: str, resolution_label: str, fallback: str) -> Tuple[int, int]:
|
||
idx = _tier_index_for_value(presets, aspect_ratio, resolution_label, fallback)
|
||
if idx is not None:
|
||
w, h, _ = _rows_for(presets, aspect_ratio, fallback)[idx]
|
||
return w, h
|
||
parsed_size = _parse_size(resolution_label)
|
||
if parsed_size is not None:
|
||
return parsed_size
|
||
w, h, _ = _rows_for(presets, aspect_ratio, fallback)[0]
|
||
return w, h
|
||
|
||
|
||
def _round_to_multiple(value: int, multiple: int) -> int:
|
||
if multiple <= 0:
|
||
return int(value)
|
||
return max(multiple, round(value / multiple) * multiple)
|
||
|
||
|
||
def _adaptive_size(width: int, height: int, target_pixels: int, divisible_by: int) -> Tuple[int, int]:
|
||
if width <= 0 or height <= 0:
|
||
return (divisible_by, divisible_by)
|
||
ratio = float(width) / float(height)
|
||
ideal_width = math.sqrt(float(target_pixels) * ratio)
|
||
ideal_height = math.sqrt(float(target_pixels) / ratio)
|
||
return (
|
||
max(divisible_by, round(ideal_width / divisible_by) * divisible_by),
|
||
max(divisible_by, round(ideal_height / divisible_by) * divisible_by),
|
||
)
|
||
|
||
|
||
def _image_dimensions(image) -> Optional[Tuple[int, int]]:
|
||
if image is None:
|
||
return None
|
||
if isinstance(image, (list, tuple)):
|
||
if not image:
|
||
return None
|
||
image = image[0]
|
||
shape = getattr(image, "shape", None)
|
||
if shape is None:
|
||
return None
|
||
dims = tuple(int(v) for v in shape)
|
||
if len(dims) == 4:
|
||
# Standard ComfyUI layout is B,H,W,C — but some custom loaders emit
|
||
# channels-first B,C,H,W tensors. Detect by which dim looks like C.
|
||
if dims[-1] in (1, 3, 4):
|
||
_, h, w, _ = dims
|
||
return (w, h)
|
||
if dims[1] in (1, 3, 4):
|
||
_, _, h, w = dims
|
||
return (w, h)
|
||
_, h, w, _ = dims
|
||
return (w, h)
|
||
if len(dims) == 3:
|
||
if dims[-1] in (1, 3, 4):
|
||
h, w, _ = dims
|
||
return (w, h)
|
||
if dims[0] in (1, 3, 4):
|
||
_, h, w = dims
|
||
return (w, h)
|
||
h, w, _ = dims
|
||
return (w, h)
|
||
if len(dims) == 2:
|
||
h, w = dims
|
||
return (w, h)
|
||
return None
|
||
|
||
|
||
def _hex_to_rgb(hex_color: str) -> Tuple[int, int, int]:
|
||
hex_color = hex_color.lstrip("#")
|
||
if len(hex_color) == 6:
|
||
return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
|
||
return (0, 0, 0)
|
||
|
||
|
||
def _resize_tensor(
|
||
tensor: torch.Tensor,
|
||
target_w: int,
|
||
target_h: int,
|
||
upscale: bool,
|
||
upscale_method: str,
|
||
small_image_mode: str,
|
||
pad_color: str,
|
||
is_mask: bool = False,
|
||
) -> torch.Tensor:
|
||
method_map = {
|
||
"lanczos": Image.LANCZOS,
|
||
"bilinear": Image.BILINEAR,
|
||
"bicubic": Image.BICUBIC,
|
||
"nearest": Image.NEAREST,
|
||
}
|
||
resize_method = method_map.get(upscale_method, Image.LANCZOS)
|
||
mask_resize_method = Image.NEAREST if upscale_method == "nearest" else resize_method
|
||
|
||
results = []
|
||
is_single_channel = len(tensor.shape) == 3
|
||
for i in range(tensor.shape[0]):
|
||
arr = tensor[i].cpu().numpy()
|
||
if is_mask:
|
||
arr = (arr * 255).astype(np.uint8)
|
||
else:
|
||
arr = (arr * 255).astype(np.uint8)
|
||
mode = "L" if is_mask or is_single_channel else "RGB"
|
||
pil_img = Image.fromarray(arr, mode=mode)
|
||
|
||
orig_w, orig_h = pil_img.width, pil_img.height
|
||
is_upscale = target_w > orig_w or target_h > orig_h
|
||
do_resize = not is_upscale or (is_upscale and upscale)
|
||
|
||
if do_resize:
|
||
if small_image_mode != "none" and (orig_w < target_w or orig_h < target_h):
|
||
target_ar = target_w / target_h
|
||
img_ar = orig_w / orig_h
|
||
if small_image_mode == "crop":
|
||
if img_ar > target_ar:
|
||
tmp_h = target_h
|
||
tmp_w = int(tmp_h * img_ar)
|
||
else:
|
||
tmp_w = target_w
|
||
tmp_h = int(tmp_w / img_ar)
|
||
pil_img = pil_img.resize((tmp_w, tmp_h), mask_resize_method if is_mask else resize_method)
|
||
left = (pil_img.width - target_w) // 2
|
||
top = (pil_img.height - target_h) // 2
|
||
pil_img = pil_img.crop((left, top, left + target_w, top + target_h))
|
||
elif small_image_mode == "pad":
|
||
pil_img.thumbnail((target_w, target_h), mask_resize_method if is_mask else resize_method)
|
||
fill = 0 if is_mask else _hex_to_rgb(pad_color)
|
||
bg = Image.new(mode, (target_w, target_h), fill)
|
||
offset = ((target_w - pil_img.width) // 2, (target_h - pil_img.height) // 2)
|
||
bg.paste(pil_img, offset)
|
||
pil_img = bg
|
||
else:
|
||
pil_img = pil_img.resize((target_w, target_h), mask_resize_method if is_mask else resize_method)
|
||
|
||
arr_out = np.array(pil_img).astype(np.float32) / 255.0
|
||
if is_mask:
|
||
if len(arr_out.shape) == 3:
|
||
arr_out = arr_out[..., 0]
|
||
results.append(arr_out)
|
||
|
||
out = torch.from_numpy(np.stack(results)).to(tensor.device)
|
||
return out
|
||
|
||
|
||
class ComfyUI_PerfectVideoResolution:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
default_model = "WAN 2.2"
|
||
config = MODEL_CONFIG[default_model]
|
||
aspect_choices = list(config["aspect_order"])
|
||
default_ar = config["fallback_aspect"]
|
||
resolution_choices = _labels_for(config["presets"], default_ar, default_ar)
|
||
default_res = resolution_choices[0]
|
||
|
||
return {
|
||
"required": {
|
||
"model": (["WAN 2.2", "MiniMax H3", "LTX"], {"default": default_model}),
|
||
"aspect_ratio": (aspect_choices, {"default": default_ar}),
|
||
"resolution": (resolution_choices, {"default": default_res}),
|
||
"divisible_by": ("INT", {"default": 16, "min": 1, "max": 128, "step": 1}),
|
||
"upscale": ("BOOLEAN", {"default": False}),
|
||
"upscale_method": (["lanczos", "bilinear", "bicubic", "nearest"], {"default": "lanczos"}),
|
||
"small_image_mode": (["none", "crop", "pad"], {"default": "none"}),
|
||
"pad_color": ("STRING", {"default": "#000000"}),
|
||
"use_image_aspect": ("BOOLEAN", {"default": True}),
|
||
},
|
||
"optional": {
|
||
"image": ("IMAGE",),
|
||
"image2": ("IMAGE",),
|
||
"mask": ("MASK",),
|
||
},
|
||
"hidden": {
|
||
"unique_id": "UNIQUE_ID",
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("INT", "INT", "IMAGE", "IMAGE", "MASK", "STRING")
|
||
RETURN_NAMES = ("width", "height", "IMAGE", "IMAGE2", "MASK", "resolution_info")
|
||
FUNCTION = "calculate"
|
||
CATEGORY = "utils"
|
||
|
||
@classmethod
|
||
def VALIDATE_INPUTS(cls, resolution=None):
|
||
return True
|
||
|
||
def calculate(
|
||
self,
|
||
model: str,
|
||
aspect_ratio: str,
|
||
resolution: str,
|
||
divisible_by: int,
|
||
upscale: bool,
|
||
upscale_method: str,
|
||
small_image_mode: str,
|
||
pad_color: str,
|
||
use_image_aspect: bool,
|
||
image=None,
|
||
image2=None,
|
||
mask=None,
|
||
unique_id=None,
|
||
):
|
||
config = MODEL_CONFIG.get(model, MODEL_CONFIG["WAN 2.2"])
|
||
presets = config["presets"]
|
||
aspect_order = config["aspect_order"]
|
||
fallback_aspect = config["fallback_aspect"]
|
||
default_divisible_by = config["default_divisible_by"]
|
||
|
||
if divisible_by <= 0:
|
||
divisible_by = default_divisible_by
|
||
|
||
# Prefer the primary image input for aspect detection, but fall back
|
||
# to image2 so a loader wired only there still drives the aspect.
|
||
image_dims = _image_dimensions(image) or _image_dimensions(image2)
|
||
resolved_aspect = aspect_ratio
|
||
resolved_resolution_label = resolution
|
||
target_w, target_h = _parse_resolution(presets, aspect_ratio, resolution, fallback_aspect)
|
||
|
||
if image_dims is not None and use_image_aspect:
|
||
image_w, image_h = image_dims
|
||
resolved_aspect = _best_aspect_ratio(image_w, image_h, aspect_order, fallback_aspect)
|
||
|
||
if model == "MiniMax H3":
|
||
idx = _tier_index_for_value(presets, aspect_ratio, resolution, fallback_aspect)
|
||
if idx is None:
|
||
parsed = _parse_size(resolution)
|
||
target_pixels = parsed[0] * parsed[1] if parsed else MINIMAX_TARGET_PIXELS[0]
|
||
else:
|
||
target_pixels = MINIMAX_TARGET_PIXELS[idx]
|
||
target_w, target_h = _adaptive_size(image_w, image_h, target_pixels, divisible_by)
|
||
resolved_resolution_label = _labels_for(presets, resolved_aspect, fallback_aspect)[
|
||
idx if idx is not None else 0
|
||
]
|
||
else:
|
||
target_w, target_h = _parse_resolution(presets, resolved_aspect, resolution, fallback_aspect)
|
||
resolved_resolution_label = _labels_for(presets, resolved_aspect, fallback_aspect)[
|
||
_tier_index_for_value(presets, resolved_aspect, resolution, fallback_aspect) or 0
|
||
]
|
||
|
||
target_w = int(_round_to_multiple(target_w, divisible_by))
|
||
target_h = int(_round_to_multiple(target_h, divisible_by))
|
||
|
||
if image is not None:
|
||
image_out = _resize_tensor(
|
||
image, target_w, target_h, upscale, upscale_method, small_image_mode, pad_color, is_mask=False
|
||
)
|
||
else:
|
||
image_out = torch.zeros((1, target_h, target_w, 3), dtype=torch.float32)
|
||
|
||
if image2 is not None:
|
||
image2_out = _resize_tensor(
|
||
image2, target_w, target_h, upscale, upscale_method, small_image_mode, pad_color, is_mask=False
|
||
)
|
||
else:
|
||
image2_out = torch.zeros((1, target_h, target_w, 3), dtype=torch.float32)
|
||
|
||
if mask is not None:
|
||
mask_out = _resize_tensor(
|
||
mask, target_w, target_h, upscale, upscale_method, small_image_mode, pad_color, is_mask=True
|
||
)
|
||
else:
|
||
mask_out = torch.ones((1, target_h, target_w), dtype=torch.float32)
|
||
|
||
approx_mb = (target_w * target_h * 3) / (1024 * 1024)
|
||
resolution_info = f"{target_w}x{target_h} | {approx_mb:.2f}MB | {target_w*target_h:,} pixels"
|
||
|
||
if unique_id and image is not None:
|
||
try:
|
||
from server import PromptServer
|
||
memory_size_mb = (image_out.numel() * image_out.element_size()) / (1024 * 1024)
|
||
PromptServer.instance.send_progress_text(
|
||
f"<tr><td>Output: </td>"
|
||
f"<td><b>{target_w}</b> x <b>{target_h}</b> | {memory_size_mb:.2f}MB | {target_w*target_h:,} pixels</td></tr>",
|
||
unique_id,
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
result = (int(target_w), int(target_h), image_out, image2_out, mask_out, resolution_info)
|
||
|
||
if image_dims is None:
|
||
return result
|
||
|
||
state = {
|
||
"model": model,
|
||
"aspect_ratio": resolved_aspect,
|
||
"resolution": resolved_resolution_label,
|
||
"source_width": image_dims[0],
|
||
"source_height": image_dims[1],
|
||
}
|
||
return {
|
||
"ui": {"perfect_video_resolution_state": [state]},
|
||
"result": result,
|
||
}
|