Update app.py

app.py

@@ -1,151 +1,12 @@
-# import all the libraries
-import math
-import numpy as np
-import scipy
-from PIL import Image
 import torch
-import torchvision.transforms as tforms
-from diffusers import DiffusionPipeline, UNet2DConditionModel, DDIMScheduler, DDIMInverseScheduler
-from diffusers.models import AutoencoderKL
-import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-# 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")
+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
+)
 
-# 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)
+model.push_to_hub("recurrentgemma-9b-it-8bit")