From 81963f887ad32512d7a11d528c0b91ae8947ed9d Mon Sep 17 00:00:00 2001 From: "remyx-ai[bot]" <289541483+remyx-ai[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:21:49 +0000 Subject: [PATCH] Add CIELAB+Sobel physical-consistency losses for inpainting --- .../physical_consistency_inpaint/README.md | 69 +++++ .../example_physical_consistency.py | 72 ++++++ .../physical_consistency.py | 243 ++++++++++++++++++ .../requirements.txt | 2 + .../test_physical_consistency.py | 115 +++++++++ 5 files changed, 501 insertions(+) create mode 100644 examples/research_projects/physical_consistency_inpaint/README.md create mode 100644 examples/research_projects/physical_consistency_inpaint/example_physical_consistency.py create mode 100644 examples/research_projects/physical_consistency_inpaint/physical_consistency.py create mode 100644 examples/research_projects/physical_consistency_inpaint/requirements.txt create mode 100644 examples/research_projects/physical_consistency_inpaint/test_physical_consistency.py diff --git a/examples/research_projects/physical_consistency_inpaint/README.md b/examples/research_projects/physical_consistency_inpaint/README.md new file mode 100644 index 000000000000..fc6f777f99fe --- /dev/null +++ b/examples/research_projects/physical_consistency_inpaint/README.md @@ -0,0 +1,69 @@ +# Physical-consistency losses for diffusion inpainting + +Adapted from **"Structured-Prior-Guided Diffusion Inpainting with Physical +Consistency for Traffic Sign Augmentation"** +(https://arxiv.org/abs/2609.02348). + +General-purpose inpainting models shift colours and deform edge structure when +applied to physically-composed objects (traffic signs, license plates, digits). +The paper adds two parameter-free, differentiable pixel-space losses on top of a +standard Stable Diffusion 1.5 inpainting backbone to hold those quantities +fixed: + +* **CIELAB chromaticity L1** — penalises hue/saturation drift in a + perceptually-uniform colour space (the `a*`, `b*` channels), while staying + indifferent to legitimate lightness changes. +* **Sobel gradient L1** — penalises edge/structure drift so digit strokes and + geometric outlines stay aligned. + +## What this project ports + +Only the two **physical-consistency loss terms** are ported here, at full +fidelity, as a drop-in `diffusers`-native module. They operate on the decoded +RGB prediction and plug into any existing inpainting reconstruction loop after +the VAE decode, optionally restricted to the inpainted region via a mask. + +The paper's three **structured-prior conditioning pathways** — a JSON-formatted +text prompt (semantic), an IP-Adapter front-view template (appearance) and a +ControlNet affine template (geometric) — are **out of scope** for this module: +they are orthogonal conditioning inputs that `diffusers` already exposes through +its existing IP-Adapter and ControlNet interfaces, and can be layered on +separately. The in-house AMAP training set and downstream detection benchmark +are likewise not reproduced. + +## Usage + +```python +from diffusers import AutoencoderKL + +from physical_consistency import ( + PhysicalConsistencyLoss, + reconstruction_loss_with_physical_consistency, +) + +physical_loss = PhysicalConsistencyLoss(lambda_cielab=1.0, lambda_sobel=0.5) + +# Inside your inpainting training step, alongside the usual latent MSE: +losses = reconstruction_loss_with_physical_consistency( + pipeline.vae, + predicted_x0_latents, # model prediction, decoded internally + target_latents, # ground-truth latents + mask=inpaint_mask, # latent-resolution mask of the edited region + physical_loss=physical_loss, +) +total = latent_mse + losses["total"] +total.backward() +``` + +Run the self-contained demo (no checkpoints/datasets required): + +```sh +pip install -r requirements.txt +python example_physical_consistency.py +``` + +## Tests + +```sh +pytest test_physical_consistency.py +``` diff --git a/examples/research_projects/physical_consistency_inpaint/example_physical_consistency.py b/examples/research_projects/physical_consistency_inpaint/example_physical_consistency.py new file mode 100644 index 000000000000..41284a50091f --- /dev/null +++ b/examples/research_projects/physical_consistency_inpaint/example_physical_consistency.py @@ -0,0 +1,72 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Minimal, self-contained demo of the physical-consistency losses. + +It composes a small diffusers ``AutoencoderKL`` with +``reconstruction_loss_with_physical_consistency`` to show exactly where the +losses plug into an inpainting training step — no checkpoints or datasets are +downloaded, so it doubles as a smoke test: + + python example_physical_consistency.py + +In a real SD1.5-inpaint / ControlNet-inpaint loop you would swap the tiny VAE +below for the pipeline's ``vae``, feed the model's predicted ``x0`` latents as +``predicted_latents``, the ground-truth latents as ``target_latents``, and add +``losses["total"]`` to the usual latent MSE before ``backward()``. +""" + +import torch +from physical_consistency import PhysicalConsistencyLoss, reconstruction_loss_with_physical_consistency + +from diffusers import AutoencoderKL + + +def build_tiny_vae() -> AutoencoderKL: + """A CPU-sized VAE mirroring the config used across the diffusers test suite.""" + return AutoencoderKL( + block_out_channels=[32, 64], + in_channels=3, + out_channels=3, + down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], + up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"], + latent_channels=4, + ) + + +def main() -> None: + torch.manual_seed(0) + vae = build_tiny_vae().eval() + + # Stand in for a training batch: ground-truth latents plus a noisier + # prediction and a mask marking the inpainted region. + target_latents = torch.randn(1, 4, 16, 16) + predicted_latents = target_latents + 0.1 * torch.randn_like(target_latents) + predicted_latents.requires_grad_(True) + mask = torch.zeros(1, 1, 16, 16) + mask[:, :, 4:12, 4:12] = 1.0 + + physical_loss = PhysicalConsistencyLoss(lambda_cielab=1.0, lambda_sobel=0.5) + losses = reconstruction_loss_with_physical_consistency( + vae, predicted_latents, target_latents, mask=mask, physical_loss=physical_loss + ) + + losses["total"].backward() + print(f"cielab: {losses['cielab'].item():.4f}") + print(f"sobel: {losses['sobel'].item():.4f}") + print(f"total: {losses['total'].item():.4f}") + print(f"grad flows back to latents: {predicted_latents.grad is not None}") + + +if __name__ == "__main__": + main() diff --git a/examples/research_projects/physical_consistency_inpaint/physical_consistency.py b/examples/research_projects/physical_consistency_inpaint/physical_consistency.py new file mode 100644 index 000000000000..c2184dfaf61f --- /dev/null +++ b/examples/research_projects/physical_consistency_inpaint/physical_consistency.py @@ -0,0 +1,243 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Physical-consistency losses for diffusion inpainting. + +Adapted from "Structured-Prior-Guided Diffusion Inpainting with Physical +Consistency for Traffic Sign Augmentation" (https://arxiv.org/abs/2609.02348). + +The paper observes that general-purpose inpainting models shift colours and +deform edge structure when applied to physically-composed objects (traffic +signs, plates, digits). It constrains the decoded prediction with two +parameter-free, differentiable pixel-space terms: + +* a **CIELAB chromaticity L1** term that penalises colour drift, measured in a + perceptually-uniform space so equal numeric errors are roughly equal + perceived errors, and +* a **Sobel gradient** term that penalises edge/structure drift. + +Only these two loss terms are ported here (the paper's full method also injects +semantic/appearance/geometric priors through a JSON prompt, an IP-Adapter +template and a ControlNet template — those are orthogonal conditioning +pathways, out of scope for this module). The terms operate on decoded RGB +images and plug into any existing diffusers reconstruction/inpainting training +loop after the VAE decode, optionally restricted to the inpainted region via a +mask. +""" + +import torch +import torch.nn.functional as F + + +# sRGB <-> linear RGB <-> CIE XYZ (D65) <-> CIELAB constants. +_SRGB_THRESHOLD = 0.04045 +_LAB_DELTA = 6.0 / 29.0 +_D65_WHITE = (0.95047, 1.0, 1.08883) +_EPS = 1e-8 + + +def _srgb_to_linear(channel: torch.Tensor) -> torch.Tensor: + """Differentiable sRGB gamma expansion for a single channel in [0, 1].""" + channel = channel.clamp(0.0, 1.0) + low = channel / 12.92 + high = ((channel + 0.055) / 1.055).clamp(min=_EPS) ** 2.4 + return torch.where(channel <= _SRGB_THRESHOLD, low, high) + + +def rgb_to_lab(image: torch.Tensor) -> torch.Tensor: + """Convert an sRGB image in [0, 1] to CIELAB. + + Args: + image: Tensor of shape ``(N, 3, H, W)`` with values in ``[0, 1]``. + + Returns: + Tensor of shape ``(N, 3, H, W)`` holding the ``L*``, ``a*`` and ``b*`` + channels. The conversion is fully differentiable, so it can sit inside + a training graph. + """ + if image.dim() != 4 or image.shape[1] != 3: + raise ValueError(f"Expected an (N, 3, H, W) RGB tensor, got shape {tuple(image.shape)}.") + + linear = _srgb_to_linear(image) + r, g, b = linear[:, 0], linear[:, 1], linear[:, 2] + + x = 0.4124564 * r + 0.3575761 * g + 0.1804375 * b + y = 0.2126729 * r + 0.7151522 * g + 0.0721750 * b + z = 0.0193339 * r + 0.1191920 * g + 0.9503041 * b + + x = x / _D65_WHITE[0] + y = y / _D65_WHITE[1] + z = z / _D65_WHITE[2] + + def _f(t: torch.Tensor) -> torch.Tensor: + cube_root = t.clamp(min=_EPS) ** (1.0 / 3.0) + linear_part = t / (3.0 * _LAB_DELTA * _LAB_DELTA) + 4.0 / 29.0 + return torch.where(t > _LAB_DELTA**3, cube_root, linear_part) + + fx, fy, fz = _f(x), _f(y), _f(z) + lightness = 116.0 * fy - 16.0 + a = 500.0 * (fx - fy) + b_star = 200.0 * (fy - fz) + return torch.stack([lightness, a, b_star], dim=1) + + +def _masked_l1(prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor: + """L1 distance, optionally averaged over a region mask. + + ``mask`` is broadcast over the channel dimension; a ``None`` mask reduces + over the whole image. + """ + diff = (prediction - target).abs() + if mask is None: + return diff.mean() + channels = prediction.shape[1] + denominator = mask.sum() * channels + _EPS + return (diff * mask).sum() / denominator + + +def cielab_chromaticity_loss( + prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor = None +) -> torch.Tensor: + """CIELAB chromaticity (``a*``, ``b*``) L1 loss between two sRGB images. + + Only the two chromaticity channels are compared, so the term penalises hue + and saturation drift while staying indifferent to lightness changes the + diffusion prior may legitimately introduce. + + Args: + prediction: Predicted sRGB image ``(N, 3, H, W)`` in ``[0, 1]``. + target: Reference sRGB image ``(N, 3, H, W)`` in ``[0, 1]``. + mask: Optional ``(N, 1, H, W)`` region mask (e.g. the inpainted area). + + Returns: + Scalar loss tensor. + """ + pred_ab = rgb_to_lab(prediction)[:, 1:] + target_ab = rgb_to_lab(target)[:, 1:] + return _masked_l1(pred_ab, target_ab, mask) + + +def _sobel_gradients(image: torch.Tensor) -> torch.Tensor: + """Return stacked horizontal/vertical Sobel responses per channel.""" + kernel_x = torch.tensor( + [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]], + dtype=image.dtype, + device=image.device, + ) + kernel_y = kernel_x.t() + channels = image.shape[1] + weight = torch.stack([kernel_x, kernel_y]).unsqueeze(1) # (2, 1, 3, 3) + weight = weight.repeat(channels, 1, 1, 1) # (2 * C, 1, 3, 3) + padded = F.pad(image, (1, 1, 1, 1), mode="replicate") + return F.conv2d(padded, weight, groups=channels) + + +def sobel_gradient_loss( + prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor = None +) -> torch.Tensor: + """Sobel edge-structure L1 loss between two sRGB images. + + Gradient magnitude is compared per channel, which keeps digit strokes and + geometric outlines aligned between the prediction and the reference. + + Args: + prediction: Predicted sRGB image ``(N, 3, H, W)`` in ``[0, 1]``. + target: Reference sRGB image ``(N, 3, H, W)`` in ``[0, 1]``. + mask: Optional ``(N, 1, H, W)`` region mask (e.g. the inpainted area). + + Returns: + Scalar loss tensor. + """ + pred_grad = _sobel_gradients(prediction) + target_grad = _sobel_gradients(target) + pred_mag = torch.sqrt(pred_grad[:, 0::2] ** 2 + pred_grad[:, 1::2] ** 2 + _EPS) + target_mag = torch.sqrt(target_grad[:, 0::2] ** 2 + target_grad[:, 1::2] ** 2 + _EPS) + # A single-channel mask broadcasts over the per-channel magnitudes; _masked_l1 + # already normalises by the channel count. + return _masked_l1(pred_mag, target_mag, mask) + + +class PhysicalConsistencyLoss(torch.nn.Module): + """Weighted sum of the CIELAB chromaticity and Sobel gradient terms. + + Args: + lambda_cielab: Weight of the CIELAB chromaticity term. + lambda_sobel: Weight of the Sobel gradient term. + """ + + def __init__(self, lambda_cielab: float = 1.0, lambda_sobel: float = 1.0): + super().__init__() + self.lambda_cielab = lambda_cielab + self.lambda_sobel = lambda_sobel + + def forward( + self, prediction: torch.Tensor, target: torch.Tensor, mask: torch.Tensor = None + ) -> dict: + """Compute the individual and combined physical-consistency terms. + + Returns a dict with ``cielab``, ``sobel`` and ``total`` scalar tensors + so a training loop can log each term while backpropagating ``total``. + """ + cielab = cielab_chromaticity_loss(prediction, target, mask) + sobel = sobel_gradient_loss(prediction, target, mask) + total = self.lambda_cielab * cielab + self.lambda_sobel * sobel + return {"cielab": cielab, "sobel": sobel, "total": total} + + +def _decode_to_images(vae, latents: torch.Tensor) -> torch.Tensor: + """Decode diffusion latents into sRGB images in ``[0, 1]`` via the VAE.""" + images = vae.decode(latents / vae.config.scaling_factor, return_dict=False)[0] + return (images / 2.0 + 0.5).clamp(0.0, 1.0) + + +def reconstruction_loss_with_physical_consistency( + vae, + predicted_latents: torch.Tensor, + target_latents: torch.Tensor, + mask: torch.Tensor = None, + physical_loss: PhysicalConsistencyLoss = None, +) -> dict: + """Decode a prediction and reference and score physical consistency. + + This is the call site that stitches the losses above onto an existing + diffusers inpainting loop: pass the model's predicted ``x0`` latents and + the ground-truth latents together with any diffusers ``AutoencoderKL`` + (SD1.5-inpaint, ControlNet-inpaint, ... all share this VAE interface). The + latents are decoded to pixel space and both terms are computed there, since + colour and edge structure are physical quantities that only exist after + decoding. + + Args: + vae: A diffusers ``AutoencoderKL`` (or any VAE exposing ``decode`` and a + ``config.scaling_factor``). + predicted_latents: Predicted ``(N, C, h, w)`` latents. + target_latents: Reference ``(N, C, h, w)`` latents. + mask: Optional latent-resolution ``(N, 1, h, w)`` mask; it is upsampled + to the decoded image resolution and used to restrict both terms to + the inpainted region. + physical_loss: Optional pre-configured :class:`PhysicalConsistencyLoss`; + a default (equal weights) instance is created when omitted. + + Returns: + The dict returned by :class:`PhysicalConsistencyLoss`. + """ + if physical_loss is None: + physical_loss = PhysicalConsistencyLoss() + + predicted_images = _decode_to_images(vae, predicted_latents) + target_images = _decode_to_images(vae, target_latents) + + if mask is not None: + mask = F.interpolate(mask, size=predicted_images.shape[-2:], mode="nearest") + + return physical_loss(predicted_images, target_images, mask) diff --git a/examples/research_projects/physical_consistency_inpaint/requirements.txt b/examples/research_projects/physical_consistency_inpaint/requirements.txt new file mode 100644 index 000000000000..ee9cc8bfeddf --- /dev/null +++ b/examples/research_projects/physical_consistency_inpaint/requirements.txt @@ -0,0 +1,2 @@ +diffusers +torch diff --git a/examples/research_projects/physical_consistency_inpaint/test_physical_consistency.py b/examples/research_projects/physical_consistency_inpaint/test_physical_consistency.py new file mode 100644 index 000000000000..14b153fb4032 --- /dev/null +++ b/examples/research_projects/physical_consistency_inpaint/test_physical_consistency.py @@ -0,0 +1,115 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import os +import sys + +import torch + +# Existing (non-new) module: the diffusers VAE the losses integrate with. +from diffusers import AutoencoderKL + + +sys.path.insert(0, os.path.dirname(__file__)) + +from physical_consistency import ( # noqa: E402 + PhysicalConsistencyLoss, + cielab_chromaticity_loss, + reconstruction_loss_with_physical_consistency, + rgb_to_lab, + sobel_gradient_loss, +) + + +def _tiny_vae() -> AutoencoderKL: + return AutoencoderKL( + block_out_channels=[32, 64], + in_channels=3, + out_channels=3, + down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], + up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"], + latent_channels=4, + ) + + +def test_rgb_to_lab_matches_reference_red(): + # Pure sRGB red converts to CIELAB ~ (53.24, 80.09, 67.20). + red = torch.tensor([1.0, 0.0, 0.0]).view(1, 3, 1, 1) + lab = rgb_to_lab(red).view(3) + assert torch.allclose(lab, torch.tensor([53.2408, 80.0925, 67.2032]), atol=1e-2) + + +def test_identical_images_have_zero_loss(): + image = torch.rand(2, 3, 32, 32) + assert cielab_chromaticity_loss(image, image).item() < 1e-4 + assert sobel_gradient_loss(image, image).item() < 1e-4 + + +def test_hue_shift_raises_chromaticity_more_than_luminance_shift(): + torch.manual_seed(0) + target = torch.rand(1, 3, 32, 32) + hue_shifted = target.clone() + hue_shifted[:, 0] = (hue_shifted[:, 0] + 0.4).clamp(0.0, 1.0) # push the red channel + darkened = (target * 0.7).clamp(0.0, 1.0) # lightness-only change + # Chromaticity term reacts to the hue shift and mostly ignores darkening. + assert cielab_chromaticity_loss(hue_shifted, target) > cielab_chromaticity_loss(darkened, target) + + +def test_edge_blur_raises_sobel_loss(): + torch.manual_seed(0) + target = torch.rand(1, 3, 48, 48) + # Blur destroys edge structure -> Sobel term should grow. + kernel = torch.ones(3, 1, 5, 5) / 25.0 + blurred = torch.nn.functional.conv2d( + torch.nn.functional.pad(target, (2, 2, 2, 2), mode="replicate"), kernel, groups=3 + ) + assert sobel_gradient_loss(blurred, target) > sobel_gradient_loss(target, target) + + +def test_mask_restricts_loss_to_region(): + torch.manual_seed(0) + target = torch.rand(1, 3, 32, 32) + prediction = target.clone() + prediction[:, :, :16, :] = torch.rand(1, 3, 16, 32) # corrupt the top half only + top_mask = torch.zeros(1, 1, 32, 32) + top_mask[:, :, :16, :] = 1.0 + bottom_mask = torch.zeros(1, 1, 32, 32) + bottom_mask[:, :, 16:, :] = 1.0 + # The bottom is untouched, so a bottom mask sees ~no error; the top does. + assert cielab_chromaticity_loss(prediction, target, top_mask) > 1e-3 + assert cielab_chromaticity_loss(prediction, target, bottom_mask) < 1e-4 + + +def test_reconstruction_helper_integrates_with_diffusers_vae(): + torch.manual_seed(0) + vae = _tiny_vae().eval() + target_latents = torch.randn(1, 4, 16, 16) + predicted_latents = (target_latents + 0.1 * torch.randn_like(target_latents)).requires_grad_(True) + mask = torch.zeros(1, 1, 16, 16) + mask[:, :, 4:12, 4:12] = 1.0 + + losses = reconstruction_loss_with_physical_consistency( + vae, + predicted_latents, + target_latents, + mask=mask, + physical_loss=PhysicalConsistencyLoss(lambda_cielab=1.0, lambda_sobel=0.5), + ) + + assert set(losses) == {"cielab", "sobel", "total"} + for value in losses.values(): + assert torch.isfinite(value) + # The wiring is differentiable end-to-end through the VAE decode. + losses["total"].backward() + assert predicted_latents.grad is not None + assert torch.isfinite(predicted_latents.grad).all()