|
|
|
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 |
|
|
|
|
|
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") |
|
|
|
|
|
|
|
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)) |
|
|
|
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 |
|
|
|
|
|
shape = (1, 4, 128, 128) |
|
w_seed = 7433 |
|
w_channel = 0 |
|
w_radius = 16 |
|
|
|
|
|
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(): |
|
|
|
|
|
|
|
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 |
|
|
|
init_latents[init_latents == float("Inf")] = 4 |
|
init_latents[init_latents == float("-Inf")] = -4 |
|
|
|
return init_latents |
|
|
|
def detect(image): |
|
|
|
curr_scheduler = pipe.scheduler |
|
pipe.scheduler = DDIMInverseScheduler.from_config(pipe.scheduler.config) |
|
|
|
|
|
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 |
|
|
|
|
|
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) |
|
|
|
|
|
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] |
|
|
|
|
|
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True) |
|
print(detect(generate("an astronaut riding a green horse"))) |
|
|
|
|
|
|
|
def manager(input, progress=gr.Progress(track_tqdm=True)): |
|
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> </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) |
|
|
