example.py 641 Bytes
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
2
3
#!/usr/bin/env python3
import tempfile

Patrick von Platen's avatar
improve  
Patrick von Platen committed
4
5
6
7
from diffusers import GaussianDDPMScheduler, UNetModel
from modeling_ddpm import DDPM


Patrick von Platen's avatar
Patrick von Platen committed
8
unet = UNetModel.from_pretrained("fusing/ddpm_dummy")
Patrick von Platen's avatar
improve  
Patrick von Platen committed
9
sampler = GaussianDDPMScheduler.from_config("fusing/ddpm_dummy")
Patrick von Platen's avatar
Patrick von Platen committed
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

# compose Diffusion Pipeline
ddpm = DDPM(unet, sampler)
# generate / sample
image = ddpm()
print(image)


# save and load with 0 extra code (handled by general `DiffusionPipeline` class)
with tempfile.TemporaryDirectory() as tmpdirname:
    ddpm.save_pretrained(tmpdirname)
    print("Model saved")
    ddpm_new = DDPM.from_pretrained(tmpdirname)
    print("Model loaded")
    print(ddpm_new)