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, AutoencoderKL 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 import random import math 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") # 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) torch.manual_seed(seed) # 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] def attack(self, img, jpeg_compression, downscale, crop, saturation, brightness, contrast): 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.Resampling.LANCZOS) if crop != 0: width, height = img.size area = width * height log_rmin = math.log(0.5) log_rmax = math.log(2.0) for _ in range(10): target_area = area * (1 - crop) aspect_ratio = math.exp(random.random() * (log_rmax - log_rmin) + log_rmin) w = int(round(math.sqrt(target_area * aspect_ratio))) h = int(round(math.sqrt(target_area / aspect_ratio))) if 0 < w <= width and 0 < h <= height: top = random.randint(0, height - h + 1) left = random.randint(0, width - w + 1) img = img.crop((left, top, left+w, top+h)) break converter = ImageEnhance.Color(img) img = converter.enhance(saturation) converter = ImageEnhance.Brightness(img) img = converter.enhance(brightness) converter = ImageEnhance.Contrast(img) img = converter.enhance(contrast) # JPEG attack mf = io.BytesIO() img.save(mf, format='JPEG', quality=jpeg_compression) mf.seek(0) img = Image.open(mf) return img def detect(self, img): # send to detection API and apply JPEG compression attack mf = io.BytesIO() img.save(mf, format='PNG') b64 = base64.b64encode(mf.getvalue()) data = { 'image': b64.decode('utf8') } headers = {} api_key = os.getenv('BZH_API_KEY') if api_key: headers['x-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'] result = "No watermark detected." rpv = 10**int(math.log10(pvalue)) if pvalue < 1e-3: result = "Watermark detected with low confidence" # (p-value<%.0e)" % rpv if pvalue < 1e-9: result = "Watermark detected with high confidence" # (p-value<%.0e)" % rpv score = min(int(-math.log10(pvalue)), 10) #print("score = ", score) return { result: score/10 } 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 brought to you by [IMATAG](https://www.imatag.com/) presents watermarking of images generated via [StableDiffusion XL Turbo](https://huggingface.co/stabilityai/sdxl-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 a demo version of [IMATAG](https://www.imatag.com/)'s in-house decoder. The watermarking system operates in zero-bit mode for improved robustness.""") gr.Markdown("""## 1. Generate Select a watermarking strength and generate images with StableDiffusion-XL Turbo from prompt and seed as usual.""") 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.Image(type="pil", width=512, height=512, sources=[], interactive=False) gr.Markdown("""## 2. Edit With these controls you may alter the generated image before detection. You may also upload your own edited image instead.""") with gr.Row(): with gr.Column(): with gr.Row(): downscale = gr.Slider(1, 3, value=1, step=0.1, label="Downscale ratio") crop = gr.Slider(0, 0.9, value=0, step=0.01, label="Random crop ratio") with gr.Row(): brightness = gr.Slider(0, 2, value=1, step=0.1, label="Brightness") contrast = gr.Slider(0, 2, value=1, step=0.1, label="Contrast") with gr.Row(): 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("Edit") with gr.Row(): attacked_image = gr.Image(type="pil", width=512, sources=['upload', 'clipboard']) gr.Markdown("""## 3. Detect Detect the watermark on the altered image. Watermark may not be detected if the image is altered too strongly. """) with gr.Row(): btn3 = gr.Button("Detect") with gr.Row(): detection_label = gr.Label(label="Detection info", show_label=False) btn1.click(fn=backend.generate, inputs=[mode, seed, inp], outputs=[watermarked_image], api_name="generate") btn2.click(fn=backend.attack, inputs=[watermarked_image, jpeg_compression, downscale, crop, saturation, brightness, contrast], outputs=[attacked_image], api_name="attack") btn3.click(fn=backend.detect, inputs=[attacked_image], outputs=[detection_label], api_name="detect") return demo if __name__ == '__main__': demo = interface() demo.launch()