Coverage for markdiffusion / watermark / tr / tr.py: 91.72%

1from ..base import BaseWatermark, BaseConfig 

2from markdiffusion.utils.media_utils import * 

3import torch 

4from typing import Dict, Union, List, Optional 

5from markdiffusion.utils.utils import set_random_seed, inherit_docstring 

6from markdiffusion.utils.diffusion_config import DiffusionConfig 

7import copy 

8import numpy as np 

9from PIL import Image 

10from markdiffusion.visualize.data_for_visualization import DataForVisualization 

11from markdiffusion.detection.tr.tr_detection import TRDetector 

12 

13class TRConfig(BaseConfig): 

14 """Config class for TR algorithm, load config file and initialize parameters.""" 

15 

16 def initialize_parameters(self) -> None: 

17 """Initialize algorithm-specific parameters.""" 

18 self.w_seed = self.config_dict['w_seed'] 

19 self.w_channel = self.config_dict['w_channel'] 

20 self.w_pattern = self.config_dict['w_pattern'] 

21 # self.w_mask_shape = self.config_dict['w_mask_shape'] 

22 self.w_radius = self.config_dict['w_radius'] 

23 self.w_pattern_const = self.config_dict['w_pattern_const'] 

24 self.threshold = self.config_dict['threshold'] 

25 self.threshold_p_value = self.config_dict.get('threshold_p_value', 0.01) 

26 

27 @property 

28 def algorithm_name(self) -> str: 

29 """Return the algorithm name.""" 

30 return 'TR' 

31 

32class TRUtils: 

33 """Utility class for TR algorithm, contains helper functions.""" 

34 

35 def __init__(self, config: TRConfig, *args, **kwargs) -> None: 

36 """ 

37 Initialize the Tree-Ring watermarking algorithm. 

38  

39 Parameters: 

40 config (TRConfig): Configuration for the Tree-Ring algorithm. 

41 """ 

42 self.config = config 

43 self.gt_patch = self._get_watermarking_pattern() 

44 self.watermarking_mask = self._get_watermarking_mask(self.config.init_latents) 

45 

46 def _circle_mask(self, size: int=64, r: int=10, x_offset: int=0, y_offset: int=0) -> np.ndarray: 

47 """Generate a circular mask.""" 

48 x0 = y0 = size // 2 

49 x0 += x_offset 

50 y0 += y_offset 

51 y, x = np.ogrid[:size, :size] 

52 y = y[::-1] 

53 

54 return ((x - x0)**2 + (y-y0)**2)<= r**2 

55 

56 def _get_watermarking_pattern(self) -> torch.Tensor: 

57 """Get the ground truth watermarking pattern.""" 

58 set_random_seed(self.config.w_seed) 

59 

60 gt_init = get_random_latents(pipe=self.config.pipe, height=self.config.image_size[0], width=self.config.image_size[1]) 

61 

62 if 'seed_ring' in self.config.w_pattern: 

63 gt_patch = gt_init 

64 

65 gt_patch_tmp = copy.deepcopy(gt_patch) 

66 for i in range(self.config.w_radius, 0, -1): 

67 tmp_mask = self._circle_mask(gt_init.shape[-1], r=i) 

68 tmp_mask = torch.tensor(tmp_mask).to(self.config.device) 

69 

70 for j in range(gt_patch.shape[1]): 

71 gt_patch[:, j, tmp_mask] = gt_patch_tmp[0, j, 0, i].item() 

72 elif 'seed_zeros' in self.config.w_pattern: 

73 gt_patch = gt_init * 0 

74 elif 'seed_rand' in self.config.w_pattern: 

75 gt_patch = gt_init 

76 elif 'rand' in self.config.w_pattern: 

77 gt_patch = torch.fft.fftshift(torch.fft.fft2(gt_init), dim=(-1, -2)) 

78 gt_patch[:] = gt_patch[0] 

79 elif 'zeros' in self.config.w_pattern: 

