From d1795e7ef20ca37f5d3d2504f941b37f282958fd Mon Sep 17 00:00:00 2001 From: GokayAI <60583610+gokay-ai@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:25:05 +0100 Subject: [PATCH] Fix SDXL watermark encode for Stability detect.py round-trip Cast images to uint8 before dwtDct encode so chroma matches what cv2.imread feeds the decoder; Fixes #15025. Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com> --- .../stable_diffusion_xl/watermark.py | 7 ++- .../test_stable_diffusion_xl_watermark.py | 63 +++++++++++++++++++ 2 files changed, 68 insertions(+), 2 deletions(-) create mode 100644 tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_watermark.py diff --git a/src/diffusers/pipelines/stable_diffusion_xl/watermark.py b/src/diffusers/pipelines/stable_diffusion_xl/watermark.py index 70d06bb6320d..876e5df37508 100644 --- a/src/diffusers/pipelines/stable_diffusion_xl/watermark.py +++ b/src/diffusers/pipelines/stable_diffusion_xl/watermark.py @@ -27,14 +27,17 @@ def apply_watermark(self, images: torch.Tensor): return images images = (255 * (images / 2 + 0.5)).cpu().permute(0, 2, 3, 1).float().numpy() + # `dwtDct` expects uint8 BGR (as Stability's detect.py reads via cv2.imread). + # Passing float32 skips OpenCV's chroma +128 offset and often fails decode after PNG save. + images = np.clip(np.round(images), 0, 255).astype(np.uint8) # Convert RGB to BGR, which is the channel order expected by the watermark encoder. images = images[:, :, :, ::-1] # Add watermark and convert BGR back to RGB - images = [self.encoder.encode(image, "dwtDct")[:, :, ::-1] for image in images] + images = [self.encoder.encode(np.ascontiguousarray(image), "dwtDct")[:, :, ::-1] for image in images] - images = np.array(images) + images = np.array(images).astype(np.float32) images = torch.from_numpy(images).permute(0, 3, 1, 2) diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_watermark.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_watermark.py new file mode 100644 index 000000000000..7946bedba292 --- /dev/null +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl_watermark.py @@ -0,0 +1,63 @@ +# Copyright 2026 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 tempfile +import unittest + +import numpy as np +import pytest +import torch + +from diffusers.image_processor import VaeImageProcessor +from diffusers.utils import is_invisible_watermark_available, is_opencv_available, load_image + +from ...testing_utils import require_torch + + +# Image from the #15025 repro: float encode fails decode; uint8 encode recovers all 48 bits. +_WATERMARK_REPRO_URL = ( + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/" + "df6c7e6873499cfe63686cb4b6edb05d6f8a5755/diffusers/sdxl_reference_adain_canny_cat.png" +) + + +def _bits_matched_via_png(pil_image, watermark_bits): + import cv2 + from imwatermark import WatermarkDecoder + + with tempfile.NamedTemporaryFile(suffix=".png") as tmp: + pil_image.save(tmp.name) + bgr = cv2.imread(tmp.name) + bits = WatermarkDecoder("bits", len(watermark_bits)).decode(bgr, "dwtDct") + return sum(int(b) == w for b, w in zip(bits, watermark_bits)) + + +@require_torch +@pytest.mark.skipif(not is_invisible_watermark_available(), reason="requires invisible-watermark") +@pytest.mark.skipif(not is_opencv_available(), reason="requires opencv") +class StableDiffusionXLWatermarkTests(unittest.TestCase): + def test_apply_watermark_decodes_after_uint8_png_roundtrip(self): + """Regression for #15025: encode on uint8 so Stability's cv2.imread detector recovers the bits.""" + from diffusers.pipelines.stable_diffusion_xl.watermark import WATERMARK_BITS, StableDiffusionXLWatermarker + + rgb = np.asarray(load_image(_WATERMARK_REPRO_URL).convert("RGB")) + images = torch.from_numpy(rgb.astype(np.float32) / 127.5 - 1.0).permute(2, 0, 1)[None] + processor = VaeImageProcessor(vae_scale_factor=8) + + watermarked = StableDiffusionXLWatermarker().apply_watermark(images.clone()) + pil = processor.postprocess(watermarked, output_type="pil")[0] + matched = _bits_matched_via_png(pil, WATERMARK_BITS) + + # detect.py treats <=33/48 as "No watermark detected"; require a clear full match. + self.assertGreaterEqual(matched, 41)