(ICME 2025) 一行代码拯救水印鲁棒性:LAWMM 如何让扩散水印“自己学会藏在哪”

from __future__ import annotations
import torch
import torch.nn.functional as F

def learn_lawmm_mask(
    inversed_steps_latents: torch.Tensor,
    sample_steps_latents: torch.Tensor,
    signature_fft: torch.Tensor,
    sample_steps=(5, 10, 15),
    threshold: float = 0.91,
    num_learning_epoch: int = 50,
lr: float = 0.02,
    show_progress: bool = False,
) -> torch.Tensor:
"""
    Learn the LAWMM binary mask (frequency-domain adaptive mask).

    Args:
        inversed_steps_latents: [T_inv, B, C, H, W] inversion latents.
        sample_steps_latents:   [T_sam, B, C, H, W] sampling latents.
        signature_fft: complex tensor, broadcastable to [B, C, H, W].
        sample_steps: which time steps to use as constraints (1-based in original impl).
        threshold: sigmoid(p) > threshold -> True.
        num_learning_epoch: optimization steps for mask logits p.
lr: Adam lrforp.
        show_progress:if True, show a tqdm progress bar (if installed).

    Returns:
        final_mask: bool tensor with shape [B, C, H, W].
"""
if inversed_steps_latents.ndim != 5or sample_steps_latents.ndim != 5:
        raise ValueError("Expected latents with shape [T, B, C, H, W].")

    device = inversed_steps_latents.device
    T_inv, B, C, H, W = inversed_steps_latents.shape

    # ---- Step indices (match original logic) ----
    step_idx_x = torch.tensor([int(i) - 1for i in sample_steps], device=device, dtype=torch.long)
    step_idx_y = (T_inv - 1) - torch.tensor([int(i) for i in sample_steps], device=device, dtype=torch.long)

if (step_idx_x.min() < 0) or (step_idx_x.max() >= sample_steps_latents.shape[0]):
        raise ValueError("sample_steps out of range for sample_steps_latents.")
if (step_idx_y.min() < 0) or (step_idx_y.max() >= inversed_steps_latents.shape[0]):
        raise ValueError("sample_steps out of range for inversed_steps_latents.")

    # ---- IMPORTANT: FFT safety => doall FFT math in float32/complex64 ----
    x_list = sample_steps_latents.index_select(0, step_idx_x).detach().float() # [K,B,C,H,W] float32
y = inversed_steps_latents.index_select(0, step_idx_y).detach().float() # [K,B,C,H,W] float32

    # ---- Broadcast signature_fft to [B,C,H,W] and ensure complex64 ----
sig = signature_fft
if not torch.is_complex(sig):
        raise TypeError("signature_fft must be a complex tensor (torch.complex64/128).")

ifsig.ndim == 2:
sig = sig[None, None, :, :] # [1,1,H,W]
    elif sig.ndim == 3:
sig = sig[None, :, :, :] # [1,C,H,W]
    elif sig.ndim != 4:
        raise ValueError("signature_fft must have shape [H,W] or [C,H,W] or [B,C,H,W].")

ifsig.shape[-2:] != (H, W):
        raise ValueError(f"signature_fft spatial size {sig.shape[-2:]} != latents {(H,W)}")
ifsig.shape[1] not in (1, C):
        raise ValueError(f"signature_fft channel dim must be 1 or {C}, got {sig.shape[1]}")
ifsig.shape[0] not in (1, B):
        raise ValueError(f"signature_fft batch dim must be 1 or {B}, got {sig.shape[0]}")

sig = sig.to(device=device)
    # Force to complex64 for stable FFT ops
ifsig.dtype != torch.complex64:
sig = sig.to(torch.complex64)
sig = sig.expand(B, C, H, W).detach() # [B,C,H,W] complex64

    # ---- Learn mask logits p in float32 (Adam + stability) ----
p = torch.rand((B, C, H, W), device=device, dtype=torch.float32, requires_grad=True)
opt = torch.optim.Adam([p], lr=float(lr))

    it = range(int(num_learning_epoch))
if show_progress:
try:
            from tqdm import tqdm
            it = tqdm(it, total=int(num_learning_epoch))
        except Exception:
            pass

for _ in it:
        loss = 0.0
        mask_soft = torch.sigmoid(p) # float32

fork in range(x_list.shape[0]):
x = x_list[k] # float32 [B,C,H,W]
            # FFT -> complex64
X = torch.fft.fft2(x)

            X_prime = mask_soft * sig + (1.0 - mask_soft) * X
            x_prime = torch.fft.ifft2(X_prime).real # float32

            # Match original style constraint:
            loss = loss + F.mse_loss(x_prime + (sig * mask_soft).real, y[k])

opt.zero_grad(set_to_none=True)
        loss.backward()
opt.step()

    final_mask = (torch.sigmoid(p) > float(threshold))
return final_mask


def apply_watermark_fft(
    last_inverse_latent: torch.Tensor,
    watermarking_mask: torch.Tensor,
    signature_fft: torch.Tensor
) -> torch.Tensor:
"""
    Replace FFT bins selected by watermarking_mask with signature_fft bins.
    last_inverse_latent: [B,C,H,W] real
    watermarking_mask : bool broadcastable to [B,C,H,W]
    signature_fft : complex broadcastable to [B,C,H,W]
"""
if last_inverse_latent.ndim != 4:
        raise ValueError("last_inverse_latent must be [B,C,H,W].")
if not torch.is_complex(signature_fft):
        raise TypeError("signature_fft must be complex.")

    B, C, H, W = last_inverse_latent.shape
    device = last_inverse_latent.device

x = last_inverse_latent.float()
X = torch.fft.fft2(x)

sig = signature_fft
ifsig.ndim == 2:
sig = sig[None, None, :, :]
    elif sig.ndim == 3:
sig = sig[None, :, :, :]
    elif sig.ndim != 4:
        raise ValueError("signature_fft must have shape [H,W] or [C,H,W] or [B,C,H,W].")

sig = sig.to(device=device)
ifsig.dtype != torch.complex64:
sig = sig.to(torch.complex64)
sig = sig.expand(B, C, H, W)

    mask = watermarking_mask
if mask.dtype != torch.bool:
        mask = mask.bool()
if mask.shape != (B, C, H, W):
        mask = mask.expand(B, C, H, W)

X = X.clone()
X[mask] = sig[mask].clone()
return torch.fft.ifft2(X).real