Revert model change

app.py

@@ -10,10 +10,9 @@ from diffusers.models import AutoencoderKL
 import gradio as gr
 
 # 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 = DiffusionPipeline.from_pretrained("common-canvas/CommonCanvas-XL-NC", torch_dtype=torch.float16)
+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")
 
@@ -31,7 +30,7 @@ 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,
-                                   512, 512,
+                                   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
@@ -45,7 +44,7 @@ def get_pattern(shape, w_seed=999999):
     return gt_patch
 
 def transform_img(image):
-    tform = tforms.Compose([tforms.Resize(512),tforms.CenterCrop(512),tforms.ToTensor()])
+    tform = tforms.Compose([tforms.Resize(1024),tforms.CenterCrop(1024),tforms.ToTensor()])
     image = tform(image)
     return 2.0 * image - 1.0
 
@@ -68,7 +67,7 @@ def get_noise():
 
     # inject watermark
     init_latents = pipe.prepare_latents(1, pipe.unet.in_channels,
-                                        512, 512,
+                                        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()
@@ -86,7 +85,7 @@ def detect(image):
 
     # 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.18215
+    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
 
@@ -146,7 +145,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="green",secondary_hue="green", f
                 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/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>')
+        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)