devingulliver commited on
Commit
5fe2738
·
verified ·
1 Parent(s): 9aff15d

Revert model change

Browse files
Files changed (1) hide show
  1. app.py +8 -9
app.py CHANGED
@@ -10,10 +10,9 @@ from diffusers.models import AutoencoderKL
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 = 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
- 512, 512,
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(512),tforms.CenterCrop(512),tforms.ToTensor()])
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
- 512, 512,
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.18215
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>&nbsp;</h1>Acknowledgements: Dendrokronos uses <a href="https://huggingface.co/common-canvas/CommonCanvas-XL-NC">CommonCanvas</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>')
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>&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>')
149
 
150
  app.queue()
151
  app.launch(show_api=False)