Spaces:
Running
on
A10G
Running
on
A10G
Revert model change
Browse files
app.py
CHANGED
@@ -10,10 +10,9 @@ from diffusers.models import AutoencoderKL
|
|
10 |
import gradio as gr
|
11 |
|
12 |
# load SDXL pipeline
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
pipe = DiffusionPipeline.from_pretrained("common-canvas/CommonCanvas-XL-NC", torch_dtype=torch.float16)
|
17 |
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
|
18 |
pipe = pipe.to("cuda")
|
19 |
|
@@ -31,7 +30,7 @@ def get_pattern(shape, w_seed=999999):
|
|
31 |
g = torch.Generator(device=pipe.device)
|
32 |
g.manual_seed(w_seed)
|
33 |
gt_init = pipe.prepare_latents(1, pipe.unet.in_channels,
|
34 |
-
|
35 |
pipe.unet.dtype, pipe.device, g)
|
36 |
gt_patch = torch.fft.fftshift(torch.fft.fft2(gt_init), dim=(-1, -2))
|
37 |
# ring pattern. paper found this to be effective
|
@@ -45,7 +44,7 @@ def get_pattern(shape, w_seed=999999):
|
|
45 |
return gt_patch
|
46 |
|
47 |
def transform_img(image):
|
48 |
-
tform = tforms.Compose([tforms.Resize(
|
49 |
image = tform(image)
|
50 |
return 2.0 * image - 1.0
|
51 |
|
@@ -68,7 +67,7 @@ def get_noise():
|
|
68 |
|
69 |
# inject watermark
|
70 |
init_latents = pipe.prepare_latents(1, pipe.unet.in_channels,
|
71 |
-
|
72 |
pipe.unet.dtype, pipe.device, None)
|
73 |
init_latents_fft = torch.fft.fftshift(torch.fft.fft2(init_latents), dim=(-1, -2))
|
74 |
init_latents_fft[w_mask] = w_key[w_mask].clone()
|
@@ -86,7 +85,7 @@ def detect(image):
|
|
86 |
|
87 |
# ddim inversion
|
88 |
img = transform_img(image).unsqueeze(0).to(pipe.unet.dtype).to(pipe.device)
|
89 |
-
image_latents = pipe.vae.encode(img).latent_dist.mode() * 0.
|
90 |
inverted_latents = pipe(prompt="", latents=image_latents, guidance_scale=1, num_inference_steps=50, output_type="latent")
|
91 |
inverted_latents = inverted_latents.images
|
92 |
|
@@ -146,7 +145,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="green",secondary_hue="green", f
|
|
146 |
det_in = gr.Image(interactive=True, sources=["upload","clipboard"], show_label=False)
|
147 |
det_btn.click(fn=manager, inputs=det_in, outputs=det_out)
|
148 |
with gr.Row():
|
149 |
-
gr.HTML('<center><h1> </h1>Acknowledgements: Dendrokronos uses <a href="https://huggingface.co/
|
150 |
|
151 |
app.queue()
|
152 |
app.launch(show_api=False)
|
|
|
10 |
import gradio as gr
|
11 |
|
12 |
# load SDXL pipeline
|
13 |
+
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
|
14 |
+
unet = UNet2DConditionModel.from_pretrained("mhdang/dpo-sdxl-text2image-v1", subfolder="unet", torch_dtype=torch.float16)
|
15 |
+
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", unet=unet, vae=vae, torch_dtype=torch.float16)
|
|
|
16 |
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
|
17 |
pipe = pipe.to("cuda")
|
18 |
|
|
|
30 |
g = torch.Generator(device=pipe.device)
|
31 |
g.manual_seed(w_seed)
|
32 |
gt_init = pipe.prepare_latents(1, pipe.unet.in_channels,
|
33 |
+
1024, 1024,
|
34 |
pipe.unet.dtype, pipe.device, g)
|
35 |
gt_patch = torch.fft.fftshift(torch.fft.fft2(gt_init), dim=(-1, -2))
|
36 |
# ring pattern. paper found this to be effective
|
|
|
44 |
return gt_patch
|
45 |
|
46 |
def transform_img(image):
|
47 |
+
tform = tforms.Compose([tforms.Resize(1024),tforms.CenterCrop(1024),tforms.ToTensor()])
|
48 |
image = tform(image)
|
49 |
return 2.0 * image - 1.0
|
50 |
|
|
|
67 |
|
68 |
# inject watermark
|
69 |
init_latents = pipe.prepare_latents(1, pipe.unet.in_channels,
|
70 |
+
1024, 1024,
|
71 |
pipe.unet.dtype, pipe.device, None)
|
72 |
init_latents_fft = torch.fft.fftshift(torch.fft.fft2(init_latents), dim=(-1, -2))
|
73 |
init_latents_fft[w_mask] = w_key[w_mask].clone()
|
|
|
85 |
|
86 |
# ddim inversion
|
87 |
img = transform_img(image).unsqueeze(0).to(pipe.unet.dtype).to(pipe.device)
|
88 |
+
image_latents = pipe.vae.encode(img).latent_dist.mode() * 0.13025
|
89 |
inverted_latents = pipe(prompt="", latents=image_latents, guidance_scale=1, num_inference_steps=50, output_type="latent")
|
90 |
inverted_latents = inverted_latents.images
|
91 |
|
|
|
145 |
det_in = gr.Image(interactive=True, sources=["upload","clipboard"], show_label=False)
|
146 |
det_btn.click(fn=manager, inputs=det_in, outputs=det_out)
|
147 |
with gr.Row():
|
148 |
+
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>')
|
149 |
|
150 |
app.queue()
|
151 |
app.launch(show_api=False)
|