diff --git a/invokeai/backend/image_util/imwatermark/vendor.py b/invokeai/backend/image_util/imwatermark/vendor.py index ef06274ff73..f96aae54521 100644 --- a/invokeai/backend/image_util/imwatermark/vendor.py +++ b/invokeai/backend/image_util/imwatermark/vendor.py @@ -7,9 +7,9 @@ # `opencv-contrib-python`. It's easier to copy the code over than complicate the installation process by # requiring an extra post-install step of removing `opencv-python` and installing `opencv-contrib-python`. -import struct -import uuid import base64 +import uuid + import cv2 import numpy as np import pywt @@ -111,14 +111,12 @@ def __init__(self, wm_type="bytes", length=0): raise NameError("%s is unsupported" % wm_type) def reconstruct_ipv4(self, bits): - ips = [str(ip) for ip in list(np.packbits(bits))] - return ".".join(ips) + packed = np.packbits(bits) + return "{}.{}.{}.{}".format(packed[0], packed[1], packed[2], packed[3]) def reconstruct_uuid(self, bits): nums = np.packbits(bits) - bstr = b"" - for i in range(16): - bstr += struct.pack(">B", nums[i]) + bstr = bytes(nums[:16]) return str(uuid.UUID(bytes=bstr)) @@ -132,10 +130,10 @@ def reconstruct_b16(self, bits): def reconstruct_bytes(self, bits): nums = np.packbits(bits) - bstr = b"" - for i in range(self._wmLen // 8): - bstr += struct.pack(">B", nums[i]) - return bstr + # Equivalent to b''.join(struct.pack(">B", nums[i]) for i in range(self._wmLen // 8)) + # Since all can be packed in one go: + end_idx = self._wmLen // 8 + return bytes(nums[:end_idx]) def reconstruct(self, bits): if len(bits) != self._wmLen: