imwatermark

flux/util.py

@@ -4,7 +4,6 @@ from dataclasses import dataclass
 import torch
 from einops import rearrange
 from huggingface_hub import hf_hub_download
-from imwatermark import WatermarkEncoder
 from PIL import ExifTags, Image
 from safetensors.torch import load_file as load_sft
 
@@ -470,47 +469,3 @@ def optionally_expand_state_dict(model: torch.nn.Module, state_dict: dict) -> di
                 state_dict[name] = expanded_state_dict_weight
 
     return state_dict
-
-
-class WatermarkEmbedder:
-    def __init__(self, watermark):
-        self.watermark = watermark
-        self.num_bits = len(WATERMARK_BITS)
-        self.encoder = WatermarkEncoder()
-        self.encoder.set_watermark("bits", self.watermark)
-
-    def __call__(self, image: torch.Tensor) -> torch.Tensor:
-        """
-        Adds a predefined watermark to the input image
-
-        Args:
-            image: ([N,] B, RGB, H, W) in range [-1, 1]
-
-        Returns:
-            same as input but watermarked
-        """
-        image = 0.5 * image + 0.5
-        squeeze = len(image.shape) == 4
-        if squeeze:
-            image = image[None, ...]
-        n = image.shape[0]
-        image_np = rearrange((255 * image).detach().cpu(), "n b c h w -> (n b) h w c").numpy()[:, :, :, ::-1]
-        # torch (b, c, h, w) in [0, 1] -> numpy (b, h, w, c) [0, 255]
-        # watermarking libary expects input as cv2 BGR format
-        for k in range(image_np.shape[0]):
-            image_np[k] = self.encoder.encode(image_np[k], "dwtDct")
-        image = torch.from_numpy(rearrange(image_np[:, :, :, ::-1], "(n b) h w c -> n b c h w", n=n)).to(
-            image.device
-        )
-        image = torch.clamp(image / 255, min=0.0, max=1.0)
-        if squeeze:
-            image = image[0]
-        image = 2 * image - 1
-        return image
-
-
-# A fixed 48-bit message that was chosen at random
-WATERMARK_MESSAGE = 0b001010101111111010000111100111001111010100101110
-# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
-WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]
-embed_watermark = WatermarkEmbedder(WATERMARK_BITS)

requirements.txt

@@ -13,5 +13,4 @@ numba
 scipy
 tqdm
 einops
-imwatermark
 https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.2.post1/flash_attn-2.7.2.post1+cu11torch2.1cxx11abiFALSE-cp310-cp310-linux_x86_64.whl