80 gt_patch = torch.fft.fftshift(torch.fft.fft2(gt_init), dim=(-1, -2)) * 0 

81 elif 'const' in self.config.w_pattern: 

82 gt_patch = torch.fft.fftshift(torch.fft.fft2(gt_init), dim=(-1, -2)) * 0 

83 gt_patch += self.config.w_pattern_const 

84 elif 'ring' in self.config.w_pattern: 

85 gt_patch = torch.fft.fftshift(torch.fft.fft2(gt_init), dim=(-1, -2)) 

86 

87 gt_patch_tmp = copy.deepcopy(gt_patch) 

88 for i in range(self.config.w_radius, 0, -1): 

89 tmp_mask = self._circle_mask(gt_init.shape[-1], r=i) 

90 tmp_mask = torch.tensor(tmp_mask).to(self.config.device) 

91 

92 for j in range(gt_patch.shape[1]): 

93 gt_patch[:, j, tmp_mask] = gt_patch_tmp[0, j, 0, i].item() 

94 

95 return gt_patch 

96 

97 def _get_watermarking_mask(self, init_latents: torch.Tensor) -> torch.Tensor: 

98 """Get the watermarking mask.""" 

99 watermarking_mask = torch.zeros(init_latents.shape, dtype=torch.bool).to(self.config.device) 

100 

101 # if self.config.w_mask_shape == 'circle': 

102 np_mask = self._circle_mask(init_latents.shape[-1], r=self.config.w_radius) 

103 torch_mask = torch.tensor(np_mask).to(self.config.device) 

104 

105 if self.config.w_channel == -1: 

106 # all channels 

107 watermarking_mask[:, :] = torch_mask 

108 else: 

109 watermarking_mask[:, self.config.w_channel] = torch_mask 

110 # elif self.config.w_mask_shape == 'square': 

111 # anchor_p = init_latents.shape[-1] // 2 

112 # if self.config.w_channel == -1: 

113 # # all channels 

114 # watermarking_mask[:, :, anchor_p-self.config.w_radius:anchor_p+self.config.w_radius, anchor_p-self.config.w_radius:anchor_p+self.config.w_radius] = True 

115 # else: 

116 # watermarking_mask[:, self.config.w_channel, anchor_p-self.config.w_radius:anchor_p+self.config.w_radius, anchor_p-self.config.w_radius:anchor_p+self.config.w_radius] = True 

117 # elif self.config.w_mask_shape == 'no': 

118 # pass 

119 # else: 

120 # raise NotImplementedError(f'w_mask_shape: {self.config.w_mask_shape}') 

121 

122 return watermarking_mask 

123 

124 def inject_watermark(self, init_latents: torch.Tensor) -> torch.Tensor: 

125 init_latents_w_fft = torch.fft.fftshift(torch.fft.fft2(init_latents), dim=(-1, -2)) 

126 target_patch = self.gt_patch 

127 

128 if not torch.is_complex(target_patch): 

129 real = target_patch.to(torch.float32) 

130 imag = torch.zeros_like(real) 

131 target_patch = torch.complex(real, imag) 

132 target_patch = target_patch.to(init_latents_w_fft.dtype) 

133 

134 init_latents_w_fft[self.watermarking_mask] = target_patch[self.watermarking_mask].clone() 

135 

136 init_latents_w = torch.fft.ifft2(torch.fft.ifftshift(init_latents_w_fft, dim=(-1, -2))).real 

137 return init_latents_w 

138 

139@inherit_docstring 

140class TR(BaseWatermark): 

141 def __init__(self, 

142 watermark_config: TRConfig, 

143 *args, **kwargs): 

144 """ 

145 Initialize the TR watermarking algorithm. 

146  

147 Parameters: 

148 watermark_config (TRConfig): Configuration instance of the Tree-Ring algorithm. 

149 """ 

150 self.config = watermark_config 

151 self.utils = TRUtils(self.config) 

152 

