expects image tensor ranging from 0 to 1, outputs bit tensor ranging from -1 to 1
(x, bits=BITS)
| 13 | |
| 14 | # convert to bit representations and back taken from https://github.com/lucidrains/bit-diffusion/blob/main/bit_diffusion/bit_diffusion.py |
| 15 | def decimal_to_bits(x, bits=BITS): |
| 16 | """expects image tensor ranging from 0 to 1, outputs bit tensor ranging from -1 to 1""" |
| 17 | device = x.device |
| 18 | |
| 19 | x = (x * 255).int().clamp(0, 255) |
| 20 | |
| 21 | mask = 2 ** torch.arange(bits - 1, -1, -1, device=device) |
| 22 | mask = rearrange(mask, "d -> d 1 1") |
| 23 | x = rearrange(x, "b c h w -> b c 1 h w") |
| 24 | |
| 25 | bits = ((x & mask) != 0).float() |
| 26 | bits = rearrange(bits, "b c d h w -> b (c d) h w") |
| 27 | bits = bits * 2 - 1 |
| 28 | return bits |
| 29 | |
| 30 | |
| 31 | def bits_to_decimal(x, bits=BITS): |