| #!/usr/bin/env python3 | |
| import torch | |
| from diffusers import ConsistencyModelPipeline, UNet2DModel | |
| device = "cpu" | |
| # Load the cd_bedroom256_lpips checkpoint. | |
| model_id_or_path = "openai/diffusers-cd_bedroom256_lpips" | |
| pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path) | |
| pipe.to(device) | |
| # Multistep sampling | |
| # Timesteps can be explicitly specified; the particular timesteps below are from the original Github repo: | |
| # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L83 | |
| for _ in range(10): | |
| image = pipe(timesteps=[17, 0]).images[0] | |
| image.show() | |