YuxinWenRick/tree-ring-watermark

About sampling and encapsulation

linju2022perfect opened this issue · 1 comments

in this code: set_random_seed(seed)
init_latents_no_w = pipe.get_random_latents()
outputs_no_w = pipe(
.......
num_inference_steps=args.num_inference_steps,
.......
)
I modified num_inference_steps=args.num_inference_steps-1,Want to get x1...
After getting it, I want to get x0 through the latent of x1 through num_inference_steps=1, but I find that the image effect of x0 is worse than that of x1, is there no way to push it out at every step?
The only way to get x1 is init_latents_no_w = pipe.get_random_latents() and num_inference_steps=args.num_inference_steps-1.......
I hope you can help me or tell me about your pipe encapsulation, I would appreciate it

Hi thanks for reaching out, and sorry for the late response. I was busy with the NeurIPS deadline recently.

If you pass args.num_inference_steps-1, you will still get x0 by doing DDIM with 49 steps (assume args.num_inference_steps=50). Also, if you pass num_inference_steps=1, the sampling function will assume the provided latent is the initial noise, and do one step from x1000 to x0. The modified stable diffusion uses the same sampling function as the original Hugging Face stable diffusion pipeline, so the usage is the same.

In order to achieve the thing you want to do, I think you need to modify the sampling function a little bit. You should still use args.num_inference_steps, but maybe provide an argument for stop step and stop the sampling at a certain step in this loop. You can also provide an argument for start step to do x1 to x0.

Let me know if you have further questions!