|
| 1 | +import torch |
| 2 | +import torch.nn.functional as F |
| 3 | +from torch import Tensor |
| 4 | + |
| 5 | +from comfy.ldm.seedvr.constants import ( |
| 6 | + CIELAB_DELTA, |
| 7 | + CIELAB_KAPPA, |
| 8 | + D65_WHITE_X, |
| 9 | + D65_WHITE_Z, |
| 10 | + WAVELET_DECOMP_LEVELS, |
| 11 | +) |
| 12 | + |
| 13 | + |
| 14 | +def wavelet_blur(image: Tensor, radius): |
| 15 | + max_safe_radius = max(1, min(image.shape[-2:]) // 8) |
| 16 | + if radius > max_safe_radius: |
| 17 | + radius = max_safe_radius |
| 18 | + |
| 19 | + num_channels = image.shape[1] |
| 20 | + |
| 21 | + kernel_vals = [ |
| 22 | + [0.0625, 0.125, 0.0625], |
| 23 | + [0.125, 0.25, 0.125], |
| 24 | + [0.0625, 0.125, 0.0625], |
| 25 | + ] |
| 26 | + kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device) |
| 27 | + kernel = kernel[None, None].repeat(num_channels, 1, 1, 1) |
| 28 | + |
| 29 | + image = F.pad(image, (radius, radius, radius, radius), mode='replicate') |
| 30 | + output = F.conv2d(image, kernel, groups=num_channels, dilation=radius) |
| 31 | + |
| 32 | + return output |
| 33 | + |
| 34 | +def wavelet_decomposition(image: Tensor, levels: int = WAVELET_DECOMP_LEVELS): |
| 35 | + high_freq = torch.zeros_like(image) |
| 36 | + |
| 37 | + for i in range(levels): |
| 38 | + radius = 2 ** i |
| 39 | + low_freq = wavelet_blur(image, radius) |
| 40 | + high_freq.add_(image).sub_(low_freq) |
| 41 | + image = low_freq |
| 42 | + |
| 43 | + return high_freq, low_freq |
| 44 | + |
| 45 | +def wavelet_reconstruction(content_feat: Tensor, style_feat: Tensor) -> Tensor: |
| 46 | + |
| 47 | + if content_feat.shape != style_feat.shape: |
| 48 | + if len(content_feat.shape) >= 3: |
| 49 | + style_feat = F.interpolate( |
| 50 | + style_feat, |
| 51 | + size=content_feat.shape[-2:], |
| 52 | + mode='bilinear', |
| 53 | + align_corners=False |
| 54 | + ) |
| 55 | + |
| 56 | + content_high_freq, content_low_freq = wavelet_decomposition(content_feat) |
| 57 | + del content_low_freq |
| 58 | + |
| 59 | + style_high_freq, style_low_freq = wavelet_decomposition(style_feat) |
| 60 | + del style_high_freq |
| 61 | + |
| 62 | + if content_high_freq.shape != style_low_freq.shape: |
| 63 | + style_low_freq = F.interpolate( |
| 64 | + style_low_freq, |
| 65 | + size=content_high_freq.shape[-2:], |
| 66 | + mode='bilinear', |
| 67 | + align_corners=False |
| 68 | + ) |
| 69 | + |
| 70 | + content_high_freq.add_(style_low_freq) |
| 71 | + |
| 72 | + return content_high_freq.clamp_(-1.0, 1.0) |
| 73 | + |
| 74 | +def _histogram_matching_channel(source: Tensor, reference: Tensor) -> Tensor: |
| 75 | + original_shape = source.shape |
| 76 | + |
| 77 | + source_flat = source.flatten() |
| 78 | + reference_flat = reference.flatten() |
| 79 | + |
| 80 | + source_sorted, source_indices = torch.sort(source_flat) |
| 81 | + reference_sorted, _ = torch.sort(reference_flat) |
| 82 | + del reference_flat |
| 83 | + |
| 84 | + n_source = len(source_sorted) |
| 85 | + n_reference = len(reference_sorted) |
| 86 | + |
| 87 | + if n_source == n_reference: |
| 88 | + matched_sorted = reference_sorted |
| 89 | + else: |
| 90 | + source_quantiles = torch.linspace(0, 1, n_source, device=source.device) |
| 91 | + ref_indices = (source_quantiles * (n_reference - 1)).long() |
| 92 | + ref_indices.clamp_(0, n_reference - 1) |
| 93 | + matched_sorted = reference_sorted[ref_indices] |
| 94 | + del source_quantiles, ref_indices, reference_sorted |
| 95 | + |
| 96 | + del source_sorted, source_flat |
| 97 | + |
| 98 | + inverse_indices = torch.argsort(source_indices) |
| 99 | + del source_indices |
| 100 | + matched_flat = matched_sorted[inverse_indices] |
| 101 | + del matched_sorted, inverse_indices |
| 102 | + |
| 103 | + return matched_flat.reshape(original_shape) |
| 104 | + |
| 105 | +def _lab_to_rgb_batch(lab: Tensor, matrix_inv: Tensor, epsilon: float, kappa: float) -> Tensor: |
| 106 | + L, a, b = lab[:, 0], lab[:, 1], lab[:, 2] |
| 107 | + |
| 108 | + fy = (L + 16.0) / 116.0 |
| 109 | + fx = a.div(500.0).add_(fy) |
| 110 | + fz = fy - b / 200.0 |
| 111 | + del L, a, b |
| 112 | + |
| 113 | + x = torch.where( |
| 114 | + fx > epsilon, |
| 115 | + torch.pow(fx, 3.0), |
| 116 | + fx.mul(116.0).sub_(16.0).div_(kappa) |
| 117 | + ) |
| 118 | + y = torch.where( |
| 119 | + fy > epsilon, |
| 120 | + torch.pow(fy, 3.0), |
| 121 | + fy.mul(116.0).sub_(16.0).div_(kappa) |
| 122 | + ) |
| 123 | + z = torch.where( |
| 124 | + fz > epsilon, |
| 125 | + torch.pow(fz, 3.0), |
| 126 | + fz.mul(116.0).sub_(16.0).div_(kappa) |
| 127 | + ) |
| 128 | + del fx, fy, fz |
| 129 | + |
| 130 | + x.mul_(D65_WHITE_X) |
| 131 | + z.mul_(D65_WHITE_Z) |
| 132 | + |
| 133 | + xyz = torch.stack([x, y, z], dim=1) |
| 134 | + del x, y, z |
| 135 | + |
| 136 | + B, _, H, W = xyz.shape |
| 137 | + xyz_flat = xyz.permute(0, 2, 3, 1).reshape(-1, 3) |
| 138 | + del xyz |
| 139 | + |
| 140 | + xyz_flat = xyz_flat.to(dtype=matrix_inv.dtype) |
| 141 | + rgb_linear_flat = torch.matmul(xyz_flat, matrix_inv.T) |
| 142 | + del xyz_flat |
| 143 | + |
| 144 | + rgb_linear = rgb_linear_flat.reshape(B, H, W, 3).permute(0, 3, 1, 2) |
| 145 | + del rgb_linear_flat |
| 146 | + |
| 147 | + mask = rgb_linear > 0.0031308 |
| 148 | + rgb = torch.where( |
| 149 | + mask, |
| 150 | + torch.pow(torch.clamp(rgb_linear, min=0.0), 1.0 / 2.4).mul_(1.055).sub_(0.055), |
| 151 | + rgb_linear * 12.92 |
| 152 | + ) |
| 153 | + del mask, rgb_linear |
| 154 | + |
| 155 | + return torch.clamp(rgb, 0.0, 1.0) |
| 156 | + |
| 157 | +def _rgb_to_lab_batch(rgb: Tensor, matrix: Tensor, epsilon: float, kappa: float) -> Tensor: |
| 158 | + mask = rgb > 0.04045 |
| 159 | + rgb_linear = torch.where( |
| 160 | + mask, |
| 161 | + torch.pow((rgb + 0.055) / 1.055, 2.4), |
| 162 | + rgb / 12.92 |
| 163 | + ) |
| 164 | + del mask |
| 165 | + |
| 166 | + B, _, H, W = rgb_linear.shape |
| 167 | + rgb_flat = rgb_linear.permute(0, 2, 3, 1).reshape(-1, 3) |
| 168 | + del rgb_linear |
| 169 | + |
| 170 | + rgb_flat = rgb_flat.to(dtype=matrix.dtype) |
| 171 | + xyz_flat = torch.matmul(rgb_flat, matrix.T) |
| 172 | + del rgb_flat |
| 173 | + |
| 174 | + xyz = xyz_flat.reshape(B, H, W, 3).permute(0, 3, 1, 2) |
| 175 | + del xyz_flat |
| 176 | + |
| 177 | + xyz[:, 0].div_(D65_WHITE_X) |
| 178 | + xyz[:, 2].div_(D65_WHITE_Z) |
| 179 | + |
| 180 | + epsilon_cubed = epsilon ** 3 |
| 181 | + mask = xyz > epsilon_cubed |
| 182 | + f_xyz = torch.where( |
| 183 | + mask, |
| 184 | + torch.pow(xyz, 1.0 / 3.0), |
| 185 | + xyz.mul(kappa).add_(16.0).div_(116.0) |
| 186 | + ) |
| 187 | + del xyz, mask |
| 188 | + |
| 189 | + L = f_xyz[:, 1].mul(116.0).sub_(16.0) |
| 190 | + a = (f_xyz[:, 0] - f_xyz[:, 1]).mul_(500.0) |
| 191 | + b = (f_xyz[:, 1] - f_xyz[:, 2]).mul_(200.0) |
| 192 | + del f_xyz |
| 193 | + |
| 194 | + return torch.stack([L, a, b], dim=1) |
| 195 | + |
| 196 | +def lab_color_transfer( |
| 197 | + content_feat: Tensor, |
| 198 | + style_feat: Tensor, |
| 199 | + luminance_weight: float = 0.8 |
| 200 | +) -> Tensor: |
| 201 | + content_feat = wavelet_reconstruction(content_feat, style_feat) |
| 202 | + |
| 203 | + if content_feat.shape != style_feat.shape: |
| 204 | + style_feat = F.interpolate( |
| 205 | + style_feat, |
| 206 | + size=content_feat.shape[-2:], |
| 207 | + mode='bilinear', |
| 208 | + align_corners=False |
| 209 | + ) |
| 210 | + |
| 211 | + device = content_feat.device |
| 212 | + original_dtype = content_feat.dtype |
| 213 | + content_feat = content_feat.float() |
| 214 | + style_feat = style_feat.float() |
| 215 | + |
| 216 | + rgb_to_xyz_matrix = torch.tensor([ |
| 217 | + [0.4124564, 0.3575761, 0.1804375], |
| 218 | + [0.2126729, 0.7151522, 0.0721750], |
| 219 | + [0.0193339, 0.1191920, 0.9503041] |
| 220 | + ], dtype=torch.float32, device=device) |
| 221 | + |
| 222 | + xyz_to_rgb_matrix = torch.tensor([ |
| 223 | + [ 3.2404542, -1.5371385, -0.4985314], |
| 224 | + [-0.9692660, 1.8760108, 0.0415560], |
| 225 | + [ 0.0556434, -0.2040259, 1.0572252] |
| 226 | + ], dtype=torch.float32, device=device) |
| 227 | + |
| 228 | + epsilon = CIELAB_DELTA |
| 229 | + kappa = CIELAB_KAPPA |
| 230 | + |
| 231 | + content_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0) |
| 232 | + style_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0) |
| 233 | + |
| 234 | + content_lab = _rgb_to_lab_batch(content_feat, rgb_to_xyz_matrix, epsilon, kappa) |
| 235 | + del content_feat |
| 236 | + |
| 237 | + style_lab = _rgb_to_lab_batch(style_feat, rgb_to_xyz_matrix, epsilon, kappa) |
| 238 | + del style_feat, rgb_to_xyz_matrix |
| 239 | + |
| 240 | + matched_a = _histogram_matching_channel(content_lab[:, 1], style_lab[:, 1]) |
| 241 | + matched_b = _histogram_matching_channel(content_lab[:, 2], style_lab[:, 2]) |
| 242 | + |
| 243 | + if luminance_weight < 1.0: |
| 244 | + matched_L = _histogram_matching_channel(content_lab[:, 0], style_lab[:, 0]) |
| 245 | + result_L = content_lab[:, 0].mul(luminance_weight).add_(matched_L.mul(1.0 - luminance_weight)) |
| 246 | + del matched_L |
| 247 | + else: |
| 248 | + result_L = content_lab[:, 0] |
| 249 | + |
| 250 | + del content_lab, style_lab |
| 251 | + |
| 252 | + result_lab = torch.stack([result_L, matched_a, matched_b], dim=1) |
| 253 | + del result_L, matched_a, matched_b |
| 254 | + |
| 255 | + result_rgb = _lab_to_rgb_batch(result_lab, xyz_to_rgb_matrix, epsilon, kappa) |
| 256 | + del result_lab, xyz_to_rgb_matrix |
| 257 | + |
| 258 | + result = result_rgb.mul_(2.0).sub_(1.0) |
| 259 | + del result_rgb |
| 260 | + |
| 261 | + result = result.to(original_dtype) |
| 262 | + |
| 263 | + return result |
| 264 | + |
| 265 | + |
| 266 | +def wavelet_color_transfer(content_feat: Tensor, style_feat: Tensor) -> Tensor: |
| 267 | + return wavelet_reconstruction(content_feat, style_feat) |
| 268 | + |
| 269 | + |
| 270 | +def adain_color_transfer(content_feat: Tensor, style_feat: Tensor, eps: float = 1e-5) -> Tensor: |
| 271 | + if content_feat.shape != style_feat.shape: |
| 272 | + style_feat = F.interpolate( |
| 273 | + style_feat, |
| 274 | + size=content_feat.shape[-2:], |
| 275 | + mode='bilinear', |
| 276 | + align_corners=False, |
| 277 | + ) |
| 278 | + |
| 279 | + original_dtype = content_feat.dtype |
| 280 | + content_feat = content_feat.float() |
| 281 | + style_feat = style_feat.float() |
| 282 | + |
| 283 | + b, c = content_feat.shape[:2] |
| 284 | + content_flat = content_feat.reshape(b, c, -1) |
| 285 | + style_flat = style_feat.reshape(b, c, -1) |
| 286 | + |
| 287 | + content_mean = content_flat.mean(dim=2).reshape(b, c, 1, 1) |
| 288 | + content_std = (content_flat.var(dim=2, correction=0) + eps).sqrt().reshape(b, c, 1, 1) |
| 289 | + style_mean = style_flat.mean(dim=2).reshape(b, c, 1, 1) |
| 290 | + style_std = (style_flat.var(dim=2, correction=0) + eps).sqrt().reshape(b, c, 1, 1) |
| 291 | + del content_flat, style_flat |
| 292 | + |
| 293 | + normalized = (content_feat - content_mean) / content_std |
| 294 | + del content_mean, content_std |
| 295 | + result = normalized * style_std + style_mean |
| 296 | + del normalized, style_mean, style_std |
| 297 | + |
| 298 | + result = result.clamp_(-1.0, 1.0) |
| 299 | + if result.dtype != original_dtype: |
| 300 | + result = result.to(original_dtype) |
| 301 | + return result |
0 commit comments