Unverified Commit ebc99a77 authored by btlorch's avatar btlorch Committed by GitHub
Browse files

Convert RGB to BGR for the SDXL watermark encoder (#7013)



* Convert channel order to BGR for the watermark encoder. Convert the watermarked BGR images back to RGB. Fixes #6292

* Revert channel order before stacking images to overcome limitations that negative strides are currently not supported

---------
Co-authored-by: default avatarSayak Paul <spsayakpaul@gmail.com>
parent fa750a15
...@@ -28,9 +28,15 @@ class StableDiffusionXLWatermarker: ...@@ -28,9 +28,15 @@ class StableDiffusionXLWatermarker:
images = (255 * (images / 2 + 0.5)).cpu().permute(0, 2, 3, 1).float().numpy() images = (255 * (images / 2 + 0.5)).cpu().permute(0, 2, 3, 1).float().numpy()
images = [self.encoder.encode(image, "dwtDct") for image in images] # Convert RGB to BGR, which is the channel order expected by the watermark encoder.
images = images[:, :, :, ::-1]
images = torch.from_numpy(np.array(images)).permute(0, 3, 1, 2) # Add watermark and convert BGR back to RGB
images = [self.encoder.encode(image, "dwtDct")[:, :, ::-1] for image in images]
images = np.array(images)
images = torch.from_numpy(images).permute(0, 3, 1, 2)
images = torch.clamp(2 * (images / 255 - 0.5), min=-1.0, max=1.0) images = torch.clamp(2 * (images / 255 - 0.5), min=-1.0, max=1.0)
return images return images
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment