Skip to content

Commit 04d1958

Browse files
feat(app): vendor in invisible-watermark
Fixes errors like `AttributeError: module 'cv2.ximgproc' has no attribute 'thinning'` which occur because there is a conflict between our own `opencv-contrib-python` dependency and the `invisible-watermark` library's `opencv-python`.
1 parent 47d7d93 commit 04d1958

File tree

3 files changed

+307
-5
lines changed

3 files changed

+307
-5
lines changed
Lines changed: 304 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,304 @@
1+
# This file is vendored from https://github.com/ShieldMnt/invisible-watermark
2+
#
3+
# `invisible-watermark` is MIT licensed as of August 23, 2025, when the code was copied into this repo.
4+
#
5+
# Why we vendored it in:
6+
# `invisible-watermark` has a dependency on `opencv-python`, which conflicts with Invoke's dependency on
7+
# `opencv-contrib-python`. It's easier to copy the code over than complicate the installation process by
8+
# requiring an extra post-install step of removing `opencv-python` and installing `opencv-contrib-python`.
9+
10+
import struct
11+
import uuid
12+
import base64
13+
import cv2
14+
import numpy as np
15+
import pywt
16+
17+
18+
class WatermarkEncoder(object):
19+
def __init__(self, content=b""):
20+
seq = np.array([n for n in content], dtype=np.uint8)
21+
self._watermarks = list(np.unpackbits(seq))
22+
self._wmLen = len(self._watermarks)
23+
self._wmType = "bytes"
24+
25+
def set_by_ipv4(self, addr):
26+
bits = []
27+
ips = addr.split(".")
28+
for ip in ips:
29+
bits += list(np.unpackbits(np.array([ip % 255], dtype=np.uint8)))
30+
self._watermarks = bits
31+
self._wmLen = len(self._watermarks)
32+
self._wmType = "ipv4"
33+
assert self._wmLen == 32
34+
35+
def set_by_uuid(self, uid):
36+
u = uuid.UUID(uid)
37+
self._wmType = "uuid"
38+
seq = np.array([n for n in u.bytes], dtype=np.uint8)
39+
self._watermarks = list(np.unpackbits(seq))
40+
self._wmLen = len(self._watermarks)
41+
42+
def set_by_bytes(self, content):
43+
self._wmType = "bytes"
44+
seq = np.array([n for n in content], dtype=np.uint8)
45+
self._watermarks = list(np.unpackbits(seq))
46+
self._wmLen = len(self._watermarks)
47+
48+
def set_by_b16(self, b16):
49+
content = base64.b16decode(b16)
50+
self.set_by_bytes(content)
51+
self._wmType = "b16"
52+
53+
def set_by_bits(self, bits=[]):
54+
self._watermarks = [int(bit) % 2 for bit in bits]
55+
self._wmLen = len(self._watermarks)
56+
self._wmType = "bits"
57+
58+
def set_watermark(self, wmType="bytes", content=""):
59+
if wmType == "ipv4":
60+
self.set_by_ipv4(content)
61+
elif wmType == "uuid":
62+
self.set_by_uuid(content)
63+
elif wmType == "bits":
64+
self.set_by_bits(content)
65+
elif wmType == "bytes":
66+
self.set_by_bytes(content)
67+
elif wmType == "b16":
68+
self.set_by_b16(content)
69+
else:
70+
raise NameError("%s is not supported" % wmType)
71+
72+
def get_length(self):
73+
return self._wmLen
74+
75+
# @classmethod
76+
# def loadModel(cls):
77+
# RivaWatermark.loadModel()
78+
79+
def encode(self, cv2Image, method="dwtDct", **configs):
80+
(r, c, channels) = cv2Image.shape
81+
if r * c < 256 * 256:
82+
raise RuntimeError("image too small, should be larger than 256x256")
83+
84+
if method == "dwtDct":
85+
embed = EmbedMaxDct(self._watermarks, wmLen=self._wmLen, **configs)
86+
return embed.encode(cv2Image)
87+
# elif method == 'dwtDctSvd':
88+
# embed = EmbedDwtDctSvd(self._watermarks, wmLen=self._wmLen, **configs)
89+
# return embed.encode(cv2Image)
90+
# elif method == 'rivaGan':
91+
# embed = RivaWatermark(self._watermarks, self._wmLen)
92+
# return embed.encode(cv2Image)
93+
else:
94+
raise NameError("%s is not supported" % method)
95+
96+
97+
class WatermarkDecoder(object):
98+
def __init__(self, wm_type="bytes", length=0):
99+
self._wmType = wm_type
100+
if wm_type == "ipv4":
101+
self._wmLen = 32
102+
elif wm_type == "uuid":
103+
self._wmLen = 128
104+
elif wm_type == "bytes":
105+
self._wmLen = length
106+
elif wm_type == "bits":
107+
self._wmLen = length
108+
elif wm_type == "b16":
109+
self._wmLen = length
110+
else:
111+
raise NameError("%s is unsupported" % wm_type)
112+
113+
def reconstruct_ipv4(self, bits):
114+
ips = [str(ip) for ip in list(np.packbits(bits))]
115+
return ".".join(ips)
116+
117+
def reconstruct_uuid(self, bits):
118+
nums = np.packbits(bits)
119+
bstr = b""
120+
for i in range(16):
121+
bstr += struct.pack(">B", nums[i])
122+
123+
return str(uuid.UUID(bytes=bstr))
124+
125+
def reconstruct_bits(self, bits):
126+
# return ''.join([str(b) for b in bits])
127+
return bits
128+
129+
def reconstruct_b16(self, bits):
130+
bstr = self.reconstruct_bytes(bits)
131+
return base64.b16encode(bstr)
132+
133+
def reconstruct_bytes(self, bits):
134+
nums = np.packbits(bits)
135+
bstr = b""
136+
for i in range(self._wmLen // 8):
137+
bstr += struct.pack(">B", nums[i])
138+
return bstr
139+
140+
def reconstruct(self, bits):
141+
if len(bits) != self._wmLen:
142+
raise RuntimeError("bits are not matched with watermark length")
143+
144+
if self._wmType == "ipv4":
145+
return self.reconstruct_ipv4(bits)
146+
elif self._wmType == "uuid":
147+
return self.reconstruct_uuid(bits)
148+
elif self._wmType == "bits":
149+
return self.reconstruct_bits(bits)
150+
elif self._wmType == "b16":
151+
return self.reconstruct_b16(bits)
152+
else:
153+
return self.reconstruct_bytes(bits)
154+
155+
def decode(self, cv2Image, method="dwtDct", **configs):
156+
(r, c, channels) = cv2Image.shape
157+
if r * c < 256 * 256:
158+
raise RuntimeError("image too small, should be larger than 256x256")
159+
160+
bits = []
161+
if method == "dwtDct":
162+
embed = EmbedMaxDct(watermarks=[], wmLen=self._wmLen, **configs)
163+
bits = embed.decode(cv2Image)
164+
# elif method == 'dwtDctSvd':
165+
# embed = EmbedDwtDctSvd(watermarks=[], wmLen=self._wmLen, **configs)
166+
# bits = embed.decode(cv2Image)
167+
# elif method == 'rivaGan':
168+
# embed = RivaWatermark(watermarks=[], wmLen=self._wmLen, **configs)
169+
# bits = embed.decode(cv2Image)
170+
else:
171+
raise NameError("%s is not supported" % method)
172+
return self.reconstruct(bits)
173+
174+
# @classmethod
175+
# def loadModel(cls):
176+
# RivaWatermark.loadModel()
177+
178+
179+
class EmbedMaxDct(object):
180+
def __init__(self, watermarks=[], wmLen=8, scales=[0, 36, 36], block=4):
181+
self._watermarks = watermarks
182+
self._wmLen = wmLen
183+
self._scales = scales
184+
self._block = block
185+
186+
def encode(self, bgr):
187+
(row, col, channels) = bgr.shape
188+
189+
yuv = cv2.cvtColor(bgr, cv2.COLOR_BGR2YUV)
190+
191+
for channel in range(2):
192+
if self._scales[channel] <= 0:
193+
continue
194+
195+
ca1, (h1, v1, d1) = pywt.dwt2(yuv[: row // 4 * 4, : col // 4 * 4, channel], "haar")
196+
self.encode_frame(ca1, self._scales[channel])
197+
198+
yuv[: row // 4 * 4, : col // 4 * 4, channel] = pywt.idwt2((ca1, (v1, h1, d1)), "haar")
199+
200+
bgr_encoded = cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
201+
return bgr_encoded
202+
203+
def decode(self, bgr):
204+
(row, col, channels) = bgr.shape
205+
206+
yuv = cv2.cvtColor(bgr, cv2.COLOR_BGR2YUV)
207+
208+
scores = [[] for i in range(self._wmLen)]
209+
for channel in range(2):
210+
if self._scales[channel] <= 0:
211+
continue
212+
213+
ca1, (h1, v1, d1) = pywt.dwt2(yuv[: row // 4 * 4, : col // 4 * 4, channel], "haar")
214+
215+
scores = self.decode_frame(ca1, self._scales[channel], scores)
216+
217+
avgScores = list(map(lambda l: np.array(l).mean(), scores))
218+
219+
bits = np.array(avgScores) * 255 > 127
220+
return bits
221+
222+
def decode_frame(self, frame, scale, scores):
223+
(row, col) = frame.shape
224+
num = 0
225+
226+
for i in range(row // self._block):
227+
for j in range(col // self._block):
228+
block = frame[
229+
i * self._block : i * self._block + self._block, j * self._block : j * self._block + self._block
230+
]
231+
232+
score = self.infer_dct_matrix(block, scale)
233+
# score = self.infer_dct_svd(block, scale)
234+
wmBit = num % self._wmLen
235+
scores[wmBit].append(score)
236+
num = num + 1
237+
238+
return scores
239+
240+
def diffuse_dct_svd(self, block, wmBit, scale):
241+
u, s, v = np.linalg.svd(cv2.dct(block))
242+
243+
s[0] = (s[0] // scale + 0.25 + 0.5 * wmBit) * scale
244+
return cv2.idct(np.dot(u, np.dot(np.diag(s), v)))
245+
246+
def infer_dct_svd(self, block, scale):
247+
u, s, v = np.linalg.svd(cv2.dct(block))
248+
249+
score = 0
250+
score = int((s[0] % scale) > scale * 0.5)
251+
return score
252+
if score >= 0.5:
253+
return 1.0
254+
else:
255+
return 0.0
256+
257+
def diffuse_dct_matrix(self, block, wmBit, scale):
258+
pos = np.argmax(abs(block.flatten()[1:])) + 1
259+
i, j = pos // self._block, pos % self._block
260+
val = block[i][j]
261+
if val >= 0.0:
262+
block[i][j] = (val // scale + 0.25 + 0.5 * wmBit) * scale
263+
else:
264+
val = abs(val)
265+
block[i][j] = -1.0 * (val // scale + 0.25 + 0.5 * wmBit) * scale
266+
return block
267+
268+
def infer_dct_matrix(self, block, scale):
269+
pos = np.argmax(abs(block.flatten()[1:])) + 1
270+
i, j = pos // self._block, pos % self._block
271+
272+
val = block[i][j]
273+
if val < 0:
274+
val = abs(val)
275+
276+
if (val % scale) > 0.5 * scale:
277+
return 1
278+
else:
279+
return 0
280+
281+
def encode_frame(self, frame, scale):
282+
"""
283+
frame is a matrix (M, N)
284+
285+
we get K (watermark bits size) blocks (self._block x self._block)
286+
287+
For i-th block, we encode watermark[i] bit into it
288+
"""
289+
(row, col) = frame.shape
290+
num = 0
291+
for i in range(row // self._block):
292+
for j in range(col // self._block):
293+
block = frame[
294+
i * self._block : i * self._block + self._block, j * self._block : j * self._block + self._block
295+
]
296+
wmBit = self._watermarks[(num % self._wmLen)]
297+
298+
diffusedBlock = self.diffuse_dct_matrix(block, wmBit, scale)
299+
# diffusedBlock = self.diffuse_dct_svd(block, wmBit, scale)
300+
frame[
301+
i * self._block : i * self._block + self._block, j * self._block : j * self._block + self._block
302+
] = diffusedBlock
303+
304+
num = num + 1

invokeai/backend/image_util/invisible_watermark.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,10 @@
66

77
import cv2
88
import numpy as np
9-
from imwatermark import WatermarkEncoder
109
from PIL import Image
1110

1211
import invokeai.backend.util.logging as logger
13-
from invokeai.app.services.config.config_default import get_config
14-
15-
config = get_config()
12+
from invokeai.backend.image_util.imwatermark.vendor import WatermarkEncoder
1613

1714

1815
class InvisibleWatermark:

pyproject.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@ dependencies = [
3838
"compel==2.1.1",
3939
"diffusers[torch]==0.33.0",
4040
"gguf",
41-
"invisible-watermark==0.2.0", # needed to install SDXL base and refiner using their repo_ids
4241
"mediapipe==0.10.14", # needed for "mediapipeface" controlnet model
4342
"numpy<2.0.0",
4443
"onnx==1.16.1",
@@ -74,6 +73,7 @@ dependencies = [
7473
"python-multipart",
7574
"requests",
7675
"semver~=3.0.1",
76+
"PyWavelets",
7777
]
7878

7979
[project.optional-dependencies]
@@ -234,6 +234,7 @@ exclude = [
234234
"invokeai/backend/image_util/mlsd/", # External code
235235
"invokeai/backend/image_util/normal_bae/", # External code
236236
"invokeai/backend/image_util/pidi/", # External code
237+
"invokeai/backend/image_util/imwatermark/", # External code
237238
]
238239

239240
[tool.ruff.lint]

0 commit comments

Comments
 (0)