expects bits from -1 to 1, outputs image tensor from 0 to 1
(x, bits=BITS)
| 29 | |
| 30 | |
| 31 | def bits_to_decimal(x, bits=BITS): |
| 32 | """expects bits from -1 to 1, outputs image tensor from 0 to 1""" |
| 33 | device = x.device |
| 34 | |
| 35 | x = (x > 0).int() |
| 36 | mask = 2 ** torch.arange(bits - 1, -1, -1, device=device, dtype=torch.int32) |
| 37 | |
| 38 | mask = rearrange(mask, "d -> d 1 1") |
| 39 | x = rearrange(x, "b (c d) h w -> b c d h w", d=8) |
| 40 | dec = reduce(x * mask, "b c d h w -> b c h w", "sum") |
| 41 | return (dec / 255).clamp(0.0, 1.0) |
| 42 | |
| 43 | |
| 44 | # modified scheduler step functions for clamping the predicted x_0 between -bit_scale and +bit_scale |