| 375 | |
| 376 | |
| 377 | class LogitNormalSampler: |
| 378 | def __init__(self, loc: float = 0.0, scale: float = 1.0): |
| 379 | """ |
| 380 | Logit-Normal sampler from the paper 'Scaling Rectified |
| 381 | Flow Transformers for High-Resolution Image Synthesis' |
| 382 | - Esser et al. (ICML 2024) |
| 383 | """ |
| 384 | self.loc = loc |
| 385 | self.scale = scale |
| 386 | |
| 387 | def __call__(self, n, device='cpu', dtype=torch.float32): |
| 388 | return torch.sigmoid(self.loc + self.scale * torch.randn(n)).to(device).to(dtype) |
| 389 | |
| 390 | |
| 391 | """ Flow Model """ |
nothing calls this directly
no outgoing calls
no test coverage detected