153 self.detector = TRDetector( 

154 watermarking_mask=self.utils.watermarking_mask, 

155 gt_patch=self.utils.gt_patch, 

156 threshold=self.config.threshold, 

157 device=self.config.device, 

158 threshold_p_value=self.config.threshold_p_value, 

159 ) 

160 

161 def _generate_watermarked_image(self, prompt: str, *args, **kwargs) -> Image.Image: 

162 """Internal method to generate a watermarked image.""" 

163 watermarked_latents = self.utils.inject_watermark(self.config.init_latents) 

164 

165 # save watermarked latents 

166 self.set_orig_watermarked_latents(watermarked_latents) 

167 

168 # Construct generation parameters 

169 generation_params = { 

170 "num_images_per_prompt": self.config.num_images, 

171 "guidance_scale": self.config.guidance_scale, 

172 "num_inference_steps": self.config.num_inference_steps, 

173 "height": self.config.image_size[0], 

174 "width": self.config.image_size[1], 

175 "latents": watermarked_latents, 

176 } 

177 

178 # Add parameters from config.gen_kwargs 

179 if hasattr(self.config, "gen_kwargs") and self.config.gen_kwargs: 

180 for key, value in self.config.gen_kwargs.items(): 

181 if key not in generation_params: 

182 generation_params[key] = value 

183 

184 # Use kwargs to override default parameters 

185 for key, value in kwargs.items(): 

186 generation_params[key] = value 

187 

188 # Ensure latents parameter is not overridden 

189 generation_params["latents"] = watermarked_latents 

190 

191 return self.config.pipe( 

192 prompt, 

193 **generation_params 

194 ).images[0] 

195 

196 def _detect_watermark_in_image(self, 

197 image: Image.Image, 

198 prompt: str = "", 

199 *args, 

200 **kwargs) -> Dict[str, float]: 

201 """Detect the watermark in the image.""" 

202 # Use config values as defaults if not explicitly provided 

203 guidance_scale_to_use = kwargs.get('guidance_scale', self.config.guidance_scale) 

204 num_steps_to_use = kwargs.get('num_inference_steps', self.config.num_inference_steps) 

205 

206 # Step 1: Get Text Embeddings 

207 do_classifier_free_guidance = (guidance_scale_to_use > 1.0) 

208 prompt_embeds, negative_prompt_embeds = self.config.pipe.encode_prompt( 

209 prompt=prompt, 

210 device=self.config.device, 

211 do_classifier_free_guidance=do_classifier_free_guidance, 

212 num_images_per_prompt=1, # TODO: Multiple image generation to be supported 

213 ) 

214 

215 if do_classifier_free_guidance: 

216 text_embeddings = torch.cat([negative_prompt_embeds, prompt_embeds]) 

217 else: 

218 text_embeddings = prompt_embeds 

219 

220 # Step 2: Preprocess Image 

221 image = transform_to_model_format(image, target_size=self.config.image_size[0]).unsqueeze(0).to(text_embeddings.dtype).to(self.config.device) 

222 

223 # Step 3: Get Image Latents 

224 image_latents = get_media_latents(pipe=self.config.pipe, media=image, sample=False, decoder_inv=kwargs.get('decoder_inv', False)) 

225 

226 # Step 4: Reverse Image Latents 

227 # Pass only known parameters to forward_diffusion, and let kwargs handle any additional parameters 

228 inversion_kwargs = {k: v for k, v in kwargs.items() if k not in ['decoder_inv', 'guidance_scale', 'num_inference_steps']} 

229 

230 reversed_latents = self.config.inversion.forward_diffusion( 

231 latents=image_latents, 

232 text_embeddings=text_embeddings, 

233 guidance_scale=guidance_scale_to_use, 

234 num_inference_steps=num_steps_to_use, 

235 **inversion_kwargs 

236 )[-1] 

237 

238 # Step 5: Evaluate Watermark 

239 if 'detector_type' in kwargs: 

240 return self.detector.eval_watermark(reversed_latents, detector_type=kwargs['detector_type']) 

241 else: 

