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