446 lines
14 KiB
Python
446 lines
14 KiB
Python
"""Canvas sizes and temporal window for tiled-fusion graphs.
|
||
|
||
Stage 1 aims at full HD in the source aspect without spatial tiling. Stage 2
|
||
is 2x that (~4K). 8K is another 2x. HD is the 720p step of the same ladder;
|
||
qHD is the 960 long-edge step.
|
||
|
||
Dimensions stay multiples of 32. ``tile_frames`` stays 8n+1, targeting 97,
|
||
stretching up to 121 only when that lets a slightly longer clip skip chunking.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
from typing import NamedTuple, Optional
|
||
|
||
from .nodes_registry import comfy_node
|
||
from .tiled_fusion_plan import plan_grid, snap_frames, temporal_plan
|
||
|
||
SPATIAL_ALIGN = 32
|
||
FRAME_STEP = 8
|
||
SLACK_PX = 200
|
||
QHD_LONG_EDGE = 960
|
||
HD_LONG_EDGE = 1280
|
||
FULLHD_LONG_EDGE = 1920
|
||
TILE_FRAMES_TARGET = 97
|
||
TILE_FRAMES_MAX = 121
|
||
OVERLAP_FRAC = 0.5
|
||
|
||
TILE_CHOICES = ("qHD", "HD", "FullHD")
|
||
OUTPUT_CHOICES = ("FullHD", "4K", "8K")
|
||
|
||
_QHD_KEYS = frozenset({"qhd", "quarterhd", "quarter", "960"})
|
||
_HD_KEYS = frozenset({"hd", "720p", "720"})
|
||
_FULLHD_KEYS = frozenset({"fullhd", "fhd", "1080p", "1080"})
|
||
_4K_KEYS = frozenset({"4k", "uhd", "2160p", "2160"})
|
||
_8K_KEYS = frozenset({"8k", "eightk", "4320p", "4320"})
|
||
|
||
|
||
class TilingSizes(NamedTuple):
|
||
qhd_width: int
|
||
qhd_height: int
|
||
hd_width: int
|
||
hd_height: int
|
||
fullhd_width: int
|
||
fullhd_height: int
|
||
overlap_frac: float
|
||
tile_frames: int
|
||
width_4k: int
|
||
height_4k: int
|
||
width_8k: int
|
||
height_8k: int
|
||
|
||
|
||
class TilingPlan(NamedTuple):
|
||
initial_canvas_width: int
|
||
initial_canvas_height: int
|
||
tile_width: int
|
||
tile_height: int
|
||
output_width: int
|
||
output_height: int
|
||
overlap_frac: float
|
||
tile_frames: int
|
||
spatial_tiles_x: int
|
||
spatial_tiles_y: int
|
||
output_tiles_x: int
|
||
output_tiles_y: int
|
||
temporal_chunks: int
|
||
canvas_label: str
|
||
tile_label: str
|
||
output_label: str
|
||
tile_clamped: bool
|
||
is_8k: bool
|
||
summary: str
|
||
|
||
|
||
def _norm_label(value: str) -> str:
|
||
return str(value or "").lower().replace(" ", "").replace("_", "").replace("-", "")
|
||
|
||
|
||
def snap_spatial(value: float, align: int = SPATIAL_ALIGN) -> int:
|
||
"""Nearest multiple of ``align``, ties round up, never below ``align``."""
|
||
value = max(float(align), float(value))
|
||
return max(align, int(math.floor(value / align + 0.5) * align))
|
||
|
||
|
||
def tile_frames_for(frame_count: int) -> int:
|
||
"""Temporal window: 97 or the clip itself, stretching to 121 to skip chunks."""
|
||
n = max(1, int(frame_count))
|
||
if n <= TILE_FRAMES_TARGET:
|
||
tf = snap_frames(n)
|
||
return min(TILE_FRAMES_TARGET, tf)
|
||
if n <= TILE_FRAMES_MAX:
|
||
covering = ((n - 1 + FRAME_STEP - 1) // FRAME_STEP) * FRAME_STEP + 1
|
||
return min(TILE_FRAMES_MAX, covering)
|
||
return TILE_FRAMES_TARGET
|
||
|
||
|
||
def _fit_long_edge(width: float, height: float, long_edge: int) -> tuple[float, float]:
|
||
long = max(width, height)
|
||
scale = long_edge / long
|
||
return width * scale, height * scale
|
||
|
||
|
||
def _keep_source_if_close(
|
||
target_w: float, target_h: float, src_w: int, src_h: int, slack: int = SLACK_PX
|
||
) -> tuple[float, float]:
|
||
if abs(src_w - target_w) <= slack and abs(src_h - target_h) <= slack:
|
||
return float(src_w), float(src_h)
|
||
return target_w, target_h
|
||
|
||
|
||
def _aligned_size(width: float, height: float) -> tuple[int, int]:
|
||
return snap_spatial(width), snap_spatial(height)
|
||
|
||
|
||
def _ladder_step(src_w: int, src_h: int, long_edge: int) -> tuple[int, int]:
|
||
tw, th = _fit_long_edge(src_w, src_h, long_edge)
|
||
tw, th = _keep_source_if_close(tw, th, src_w, src_h)
|
||
return _aligned_size(tw, th)
|
||
|
||
|
||
def compute_tiling_sizes(width: int, height: int, frame_count: int) -> TilingSizes:
|
||
"""Return the qHD / HD / full-HD / 4K / 8K ladder plus overlap and tile_frames."""
|
||
src_w = max(1, int(width))
|
||
src_h = max(1, int(height))
|
||
|
||
qhd_w, qhd_h = _ladder_step(src_w, src_h, QHD_LONG_EDGE)
|
||
hd_w, hd_h = _ladder_step(src_w, src_h, HD_LONG_EDGE)
|
||
fhd_w, fhd_h = _ladder_step(src_w, src_h, FULLHD_LONG_EDGE)
|
||
|
||
k4_w, k4_h = _keep_source_if_close(fhd_w * 2, fhd_h * 2, src_w, src_h)
|
||
k4_w, k4_h = _aligned_size(k4_w, k4_h)
|
||
|
||
k8_w, k8_h = _keep_source_if_close(k4_w * 2, k4_h * 2, src_w, src_h)
|
||
k8_w, k8_h = _aligned_size(k8_w, k8_h)
|
||
|
||
return TilingSizes(
|
||
qhd_width=qhd_w,
|
||
qhd_height=qhd_h,
|
||
hd_width=hd_w,
|
||
hd_height=hd_h,
|
||
fullhd_width=fhd_w,
|
||
fullhd_height=fhd_h,
|
||
overlap_frac=OVERLAP_FRAC,
|
||
tile_frames=tile_frames_for(frame_count),
|
||
width_4k=k4_w,
|
||
height_4k=k4_h,
|
||
width_8k=k8_w,
|
||
height_8k=k8_h,
|
||
)
|
||
|
||
|
||
def _pick_tile_or_canvas(sizes: TilingSizes, choice: str) -> tuple[int, int, str]:
|
||
key = _norm_label(choice)
|
||
if key in _QHD_KEYS:
|
||
return sizes.qhd_width, sizes.qhd_height, "qHD"
|
||
if key in _HD_KEYS:
|
||
return sizes.hd_width, sizes.hd_height, "HD"
|
||
if key in _FULLHD_KEYS:
|
||
return sizes.fullhd_width, sizes.fullhd_height, "FullHD"
|
||
raise ValueError(
|
||
f"Unknown tile/canvas size {choice!r}; expected qHD, HD, or FullHD."
|
||
)
|
||
|
||
|
||
def _pick_output(sizes: TilingSizes, choice: str) -> tuple[int, int, str]:
|
||
key = _norm_label(choice)
|
||
if key in _8K_KEYS:
|
||
return sizes.width_8k, sizes.height_8k, "8K"
|
||
if key in _4K_KEYS:
|
||
return sizes.width_4k, sizes.height_4k, "4K"
|
||
if key in _FULLHD_KEYS:
|
||
return sizes.fullhd_width, sizes.fullhd_height, "FullHD"
|
||
raise ValueError(f"Unknown output size {choice!r}; expected FullHD, 4K, or 8K.")
|
||
|
||
|
||
def _latent_wh(width: int, height: int) -> tuple[int, int]:
|
||
return max(1, int(width) // SPATIAL_ALIGN), max(1, int(height) // SPATIAL_ALIGN)
|
||
|
||
|
||
def spatial_tile_grid(
|
||
canvas_w: int,
|
||
canvas_h: int,
|
||
tile_w: int,
|
||
tile_h: int,
|
||
overlap_frac: float = OVERLAP_FRAC,
|
||
) -> tuple[int, int]:
|
||
"""Number of spatial tiles covering the canvas, matching the sampler grid."""
|
||
cw, ch = _latent_wh(canvas_w, canvas_h)
|
||
tw, th = _latent_wh(tile_w, tile_h)
|
||
tw, th = min(tw, cw), min(th, ch)
|
||
min_ov_x = max(2, int(tw * overlap_frac))
|
||
min_ov_y = max(2, int(th * overlap_frac))
|
||
return (
|
||
len(plan_grid(cw, tw, min_ov_x)),
|
||
len(plan_grid(ch, th, min_ov_y)),
|
||
)
|
||
|
||
|
||
def temporal_chunk_count(frame_count: int, tile_frames: int) -> int:
|
||
"""Number of temporal windows, matching the sampler / guide handshake."""
|
||
n = max(1, int(frame_count))
|
||
tf = max(1, int(tile_frames))
|
||
n_lat = (n - 1) // FRAME_STEP + 1
|
||
tf_lat = min(n_lat, (tf - 1) // FRAME_STEP + 1)
|
||
return len(temporal_plan(n_lat, tf_lat))
|
||
|
||
|
||
def format_summary(
|
||
*,
|
||
src_w: int,
|
||
src_h: int,
|
||
frame_count: int,
|
||
plan_fields: dict,
|
||
) -> str:
|
||
sx, sy = plan_fields["spatial_tiles_x"], plan_fields["spatial_tiles_y"]
|
||
ox, oy = plan_fields["output_tiles_x"], plan_fields["output_tiles_y"]
|
||
chunks = plan_fields["temporal_chunks"]
|
||
tile_note = " (clamped to canvas)" if plan_fields["tile_clamped"] else ""
|
||
if (sx, sy) == (ox, oy):
|
||
spatial = f"spatial {sx}×{sy} tiles"
|
||
else:
|
||
spatial = f"spatial {sx}×{sy} on canvas, {ox}×{oy} on output"
|
||
chunk_word = "chunk" if chunks == 1 else "chunks"
|
||
return "\n".join(
|
||
(
|
||
f"source {src_w}×{src_h} × {frame_count}f",
|
||
f"canvas {plan_fields['initial_canvas_width']}×{plan_fields['initial_canvas_height']}"
|
||
f" {plan_fields['canvas_label']}",
|
||
f"tile {plan_fields['tile_width']}×{plan_fields['tile_height']}"
|
||
f" {plan_fields['tile_label']}{tile_note}",
|
||
f"output {plan_fields['output_width']}×{plan_fields['output_height']}"
|
||
f" {plan_fields['output_label']}",
|
||
f"overlap {plan_fields['overlap_frac']:g}",
|
||
f"tile_frames {plan_fields['tile_frames']}",
|
||
spatial,
|
||
f"temporal {chunks} {chunk_word}",
|
||
)
|
||
)
|
||
|
||
|
||
def resolve_tiling_plan(
|
||
width: int,
|
||
height: int,
|
||
frame_count: int,
|
||
tile_size: str = "FullHD",
|
||
initial_canvas_size: str = "FullHD",
|
||
output_size: str = "FullHD",
|
||
) -> TilingPlan:
|
||
"""Pick canvas / tile / output from the ladder and count tiles and chunks."""
|
||
src_w = max(1, int(width))
|
||
src_h = max(1, int(height))
|
||
frames = max(1, int(frame_count))
|
||
sizes = compute_tiling_sizes(src_w, src_h, frames)
|
||
|
||
canvas_w, canvas_h, canvas_label = _pick_tile_or_canvas(sizes, initial_canvas_size)
|
||
tile_w, tile_h, tile_label = _pick_tile_or_canvas(sizes, tile_size)
|
||
out_w, out_h, out_label = _pick_output(sizes, output_size)
|
||
|
||
clamped_w, clamped_h = min(tile_w, canvas_w), min(tile_h, canvas_h)
|
||
tile_clamped = (clamped_w, clamped_h) != (tile_w, tile_h)
|
||
tile_w, tile_h = clamped_w, clamped_h
|
||
|
||
sx, sy = spatial_tile_grid(canvas_w, canvas_h, tile_w, tile_h, sizes.overlap_frac)
|
||
ox, oy = spatial_tile_grid(out_w, out_h, tile_w, tile_h, sizes.overlap_frac)
|
||
chunks = temporal_chunk_count(frames, sizes.tile_frames)
|
||
|
||
fields = dict(
|
||
initial_canvas_width=canvas_w,
|
||
initial_canvas_height=canvas_h,
|
||
tile_width=tile_w,
|
||
tile_height=tile_h,
|
||
output_width=out_w,
|
||
output_height=out_h,
|
||
overlap_frac=sizes.overlap_frac,
|
||
tile_frames=sizes.tile_frames,
|
||
spatial_tiles_x=sx,
|
||
spatial_tiles_y=sy,
|
||
output_tiles_x=ox,
|
||
output_tiles_y=oy,
|
||
temporal_chunks=chunks,
|
||
canvas_label=canvas_label,
|
||
tile_label=tile_label,
|
||
output_label=out_label,
|
||
tile_clamped=tile_clamped,
|
||
is_8k=out_label == "8K",
|
||
)
|
||
summary = format_summary(
|
||
src_w=src_w, src_h=src_h, frame_count=frames, plan_fields=fields
|
||
)
|
||
return TilingPlan(summary=summary, **fields)
|
||
|
||
|
||
def _dims_from_image(image) -> tuple[int, int, int]:
|
||
"""Comfy IMAGE is [frames, height, width, channels]."""
|
||
frames, height, width = (
|
||
int(image.shape[0]),
|
||
int(image.shape[1]),
|
||
int(image.shape[2]),
|
||
)
|
||
return width, height, frames
|
||
|
||
|
||
@comfy_node(
|
||
name="LTXVGetTilingSizes",
|
||
description="Get Tiling Sizes",
|
||
)
|
||
class LTXVGetTilingSizes:
|
||
"""Size calculator in the spirit of Comfy's Get Image Size, for tiled fusion."""
|
||
|
||
CATEGORY = "ltxvideo"
|
||
FUNCTION = "get_sizes"
|
||
OUTPUT_NODE = True
|
||
RETURN_TYPES = (
|
||
"INT",
|
||
"INT",
|
||
"INT",
|
||
"INT",
|
||
"INT",
|
||
"INT",
|
||
"FLOAT",
|
||
"INT",
|
||
"BOOLEAN",
|
||
"STRING",
|
||
)
|
||
RETURN_NAMES = (
|
||
"initial_canvas_width",
|
||
"initial_canvas_height",
|
||
"tile_width",
|
||
"tile_height",
|
||
"output_width",
|
||
"output_height",
|
||
"overlap_frac",
|
||
"tile_frames",
|
||
"is_8k",
|
||
"summary",
|
||
)
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"width": (
|
||
"INT",
|
||
{
|
||
"default": 1920,
|
||
"min": 1,
|
||
"max": 16384,
|
||
"tooltip": "Source width in pixels (from Load Video / Get Image Size).",
|
||
},
|
||
),
|
||
"height": (
|
||
"INT",
|
||
{
|
||
"default": 1080,
|
||
"min": 1,
|
||
"max": 16384,
|
||
"tooltip": "Source height in pixels.",
|
||
},
|
||
),
|
||
"frame_count": (
|
||
"INT",
|
||
{
|
||
"default": 97,
|
||
"min": 1,
|
||
"max": 10000,
|
||
"tooltip": "Source frame count. tile_frames is 8n+1, 97 or less, "
|
||
"or up to 121 when that avoids temporal chunking.",
|
||
},
|
||
),
|
||
"tile_size": (
|
||
list(TILE_CHOICES),
|
||
{
|
||
"default": "FullHD",
|
||
"tooltip": "Spatial tile long-edge. qHD=960, HD=1280, FullHD=1920. "
|
||
"Clamped to the initial canvas so a FullHD tile on a smaller "
|
||
"canvas does not force tiling.",
|
||
},
|
||
),
|
||
"initial_canvas_size": (
|
||
list(TILE_CHOICES),
|
||
{
|
||
"default": "FullHD",
|
||
"tooltip": "Stage-1 canvas long-edge (same choices as tile size). "
|
||
"Source within 200px of FullHD is kept instead of rescaled.",
|
||
},
|
||
),
|
||
"output_size": (
|
||
list(OUTPUT_CHOICES),
|
||
{
|
||
"default": "FullHD",
|
||
"tooltip": "Final canvas. FullHD is the 1920 long-edge step; "
|
||
"4K is 2× that; 8K is 4×. Single-stage graphs use this as the "
|
||
"canvas. Two-stage graphs use initial canvas for stage 1 and "
|
||
"this for stage 2.",
|
||
},
|
||
),
|
||
},
|
||
"optional": {
|
||
"image": (
|
||
"IMAGE",
|
||
{
|
||
"tooltip": "Optional video/image batch. When connected, width, "
|
||
"height, and frame_count are taken from it (Get Image Size style)."
|
||
},
|
||
),
|
||
},
|
||
}
|
||
|
||
def get_sizes(
|
||
self,
|
||
width,
|
||
height,
|
||
frame_count,
|
||
tile_size="FullHD",
|
||
initial_canvas_size="FullHD",
|
||
output_size="FullHD",
|
||
image: Optional[object] = None,
|
||
**_kwargs,
|
||
):
|
||
if image is not None:
|
||
width, height, frame_count = _dims_from_image(image)
|
||
plan = resolve_tiling_plan(
|
||
width,
|
||
height,
|
||
frame_count,
|
||
tile_size=tile_size,
|
||
initial_canvas_size=initial_canvas_size,
|
||
output_size=output_size,
|
||
)
|
||
return {
|
||
"ui": {"text": [plan.summary]},
|
||
"result": (
|
||
plan.initial_canvas_width,
|
||
plan.initial_canvas_height,
|
||
plan.tile_width,
|
||
plan.tile_height,
|
||
plan.output_width,
|
||
plan.output_height,
|
||
plan.overlap_frac,
|
||
plan.tile_frames,
|
||
plan.is_8k,
|
||
plan.summary,
|
||
),
|
||
}
|