Skip to content

Commit 1d7adf1

Browse files
Improve unconditional diffusers example (huggingface#414)
* use gpu and improve * Update README.md Co-authored-by: Patrick von Platen <[email protected]> * Update README.md Co-authored-by: Patrick von Platen <[email protected]> Co-authored-by: Patrick von Platen <[email protected]>
1 parent f4a785c commit 1d7adf1

File tree

1 file changed

+6
-2
lines changed

1 file changed

+6
-2
lines changed

README.md

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -279,18 +279,22 @@ image.save("squirrel.png")
279279
- [Unconditional Diffusion with discrete scheduler](https://huggingface.co/google/ddpm-celebahq-256)
280280
```python
281281
# !pip install diffusers
282+
from torch import autocast
282283
from diffusers import DDPMPipeline, DDIMPipeline, PNDMPipeline
283284

284285
model_id = "google/ddpm-celebahq-256"
286+
device = "cuda"
285287

286288
# load model and scheduler
287289
ddpm = DDPMPipeline.from_pretrained(model_id) # you can replace DDPMPipeline with DDIMPipeline or PNDMPipeline for faster inference
290+
ddpm.to(device)
288291

289292
# run pipeline in inference (sample random noise and denoise)
290-
image = ddpm().images
293+
with autocast("cuda"):
294+
image = ddpm().images[0]
291295

292296
# save image
293-
image[0].save("ddpm_generated_image.png")
297+
image.save("ddpm_generated_image.png")
294298
```
295299
- [Unconditional Latent Diffusion](https://huggingface.co/CompVis/ldm-celebahq-256)
296300
- [Unconditional Diffusion with continous scheduler](https://huggingface.co/google/ncsnpp-ffhq-1024)

0 commit comments

Comments
 (0)