facok
e7a836db4d
Two bugs caused sharpness edits to shift color even when calibrated with lcs_data: 1. Calibration order was wrong: DC removal before LCS color removal operated in the wrong space (LCS basis has large DC components), producing incorrect color projection and incomplete removal. Fix: remove color first (in raw space), then DC. 2. Intervention DC removal broke orthogonality: subtracting mean(pc1_dir) introduced a color component because the all-ones vector projects onto LCS basis. Fix: re-orthogonalize against LCS basis after DC removal. Also: store lcs_basis in SharpnessData/safetensors so intervention can re-orthogonalize; bump cache name to _grating to force re-calibration with the corrected pipeline.
214 lines
8.6 KiB
Python
214 lines
8.6 KiB
Python
"""Sharpness subspace calibration via sinusoidal grating stimuli.
|
||
|
||
Replaces the previous Gaussian blur approach with narrowband frequency
|
||
gratings, which achieve higher linearity (R²=0.94 vs 0.88) because each
|
||
stimulus contains a single spatial frequency — a purer probe of the VAE's
|
||
frequency encoding axis.
|
||
|
||
The two methods discover the same 1D subspace (|cos|=0.986, 9.7° apart),
|
||
but grating stimuli yield a cleaner PC1 direction.
|
||
"""
|
||
|
||
import math
|
||
from dataclasses import dataclass
|
||
from typing import List, Optional, Tuple
|
||
import warnings
|
||
|
||
import torch
|
||
import comfy.utils
|
||
|
||
from .patchify import patchify
|
||
from .lcs_data import LCSData
|
||
|
||
|
||
@dataclass
|
||
class SharpnessData:
|
||
"""Calibration data for the sharpness subspace.
|
||
|
||
Produced by PCA on FLUX VAE-encoded sinusoidal gratings at varying
|
||
spatial frequencies. PC1 captures ~94% of variance with R²=0.94
|
||
linearity vs log₂(frequency).
|
||
"""
|
||
|
||
basis: torch.Tensor # [64, K] PCA basis (columns), K typically 1-2
|
||
mean: torch.Tensor # [64] PCA mean (in color-removed space if lcs_data was used)
|
||
sign: float # +1 or -1: ensures positive strength = sharper
|
||
lcs_basis: Optional[torch.Tensor] = None # [64, 3] LCS basis used during calibration (for re-orthogonalization)
|
||
|
||
def to(self, device, dtype=None):
|
||
"""Move all tensors to device/dtype."""
|
||
kw = {"device": device}
|
||
if dtype is not None:
|
||
kw["dtype"] = dtype
|
||
return SharpnessData(
|
||
basis=self.basis.to(**kw),
|
||
mean=self.mean.to(**kw),
|
||
sign=self.sign,
|
||
lcs_basis=self.lcs_basis.to(**kw) if self.lcs_basis is not None else None,
|
||
)
|
||
|
||
|
||
def _generate_grating_batch(
|
||
indices: List[int],
|
||
angles: torch.Tensor,
|
||
phases: torch.Tensor,
|
||
frequencies: Tuple[float, ...],
|
||
coord_x: torch.Tensor,
|
||
coord_y: torch.Tensor,
|
||
) -> torch.Tensor:
|
||
"""Generate a batch of sinusoidal grating stimuli by flat index.
|
||
|
||
Each flat index maps to (orientation, frequency) via divmod.
|
||
Returns [len(indices), 3, H, W] tensor.
|
||
"""
|
||
num_freqs = len(frequencies)
|
||
batch = []
|
||
for idx in indices:
|
||
ori = idx // num_freqs
|
||
freq = frequencies[idx % num_freqs]
|
||
angle = angles[ori].item()
|
||
phase = phases[ori].item()
|
||
cos_a, sin_a = math.cos(angle), math.sin(angle)
|
||
coord = coord_x * cos_a + coord_y * sin_a
|
||
grating = 0.5 + 0.3 * torch.sin(2 * math.pi * freq * coord + phase)
|
||
batch.append(grating.unsqueeze(0).expand(3, -1, -1))
|
||
return torch.stack(batch, dim=0)
|
||
|
||
|
||
def calibrate_sharpness(vae, num_samples: int = 64, image_size: int = 512,
|
||
frequencies: Tuple[float, ...] = (1, 2, 4, 8, 16, 32, 64),
|
||
batch_size: int = 8,
|
||
lcs_data: LCSData = None,
|
||
# Legacy parameter — accepted but ignored
|
||
blur_levels: Optional[Tuple[float, ...]] = None,
|
||
) -> SharpnessData:
|
||
"""Compute sharpness subspace data (PCA basis, mean, sign) from FLUX VAE.
|
||
|
||
Generates sinusoidal gratings at varying spatial frequencies (one pure
|
||
frequency per stimulus), VAE-encodes them, and runs PCA to find the
|
||
sharpness/frequency direction in 64D patch space.
|
||
|
||
Args:
|
||
vae: ComfyUI VAE object
|
||
num_samples: Number of orientations (each combined with all frequencies)
|
||
image_size: Size of generated images
|
||
frequencies: Spatial frequencies in cycles/image
|
||
batch_size: Batch size for VAE encoding
|
||
lcs_data: Optional LCS data for removing color component during calibration.
|
||
When provided, the sharpness PC1 will be orthogonal to the color subspace,
|
||
preventing color shifts during intervention.
|
||
|
||
Returns: SharpnessData
|
||
"""
|
||
if blur_levels is not None:
|
||
warnings.warn(
|
||
"blur_levels is deprecated and ignored; calibration now uses sinusoidal gratings",
|
||
DeprecationWarning, stacklevel=2,
|
||
)
|
||
|
||
n_freqs = len(frequencies)
|
||
total_images = num_samples * n_freqs
|
||
|
||
print(f"\n[LCS Sharpness Calibration] Starting: {num_samples} orientations × {n_freqs} frequencies = {total_images} stimuli")
|
||
print(f"[LCS Sharpness Calibration] Frequencies: {list(frequencies)} cycles/image")
|
||
|
||
# Pre-compute shared state for grating generation
|
||
gen = torch.Generator().manual_seed(42)
|
||
angles = torch.rand(num_samples, generator=gen) * math.pi # [0, π)
|
||
phases = torch.rand(num_samples, generator=gen) * 2 * math.pi # [0, 2π)
|
||
y_coords = torch.linspace(-0.5, 0.5, image_size).unsqueeze(1)
|
||
x_coords = torch.linspace(-0.5, 0.5, image_size).unsqueeze(0)
|
||
coord_y = y_coords.expand(image_size, image_size)
|
||
coord_x = x_coords.expand(image_size, image_size)
|
||
|
||
# Build frequency labels for all stimuli (flat index → frequency)
|
||
freq_labels = [frequencies[idx % n_freqs] for idx in range(total_images)]
|
||
freq_labels_t = torch.tensor(freq_labels, dtype=torch.float32)
|
||
log_freq = torch.log2(freq_labels_t.clamp(min=0.5))
|
||
|
||
# Generate stimuli lazily per batch and VAE encode
|
||
vectors = []
|
||
pbar = comfy.utils.ProgressBar(total_images)
|
||
|
||
for batch_start in range(0, total_images, batch_size):
|
||
batch_end = min(batch_start + batch_size, total_images)
|
||
indices = list(range(batch_start, batch_end))
|
||
batch = _generate_grating_batch(indices, angles, phases, frequencies, coord_x, coord_y)
|
||
actual_batch = batch.shape[0]
|
||
|
||
# Convert BCHW → BHWC for ComfyUI VAE
|
||
imgs_bhwc = batch.permute(0, 2, 3, 1).contiguous().cpu()
|
||
|
||
# VAE encode — try batch first, fall back to per-image for video VAEs
|
||
latent = vae.encode(imgs_bhwc)
|
||
patches, _, _, _ = patchify(latent)
|
||
avg = patches.mean(dim=1).cpu()
|
||
|
||
if avg.shape[0] == actual_batch:
|
||
vectors.extend(avg.unbind(0))
|
||
else:
|
||
# Video VAE: batch not fully supported, encode one by one
|
||
vectors.extend(avg.unbind(0))
|
||
for k in range(1, actual_batch):
|
||
single = imgs_bhwc[k:k+1]
|
||
lat = vae.encode(single)
|
||
p, _, _, _ = patchify(lat)
|
||
vectors.append(p.mean(dim=1).cpu().squeeze(0))
|
||
|
||
pbar.update(actual_batch)
|
||
|
||
# Stack all vectors: [N, 64]
|
||
X = torch.stack(vectors, dim=0).float()
|
||
print(f"[LCS Sharpness Calibration] Collected {X.shape[0]} vectors of dimension {X.shape[1]}")
|
||
|
||
# Remove LCS color component FIRST, in the raw space where LCS was calibrated.
|
||
# This must happen before per-vector DC removal, because the LCS basis has
|
||
# significant DC components (PC1 ≈ brightness). Doing DC removal first would
|
||
# shift vectors into a different space where B^T(x - mu) is incorrect.
|
||
if lcs_data is not None:
|
||
print("[LCS Sharpness Calibration] Removing LCS color component...")
|
||
lcs_mean = lcs_data.mean.to(X.device, X.dtype)
|
||
lcs_basis = lcs_data.basis.to(X.device, X.dtype)
|
||
# Project out color: X' = X - B B^T (X - mu)
|
||
centered = X - lcs_mean
|
||
lcs_coords = centered @ lcs_basis # [N, 3]
|
||
X = X - lcs_coords @ lcs_basis.T
|
||
print("[LCS Sharpness Calibration] Color component removed")
|
||
|
||
# Remove per-vector DC AFTER color removal.
|
||
# VAE encoding shifts the latent mean depending on stimulus content.
|
||
# Per-vector zero-mean forces PCA to find patterns in the relative channel
|
||
# structure, not in the absolute level.
|
||
X = X - X.mean(dim=1, keepdim=True)
|
||
|
||
# Step 3: PCA
|
||
print("[LCS Sharpness Calibration] Computing PCA...")
|
||
mean = X.mean(dim=0) # [64]
|
||
X_centered = X - mean
|
||
U, S, Vh = torch.linalg.svd(X_centered, full_matrices=False)
|
||
# Top 2 components
|
||
basis = Vh[:2].T # [64, 2]
|
||
|
||
# Variance explained
|
||
total_var = (S ** 2).sum()
|
||
explained = (S[:2] ** 2) / total_var
|
||
print(f"[LCS Sharpness Calibration] PC1: {explained[0]:.1%}, PC2: {explained[1]:.1%} ({(explained[0]+explained[1]):.1%} total)")
|
||
|
||
# Step 4: Determine sign convention
|
||
# Project all vectors onto PC1
|
||
pc1_scores = X_centered @ basis[:, 0] # [N]
|
||
|
||
# Correlate PC1 score with log₂(frequency)
|
||
# Higher frequency = sharper → if positive correlation, sign = +1
|
||
correlation = torch.corrcoef(torch.stack([pc1_scores, log_freq]))[0, 1]
|
||
sign = 1.0 if correlation > 0 else -1.0
|
||
print(f"[LCS Sharpness Calibration] PC1-frequency correlation: {correlation:.3f} → sign = {sign:+.0f}")
|
||
print(f"[LCS Sharpness Calibration] Complete! Basis shape: {basis.shape}")
|
||
|
||
return SharpnessData(
|
||
basis=basis,
|
||
mean=mean,
|
||
sign=sign,
|
||
lcs_basis=lcs_data.basis.clone() if lcs_data is not None else None,
|
||
)
|