Spaces:
Sleeping
Sleeping
File size: 5,712 Bytes
9d3899c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """Mitigation tests.
These exercise the real mitigation code path (text-encoder dropout averaging +
the detect -> mitigate -> block decision) against a tiny toy pipeline -- no
Stable Diffusion download. torch is required, so the module is skipped where
torch is absent (e.g. the minimal CI install); the memguard env has torch, so it
runs there.
The toy UNet emits ``mean(tanh(embedding)) * base``, so the memorization signal
``||eps_c - eps_uc||`` is monotonic in the gap between the conditional and
unconditional embeddings. The toy text encoder models attention dropout as a
pull of the conditional embedding toward the unconditional one, so averaging
dropout draws (the mitigation) reduces that gap.
"""
import types
import pytest
torch = pytest.importorskip("torch")
from memguard.detector import MemorizationDetector # noqa: E402
from memguard.mitigation import mitigate_prompt_embeddings # noqa: E402
from memguard.pipeline import GuardedDiffusionPipeline # noqa: E402
C, S, SEQ, DIM = 4, 8, 4, 6
SD1_CAL = (1.04, 0.96) # so the toy score spans the 0.9 threshold like SD1.4
class _Unet:
def __init__(self):
self.config = types.SimpleNamespace(in_channels=C, sample_size=S)
self.dtype = torch.float32
self.base = torch.ones(1, C, S, S)
def __call__(self, x, t, encoder_hidden_states=None, return_dict=True):
a = torch.tanh(encoder_hidden_states).mean(dim=(1, 2)) # [B]
out = a.view(-1, 1, 1, 1) * self.base
return (out,) if not return_dict else types.SimpleNamespace(sample=out)
class _Scheduler:
def __init__(self):
self.alphas_cumprod = torch.linspace(0.9999, 0.001, 1000)
self.init_noise_sigma = 1.0
self.timesteps = torch.arange(49, -1, -1)
def set_timesteps(self, n, device=None):
self.timesteps = torch.arange(n - 1, -1, -1)
def scale_model_input(self, x, t):
return x
class _TextEncoder(torch.nn.Module):
"""Minimal CLIP-like encoder exposing a float ``dropout`` attribute."""
def __init__(self):
super().__init__()
self.dropout = 0.0
class _Pipe:
"""Minimal stand-in implementing just the surface mitigation/metric use."""
def __init__(self, cond_val=0.5, uncond_val=0.1):
self.unet = _Unet()
self.scheduler = _Scheduler()
self.text_encoder = _TextEncoder().eval() # inference pipelines keep this in eval
self._execution_device = "cpu"
self._cond_val = cond_val
self._uncond_val = uncond_val
self.calls = []
def encode_prompt(self, prompt, device, num_images_per_prompt=1,
do_classifier_free_guidance=True, negative_prompt=None):
te = self.text_encoder
cv = self._cond_val
if te.training and te.dropout > 0:
# Attention dropout pulls the conditional embedding toward the
# unconditional one (breaking the memorization trigger), with noise.
frac = min(1.0, te.dropout * 3.0) # p=0.3 -> ~0.9 pull
cv = (self._cond_val - (self._cond_val - self._uncond_val) * frac
+ float(torch.randn(1)) * 0.02)
cond = torch.full((1, SEQ, DIM), cv)
uncond = torch.full((1, SEQ, DIM), self._uncond_val)
return cond, uncond
def __call__(self, prompt=None, *, prompt_embeds=None, negative_prompt_embeds=None,
num_inference_steps=50, guidance_scale=7.5, latents=None, **kwargs):
self.calls.append("embeds" if prompt_embeds is not None else "prompt")
return types.SimpleNamespace(images=[object()])
def test_dropout_mitigation_reduces_signal_and_restores_state():
torch.manual_seed(0)
pipe = _Pipe(cond_val=0.5, uncond_val=0.1)
latents = torch.ones(1, C, S, S)
out = mitigate_prompt_embeddings(
pipe, "memorized prompt", latents=latents,
num_inference_steps=50, num_samples=10, dropout_p=0.3,
threshold=0.9, calibration=SD1_CAL,
)
assert out["samples"] == 10
assert out["signal_after"] < out["signal_before"]
assert out["passed"] is True
assert torch.isfinite(out["prompt_embeds"]).all()
# the text encoder is returned to its original state (dropout off, eval mode)
assert pipe.text_encoder.dropout == 0.0
assert pipe.text_encoder.training is False
def test_guarded_mitigates_when_possible():
torch.manual_seed(0)
pipe = _Pipe(cond_val=0.5, uncond_val=0.1) # strong signal -> memorized
guarded = GuardedDiffusionPipeline(
pipe, MemorizationDetector(threshold=0.9, calibration=SD1_CAL), mitigate=True,
)
res = guarded.generate("memorized prompt", num_inference_steps=5, seed=0)
assert res["mitigated"] is True
assert res["blocked"] is False
assert res["memorized"] is False
assert res["image"] is not None
assert res["signal_after"] < res["signal_before"]
def test_guarded_blocks_without_mitigation():
pipe = _Pipe(cond_val=0.5, uncond_val=0.1)
guarded = GuardedDiffusionPipeline(
pipe, MemorizationDetector(threshold=0.9), mitigate=False,
)
res = guarded.generate("memorized prompt", num_inference_steps=5, seed=0)
assert res["memorized"] is True
assert res["blocked"] is True
assert res["image"] is None
assert res["mitigated"] is False
def test_guarded_passes_benign():
pipe = _Pipe(cond_val=0.12, uncond_val=0.1) # tiny gap -> not memorized
guarded = GuardedDiffusionPipeline(pipe, MemorizationDetector(threshold=0.9))
res = guarded.generate("benign prompt", num_inference_steps=5, seed=0)
assert res["memorized"] is False
assert res["blocked"] is False
assert res["mitigated"] is False
assert res["image"] is not None
|