912 lines
34 KiB
Python
912 lines
34 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import math
|
|
from typing import Iterable, Optional, Sequence
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def _sigma_scalar(timestep: torch.Tensor | float | int) -> float:
|
|
if isinstance(timestep, (float, int)):
|
|
return float(timestep)
|
|
if timestep.numel() == 0:
|
|
return 0.0
|
|
return float(timestep.detach().flatten()[0].cpu().item())
|
|
|
|
|
|
def clone_latent(latent: dict) -> dict:
|
|
return dict(latent)
|
|
|
|
|
|
def _stage_cpu_tensor(x: torch.Tensor, dtype: torch.dtype, pin_memory: bool) -> torch.Tensor:
|
|
y = x.detach().to(device="cpu", dtype=dtype).contiguous().clone()
|
|
if pin_memory:
|
|
try:
|
|
y = y.pin_memory()
|
|
except Exception:
|
|
pass
|
|
return y
|
|
|
|
|
|
def _round_to_multiple(v: int, multiple: int) -> int:
|
|
if multiple <= 1:
|
|
return max(1, int(v))
|
|
return max(multiple, int(round(v / multiple) * multiple))
|
|
|
|
|
|
def resize_4d_tensor_nearest(x: torch.Tensor, size_hw: tuple[int, int]) -> torch.Tensor:
|
|
if tuple(x.shape[-2:]) == tuple(size_hw):
|
|
return x
|
|
return F.interpolate(x, size=size_hw, mode="nearest")
|
|
|
|
|
|
def _interp_mode(src_hw: tuple[int, int], dst_hw: tuple[int, int]) -> str:
|
|
src_h, src_w = src_hw
|
|
dst_h, dst_w = dst_hw
|
|
return "area" if (dst_h < src_h or dst_w < src_w) else "bilinear"
|
|
|
|
|
|
def resize_4d_tensor(x: torch.Tensor, size_hw: tuple[int, int]) -> torch.Tensor:
|
|
"""Resize BCHW tensor with sane defaults for latent tensors."""
|
|
if tuple(x.shape[-2:]) == tuple(size_hw):
|
|
return x
|
|
mode = _interp_mode(tuple(x.shape[-2:]), tuple(size_hw))
|
|
if mode == "area":
|
|
return F.interpolate(x, size=size_hw, mode=mode)
|
|
return F.interpolate(x, size=size_hw, mode=mode, align_corners=False)
|
|
|
|
|
|
def resize_mask(mask: Optional[torch.Tensor], size_hw: tuple[int, int], batch: int, channels: int) -> Optional[torch.Tensor]:
|
|
if mask is None:
|
|
return None
|
|
|
|
m = mask
|
|
if m.ndim == 2:
|
|
m = m.unsqueeze(0).unsqueeze(0)
|
|
elif m.ndim == 3:
|
|
m = m.unsqueeze(1)
|
|
elif m.ndim == 4 and m.shape[1] != 1:
|
|
# Unexpected format: collapse to a single channel mask.
|
|
m = m.mean(dim=1, keepdim=True)
|
|
|
|
m = m.float()
|
|
mode = _interp_mode(tuple(m.shape[-2:]), tuple(size_hw))
|
|
if mode == "area":
|
|
m = F.interpolate(m, size=size_hw, mode=mode)
|
|
else:
|
|
m = F.interpolate(m, size=size_hw, mode=mode, align_corners=False)
|
|
|
|
if m.shape[0] < batch:
|
|
repeat = math.ceil(batch / max(1, m.shape[0]))
|
|
m = m.repeat(repeat, 1, 1, 1)[:batch]
|
|
elif m.shape[0] > batch:
|
|
m = m[:batch]
|
|
|
|
m = m.expand(batch, channels, size_hw[0], size_hw[1]).contiguous()
|
|
return m
|
|
|
|
|
|
def latent_target_hw_from_megapixels(
|
|
latent_bchw: torch.Tensor,
|
|
target_megapixels: float,
|
|
latent_multiple: int = 2,
|
|
) -> tuple[int, int]:
|
|
"""
|
|
Compute a target latent spatial size that preserves aspect ratio while targeting a pixel-space megapixel count.
|
|
Latent tensors are expected to be 8x downsampled in H/W for standard image models used by ComfyUI.
|
|
"""
|
|
_, _, h_lat, w_lat = latent_bchw.shape
|
|
h_px = h_lat * 8
|
|
w_px = w_lat * 8
|
|
aspect = w_px / max(1.0, float(h_px))
|
|
target_px_area = max(target_megapixels, 0.01) * 1_000_000.0
|
|
|
|
h_target_px = math.sqrt(target_px_area / max(aspect, 1e-8))
|
|
w_target_px = h_target_px * aspect
|
|
|
|
# Keep latent compatibility and avoid absurdly small sizes.
|
|
h_target_lat = _round_to_multiple(int(round(h_target_px / 8.0)), latent_multiple)
|
|
w_target_lat = _round_to_multiple(int(round(w_target_px / 8.0)), latent_multiple)
|
|
|
|
return (h_target_lat, w_target_lat)
|
|
|
|
|
|
def resize_latent_dict(latent: dict, size_hw: tuple[int, int]) -> dict:
|
|
out = clone_latent(latent)
|
|
out["samples"] = resize_4d_tensor(latent["samples"], size_hw)
|
|
if "noise_mask" in latent:
|
|
nm = resize_mask(latent.get("noise_mask"), size_hw, out["samples"].shape[0], 1)
|
|
if nm is not None:
|
|
out["noise_mask"] = nm[:, :1, :, :].squeeze(1)
|
|
return out
|
|
|
|
|
|
def build_nested_noise(
|
|
lowres_noise: torch.Tensor,
|
|
target_shape: tuple[int, int, int, int],
|
|
seed: int,
|
|
hf_strength: float,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Lift low-res noise to high-res while keeping the coarse field aligned.
|
|
|
|
The implementation uses:
|
|
base = nearest-upsampled low-res noise
|
|
hf = random noise with its low-frequency projection removed
|
|
so that the high-res branch shares the same coarse stochastic layout but can still invent fine detail.
|
|
"""
|
|
batch, channels, target_h, target_w = target_shape
|
|
device = lowres_noise.device
|
|
dtype = lowres_noise.dtype
|
|
|
|
base = resize_4d_tensor_nearest(lowres_noise.float(), (target_h, target_w))
|
|
if hf_strength <= 0.0:
|
|
return base.to(device=device, dtype=dtype)
|
|
|
|
gen = torch.Generator(device="cpu")
|
|
gen.manual_seed(int(seed) & 0xFFFFFFFFFFFFFFFF)
|
|
hf = torch.randn((batch, channels, target_h, target_w), generator=gen, dtype=torch.float32, device="cpu")
|
|
|
|
low_hw = tuple(lowres_noise.shape[-2:])
|
|
hf_low = resize_4d_tensor(hf, low_hw)
|
|
hf_low_up = resize_4d_tensor_nearest(hf_low, (target_h, target_w))
|
|
hf = hf - hf_low_up
|
|
|
|
std = hf.std(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
|
|
hf = (hf / std).to(device=device, non_blocking=True)
|
|
|
|
out = base + float(hf_strength) * hf
|
|
out_std = out.std(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
|
|
out = out / out_std
|
|
|
|
return out.to(device=device, dtype=dtype)
|
|
|
|
|
|
def lowpass_latent(x: torch.Tensor, cutoff: float) -> torch.Tensor:
|
|
"""
|
|
Cheap, sampler-friendly low-pass filter implemented as area downsample + bilinear upsample.
|
|
|
|
cutoff is a retained-resolution fraction in (0, 1].
|
|
Example: cutoff=0.33 means keep roughly one third of the spatial resolution for the coarse field.
|
|
"""
|
|
cutoff = float(max(0.01, min(1.0, cutoff)))
|
|
h, w = x.shape[-2:]
|
|
if cutoff >= 0.999:
|
|
return x
|
|
coarse_h = max(1, int(round(h * cutoff)))
|
|
coarse_w = max(1, int(round(w * cutoff)))
|
|
y = resize_4d_tensor(x, (coarse_h, coarse_w))
|
|
y = resize_4d_tensor(y, (h, w))
|
|
return y
|
|
|
|
|
|
def residual_lock_multiband(
|
|
base_denoised: torch.Tensor,
|
|
anchor_denoised: torch.Tensor,
|
|
low_strength: float,
|
|
mid_strength: float,
|
|
low_cutoff: float,
|
|
mid_cutoff: float,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Preserve the highest-frequency residual while independently blending:
|
|
- a low / coarse band
|
|
- a mid-frequency structural band
|
|
toward the planner anchor.
|
|
"""
|
|
low_strength = float(max(0.0, min(1.0, low_strength)))
|
|
mid_strength = float(max(0.0, min(1.0, mid_strength)))
|
|
mid_cutoff = float(max(low_cutoff, mid_cutoff))
|
|
|
|
base_mid_lp = lowpass_latent(base_denoised, mid_cutoff)
|
|
anchor_mid_lp = lowpass_latent(anchor_denoised, mid_cutoff)
|
|
base_low = lowpass_latent(base_denoised, low_cutoff)
|
|
anchor_low = lowpass_latent(anchor_denoised, low_cutoff)
|
|
base_mid = base_mid_lp - base_low
|
|
anchor_mid = anchor_mid_lp - anchor_low
|
|
base_high = base_denoised - base_mid_lp
|
|
return base_high + torch.lerp(base_low, anchor_low, low_strength) + torch.lerp(base_mid, anchor_mid, mid_strength)
|
|
|
|
|
|
def _base_grid(batch: int, height: int, width: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
|
yy, xx = torch.meshgrid(
|
|
torch.linspace(-1.0, 1.0, height, device=device, dtype=dtype),
|
|
torch.linspace(-1.0, 1.0, width, device=device, dtype=dtype),
|
|
indexing="ij",
|
|
)
|
|
return torch.stack([xx, yy], dim=-1).unsqueeze(0).expand(batch, -1, -1, -1).contiguous()
|
|
|
|
|
|
def warp_4d_tensor(x: torch.Tensor, flow_xy_norm: torch.Tensor, mode: str = "bilinear") -> torch.Tensor:
|
|
if x.ndim != 4:
|
|
raise ValueError(f"Expected BCHW tensor, got shape {tuple(x.shape)}")
|
|
batch, _, height, width = x.shape
|
|
expected = (batch, height, width, 2)
|
|
if tuple(flow_xy_norm.shape) != expected:
|
|
raise ValueError(f"Expected flow shape {expected}, got {tuple(flow_xy_norm.shape)}")
|
|
base_grid = _base_grid(batch, height, width, x.device, x.dtype)
|
|
sample_grid = (base_grid + flow_xy_norm).clamp(-1.25, 1.25)
|
|
return F.grid_sample(
|
|
x,
|
|
sample_grid,
|
|
mode=mode,
|
|
padding_mode="border",
|
|
align_corners=True,
|
|
)
|
|
|
|
|
|
def _expand_mask_channels(mask: torch.Tensor, like: torch.Tensor) -> torch.Tensor:
|
|
if mask.ndim != 4:
|
|
raise ValueError(f"Expected mask tensor BCHW, got shape {tuple(mask.shape)}")
|
|
if mask.shape[0] < like.shape[0]:
|
|
repeat = math.ceil(like.shape[0] / max(1, mask.shape[0]))
|
|
mask = mask.repeat(repeat, 1, 1, 1)[: like.shape[0]]
|
|
elif mask.shape[0] > like.shape[0]:
|
|
mask = mask[: like.shape[0]]
|
|
if mask.shape[1] == like.shape[1]:
|
|
return mask
|
|
if mask.shape[1] == 1:
|
|
return mask.expand(like.shape[0], like.shape[1], like.shape[-2], like.shape[-1]).contiguous()
|
|
return mask.mean(dim=1, keepdim=True).expand(like.shape[0], like.shape[1], like.shape[-2], like.shape[-1]).contiguous()
|
|
|
|
|
|
def _latent_activity_map(x: torch.Tensor, mask_1ch: torch.Tensor) -> torch.Tensor:
|
|
mass = mask_1ch.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
|
|
mean = (x * mask_1ch).sum(dim=(-2, -1), keepdim=True) / mass
|
|
centered = x - mean
|
|
activity = centered.square().mean(dim=1, keepdim=True)
|
|
activity = activity * mask_1ch
|
|
activity = activity + mask_1ch * 1e-8
|
|
return activity
|
|
|
|
|
|
def _masked_spatial_stats(x: torch.Tensor, mask_1ch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
batch, _, height, width = x.shape
|
|
yy, xx = torch.meshgrid(
|
|
torch.linspace(-1.0, 1.0, height, device=x.device, dtype=x.dtype),
|
|
torch.linspace(-1.0, 1.0, width, device=x.device, dtype=x.dtype),
|
|
indexing="ij",
|
|
)
|
|
xx = xx.view(1, 1, height, width).expand(batch, -1, -1, -1)
|
|
yy = yy.view(1, 1, height, width).expand(batch, -1, -1, -1)
|
|
|
|
activity = _latent_activity_map(x, mask_1ch)
|
|
mass = activity.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
|
|
cx = (activity * xx).sum(dim=(-2, -1), keepdim=True) / mass
|
|
cy = (activity * yy).sum(dim=(-2, -1), keepdim=True) / mass
|
|
|
|
dx = xx - cx
|
|
dy = yy - cy
|
|
rx = torch.sqrt((activity * dx.square()).sum(dim=(-2, -1), keepdim=True) / mass).clamp_min(1e-4)
|
|
ry = torch.sqrt((activity * dy.square()).sum(dim=(-2, -1), keepdim=True) / mass).clamp_min(1e-4)
|
|
return cx, cy, rx, ry
|
|
|
|
|
|
def estimate_latent_compaction_flow(
|
|
anchor_low: torch.Tensor,
|
|
base_low: torch.Tensor,
|
|
mask_1ch: torch.Tensor,
|
|
strength: float,
|
|
radial_strength: float,
|
|
anisotropy: float,
|
|
translation_strength: float,
|
|
max_shift_px: float,
|
|
) -> torch.Tensor:
|
|
batch, _, height, width = base_low.shape
|
|
anchor_cx, anchor_cy, anchor_rx, anchor_ry = _masked_spatial_stats(anchor_low, mask_1ch)
|
|
base_cx, base_cy, base_rx, base_ry = _masked_spatial_stats(base_low, mask_1ch)
|
|
|
|
ratio_x = (base_rx / anchor_rx.clamp_min(1e-5)).clamp(0.85, 1.25)
|
|
ratio_y = (base_ry / anchor_ry.clamp_min(1e-5)).clamp(0.85, 1.25)
|
|
outward_x = (ratio_x - 1.0).clamp(min=0.0)
|
|
outward_y = (ratio_y - 1.0).clamp(min=0.0)
|
|
|
|
yy, xx = torch.meshgrid(
|
|
torch.linspace(-1.0, 1.0, height, device=base_low.device, dtype=base_low.dtype),
|
|
torch.linspace(-1.0, 1.0, width, device=base_low.device, dtype=base_low.dtype),
|
|
indexing="ij",
|
|
)
|
|
xx = xx.view(1, 1, height, width).expand(batch, -1, -1, -1)
|
|
yy = yy.view(1, 1, height, width).expand(batch, -1, -1, -1)
|
|
|
|
dx = xx - base_cx
|
|
dy = yy - base_cy
|
|
ex = dx / base_rx.clamp_min(1e-5)
|
|
ey = dy / base_ry.clamp_min(1e-5)
|
|
radius = torch.sqrt(ex.square() + ey.square() + 1e-8)
|
|
edge_envelope = torch.clamp(radius / 1.25, 0.0, 1.0)
|
|
|
|
smooth_mask = lowpass_latent(mask_1ch, 0.35).clamp(0.0, 1.0)
|
|
mean_outward = 0.5 * (outward_x + outward_y)
|
|
axis_x = mean_outward * float(radial_strength) + (outward_x - mean_outward) * float(anisotropy)
|
|
axis_y = mean_outward * float(radial_strength) + (outward_y - mean_outward) * float(anisotropy)
|
|
|
|
shift_x = ex * edge_envelope * smooth_mask * float(strength) * axis_x
|
|
shift_y = ey * edge_envelope * smooth_mask * float(strength) * axis_y
|
|
|
|
trans_x = (base_cx - anchor_cx) * smooth_mask * float(strength) * float(translation_strength)
|
|
trans_y = (base_cy - anchor_cy) * smooth_mask * float(strength) * float(translation_strength)
|
|
|
|
max_shift_norm_x = 2.0 * float(max_shift_px) / max(width - 1, 1)
|
|
max_shift_norm_y = 2.0 * float(max_shift_px) / max(height - 1, 1)
|
|
|
|
shift_x = (shift_x + trans_x).clamp(-max_shift_norm_x, max_shift_norm_x)
|
|
shift_y = (shift_y + trans_y).clamp(-max_shift_norm_y, max_shift_norm_y)
|
|
return torch.stack([shift_x[:, 0], shift_y[:, 0]], dim=-1)
|
|
|
|
|
|
def _weighted_channel_stats(x: torch.Tensor, mask_1ch: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
|
mass = mask_1ch.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
|
|
mean = (x * mask_1ch).sum(dim=(-2, -1), keepdim=True) / mass
|
|
var = (((x - mean).square()) * mask_1ch).sum(dim=(-2, -1), keepdim=True) / mass
|
|
std = torch.sqrt(var.clamp_min(1e-8))
|
|
return mean, std
|
|
|
|
|
|
def _weighted_global_std(x: torch.Tensor, mask_1ch: torch.Tensor, mean: Optional[torch.Tensor] = None) -> torch.Tensor:
|
|
if mean is None:
|
|
mean, _ = _weighted_channel_stats(x, mask_1ch)
|
|
mass = mask_1ch.sum(dim=(-2, -1), keepdim=True).clamp_min(1e-6)
|
|
denom = mass * float(max(1, x.shape[1]))
|
|
var = (((x - mean).square()) * mask_1ch).sum(dim=(1, 2, 3), keepdim=True) / denom
|
|
return torch.sqrt(var.clamp_min(1e-8))
|
|
|
|
|
|
def _bounded_ratio(target: torch.Tensor, source: torch.Tensor, gain_cap: float) -> torch.Tensor:
|
|
gain_cap = float(max(1.0, gain_cap))
|
|
lo = 1.0 / gain_cap
|
|
hi = gain_cap
|
|
return (target / source.clamp_min(1e-5)).clamp(lo, hi)
|
|
|
|
|
|
def tether_latent_low_frequency_energy(
|
|
base_low: torch.Tensor,
|
|
anchor_low: torch.Tensor,
|
|
mask_1ch: torch.Tensor,
|
|
energy_tether: float,
|
|
channel_tether: float,
|
|
gain_cap: float,
|
|
) -> torch.Tensor:
|
|
energy_tether = float(max(0.0, min(1.0, energy_tether)))
|
|
channel_tether = float(max(0.0, min(1.0, channel_tether)))
|
|
if energy_tether <= 0.0 and channel_tether <= 0.0:
|
|
return base_low
|
|
|
|
base_mean, _ = _weighted_channel_stats(base_low, mask_1ch)
|
|
anchor_mean, anchor_std = _weighted_channel_stats(anchor_low, mask_1ch)
|
|
|
|
regulated = base_low
|
|
if energy_tether > 0.0:
|
|
base_global_std = _weighted_global_std(regulated, mask_1ch, mean=base_mean)
|
|
anchor_global_std = _weighted_global_std(anchor_low, mask_1ch, mean=anchor_mean)
|
|
global_gain = _bounded_ratio(anchor_global_std, base_global_std, gain_cap)
|
|
global_gain = torch.lerp(torch.ones_like(global_gain), global_gain, energy_tether)
|
|
regulated = base_mean + (regulated - base_mean) * global_gain
|
|
|
|
if channel_tether > 0.0:
|
|
regulated_mean, regulated_std = _weighted_channel_stats(regulated, mask_1ch)
|
|
channel_gain = _bounded_ratio(anchor_std, regulated_std, gain_cap)
|
|
channel_gain = torch.lerp(torch.ones_like(channel_gain), channel_gain, channel_tether)
|
|
regulated = regulated_mean + (regulated - regulated_mean) * channel_gain
|
|
|
|
return regulated
|
|
|
|
|
|
def restore_latent_low_frequency_stats(
|
|
warped_low: torch.Tensor,
|
|
anchor_low: torch.Tensor,
|
|
mask_1ch: torch.Tensor,
|
|
anchor_mix: float,
|
|
mean_anchor_mix: float,
|
|
contrast_restore: float,
|
|
) -> torch.Tensor:
|
|
warped_mean, warped_std = _weighted_channel_stats(warped_low, mask_1ch)
|
|
anchor_mean, anchor_std = _weighted_channel_stats(anchor_low, mask_1ch)
|
|
|
|
mean_target = torch.lerp(warped_mean, anchor_mean, float(mean_anchor_mix))
|
|
contrast_gain = torch.lerp(
|
|
torch.ones_like(anchor_std),
|
|
anchor_std / warped_std.clamp_min(1e-5),
|
|
float(contrast_restore),
|
|
)
|
|
restored = mean_target + (warped_low - warped_mean) * contrast_gain
|
|
return torch.lerp(restored, anchor_low, float(anchor_mix))
|
|
|
|
|
|
def latent_manifold_compand(
|
|
base_denoised: torch.Tensor,
|
|
anchor_denoised: torch.Tensor,
|
|
mask: Optional[torch.Tensor],
|
|
strength: float,
|
|
cutoff: float,
|
|
radial_strength: float,
|
|
anisotropy: float,
|
|
translation_strength: float,
|
|
anchor_mix: float,
|
|
mean_anchor_mix: float,
|
|
contrast_restore: float,
|
|
energy_tether: float,
|
|
channel_tether: float,
|
|
energy_gain_cap: float,
|
|
max_shift_px: float,
|
|
) -> torch.Tensor:
|
|
strength = float(max(0.0, min(1.0, strength)))
|
|
if strength <= 0.0:
|
|
return base_denoised
|
|
|
|
cutoff = float(max(0.05, min(1.0, cutoff)))
|
|
work_dtype = base_denoised.dtype
|
|
compute_dtype = torch.float32
|
|
|
|
base = base_denoised.to(dtype=compute_dtype)
|
|
anchor = anchor_denoised.to(device=base.device, dtype=compute_dtype)
|
|
if tuple(anchor.shape[-2:]) != tuple(base.shape[-2:]):
|
|
anchor = resize_4d_tensor(anchor, tuple(base.shape[-2:]))
|
|
|
|
if mask is None:
|
|
mask_1ch = torch.ones((base.shape[0], 1, base.shape[-2], base.shape[-1]), device=base.device, dtype=compute_dtype)
|
|
else:
|
|
if mask.ndim == 3:
|
|
mask = mask.unsqueeze(1)
|
|
elif mask.ndim == 4 and mask.shape[1] != 1:
|
|
mask = mask.mean(dim=1, keepdim=True)
|
|
mask_1ch = _expand_mask_channels(mask.to(device=base.device, dtype=compute_dtype), base)[:, :1].clamp(0.0, 1.0)
|
|
|
|
base_low = lowpass_latent(base, cutoff)
|
|
anchor_low = lowpass_latent(anchor, cutoff)
|
|
base_high = base - base_low
|
|
|
|
regulated_low = tether_latent_low_frequency_energy(
|
|
base_low,
|
|
anchor_low,
|
|
mask_1ch,
|
|
energy_tether=float(max(0.0, min(1.0, energy_tether))) * strength,
|
|
channel_tether=float(max(0.0, min(1.0, channel_tether))) * strength,
|
|
gain_cap=float(max(1.0, energy_gain_cap)),
|
|
)
|
|
|
|
flow = estimate_latent_compaction_flow(
|
|
anchor_low=anchor_low,
|
|
base_low=regulated_low,
|
|
mask_1ch=mask_1ch,
|
|
strength=strength,
|
|
radial_strength=radial_strength,
|
|
anisotropy=anisotropy,
|
|
translation_strength=translation_strength,
|
|
max_shift_px=max_shift_px,
|
|
)
|
|
warped_low = warp_4d_tensor(regulated_low, flow, mode="bilinear")
|
|
restored_low = restore_latent_low_frequency_stats(
|
|
warped_low,
|
|
anchor_low,
|
|
mask_1ch,
|
|
anchor_mix=float(max(0.0, min(1.0, anchor_mix))) * strength,
|
|
mean_anchor_mix=float(max(0.0, min(1.0, mean_anchor_mix))) * strength,
|
|
contrast_restore=float(max(0.0, min(1.0, contrast_restore))) * strength,
|
|
)
|
|
corrected = base_high + restored_low
|
|
return corrected.to(dtype=work_dtype)
|
|
|
|
|
|
def clamp(value: float, lo: float, hi: float) -> float:
|
|
return float(max(lo, min(hi, value)))
|
|
|
|
|
|
def sigma_progress(planner_sigmas: Optional[Sequence[float]], idx: int) -> Optional[float]:
|
|
if planner_sigmas is None:
|
|
return None
|
|
|
|
if len(planner_sigmas) == 0:
|
|
return None
|
|
|
|
idx = max(0, min(int(idx), len(planner_sigmas) - 1))
|
|
sigma_hi = max(float(planner_sigmas[0]), 1e-6)
|
|
sigma_lo = max(float(planner_sigmas[-1]), 1e-6)
|
|
sigma = max(float(planner_sigmas[idx]), 1e-6)
|
|
|
|
hi = math.log(sigma_hi)
|
|
lo = math.log(sigma_lo)
|
|
denom = hi - lo
|
|
if abs(denom) < 1e-6:
|
|
return None
|
|
|
|
cur = math.log(sigma)
|
|
return clamp((hi - cur) / denom, 0.0, 1.0)
|
|
|
|
|
|
def schedule_curve(progress: float, mode: str, power: float = 2.0, hold: float = 0.0) -> float:
|
|
p = clamp(progress, 0.0, 1.0)
|
|
power = max(1e-6, float(power))
|
|
hold = clamp(hold, 0.0, 0.95)
|
|
|
|
if mode == "flat":
|
|
return 0.0
|
|
if mode == "linear":
|
|
return p
|
|
if mode == "cosine":
|
|
return 0.5 - 0.5 * math.cos(math.pi * p)
|
|
if mode == "smoothstep":
|
|
return p * p * (3.0 - 2.0 * p)
|
|
if mode == "smootherstep":
|
|
return p * p * p * (p * (p * 6.0 - 15.0) + 10.0)
|
|
if mode == "ease_in":
|
|
return p ** power
|
|
if mode == "ease_out":
|
|
return 1.0 - (1.0 - p) ** power
|
|
if mode == "ease_in_out":
|
|
if p < 0.5:
|
|
return 0.5 * ((2.0 * p) ** power)
|
|
return 1.0 - 0.5 * ((2.0 * (1.0 - p)) ** power)
|
|
if mode == "hold_then_drop":
|
|
if p <= hold:
|
|
return 0.0
|
|
u = (p - hold) / max(1e-6, 1.0 - hold)
|
|
return u ** power
|
|
if mode == "fast_drop":
|
|
return p ** (1.0 / power)
|
|
return p
|
|
|
|
|
|
def schedule_value(start: float, end: float, progress: float, mode: str, power: float = 2.0, hold: float = 0.0) -> float:
|
|
t = schedule_curve(progress, mode, power=power, hold=hold)
|
|
return float(start + (end - start) * t)
|
|
|
|
|
|
@dataclass
|
|
class TrajectoryRecorder:
|
|
store_dtype: torch.dtype = torch.float16
|
|
capture_noisy_xt: bool = False
|
|
pin_memory: bool = False
|
|
|
|
def __post_init__(self) -> None:
|
|
self.x0_steps: list[torch.Tensor] = []
|
|
self.xt_steps: list[torch.Tensor] = []
|
|
self.step_sigmas: list[float] = []
|
|
|
|
def callback(self, step: int, x0: torch.Tensor, x: torch.Tensor, total_steps: int) -> None:
|
|
del step, total_steps
|
|
self.x0_steps.append(_stage_cpu_tensor(x0, self.store_dtype, self.pin_memory))
|
|
if self.capture_noisy_xt:
|
|
self.xt_steps.append(_stage_cpu_tensor(x, self.store_dtype, self.pin_memory))
|
|
|
|
|
|
class ScaleLockedCFGGuider:
|
|
"""
|
|
A custom ComfyUI guider that applies the scale lock in denoised-latent space.
|
|
|
|
It subclasses the runtime behavior of comfy.samplers.CFGGuider at usage time
|
|
(the actual inheritance happens in nodes.py where comfy is available).
|
|
This standalone class only implements the extra logic/state we need.
|
|
"""
|
|
|
|
def _init_scale_lock_state(
|
|
self,
|
|
*,
|
|
model,
|
|
anchors_x0_cpu: Iterable[torch.Tensor],
|
|
planner_sigmas: Optional[Iterable[float]],
|
|
lock_strength: float,
|
|
lock_strength_start: float,
|
|
lock_strength_end: float,
|
|
cutoff: float,
|
|
mid_cutoff: float,
|
|
mid_strength: float,
|
|
schedule: str,
|
|
schedule_power: float,
|
|
schedule_hold: float,
|
|
mid_strength_start: float,
|
|
mid_strength_end: float,
|
|
mid_schedule: str,
|
|
mid_schedule_power: float,
|
|
mid_schedule_hold: float,
|
|
spatial_mask: Optional[torch.Tensor],
|
|
manifold_enabled: bool = False,
|
|
manifold_strength: float = 0.0,
|
|
manifold_strength_start: float = 1.0,
|
|
manifold_strength_end: float = 0.0,
|
|
manifold_schedule: str = "ease_out",
|
|
manifold_schedule_power: float = 2.0,
|
|
manifold_schedule_hold: float = 0.0,
|
|
manifold_cutoff: float = 0.18,
|
|
manifold_radial_strength: float = 1.0,
|
|
manifold_anisotropy: float = 0.15,
|
|
manifold_translation_strength: float = 1.0,
|
|
manifold_anchor_mix: float = 0.18,
|
|
manifold_mean_anchor_mix: float = 0.12,
|
|
manifold_contrast_restore: float = 0.10,
|
|
manifold_energy_tether: float = 0.0,
|
|
manifold_channel_tether: float = 0.0,
|
|
manifold_energy_gain_cap: float = 1.75,
|
|
manifold_max_shift_px: float = 3.0,
|
|
manifold_spatial_mask: Optional[torch.Tensor] = None,
|
|
) -> None:
|
|
init_scale_lock_state(
|
|
self,
|
|
model=model,
|
|
anchors_x0_cpu=anchors_x0_cpu,
|
|
planner_sigmas=planner_sigmas,
|
|
lock_strength=lock_strength,
|
|
lock_strength_start=lock_strength_start,
|
|
lock_strength_end=lock_strength_end,
|
|
cutoff=cutoff,
|
|
mid_cutoff=mid_cutoff,
|
|
mid_strength=mid_strength,
|
|
schedule=schedule,
|
|
schedule_power=schedule_power,
|
|
schedule_hold=schedule_hold,
|
|
mid_strength_start=mid_strength_start,
|
|
mid_strength_end=mid_strength_end,
|
|
mid_schedule=mid_schedule,
|
|
mid_schedule_power=mid_schedule_power,
|
|
mid_schedule_hold=mid_schedule_hold,
|
|
spatial_mask=spatial_mask,
|
|
manifold_enabled=manifold_enabled,
|
|
manifold_strength=manifold_strength,
|
|
manifold_strength_start=manifold_strength_start,
|
|
manifold_strength_end=manifold_strength_end,
|
|
manifold_schedule=manifold_schedule,
|
|
manifold_schedule_power=manifold_schedule_power,
|
|
manifold_schedule_hold=manifold_schedule_hold,
|
|
manifold_cutoff=manifold_cutoff,
|
|
manifold_radial_strength=manifold_radial_strength,
|
|
manifold_anisotropy=manifold_anisotropy,
|
|
manifold_translation_strength=manifold_translation_strength,
|
|
manifold_anchor_mix=manifold_anchor_mix,
|
|
manifold_mean_anchor_mix=manifold_mean_anchor_mix,
|
|
manifold_contrast_restore=manifold_contrast_restore,
|
|
manifold_energy_tether=manifold_energy_tether,
|
|
manifold_channel_tether=manifold_channel_tether,
|
|
manifold_energy_gain_cap=manifold_energy_gain_cap,
|
|
manifold_max_shift_px=manifold_max_shift_px,
|
|
manifold_spatial_mask=manifold_spatial_mask,
|
|
)
|
|
|
|
def _slrd_resolve_step_index(self, timestep: torch.Tensor | float | int) -> int:
|
|
return resolve_scale_lock_step_index(self, timestep)
|
|
|
|
def _slrd_strength_for_step(self, idx: int) -> float:
|
|
return scale_lock_strength_for_step(self, idx)
|
|
|
|
def _slrd_anchor_for(self, idx: int, like: torch.Tensor) -> torch.Tensor:
|
|
return scale_lock_anchor_for(self, idx, like)
|
|
|
|
def _slrd_mask_for(self, like: torch.Tensor) -> Optional[torch.Tensor]:
|
|
return scale_lock_mask_for(self, like)
|
|
|
|
def _slrd_manifold_strength_for_step(self, idx: int) -> float:
|
|
return scale_lock_manifold_strength_for_step(self, idx)
|
|
|
|
def _slrd_manifold_mask_for(self, like: torch.Tensor) -> Optional[torch.Tensor]:
|
|
return scale_lock_manifold_mask_for(self, like)
|
|
|
|
|
|
def init_scale_lock_state(
|
|
guider,
|
|
*,
|
|
model,
|
|
anchors_x0_cpu: Iterable[torch.Tensor],
|
|
planner_sigmas: Optional[Iterable[float]],
|
|
lock_strength: float,
|
|
lock_strength_start: float,
|
|
lock_strength_end: float,
|
|
cutoff: float,
|
|
mid_cutoff: float,
|
|
mid_strength: float,
|
|
schedule: str,
|
|
schedule_power: float,
|
|
schedule_hold: float,
|
|
mid_strength_start: float,
|
|
mid_strength_end: float,
|
|
mid_schedule: str,
|
|
mid_schedule_power: float,
|
|
mid_schedule_hold: float,
|
|
spatial_mask: Optional[torch.Tensor],
|
|
manifold_enabled: bool = False,
|
|
manifold_strength: float = 0.0,
|
|
manifold_strength_start: float = 1.0,
|
|
manifold_strength_end: float = 0.0,
|
|
manifold_schedule: str = "ease_out",
|
|
manifold_schedule_power: float = 2.0,
|
|
manifold_schedule_hold: float = 0.0,
|
|
manifold_cutoff: float = 0.18,
|
|
manifold_radial_strength: float = 1.0,
|
|
manifold_anisotropy: float = 0.15,
|
|
manifold_translation_strength: float = 1.0,
|
|
manifold_anchor_mix: float = 0.18,
|
|
manifold_mean_anchor_mix: float = 0.12,
|
|
manifold_contrast_restore: float = 0.10,
|
|
manifold_energy_tether: float = 0.0,
|
|
manifold_channel_tether: float = 0.0,
|
|
manifold_energy_gain_cap: float = 1.75,
|
|
manifold_max_shift_px: float = 3.0,
|
|
manifold_spatial_mask: Optional[torch.Tensor] = None,
|
|
) -> None:
|
|
guider._slrd_model = model
|
|
guider._slrd_anchors_x0_cpu = list(anchors_x0_cpu)
|
|
guider._slrd_planner_sigmas = [float(x) for x in planner_sigmas] if planner_sigmas is not None else None
|
|
guider._slrd_lock_strength = float(lock_strength)
|
|
guider._slrd_lock_strength_start = float(lock_strength_start)
|
|
guider._slrd_lock_strength_end = float(lock_strength_end)
|
|
guider._slrd_cutoff = float(cutoff)
|
|
guider._slrd_mid_cutoff = float(max(cutoff, mid_cutoff))
|
|
guider._slrd_mid_strength = float(mid_strength)
|
|
guider._slrd_schedule = schedule
|
|
guider._slrd_schedule_power = float(schedule_power)
|
|
guider._slrd_schedule_hold = float(schedule_hold)
|
|
guider._slrd_mid_strength_start = float(mid_strength_start)
|
|
guider._slrd_mid_strength_end = float(mid_strength_end)
|
|
guider._slrd_mid_schedule = mid_schedule
|
|
guider._slrd_mid_schedule_power = float(mid_schedule_power)
|
|
guider._slrd_mid_schedule_hold = float(mid_schedule_hold)
|
|
guider._slrd_seen_sigmas = []
|
|
guider._slrd_prev_match_idx = 0
|
|
guider._slrd_spatial_mask = spatial_mask
|
|
guider._slrd_last_sigma = None
|
|
guider._slrd_manifold_enabled = bool(manifold_enabled)
|
|
guider._slrd_manifold_strength = float(manifold_strength)
|
|
guider._slrd_manifold_strength_start = float(manifold_strength_start)
|
|
guider._slrd_manifold_strength_end = float(manifold_strength_end)
|
|
guider._slrd_manifold_schedule = manifold_schedule
|
|
guider._slrd_manifold_schedule_power = float(manifold_schedule_power)
|
|
guider._slrd_manifold_schedule_hold = float(manifold_schedule_hold)
|
|
guider._slrd_manifold_cutoff = float(max(0.05, min(1.0, manifold_cutoff)))
|
|
guider._slrd_manifold_radial_strength = float(manifold_radial_strength)
|
|
guider._slrd_manifold_anisotropy = float(manifold_anisotropy)
|
|
guider._slrd_manifold_translation_strength = float(manifold_translation_strength)
|
|
guider._slrd_manifold_anchor_mix = float(manifold_anchor_mix)
|
|
guider._slrd_manifold_mean_anchor_mix = float(manifold_mean_anchor_mix)
|
|
guider._slrd_manifold_contrast_restore = float(manifold_contrast_restore)
|
|
guider._slrd_manifold_energy_tether = float(max(0.0, min(1.0, manifold_energy_tether)))
|
|
guider._slrd_manifold_channel_tether = float(max(0.0, min(1.0, manifold_channel_tether)))
|
|
guider._slrd_manifold_energy_gain_cap = float(max(1.0, manifold_energy_gain_cap))
|
|
guider._slrd_manifold_max_shift_px = float(manifold_max_shift_px)
|
|
guider._slrd_manifold_spatial_mask = manifold_spatial_mask if manifold_spatial_mask is not None else spatial_mask
|
|
|
|
|
|
def resolve_scale_lock_step_index(guider, timestep: torch.Tensor | float | int) -> int:
|
|
sigma = _sigma_scalar(timestep)
|
|
|
|
if guider._slrd_planner_sigmas:
|
|
best_idx = 0
|
|
best_dist = float("inf")
|
|
for i, planner_sigma in enumerate(guider._slrd_planner_sigmas):
|
|
dist = abs(planner_sigma - sigma)
|
|
if dist < best_dist:
|
|
best_idx = i
|
|
best_dist = dist
|
|
best_idx = max(best_idx, guider._slrd_prev_match_idx)
|
|
best_idx = min(best_idx, len(guider._slrd_anchors_x0_cpu) - 1)
|
|
guider._slrd_prev_match_idx = best_idx
|
|
return best_idx
|
|
|
|
if guider._slrd_last_sigma is None:
|
|
guider._slrd_last_sigma = sigma
|
|
step_index = 0
|
|
guider._slrd_seen_sigmas = [sigma]
|
|
else:
|
|
tol = 1e-6 * max(1.0, abs(guider._slrd_last_sigma), abs(sigma))
|
|
if abs(sigma - guider._slrd_last_sigma) > tol:
|
|
guider._slrd_last_sigma = sigma
|
|
guider._slrd_seen_sigmas.append(sigma)
|
|
unique_sigmas = []
|
|
for seen_sigma in guider._slrd_seen_sigmas:
|
|
if all(
|
|
abs(seen_sigma - unique_sigma) > (1e-6 * max(1.0, abs(seen_sigma), abs(unique_sigma)))
|
|
for unique_sigma in unique_sigmas
|
|
):
|
|
unique_sigmas.append(seen_sigma)
|
|
step_index = max(0, len(unique_sigmas) - 1)
|
|
step_index = max(step_index, guider._slrd_prev_match_idx)
|
|
guider._slrd_prev_match_idx = step_index
|
|
if not guider._slrd_anchors_x0_cpu:
|
|
return 0
|
|
return min(step_index, len(guider._slrd_anchors_x0_cpu) - 1)
|
|
|
|
|
|
def scale_lock_strength_for_step(guider, idx: int) -> float:
|
|
return scale_lock_strengths_for_step(guider, idx)[0]
|
|
|
|
|
|
def scale_lock_progress_for_step(guider, idx: int) -> float:
|
|
sigma_based = sigma_progress(getattr(guider, "_slrd_planner_sigmas", None), idx)
|
|
if sigma_based is not None:
|
|
return sigma_based
|
|
total = max(1, len(guider._slrd_anchors_x0_cpu) - 1)
|
|
return clamp(idx / total, 0.0, 1.0)
|
|
|
|
|
|
def _scheduled_strength(
|
|
base_strength: float,
|
|
start: float,
|
|
end: float,
|
|
progress: float,
|
|
mode: str,
|
|
power: float,
|
|
hold: float,
|
|
) -> float:
|
|
scheduled = schedule_value(start, end, progress, mode, power=power, hold=hold)
|
|
return clamp(base_strength * scheduled, 0.0, 1.0)
|
|
|
|
|
|
def scale_lock_strengths_for_step(guider, idx: int) -> tuple[float, float]:
|
|
progress = scale_lock_progress_for_step(guider, idx)
|
|
low_strength = _scheduled_strength(
|
|
guider._slrd_lock_strength,
|
|
guider._slrd_lock_strength_start,
|
|
guider._slrd_lock_strength_end,
|
|
progress,
|
|
guider._slrd_schedule,
|
|
guider._slrd_schedule_power,
|
|
guider._slrd_schedule_hold,
|
|
)
|
|
if getattr(guider, "_slrd_mid_schedule", "linked") == "linked":
|
|
mid_strength = clamp(low_strength * guider._slrd_mid_strength, 0.0, 1.0)
|
|
else:
|
|
mid_strength = _scheduled_strength(
|
|
guider._slrd_mid_strength,
|
|
guider._slrd_mid_strength_start,
|
|
guider._slrd_mid_strength_end,
|
|
progress,
|
|
guider._slrd_mid_schedule,
|
|
guider._slrd_mid_schedule_power,
|
|
guider._slrd_mid_schedule_hold,
|
|
)
|
|
return low_strength, mid_strength
|
|
|
|
|
|
def scale_lock_manifold_strength_for_step(guider, idx: int) -> float:
|
|
if not getattr(guider, "_slrd_manifold_enabled", False):
|
|
return 0.0
|
|
progress = scale_lock_progress_for_step(guider, idx)
|
|
return _scheduled_strength(
|
|
guider._slrd_manifold_strength,
|
|
guider._slrd_manifold_strength_start,
|
|
guider._slrd_manifold_strength_end,
|
|
progress,
|
|
guider._slrd_manifold_schedule,
|
|
guider._slrd_manifold_schedule_power,
|
|
guider._slrd_manifold_schedule_hold,
|
|
)
|
|
|
|
|
|
def scale_lock_anchor_for(guider, idx: int, like: torch.Tensor) -> torch.Tensor:
|
|
anchor = guider._slrd_anchors_x0_cpu[idx].to(device=like.device, dtype=like.dtype, non_blocking=True)
|
|
if tuple(anchor.shape[-2:]) != tuple(like.shape[-2:]):
|
|
anchor = resize_4d_tensor(anchor, tuple(like.shape[-2:]))
|
|
return anchor
|
|
|
|
|
|
def scale_lock_mask_for(guider, like: torch.Tensor) -> Optional[torch.Tensor]:
|
|
if guider._slrd_spatial_mask is None:
|
|
return None
|
|
mask = guider._slrd_spatial_mask.to(device=like.device, dtype=like.dtype, non_blocking=True)
|
|
if tuple(mask.shape[-2:]) != tuple(like.shape[-2:]):
|
|
mask = resize_4d_tensor(mask, tuple(like.shape[-2:]))
|
|
if mask.shape[0] < like.shape[0]:
|
|
repeat = math.ceil(like.shape[0] / max(1, mask.shape[0]))
|
|
mask = mask.repeat(repeat, 1, 1, 1)[: like.shape[0]]
|
|
elif mask.shape[0] > like.shape[0]:
|
|
mask = mask[: like.shape[0]]
|
|
return mask
|
|
|
|
|
|
def scale_lock_manifold_mask_for(guider, like: torch.Tensor) -> Optional[torch.Tensor]:
|
|
manifold_mask = getattr(guider, "_slrd_manifold_spatial_mask", None)
|
|
if manifold_mask is None:
|
|
return None
|
|
mask = manifold_mask.to(device=like.device, dtype=like.dtype, non_blocking=True)
|
|
if tuple(mask.shape[-2:]) != tuple(like.shape[-2:]):
|
|
mask = resize_4d_tensor(mask, tuple(like.shape[-2:]))
|
|
if mask.shape[0] < like.shape[0]:
|
|
repeat = math.ceil(like.shape[0] / max(1, mask.shape[0]))
|
|
mask = mask.repeat(repeat, 1, 1, 1)[: like.shape[0]]
|
|
elif mask.shape[0] > like.shape[0]:
|
|
mask = mask[: like.shape[0]]
|
|
return mask
|
|
|
|
|
|
|
|
|