memef4rmer/edit_anything

1import numpy as np2import torch3 4 5def torch_samps_to_imgs(imgs, uncenter=True):6    if uncenter:7        imgs = (imgs + 1) / 2  # [-1, 1] -> [0, 1]8    imgs = (imgs * 255).clamp(0, 255)9    imgs = imgs.to(torch.uint8)10    imgs = imgs.permute(0, 2, 3, 1)11    imgs = imgs.cpu().numpy()12    return imgs13 14 15def imgs_to_torch(imgs):16    assert imgs.dtype == np.uint817    assert len(imgs.shape) == 4 and imgs.shape[-1] == 3, "expect (N, H, W, C)"18    _, H, W, _ = imgs.shape19 20    imgs = imgs.transpose(0, 3, 1, 2)21    imgs = (imgs / 255).astype(np.float32)22    imgs = (imgs * 2) - 123    imgs = torch.as_tensor(imgs)24    H, W = [_l - (_l % 32) for _l in (H, W)]25    imgs = torch.nn.functional.interpolate(imgs, (H, W), mode="bilinear")26    return imgs27 28 29def test_encode_decode():30    import imageio31    from run_img_sampling import ScoreAdapter, SD32    from vis import _draw33 34    fname = "~/clean.png"35    raw = imageio.imread(fname)36    raw = imgs_to_torch(raw[np.newaxis, ...])37 38    model: ScoreAdapter = SD().run()39    raw = raw.to(model.device)40    zs = model.encode(raw)41    img = model.decode(zs)42    img = torch_samps_to_imgs(img)43    _draw(44        [imageio.imread(fname), img.squeeze(0)],45    )46 47 48def test():49    test_encode_decode()50 51 52if __name__ == "__main__":53    test()54