| 9 | |
| 10 | |
| 11 | class SD3RectFlow: |
| 12 | |
| 13 | def __init__( |
| 14 | self, |
| 15 | logit_mean: float = 0.0, |
| 16 | logid_std: float = 1.0, |
| 17 | base_height: int = 256, |
| 18 | base_width: int = 256, |
| 19 | base_frames: int = 1, |
| 20 | base_scale: float = 1, |
| 21 | ) -> None: |
| 22 | self.t_scale = 1000 |
| 23 | self.logit_mean = logit_mean |
| 24 | self.logid_std = logid_std |
| 25 | self.base_height = base_height |
| 26 | self.base_width = base_width |
| 27 | self.base_frames = base_frames |
| 28 | self.base_scale = base_scale |
| 29 | |
| 30 | def sample_t_and_sigma( |
| 31 | self, |
| 32 | batch_size: int, |
| 33 | frames: int, |
| 34 | height: int, |
| 35 | width: int, |
| 36 | sample_type: str = 'logitnorm', |
| 37 | device=None, |
| 38 | ): |
| 39 | if sample_type == 'logitnorm': |
| 40 | t = self.logitnorm_sample_t( |
| 41 | batch_size, self.logit_mean, self.logid_std, device=device) |
| 42 | elif sample_type == 'uniform': |
| 43 | t = torch.rand((batch_size, ), device=device) |
| 44 | sigma = self.sigma_shift(t, frames, height, width) |
| 45 | # NOTE in the training process, the t is also shifted according to resolution |
| 46 | return sigma |
| 47 | |
| 48 | @staticmethod |
| 49 | def logitnorm_sample_t( |
| 50 | batch_size: int, |
| 51 | logit_mean: float = 0.0, |
| 52 | logid_std: float = 1.0, |
| 53 | device=None, |
| 54 | ): |
| 55 | """ |
| 56 | NOTE the returned t is in range [0, 1], before sending into transformer, you need to scale it by num_train_steps = 1000 |
| 57 | """ |
| 58 | t = torch.normal( |
| 59 | mean=logit_mean, |
| 60 | std=logid_std, |
| 61 | size=(batch_size, ), |
| 62 | device=device or 'cpu') |
| 63 | t = F.sigmoid(t) |
| 64 | return t |
| 65 | |
| 66 | def sigma_shift( |
| 67 | self, |
| 68 | t, |
nothing calls this directly
no outgoing calls
no test coverage detected