Update app.py

app.py

@@ -31,7 +31,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,
-                                   1024, 1024,
+                                   512, 512,
                                    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 +45,7 @@ def get_pattern(shape, w_seed=999999):
     return gt_patch
 
 def transform_img(image):
-    tform = tforms.Compose([tforms.Resize(1024),tforms.CenterCrop(1024),tforms.ToTensor()])
+    tform = tforms.Compose([tforms.Resize(512),tforms.CenterCrop(512),tforms.ToTensor()])
     image = tform(image)
     return 2.0 * image - 1.0
 
@@ -68,7 +68,7 @@ def get_noise():
 
     # inject watermark
     init_latents = pipe.prepare_latents(1, pipe.unet.in_channels,
-                                        1024, 1024,
+                                        512, 512,
                                         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 +86,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.13025
+    image_latents = pipe.vae.encode(img).latent_dist.mode() * 0.18215
     inverted_latents = pipe(prompt="", latents=image_latents, guidance_scale=1, num_inference_steps=50, output_type="latent")
     inverted_latents = inverted_latents.images