ColossalAI/examples/images/diffusion/scripts/tests/test_watermark.py