use packaged VAEs

app.py

@@ -7,7 +7,7 @@ import numpy as np
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoencoderKL
 import torchvision.transforms as transforms
 
 from copy import deepcopy
@@ -26,42 +26,35 @@ class BZHStableSignatureDemo(object):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+
         self.pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16").to("cuda")
-        try:
-            print("self.pipe.watermark = ", self.pipe.watermark)
-        except:
-            print("no self.pipe.watermark")
-
-        # load the patched VQ-VAEs
-        sd1 = deepcopy(self.pipe.vae.state_dict()) # save initial state dict
-        self.decoders = decoders = OrderedDict([("no watermark", sd1)])
-        for name, patched_decoder_ckpt in (
-                ("weak", "models/checkpoint_000.pth.50000"),
-                ("medium", "models/checkpoint_000.pth.150000"),
-                ("strong", "models/checkpoint_000.pth.500000"),
-                ("extreme", "models/checkpoint_000.pth.1500000")):
-            sd2 = torch.load(patched_decoder_ckpt)['ldm_decoder']
-            msg = self.pipe.vae.load_state_dict(sd2, strict=False)
-            print(f"loaded LDM decoder state_dict with message\n{msg}")
-            print("you should check that the decoder keys are correctly matched")
-            decoders[name] = sd2
+
+        # disable invisible-watermark
+        self.pipe.watermark = None
+
+        # save the original VAE
+        decoders = OrderedDict([("no watermark", self.pipe.vae)])
+
+        # load the patched VAEs
+        for name in ("weak", "medium", "strong", "extreme"):
+            vae = AutoencoderKL.from_pretrained(f"imatag/stable-signature-bzh-sdxl-vae-{name}", torch_dtype=torch.float16).to("cuda")
+            decoders[name] = vae
+
         self.decoders = decoders
 
     def generate(self, mode, seed, prompt):
         generator = torch.Generator(device=device)
-        #if seed:
         torch.manual_seed(seed)
 
-        # load the patched VAE decoder
-        sd = self.decoders[mode]
-        self.pipe.vae.load_state_dict(sd, strict=False)
+        # load the patched VAE
+        vae = self.decoders[mode]
+        self.pipe.vae = vae
 
         output = self.pipe(prompt, num_inference_steps=4, guidance_scale=0.0, output_type="pil")
-        return output.images[0] #{ "background": output.images[0], "layers": [], "composite": None }
+        return output.images[0]
 
     def attack_detect(self, img, jpeg_compression, downscale, crop, saturation):
 
-        #img = img_edit["composite"]
         img = img.convert("RGB")
         
         # attack
@@ -69,6 +62,7 @@ class BZHStableSignatureDemo(object):
             size = img.size
             size = (int(size[0] / downscale), int(size[1] / downscale))
             img = img.resize(size, Image.Resampling.LANCZOS)
+
         if crop != 0:
             width, height = img.size
             area = width * height
@@ -108,17 +102,15 @@ class BZHStableSignatureDemo(object):
 
         mf.seek(0)
         img0 = Image.open(mf) # reload to show JPEG attack
-        #result = "resolution = %dx%d  p-value = %e" % (img.size[0], img.size[1], pvalue))
+
         result = "No watermark detected."
-        chances = int(1 / pvalue + 1)
         rpv = 10**int(math.log10(pvalue))
         if pvalue < 1e-3:
-            result = "Watermark detected with low confidence (p-value<%.0e)" % rpv # (< 1/%d chances of being wrong)" % chances
+            result = "Watermark detected with low confidence (p-value<%.0e)" % rpv
         if pvalue < 1e-9:
-            result = "Watermark detected with high confidence (p-value<%.0e)" % rpv # (< 1/%d chances of being wrong)" % chances
+            result = "Watermark detected with high confidence (p-value<%.0e)" % rpv
         return (img0, result)
 
-
 def interface():
     prompt = "sailing ship in storm by Rembrandt"