from functools import lru_cache import gc import glob import math import os import re import comfy.samplers import folder_paths import torch try: import torch.nn as nn import torch.nn.functional as F except Exception: # pragma: no cover - import-time fallback for the test shim nn = None F = None try: from comfy.ldm.minimax.model import FRAME_PER_TOKEN except Exception: # pragma: no cover - import-time fallback for the test shim FRAME_PER_TOKEN = (1, 4, 4, 4, 4) LATENTS_MEAN = [ 0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075, -0.2727581858634949, -1.3675414323806763, -0.2553254961967468, -0.26907554268836975, -0.5376840829849243, -0.0464097298681736, 0.6657370328903198, 0.19690127670764923, -0.5460608005523682, -0.4035342037677765, -0.23683024942874908, 0.25928452610969543, -0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279, -0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264, ] LATENTS_STD = [ 1.2223774194717407, 1.2767263650894165, 1.6831774711608887, 1.7549455165863037, 1.5636216402053833, 2.194143533706665, 0.9653137922286987, 1.0569885969161987, 0.841948926448822, 0.7729952931404114, 1.8955937623977661, 0.946841835975647, 0.7996809482574463, 0.44988900423049927, 0.7197399735450745, 0.6936293244361877, 2.961095094680786, 2.7694199085235596, 3.0496184825897217, 2.1088054180265264, 3.276226282119751, 3.1627357006073, 2.2816812992095947, 2.6127843856811523, ] _LATENT_UPSCALE_FOLDER = "latent_upscale_models" MP_UNIT = 1024 * 1024 RES_MULTIPLE = 32 def _models_dir(): try: if _LATENT_UPSCALE_FOLDER not in folder_paths.folder_names_and_paths: folder_paths.add_model_folder_path( _LATENT_UPSCALE_FOLDER, os.path.join(folder_paths.models_dir, _LATENT_UPSCALE_FOLDER), ) return folder_paths.get_folder_paths(_LATENT_UPSCALE_FOLDER)[0] except Exception: return os.path.join(getattr(folder_paths, "models_dir", ""), _LATENT_UPSCALE_FOLDER) def _scan_models(): try: model_dir = _models_dir() files = [] for ext in ("*.pth", "*.safetensors"): files.extend(glob.glob(os.path.join(model_dir, ext))) names = sorted(os.path.basename(path) for path in files) return ["none"] + names if names else [f"(no upscale models found in: {model_dir})"] except Exception: return ["none"] def _make_norm_tensors(device, dtype): mean = torch.tensor(LATENTS_MEAN, dtype=dtype, device=device).view(1, -1, 1, 1, 1) std = torch.tensor(LATENTS_STD, dtype=dtype, device=device).view(1, -1, 1, 1, 1) return mean, std if nn is not None: def _normalization(channels): return nn.GroupNorm(32, channels) def _zero_module(module): for p in module.parameters(): p.detach().zero_() return module class _AttnBlock3D(nn.Module): def __init__(self, in_channels): super().__init__() self.norm = _normalization(in_channels) self.q = nn.Conv3d(in_channels, in_channels, 1) self.k = nn.Conv3d(in_channels, in_channels, 1) self.v = nn.Conv3d(in_channels, in_channels, 1) self.proj_out = nn.Conv3d(in_channels, in_channels, 1) def forward(self, x): h = self.norm(x) b, c, t, hh, w = h.shape q = self.q(h).flatten(2).transpose(1, 2) k = self.k(h).flatten(2).transpose(1, 2) v = self.v(h).flatten(2).transpose(1, 2) h = F.scaled_dot_product_attention(q, k, v) h = h.transpose(1, 2).view(b, c, t, hh, w) return x + self.proj_out(h) class _ResBlockEmb3D(nn.Module): def __init__(self, channels, emb_channels, dropout=0, out_channels=None): super().__init__() self.out_channels = out_channels or channels self.in_layers = nn.Sequential( _normalization(channels), nn.SiLU(), nn.Conv3d(channels, self.out_channels, 3, padding=1), ) self.emb_layers = nn.Sequential( nn.SiLU(), nn.Linear(emb_channels, 2 * self.out_channels), ) self.out_norm = _normalization(self.out_channels) self.out_layers = nn.Sequential( nn.SiLU(), nn.Dropout(p=dropout), _zero_module(nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1)), ) self.skip = ( nn.Conv3d(channels, self.out_channels, 1) if self.out_channels != channels else nn.Identity() ) def forward(self, x, emb): h = self.in_layers(x) emb_out = self.emb_layers(emb).type(h.dtype) while len(emb_out.shape) < len(h.shape): emb_out = emb_out[..., None] scale, shift = torch.chunk(emb_out, 2, dim=1) h = self.out_norm(h) * (1 + scale) + shift h = self.out_layers(h) return self.skip(x) + h class _TemporalConv(nn.Module): def __init__(self, channels, kernel_size=5): super().__init__() padding = kernel_size // 2 self.norm = _normalization(channels) self.dwconv = nn.Conv3d( channels, channels, kernel_size=(kernel_size, 1, 1), padding=(padding, 0, 0), groups=channels, ) self.pwconv = nn.Conv3d(channels, channels, kernel_size=1) nn.init.zeros_(self.pwconv.weight) nn.init.zeros_(self.pwconv.bias) def forward(self, x): identity = x h = self.norm(x) h = F.silu(h) h = self.dwconv(h) h = self.pwconv(h) return identity + h class _LatentResizer3D(nn.Module): def __init__(self, in_channels=24, in_blocks=12, out_blocks=12, channels=512, dropout=0.1, attn=False, temporal_every=2, temporal_kernel=5): super().__init__() self.conv_in = nn.Conv3d(in_channels, channels, 3, padding=1) embed_dim = 64 self.embed = nn.Sequential( nn.Linear(1, embed_dim), nn.SiLU(), nn.Linear(embed_dim, embed_dim)) self.in_blocks = nn.ModuleList() for b in range(in_blocks): if (b == 1 or b == in_blocks - 1) and attn: self.in_blocks.append(_AttnBlock3D(channels)) self.in_blocks.append(_ResBlockEmb3D(channels, embed_dim, dropout)) if temporal_every > 0 and b % temporal_every == 0: self.in_blocks.append(_TemporalConv(channels, temporal_kernel)) self.out_blocks = nn.ModuleList() for b in range(out_blocks): if (b == 1 or b == out_blocks - 1) and attn: self.out_blocks.append(_AttnBlock3D(channels)) self.out_blocks.append(_ResBlockEmb3D(channels, embed_dim, dropout)) if temporal_every > 0 and b % temporal_every == 0: self.out_blocks.append(_TemporalConv(channels, temporal_kernel)) self.norm_out = _normalization(channels) self.conv_out = nn.Conv3d(channels, in_channels, 3, padding=1) def forward(self, x, scale=None, target_size=None): if target_size is not None: size = target_size elif scale is not None: size = tuple(int(round(s * scale)) for s in x.shape[-3:]) else: return x if size == x.shape[-3:]: return x scale_emb = torch.tensor( [scale - 1 if scale is not None else 0.0], dtype=x.dtype, device=x.device).unsqueeze(0) emb = self.embed(scale_emb) x = self.conv_in(x) for block in self.in_blocks: if isinstance(block, _ResBlockEmb3D): x = block(x, emb.expand(x.shape[0], -1)) else: x = block(x) x = F.interpolate(x, size=size, mode="trilinear", align_corners=False) for block in self.out_blocks: if isinstance(block, _ResBlockEmb3D): x = block(x, emb.expand(x.shape[0], -1)) else: x = block(x) x = self.norm_out(x) x = F.silu(x) x = self.conv_out(x) return x else: # pragma: no cover - import-time fallback for the test shim _LatentResizer3D = None _MODEL_CACHE = {} def _load_raw_sd(path): if path.endswith(".safetensors"): from safetensors.torch import load_file sd = load_file(path, device="cpu") else: sd = torch.load(path, map_location="cpu", weights_only=False) if isinstance(sd, dict) and "model" in sd: sd = sd["model"] float8 = getattr(torch, "float8_e4m3fn", None) if float8 is not None: sd = {k: v.to(torch.float16) if getattr(v, "dtype", None) == float8 else v for k, v in sd.items()} return sd def _extract_upscaler_sd(sd): if any(k.startswith("upscaler.") for k in sd): return {k[len("upscaler."):]: v for k, v in sd.items() if k.startswith("upscaler.")} return sd def _detect_arch(sd): cfg = { "in_channels": 24, "in_blocks": 12, "out_blocks": 12, "channels": 512, "dropout": 0.1, "attn": False, "temporal_every": 2, "temporal_kernel": 5, } conv_key = "conv_in.weight" if conv_key in sd: cfg["in_channels"] = sd[conv_key].shape[1] cfg["channels"] = sd[conv_key].shape[0] in_ids, out_ids = set(), set() temporal_in_indices, temporal_out_indices = set(), set() for k in sd.keys(): m = re.match(r"in_blocks\.(\d+)\.in_layers\.", k) if m: in_ids.add(int(m.group(1))) m = re.match(r"out_blocks\.(\d+)\.in_layers\.", k) if m: out_ids.add(int(m.group(1))) m = re.match(r"in_blocks\.(\d+)\.dwconv\.weight", k) if m: temporal_in_indices.add(int(m.group(1))) m = re.match(r"out_blocks\.(\d+)\.dwconv\.weight", k) if m: temporal_out_indices.add(int(m.group(1))) if in_ids: cfg["in_blocks"] = len(in_ids) if out_ids: cfg["out_blocks"] = len(out_ids) if temporal_in_indices or temporal_out_indices: cfg["temporal_every"] = 2 for k in sd.keys(): if "dwconv.weight" in k and k.endswith("dwconv.weight"): cfg["temporal_kernel"] = sd[k].shape[2] break else: cfg["temporal_every"] = 0 cfg["attn"] = False return cfg def load_upscale_model(name, device, precision): if _LatentResizer3D is None: raise RuntimeError("latent upscaler requires torch.nn") cache_key = f"{name}::{device}::{precision}" if cache_key in _MODEL_CACHE: return _MODEL_CACHE[cache_key].to(device) path = os.path.join(_models_dir(), name) if not os.path.exists(path): raise FileNotFoundError(f"Model file not found: {path}") raw_sd = _load_raw_sd(path) up_sd = _extract_upscaler_sd(raw_sd) cfg = _detect_arch(up_sd) if cfg["in_channels"] != 24: raise ValueError( f"Checkpoint '{name}' is not an H3 latent upscaler (expected 24 input channels, got {cfg['in_channels']})." ) model = _LatentResizer3D( in_channels=cfg["in_channels"], in_blocks=cfg["in_blocks"], out_blocks=cfg["out_blocks"], channels=cfg["channels"], dropout=cfg["dropout"], attn=cfg["attn"], temporal_every=cfg["temporal_every"], temporal_kernel=cfg["temporal_kernel"], ) model.load_state_dict(up_sd, strict=True) dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}.get(precision, torch.float32) model = model.to(device).eval().requires_grad_(False) if dtype != torch.float32: model = model.to(dtype) _MODEL_CACHE[cache_key] = model return model def unload_upscale_model(name, device, precision): cache_key = f"{name}::{device}::{precision}" model = _MODEL_CACHE.get(cache_key) if model is not None and str(next(model.parameters()).device) != "cpu": model.to("cpu") if str(device) == "cuda" and hasattr(torch, "cuda"): try: torch.cuda.empty_cache() except Exception: pass def _compute_upscale_target(width, height, h_in, w_in): ds = 16 w_px = float(width) h_px = float(height) eff = (w_px / (w_in * ds) + h_px / (h_in * ds)) / 2.0 w_px_f = round(w_px / ds) * ds h_px_f = round(h_px / ds) * ds w_out = max(1, int(w_px_f // ds)) h_out = max(1, int(h_px_f // ds)) return h_out, w_out, eff def _scale_to_megapixels(w, h, mp, multiple=RES_MULTIPLE): if not mp or mp <= 0 or w <= 0 or h <= 0: return int(h), int(w) multiple = max(1, int(multiple)) scale = math.sqrt((float(mp) * MP_UNIT) / float(w * h)) nw = max(multiple, int(round(w * scale / multiple)) * multiple) nh = max(multiple, int(round(h * scale / multiple)) * multiple) return nh, nw def _resolve_target_size(param, h_in, w_in): width = int(param.get("width", 0) or 0) height = int(param.get("height", 0) or 0) megapixels = float(param.get("megapixels", 0.0) or 0.0) if width > 0 and height > 0: return height, width if megapixels > 0: return _scale_to_megapixels(w_in, h_in, megapixels) return int(h_in), int(w_in) def _frames_for_tokens(n): return sum(FRAME_PER_TOKEN[i % 5] for i in range(int(n))) def _snap_frame_boundary(f, max_tokens, phase=5): best_k, best_f, best_d = 0, 0, abs(int(f)) for k in range(0, int(max_tokens) + 1, phase): acc = _frames_for_tokens(k) d = abs(acc - f) if d < best_d: best_k, best_f, best_d = k, acc, d return best_k, best_f def _temporal_segments(token_count, chunk_length, overlap): token_count = int(token_count) chunk_length = int(chunk_length) overlap = int(overlap) if token_count <= 0: return [(0, 0, 0, 0)] if chunk_length <= 0: raise ValueError("chunk_length must be positive") if overlap < 0: raise ValueError("temporal_overlap must be non-negative") if chunk_length <= overlap: raise ValueError("temporal_overlap must be smaller than chunk_length") frame_count = _frames_for_tokens(token_count) if frame_count <= chunk_length: return [(0, 0, token_count, frame_count)] hop = chunk_length - overlap bounds = [] prev_end_k = 0 i = 0 while True: s = i * hop e = min(s + chunk_length, frame_count) if i == 0: k0, f0 = 0, 0 else: k0, f0 = _snap_frame_boundary(s, token_count, phase=5) if k0 > prev_end_k: k0, f0 = prev_end_k, _frames_for_tokens(prev_end_k) if e >= frame_count: k1, f1 = token_count, frame_count else: k1, f1 = _snap_frame_boundary(e, token_count, phase=5) if k1 <= k0: k1 = min(token_count, k0 + 5) f1 = _frames_for_tokens(k1) if k1 >= token_count: k1, f1 = token_count, frame_count bounds.append((k0, f0, k1, f1)) if k1 >= token_count: break prev_end_k = k1 i += 1 return bounds def _temporal_blend_weights(length, anchor_strength): length = int(length) if length <= 0: return None anchor_strength = float(anchor_strength) anchor_strength = min(1.0, max(0.0, anchor_strength)) start = max(0.0, 1.0 - anchor_strength) return torch.linspace(start, 1.0, length) def _spatial_blend_weights(t, overlap_mode, overlap_blend="linear"): if overlap_blend == "overwrite": return torch.ones_like(t) if overlap_mode == "later" else torch.zeros_like(t) if overlap_blend == "midpoint": base = (t >= 0.5).to(t.dtype) elif overlap_blend == "smoothstep": base = t * t * (3.0 - 2.0 * t) else: base = t return base if overlap_mode == "later" else 1.0 - base def _grid_1d(size, tile, overlap, min_tile): if size <= tile: return [0], [size], [0] step = tile - overlap n = math.ceil((size - overlap) / step) if (n - 1) * step + tile < size: n += 1 rows = [i * step for i in range(n)] tiles = [min(tile, size - r) for r in rows] if min_tile > 0 and n >= 2: edge = size - rows[-1] if edge < min_tile: new_last = size - min_tile if rows[-2] < new_last < rows[-2] + tiles[-2]: rows[-1] = new_last tiles[-1] = size - new_last overlaps = [0] * n for i in range(1, n): overlaps[i] = max(0, rows[i - 1] + tiles[i - 1] - rows[i]) return rows, tiles, overlaps def _compute_spatial_grid(h, w, th, tw, ol_h, ol_w, min_th=0, min_tw=0): if th <= 0 or tw <= 0: raise ValueError("tile dimensions must be positive") if ol_h >= th or ol_w >= tw: raise ValueError("overlap must be smaller than the tile size") if min_th < 0 or min_tw < 0: raise ValueError("minimum tile size must be non-negative") if min_th > th or min_tw > tw: raise ValueError("minimum tile size must not exceed the tile size") rows, trows, row_ovl = _grid_1d(h, th, ol_h, min_th) cols, tcols, col_ovl = _grid_1d(w, tw, ol_w, min_tw) return rows, cols, trows, tcols, row_ovl, col_ovl def _shrink_model_tile_param(param): def _halve_32(v): v = int(v or 0) if v <= 0: return 0 v = (v // 2 // 32) * 32 return max(0, v) next_param = dict(param) mode = str(next_param.get("tile_size_mode") or "specific_size") if mode == "rows_cols": rows = min(9, max(1, int(next_param.get("grid_rows", 2) or 2) + 1)) cols = min(9, max(1, int(next_param.get("grid_cols", 2) or 2) + 1)) if rows == int(next_param.get("grid_rows", 2) or 2) and cols == int(next_param.get("grid_cols", 2) or 2): return None next_param["grid_rows"] = rows next_param["grid_cols"] = cols next_param["spatial_w_overlap"] = _halve_32(next_param.get("spatial_w_overlap", 0)) next_param["spatial_h_overlap"] = _halve_32(next_param.get("spatial_h_overlap", 0)) next_param["fade_width"] = _halve_32(next_param.get("fade_width", 0)) next_param["fade_height"] = _halve_32(next_param.get("fade_height", 0)) next_param["min_tile_size"] = _halve_32(next_param.get("min_tile_size", 0)) return next_param tile_w = int(next_param.get("tile_width", 0) or 0) tile_h = int(next_param.get("tile_height", 0) or 0) new_w = _halve_32(tile_w) new_h = _halve_32(tile_h) if new_w < 32 or new_h < 32: return None if new_w >= tile_w and new_h >= tile_h: return None next_param["tile_width"] = new_w next_param["tile_height"] = new_h next_param["overlap"] = _halve_32(next_param.get("overlap", 0)) next_param["fade_width"] = _halve_32(next_param.get("fade_width", 0)) next_param["fade_height"] = _halve_32(next_param.get("fade_height", 0)) next_param["min_tile_size"] = min( _halve_32(next_param.get("min_tile_size", 0)), new_w, new_h, ) return next_param def _shrink_temporal_param(param): def _lower_17(v): v = int(v or 0) if v <= 0: return 0 v = (v // 2 // 17) * 17 return max(0, v) next_param = dict(param) chunk_length = int(next_param.get("chunk_length", 0) or 0) temporal_overlap = int(next_param.get("temporal_overlap", 0) or 0) if chunk_length <= 17: return None new_chunk_length = _lower_17(chunk_length) if new_chunk_length < 17: new_chunk_length = chunk_length - 17 if new_chunk_length < 17: return None if new_chunk_length >= chunk_length: new_chunk_length = chunk_length - 17 if new_chunk_length < 17: return None new_overlap = min(_lower_17(temporal_overlap), max(0, new_chunk_length - 17)) if new_overlap >= new_chunk_length: new_overlap = max(0, new_chunk_length - 17) if new_overlap >= new_chunk_length: new_overlap = 0 if new_overlap >= new_chunk_length: return None next_param["chunk_length"] = new_chunk_length next_param["temporal_overlap"] = new_overlap return next_param def _upscale_video_model_core(video, param): model_name = param["model_name"] device = param.get("device", "cuda") precision = param.get("precision", "fp16") orig_dtype = video.dtype dev = torch.device(device if (device == "cpu" or torch.cuda.is_available()) else "cpu") compute_dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[precision] _, c, t, h_in, w_in = video.shape h_out, w_out = _resolve_target_size(param, h_in, w_in) eff = (w_out / float(w_in) + h_out / float(h_in)) / 2.0 if w_in and h_in else 1.0 if eff < 1.0 and (w_out < w_in or h_out < h_in): raise ValueError("This model only supports upscaling (effective scale >= 1.0).") if w_out == w_in and h_out == h_in: return video, h_in, w_in if str(model_name).startswith("("): raise ValueError("Please place H3 upscale model files into the latent_upscale_models directory") s = video.to(device=dev, dtype=compute_dtype, copy=True) model = load_upscale_model(model_name, dev, precision) norm_mean, norm_std = _make_norm_tensors(dev, compute_dtype) try: with torch.inference_mode(): s = s.sub(norm_mean).div(norm_std) out = model(s, scale=eff, target_size=(t, h_out, w_out)) del s out = out.mul(norm_std).add(norm_mean) out = out.to(device="cpu", dtype=orig_dtype) return out, h_out, w_out finally: unload_upscale_model(model_name, dev, precision) try: gc.collect() except Exception: pass def upscale_video_model(video, param): try: return _upscale_video_model_tiled(video, param) except RuntimeError as exc: if "out of memory" not in str(exc).lower(): raise smaller = _shrink_model_tile_param(param) if smaller is None: raise try: gc.collect() except Exception: pass if torch.cuda.is_available(): try: torch.cuda.empty_cache() except Exception: pass return upscale_video_model(video, smaller) def _upscale_video_model_tiled(video, param): _, _, _, h_in, w_in = video.shape h_out, w_out = _resolve_target_size(param, h_in, w_in) mode = str(param.get("tile_size_mode") or "specific_size") tile_w = int(param.get("tile_width", 0) or 0) tile_h = int(param.get("tile_height", 0) or 0) overlap = max(0, int(param.get("overlap", 0) or 0)) fade_w = max(0, int(param.get("fade_width", 0) or 0)) fade_h = max(0, int(param.get("fade_height", 0) or 0)) overlap_mode = str(param.get("overlap_mode", "earlier")) overlap_blend = str(param.get("overlap_blend", "linear")) grid_rows = max(1, int(param.get("grid_rows", 2) or 2)) grid_cols = max(1, int(param.get("grid_cols", 2) or 2)) spatial_w_overlap = max(0, int(param.get("spatial_w_overlap", overlap) or overlap)) spatial_h_overlap = max(0, int(param.get("spatial_h_overlap", overlap) or overlap)) min_tile_size = max(0, int(param.get("min_tile_size", 0) or 0)) if mode == "rows_cols": tile_w = max(32, int(round(w_out / float(grid_cols) / 32.0)) * 32) tile_h = max(32, int(round(h_out / float(grid_rows) / 32.0)) * 32) min_tile_size = min(min_tile_size, tile_w, tile_h) rows, cols, trows, tcols, row_ovl, col_ovl = _compute_spatial_grid( h_out, w_out, tile_h, tile_w, spatial_h_overlap, spatial_w_overlap, min_tile_size, min_tile_size ) else: if tile_w <= 0 or tile_h <= 0 or (tile_w >= w_out and tile_h >= h_out): return _upscale_video_model_core(video, param) for name, value in (("tile_width", tile_w), ("tile_height", tile_h), ("overlap", overlap), ("fade_width", fade_w), ("fade_height", fade_h), ("min_tile_size", min_tile_size)): if value % 32 != 0: raise ValueError(f"'{name}' must be a multiple of 32 pixels; got {value}.") min_tile_size = min(min_tile_size, tile_w, tile_h) rows, cols, trows, tcols, row_ovl, col_ovl = _compute_spatial_grid( h_out, w_out, tile_h, tile_w, overlap, overlap, min_tile_size, min_tile_size ) # If the requested tile is not smaller than the target on either axis, # the tiled path would just duplicate work. if len(rows) == 1 and len(cols) == 1: return _upscale_video_model_core(video, param) scale_h = h_out / float(h_in) scale_w = w_out / float(w_in) orig_dtype = video.dtype out = torch.zeros((video.shape[0], video.shape[1], video.shape[2], h_out, w_out), device="cpu", dtype=orig_dtype) for i, r0 in enumerate(rows): tr = trows[i] ovh = row_ovl[i] src_r0 = max(0, int(round(r0 / scale_h))) src_r1 = min(h_in, max(src_r0 + 1, int(round((r0 + tr) / scale_h)))) for j, c0 in enumerate(cols): tc = tcols[j] ovw = col_ovl[j] src_c0 = max(0, int(round(c0 / scale_w))) src_c1 = min(w_in, max(src_c0 + 1, int(round((c0 + tc) / scale_w)))) tile_video = video[:, :, :, src_r0:src_r1, src_c0:src_c1].contiguous() tile_param = dict(param) tile_param["width"] = int(tc) tile_param["height"] = int(tr) tile_param["megapixels"] = 0.0 tile_out, tile_h_out, tile_w_out = _upscale_video_model_core(tile_video, tile_param) if tile_h_out != tr or tile_w_out != tc: flat = tile_out.permute(0, 2, 1, 3, 4).reshape(-1, tile_out.shape[1], tile_h_out, tile_w_out) flat = F.interpolate(flat, size=(tr, tc), mode="bilinear", align_corners=False) tile_out = flat.reshape(tile_out.shape[0], tile_out.shape[2], tile_out.shape[1], tr, tc).permute(0, 2, 1, 3, 4).contiguous() region = out[:, :, :, r0:r0 + tr, c0:c0 + tc] base_region = region.clone() region.copy_(tile_out.to(dtype=orig_dtype)) if j > 0 and ovw > 0: t = torch.linspace(0.0, 1.0, ovw, device=region.device, dtype=region.dtype) w = _spatial_blend_weights(t, overlap_mode, overlap_blend) if fade_w > 0: w = w.clone() w[:min(fade_w, ovw)] = 0.0 region[:, :, :, :, :ovw] = ( base_region[:, :, :, :, :ovw] * (1.0 - w[None, None, None, None, :]) + tile_out[:, :, :, :, :ovw].to(dtype=orig_dtype) * w[None, None, None, None, :] ) if i > 0 and ovh > 0: t = torch.linspace(0.0, 1.0, ovh, device=region.device, dtype=region.dtype) w = _spatial_blend_weights(t, overlap_mode, overlap_blend) if fade_h > 0: w = w.clone() w[:min(fade_h, ovh)] = 0.0 region[:, :, :, :ovh, :] = ( base_region[:, :, :, :ovh, :] * (1.0 - w[None, None, None, :, None]) + tile_out[:, :, :, :ovh, :].to(dtype=orig_dtype) * w[None, None, None, :, None] ) out[:, :, :, r0:r0 + tr, c0:c0 + tc] = region return out, h_out, w_out def upscale_video_interp(video, param): method = str(param.get("method") or "bilinear") _, c, t, h_in, w_in = video.shape h_out, w_out = _resolve_target_size(param, h_in, w_in) if h_out == h_in and w_out == w_in: return video, h_in, w_in video_bt = video.permute(0, 2, 1, 3, 4).reshape(-1, c, h_in, w_in) up = F.interpolate(video_bt, size=(h_out, w_out), mode=method) up = up.reshape(video.shape[0], t, c, h_out, w_out).permute(0, 2, 1, 3, 4).contiguous() return up, h_out, w_out def _upscale_video_temporal_chunks(video, param, upscaler): chunk_length = int(param.get("chunk_length", 0) or 0) temporal_overlap = int(param.get("temporal_overlap", 0) or 0) anchor_strength = float(param.get("anchor_strength", 0.999) or 0.999) t = int(video.shape[2]) frame_count = _frames_for_tokens(t) if chunk_length <= 0 or frame_count <= chunk_length: return upscaler(video, param) bounds = _temporal_segments(t, chunk_length, temporal_overlap) if len(bounds) <= 1: return upscaler(video, param) orig_dtype = video.dtype out = None out_h = out_w = None for i, (k0, f0, k1, f1) in enumerate(bounds): chunk = video[:, :, k0:k1].contiguous() try: chunk_out, chunk_h, chunk_w = upscaler(chunk, param) except RuntimeError as exc: if "out of memory" not in str(exc).lower(): raise smaller = _shrink_temporal_param(param) if smaller is None: raise try: gc.collect() except Exception: pass if torch.cuda.is_available(): try: torch.cuda.empty_cache() except Exception: pass return _upscale_video_temporal_chunks(video, smaller, upscaler) chunk_out = chunk_out.to(device="cpu", dtype=orig_dtype) if out is None: out_h, out_w = chunk_h, chunk_w out = torch.zeros((video.shape[0], video.shape[1], t, out_h, out_w), device="cpu", dtype=orig_dtype) out[:, :, k0:k1] = chunk_out continue ov = min(temporal_overlap, k1 - k0, t - k0) if ov <= 0: out[:, :, k0:k1] = chunk_out continue prev_region = out[:, :, k0:k0 + ov] new_region = chunk_out[:, :, :ov] w = _temporal_blend_weights(ov, anchor_strength).to(device=prev_region.device, dtype=prev_region.dtype) out[:, :, k0:k0 + ov] = ( prev_region * (1.0 - w[None, None, :, None, None]) + new_region * w[None, None, :, None, None] ) out[:, :, k0 + ov:k1] = chunk_out[:, :, ov:] return out, out_h, out_w def upscale_latent_video(video, param): mode = str(param.get("mode") or "off") if mode == "off": return video, video.shape[-2], video.shape[-1] if mode == "model": return _upscale_video_temporal_chunks(video, param, upscale_video_model) return _upscale_video_temporal_chunks(video, param, upscale_video_interp) class H3LatentUpscaleParams: CATEGORY = "Dumas/MiniMax" FUNCTION = "build" RETURN_TYPES = ("DUMAS_H3_LATENT_UPSCALE_PARAM",) RETURN_NAMES = ("latent_upscale_param",) @classmethod def INPUT_TYPES(cls): return { "required": { "mode": (["off", "model", "interp"], {"default": "off", "tooltip": "Latent refinement mode. 'off' skips the stage, 'model' uses the H3 latent upscaler model, 'interp' uses model-free interpolation."}), "model_name": (_scan_models(), { "default": "none", "tooltip": "H3 latent upscale checkpoint from models/latent_upscale_models, used when mode = model."}), "method": (["nearest-exact", "bilinear", "area", "bicubic"], {"default": "bilinear", "tooltip": "Interpolation method used when mode = interp."}), "width": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 32, "tooltip": "Explicit target width for the latent refinement stage. Leave at 0 to let megapixels choose the size instead."}), "height": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 32, "tooltip": "Explicit target height for the latent refinement stage. Leave at 0 to let megapixels choose the size instead."}), "device": (["cuda", "cpu"], {"default": "cuda", "tooltip": "Device used by the H3 latent upscaler model when mode = model."}), "precision": (["fp16", "fp32", "bf16"], {"default": "fp16", "tooltip": "Computation precision used by the H3 latent upscaler model when mode = model."}), "sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler_ancestral", "tooltip": "Sampler used for the latent refinement pass. Default matches the current H3 preference."}), "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "simple", "tooltip": "Scheduler used for the latent refinement pass."}), "steps": ("INT", {"default": 2, "min": 1, "max": 50, "step": 1, "tooltip": "Number of refinement steps applied after the latent upscaler stage."}), "denoise": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "How much the refinement pass may rewrite the upscaled latent. Lower = safer, higher = freer."}), "megapixels": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 4.0, "step": 0.01, "tooltip": "Primary target size for the latent refinement stage. If width and height are both set, they win; otherwise the node scales the current shot to this pixel budget while preserving aspect ratio. 0 keeps the incoming latent size."}), "tile_width": ("INT", {"default": 512, "min": 32, "max": 4096, "step": 32, "tooltip": "Spatial tile width for the refinement stage in pixels. 512 matches the upstream latent-split default."}), "tile_height": ("INT", {"default": 512, "min": 32, "max": 4096, "step": 32, "tooltip": "Spatial tile height for the refinement stage in pixels. 512 matches the upstream latent-split default."}), "overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 32, "tooltip": "Pixel overlap between neighbouring spatial tiles. 64 matches the upstream latent-split default."}), "fade_width": ("INT", {"default": 32, "min": 0, "max": 4096, "step": 32, "tooltip": "Width in pixels of the freeze-to-free transition inside each overlap strip. 0 freezes the whole strip; 32 matches the upstream default."}), "fade_height": ("INT", {"default": 32, "min": 0, "max": 4096, "step": 32, "tooltip": "Height in pixels of the freeze-to-free transition inside each overlap strip. 0 freezes the whole strip; 32 matches the upstream default."}), "overlap_mode": (["earlier", "later"], {"default": "earlier", "tooltip": "Which tile wins the overlap band when stitching the spatial batches back together."}), "overlap_blend": (["linear", "smoothstep", "overwrite", "midpoint"], {"default": "linear", "tooltip": "How overlap bands are blended when the spatial batches are stitched back together."}), "tile_size_mode": (["specific_size", "rows_cols"], {"default": "specific_size", "tooltip": "How the spatial tile size is determined. 'specific_size' uses tile_width/tile_height. 'rows_cols' solves equal-size tiles from the target width/height and grid counts."}), "grid_rows": ("INT", {"default": 2, "min": 1, "max": 9, "step": 1, "tooltip": "Number of tile rows used when tile_size_mode = rows_cols."}), "grid_cols": ("INT", {"default": 2, "min": 1, "max": 9, "step": 1, "tooltip": "Number of tile columns used when tile_size_mode = rows_cols."}), "spatial_w_overlap": ("INT", {"default": 128, "min": 0, "max": 4096, "step": 32, "tooltip": "Desired horizontal overlap when tile_size_mode = rows_cols. The node solves the actual equal-tile overlap from the requested target size."}), "spatial_h_overlap": ("INT", {"default": 128, "min": 0, "max": 4096, "step": 32, "tooltip": "Desired vertical overlap when tile_size_mode = rows_cols. The node solves the actual equal-tile overlap from the requested target size."}), "min_tile_size": ("INT", {"default": 256, "min": 0, "max": 4096, "step": 32, "tooltip": "Minimum leftover edge tile size allowed when the spatial grid is solved."}), "masked_area_noise": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "How much noise is allowed into the masked overlap band during spatial refinement."}), "brightness_match": ("BOOLEAN", {"default": False, "tooltip": "Match each sampled tile's brightness baseline to the source tile region after sampling."}), "dynamic_fade": (["off", "narrowing", "widening"], {"default": "off", "tooltip": "Temporal fade schedule over each tile's sampling. Off keeps the fade fixed; narrowing shrinks it over steps; widening grows it over steps."}), "dynamic_fade_min": ("INT", {"default": 32, "min": 0, "max": 4096, "step": 32, "tooltip": "Minimum fade width used by dynamic_fade when it is enabled."}), "chunk_length": ("INT", {"default": 85, "min": 17, "max": 100000, "step": 17, "tooltip": "Temporal chunk length for the latent upscale stage. 85 matches the practical upstream example and helps lower peak VRAM."}), "temporal_overlap": ("INT", {"default": 17, "min": 0, "max": 100000, "step": 17, "tooltip": "Temporal overlap between latent chunks. 17 matches the upstream split example and reduces seam risk."}), "resize_conditioning": ("BOOLEAN", {"default": False, "tooltip": "Reserved for upstream split compatibility. Leave OFF unless you need the original fallback behavior."}), "anchor_strength": ("FLOAT", {"default": 0.999, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Reserved for upstream split compatibility. 0.999 matches the original anchor default."}), } } def build(self, mode, model_name, method, width, height, device, precision, sampler_name, scheduler, steps, denoise, megapixels, tile_width, tile_height, overlap, fade_width, fade_height, overlap_mode, overlap_blend, tile_size_mode, grid_rows, grid_cols, spatial_w_overlap, spatial_h_overlap, min_tile_size, masked_area_noise, brightness_match, dynamic_fade, dynamic_fade_min, chunk_length, temporal_overlap, resize_conditioning, anchor_strength): width = int(width) height = int(height) steps = int(steps) tile_width = int(tile_width) tile_height = int(tile_height) overlap = int(overlap) fade_width = int(fade_width) fade_height = int(fade_height) grid_rows = int(grid_rows) grid_cols = int(grid_cols) spatial_w_overlap = int(spatial_w_overlap) spatial_h_overlap = int(spatial_h_overlap) min_tile_size = int(min_tile_size) masked_area_noise = float(masked_area_noise) dynamic_fade_min = int(dynamic_fade_min) chunk_length = int(chunk_length) temporal_overlap = int(temporal_overlap) anchor_strength = float(anchor_strength) if mode == "off": return ({ "mode": "off", "width": width, "height": height, "sampler_name": sampler_name, "scheduler": scheduler, "steps": steps, "denoise": float(denoise), "refine_denoise": float(denoise), "megapixels": float(megapixels), "tile_width": tile_width, "tile_height": tile_height, "overlap": overlap, "fade_width": fade_width, "fade_height": fade_height, "overlap_mode": overlap_mode, "overlap_blend": overlap_blend, "tile_size_mode": tile_size_mode, "grid_rows": grid_rows, "grid_cols": grid_cols, "spatial_w_overlap": spatial_w_overlap, "spatial_h_overlap": spatial_h_overlap, "min_tile_size": min_tile_size, "masked_area_noise": masked_area_noise, "brightness_match": bool(brightness_match), "dynamic_fade": dynamic_fade, "dynamic_fade_min": dynamic_fade_min, "chunk_length": chunk_length, "temporal_overlap": temporal_overlap, "resize_conditioning": bool(resize_conditioning), "anchor_strength": anchor_strength, },) if chunk_length % 17 != 0: raise ValueError("chunk_length must be a multiple of 17 pixels") if temporal_overlap % 17 != 0: raise ValueError("temporal_overlap must be a multiple of 17 pixels") if temporal_overlap >= chunk_length: raise ValueError("temporal_overlap must be smaller than chunk_length") if width > 0: width = int(round(width / 32.0)) * 32 if height > 0: height = int(round(height / 32.0)) * 32 return ({ "mode": mode, "model_name": model_name, "method": method, "width": width, "height": height, "device": device, "precision": precision, "sampler_name": sampler_name, "scheduler": scheduler, "steps": steps, "denoise": float(denoise), "refine_denoise": float(denoise), "megapixels": float(megapixels), "tile_width": tile_width, "tile_height": tile_height, "overlap": overlap, "fade_width": fade_width, "fade_height": fade_height, "overlap_mode": overlap_mode, "overlap_blend": overlap_blend, "tile_size_mode": tile_size_mode, "grid_rows": grid_rows, "grid_cols": grid_cols, "spatial_w_overlap": spatial_w_overlap, "spatial_h_overlap": spatial_h_overlap, "min_tile_size": min_tile_size, "masked_area_noise": masked_area_noise, "brightness_match": bool(brightness_match), "dynamic_fade": dynamic_fade, "dynamic_fade_min": dynamic_fade_min, "chunk_length": chunk_length, "temporal_overlap": temporal_overlap, "resize_conditioning": bool(resize_conditioning), "anchor_strength": anchor_strength, },) NODE_CLASS_MAPPINGS = {"DumasH3LatentUpscaleParams": H3LatentUpscaleParams} NODE_DISPLAY_NAME_MAPPINGS = {"DumasH3LatentUpscaleParams": "Dumas H3 Latent Upscale Params"} __all__ = [ "NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "upscale_latent_video", ]