Upload app.py

app.py

@@ -1,12 +1,151 @@
+# import all the libraries
+import math
+import numpy as np
+import scipy
+from PIL import Image
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+import torchvision.transforms as tforms
+from diffusers import DiffusionPipeline, UNet2DConditionModel, DDIMScheduler, DDIMInverseScheduler
+from diffusers.models import AutoencoderKL
+import gradio as gr
 
-quant_config = BitsAndBytesConfig(load_in_8bit=True, llm_int8_skip_modules=["temporal_block"])
-tokenizer = AutoTokenizer.from_pretrained("alpindale/recurrentgemma-9b-it")
-model = AutoModelForCausalLM.from_pretrained(
-    "alpindale/recurrentgemma-9b-it",
-    device_map="auto", torch_dtype=torch.float16,
-    quantization_config=quant_config
-)
+# load SDXL pipeline
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
+unet = UNet2DConditionModel.from_pretrained("mhdang/dpo-sdxl-text2image-v1", subfolder="unet", torch_dtype=torch.float16)
+pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", unet=unet, vae=vae, torch_dtype=torch.float16)
+pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+pipe = pipe.to("cuda")
 
-model.push_to_hub("recurrentgemma-9b-it-8bit")
+# watermarking helper functions. paraphrased from the reference impl of arXiv:2305.20030
+
+def circle_mask(size=128, r=16, x_offset=0, y_offset=0):
+    x0 = y0 = size // 2
+    x0 += x_offset
+    y0 += y_offset
+    y, x = np.ogrid[:size, :size]
+    y = y[::-1]
+    return ((x - x0)**2 + (y-y0)**2)<= r**2
+
+def get_pattern(shape, w_seed=999999):
+    g = torch.Generator(device=pipe.device)
+    g.manual_seed(w_seed)
+    gt_init = pipe.prepare_latents(1, pipe.unet.in_channels,
+                                   1024, 1024,
+                                   pipe.unet.dtype, pipe.device, g)
+    gt_patch = torch.fft.fftshift(torch.fft.fft2(gt_init), dim=(-1, -2))
+    # ring pattern. paper found this to be effective
+    gt_patch_tmp = gt_patch.clone().detach()
+    for i in range(shape[-1] // 2, 0, -1):
+        tmp_mask = circle_mask(gt_init.shape[-1], r=i)
+        tmp_mask = torch.tensor(tmp_mask)
+        for j in range(gt_patch.shape[1]):
+            gt_patch[:, j, tmp_mask] = gt_patch_tmp[0, j, 0, i].item()
+
+    return gt_patch
+
+def transform_img(image):
+    tform = tforms.Compose([tforms.Resize(1024),tforms.CenterCrop(1024),tforms.ToTensor()])
+    image = tform(image)
+    return 2.0 * image - 1.0
+
+# hyperparameters
+shape = (1, 4, 128, 128)
+w_seed = 7433 # TREE :)
+w_channel = 0
+w_radius = 16 # the suggested r from section 4.4 of paper
+
+# get w_key and w_mask
+np_mask = circle_mask(shape[-1], r=w_radius)
+torch_mask = torch.tensor(np_mask).to(pipe.device)
+w_mask = torch.zeros(shape, dtype=torch.bool).to(pipe.device)
+w_mask[:, w_channel] = torch_mask
+w_key = get_pattern(shape, w_seed=w_seed).to(pipe.device)
+
+
+def get_noise():
+    # moved w_key and w_mask to globals
+
+    # inject watermark
+    init_latents = pipe.prepare_latents(1, pipe.unet.in_channels,
+                                        1024, 1024,
+                                        pipe.unet.dtype, pipe.device, None)
+    init_latents_fft = torch.fft.fftshift(torch.fft.fft2(init_latents), dim=(-1, -2))
+    init_latents_fft[w_mask] = w_key[w_mask].clone()
+    init_latents = torch.fft.ifft2(torch.fft.ifftshift(init_latents_fft, dim=(-1, -2))).real
+    # hot fix to prevent out of bounds values. will "properly" fix this later
+    init_latents[init_latents == float("Inf")] = 4
+    init_latents[init_latents == float("-Inf")] = -4
+
+    return init_latents
+
+def detect(image):
+    # invert scheduler
+    curr_scheduler = pipe.scheduler
+    pipe.scheduler = DDIMInverseScheduler.from_config(pipe.scheduler.config)
+
+    # ddim inversion
+    img = transform_img(image).unsqueeze(0).to(pipe.unet.dtype).to(pipe.device)
+    image_latents = pipe.vae.encode(img).latent_dist.mode() * 0.13025
+    inverted_latents = pipe(prompt="", latents=image_latents, guidance_scale=1, num_inference_steps=50, output_type="latent")
+    inverted_latents = inverted_latents.images
+
+    # calculate p-value instead of detection threshold. more rigorous, plus we can do a non-boolean output
+    inverted_latents_fft = torch.fft.fftshift(torch.fft.fft2(inverted_latents), dim=(-1, -2))[w_mask].flatten()
+    target = w_key[w_mask].flatten()
+    inverted_latents_fft = torch.concatenate([inverted_latents_fft.real, inverted_latents_fft.imag])
+    target = torch.concatenate([target.real, target.imag])
+
+    sigma = inverted_latents_fft.std()
+    lamda = (target ** 2 / sigma ** 2).sum().item()
+    x = (((inverted_latents_fft - target) / sigma) ** 2).sum().item()
+    p_value = scipy.stats.ncx2.cdf(x=x, df=len(target), nc=lamda)
+
+    # revert scheduler
+    pipe.scheduler = curr_scheduler
+
+    if p_value == 0:
+        return 1.0
+    else:
+        return max(0.0, 1-1/math.log(5/p_value,10))
+
+def generate(prompt):
+    return pipe(prompt=prompt, negative_prompt="monochrome", num_inference_steps=50, latents=get_noise()).images[0]
+
+# optimize for speed
+pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+print(detect(generate("an astronaut riding a green horse"))) # warmup after jit
+
+# actual gradio demo
+
+def manager(input, progress=gr.Progress(track_tqdm=True)): # to prevent the queue from overloading
+    if type(input) == str:
+        return generate(input)
+    elif type(input) == np.ndarray:
+        image = Image.fromarray(input)
+        percent = detect(image)
+        return {"watermarked": percent, "not_watermarked": 1.0-percent}
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="green",secondary_hue="green", font=gr.themes.GoogleFont("Fira Sans"))) as app:
+    with gr.Row():
+        gr.HTML('<center><p>Bad actors are using generative AI to destroy the livelihoods of real artists. We need transparency now.</p><h1><span style="font-size:1.5em">Introducing Dendrokronos 🌳</span></h1></center>')
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("# Generate\nType a prompt and hit Go. Dendrokronos will generate an invisibly-watermarked image.  \nYou can click the download button to save the finished image. Try it with the detector.")
+            with gr.Group():
+                with gr.Row():
+                    gen_in = gr.Textbox(max_lines=1, placeholder='try "a majestic tree at sunset, oil painting"', show_label=False, scale=4)
+                    gen_btn = gr.Button("Go", variant="primary", scale=0)
+            gen_out = gr.Image(interactive=False, show_label=False)
+            gen_btn.click(fn=manager, inputs=gen_in, outputs=gen_out)
+        with gr.Column():
+            gr.Markdown("# Detect\nUpload an image and hit Detect. Dendrokronos will predict the probability it was watermarked.  \nNote: Dendrokronos can only detect its own watermark. It won't detect other AIs, such as DALL-E.")
+            det_out = gr.Label(show_label=False)
+            with gr.Group():
+                det_btn = gr.Button("Detect", variant="primary")
+                det_in = gr.Image(interactive=True, sources=["upload","clipboard"], show_label=False)
+            det_btn.click(fn=manager, inputs=det_in, outputs=det_out)
+    with gr.Row():
+        gr.HTML('<center><h1>&nbsp;</h1>Acknowledgements: Dendrokronos uses <a href="https://huggingface.co/mhdang/dpo-sdxl-text2image-v1">SDXL DPO 1.0</a> for the underlying image generation and <a href="https://arxiv.org/abs/2305.20030">an algorithm by UMD researchers</a> for the watermark technology.<br />Dendrokronos is a project by Devin Gulliver.</center>')
+
+app.queue()
+app.launch(show_api=False)