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

Patrick von Platen's avatar
improve  
Patrick von Platen committed
5
6
from modeling_ddpm import DDPM

Patrick von Platen's avatar
up  
Patrick von Platen committed
7
8
model_id = sys.argv[1]

Patrick von Platen's avatar
up  
Patrick von Platen committed
9
ddpm = DDPM.from_pretrained(model_id)
Patrick von Platen's avatar
Patrick von Platen committed
10
11
image = ddpm()

Patrick von Platen's avatar
up  
Patrick von Platen committed
12
13
14
15
16
17
18
import PIL.Image
import numpy as np
image_processed = image.cpu().permute(0, 2, 3, 1)
image_processed = (image_processed + 1.0) * 127.5
image_processed = image_processed.numpy().astype(np.uint8)
image_pil = PIL.Image.fromarray(image_processed[0])
image_pil.save("test.png")
Patrick von Platen's avatar
Patrick von Platen committed
19

Patrick von Platen's avatar
up  
Patrick von Platen committed
20
import ipdb; ipdb.set_trace()