import gradio as gr import os import torch import numpy as np device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') from diffusers import DiffusionPipeline import torchvision.transforms as transforms from copy import deepcopy from collections import OrderedDict import requests import json from PIL import Image, ImageEnhance import base64 import io 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") # 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 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) output = self.pipe(prompt, num_inference_steps=4, guidance_scale=0.0, output_type="pil") return { "background": output.images[0], "layers": [], "composite": None } def attack_detect(self, img_edit, jpeg_compression, downscale, saturation): img = img_edit["composite"] img = img.convert("RGB") # attack if downscale != 1: size = img.size size = (int(size[0] / downscale), int(size[1] / downscale)) img = img.resize(size, Image.BICUBIC) converter = ImageEnhance.Color(img) img = converter.enhance(saturation) # send to detection API and apply JPEG compression attack mf = io.BytesIO() img.save(mf, format='JPEG', quality=jpeg_compression) # includes JPEG attack b64 = base64.b64encode(mf.getvalue()) data = { 'image': b64.decode('utf8') } headers = {} api_key = os.environ.get('BZH_API_KEY', None) if api_key: headers['BZH_API_KEY'] = api_key response = requests.post('https://bzh.imatag.com/bzh/api/v1.0/detect', json=data, headers=headers) response.raise_for_status() data = response.json() pvalue = data['p-value'] 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) if pvalue < 1e-3: result = "Weak watermark detected" # (< 1/%d chances of being wrong)" % chances if pvalue < 1e-9: result = "Strong watermark detected" # (< 1/%d chances of being wrong)" % chances return (img0, result) def interface(): prompt = "sailing ship in storm by Rembrandt" backend = BZHStableSignatureDemo() decoders = list(backend.decoders.keys()) with gr.Blocks() as demo: gr.Markdown("""# Watermarked SDXL-Turbo demo This demo presents watermarking of images generated via StableDiffusion XL Turbo. Using the method presented in [StableSignature](https://ai.meta.com/blog/stable-signature-watermarking-generative-ai/), the VAE decoder of StableDiffusion is fine-tuned to produce images including a specific invisible watermark. We combined this method with our in-house decoder which operates in zero-bit mode for improved robustness.""") with gr.Row(): inp = gr.Textbox(label="Prompt", value=prompt) seed = gr.Number(label="Seed", precision=0) mode = gr.Dropdown(choices=decoders, label="Watermark strength", value="medium") with gr.Row(): btn1 = gr.Button("Generate") with gr.Row(): watermarked_image = gr.ImageEditor(type="pil", width=512, height=512) with gr.Column(): downscale = gr.Slider(1, 3, value=1, step=0.1, label="Downscale ratio") saturation = gr.Slider(0, 2, value=1, step=0.1, label="Color saturation") jpeg_compression = gr.Slider(value=100, step=5, label="JPEG quality") btn2 = gr.Button("Attack & Detect") with gr.Row(): attacked_image = gr.Image(type="pil", width=256) detection_label = gr.Label(label="Detection info") btn1.click(fn=backend.generate, inputs=[mode, seed, inp], outputs=[watermarked_image], api_name="generate") btn2.click(fn=backend.attack_detect, inputs=[watermarked_image, jpeg_compression, downscale, saturation], outputs=[attacked_image, detection_label], api_name="detect") return demo if __name__ == '__main__': demo = interface() demo.launch()