Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import torch | |
| def torch_samps_to_imgs(imgs, uncenter=True): | |
| if uncenter: | |
| imgs = (imgs + 1) / 2 # [-1, 1] -> [0, 1] | |
| imgs = (imgs * 255).clamp(0, 255) | |
| imgs = imgs.to(torch.uint8) | |
| imgs = imgs.permute(0, 2, 3, 1) | |
| imgs = imgs.cpu().numpy() | |
| return imgs | |
| def imgs_to_torch(imgs): | |
| assert imgs.dtype == np.uint8 | |
| assert len(imgs.shape) == 4 and imgs.shape[-1] == 3, "expect (N, H, W, C)" | |
| _, H, W, _ = imgs.shape | |
| imgs = imgs.transpose(0, 3, 1, 2) | |
| imgs = (imgs / 255).astype(np.float32) | |
| imgs = (imgs * 2) - 1 | |
| imgs = torch.as_tensor(imgs) | |
| H, W = [_l - (_l % 32) for _l in (H, W)] | |
| imgs = torch.nn.functional.interpolate(imgs, (H, W), mode="bilinear") | |
| return imgs | |
| def test_encode_decode(): | |
| import imageio | |
| from run_img_sampling import ScoreAdapter, SD | |
| from vis import _draw | |
| fname = "~/clean.png" | |
| raw = imageio.imread(fname) | |
| raw = imgs_to_torch(raw[np.newaxis, ...]) | |
| model: ScoreAdapter = SD().run() | |
| raw = raw.to(model.device) | |
| zs = model.encode(raw) | |
| img = model.decode(zs) | |
| img = torch_samps_to_imgs(img) | |
| _draw( | |
| [imageio.imread(fname), img.squeeze(0)], | |
| ) | |
| def test(): | |
| test_encode_decode() | |
| if __name__ == "__main__": | |
| test() | |