242 return self.detector.eval_watermark(reversed_latents) 

243 

244 def get_data_for_visualize(self, 

245 image: Image.Image, 

246 prompt: str="", 

247 guidance_scale: Optional[float]=None, 

248 decoder_inv: bool=False, 

249 *args, 

250 **kwargs) -> DataForVisualization: 

251 """Get data for visualization including detection inversion - similar to GS logic.""" 

252 # Use config values as defaults if not explicitly provided 

253 guidance_scale_to_use = guidance_scale if guidance_scale is not None else self.config.guidance_scale 

254 

255 # Step 1: Generate watermarked latents (generation process) 

256 set_random_seed(self.config.gen_seed) 

257 watermarked_latents = self.utils.inject_watermark(self.config.init_latents) 

258 

259 # Step 2: Generate actual watermarked image using the same process as _generate_watermarked_image 

260 generation_params = { 

261 "num_images_per_prompt": self.config.num_images, 

262 "guidance_scale": self.config.guidance_scale, 

263 "num_inference_steps": self.config.num_inference_steps, 

264 "height": self.config.image_size[0], 

265 "width": self.config.image_size[1], 

266 "latents": watermarked_latents, 

267 } 

268 

269 # Add parameters from config.gen_kwargs 

270 if hasattr(self.config, "gen_kwargs") and self.config.gen_kwargs: 

271 for key, value in self.config.gen_kwargs.items(): 

272 if key not in generation_params: 

273 generation_params[key] = value 

274 

275 # Generate the actual watermarked image 

276 watermarked_image = self.config.pipe( 

277 prompt, 

278 **generation_params 

279 ).images[0] 

280 

281 # Step 3: Perform watermark detection to get inverted latents (detection process) 

282 inverted_latents = None 

283 try: 

284 # Get Text Embeddings for detection 

285 do_classifier_free_guidance = (guidance_scale_to_use > 1.0) 

286 prompt_embeds, negative_prompt_embeds = self.config.pipe.encode_prompt( 

287 prompt=prompt, 

288 device=self.config.device, 

289 do_classifier_free_guidance=do_classifier_free_guidance, 

290 num_images_per_prompt=1, # TODO: Multiple image generation to be supported 

291 ) 

292 

293 if do_classifier_free_guidance: 

294 text_embeddings = torch.cat([negative_prompt_embeds, prompt_embeds]) 

295 else: 

296 text_embeddings = prompt_embeds 

297 

298 # Preprocess watermarked image for detection 

299 processed_image = transform_to_model_format( 

300 watermarked_image, 

301 target_size=self.config.image_size[0] 

302 ).unsqueeze(0).to(text_embeddings.dtype).to(self.config.device) 

303 

304 # Get Image Latents 

305 image_latents = get_media_latents( 

306 pipe=self.config.pipe, 

307 media=processed_image, 

308 sample=False, 

309 decoder_inv=decoder_inv 

310 ) 

311 

312 # Reverse Image Latents to get inverted noise 

313 inversion_kwargs = {k: v for k, v in kwargs.items() if k not in ['prompt', 'decoder_inv', 'guidance_scale', 'num_inference_steps']} 

314 

315 reversed_latents_list = self.config.inversion.forward_diffusion( 

316 latents=image_latents, 

317 text_embeddings=text_embeddings, 

318 guidance_scale=guidance_scale_to_use, 

319 num_inference_steps=self.config.num_inference_steps, 

320 **inversion_kwargs 

321 ) 

322 

323 inverted_latents = reversed_latents_list[-1] 

324 

325 except Exception as e: 

326 print(f"Warning: Could not perform inversion for visualization: {e}") 

327 inverted_latents = None 

328 

329 # Step 4: Prepare visualization data  

330 return DataForVisualization( 

331 config=self.config, 

332 utils=self.utils, 

333 reversed_latents=reversed_latents_list, 

334 orig_watermarked_latents=self.orig_watermarked_latents, 

335 image=image, 

336 ) 

337# try tr