| 31 | |
| 32 | |
| 33 | class VOSR(nn.Module): |
| 34 | def __init__( |
| 35 | self, |
| 36 | time_dist = ['lognorm', -0.4, 1.0], |
| 37 | cfg_ratio = 0.10, |
| 38 | cfg_scale = 2.0, |
| 39 | interp_type = 'lin', |
| 40 | u_weight = 1, |
| 41 | a = 1.0, |
| 42 | b = 1.0, |
| 43 | accelerator = None, |
| 44 | t_start: float = 0.0, |
| 45 | t_end: float = 1.0, |
| 46 | args=None, |
| 47 | ): |
| 48 | super().__init__() |
| 49 | self.time_dist = time_dist |
| 50 | self.interp_type = interp_type |
| 51 | self.cfg_ratio = cfg_ratio |
| 52 | self.cfg_scale = cfg_scale |
| 53 | self.args = args |
| 54 | self.u_weight = u_weight |
| 55 | self.device = accelerator.device |
| 56 | self.a = a |
| 57 | self.b = b |
| 58 | self.t_start = t_start |
| 59 | self.t_end = t_end |
| 60 | |
| 61 | |
| 62 | |
| 63 | def interpolate(self, cond, uncond, alpha, interp_type='linear'): |
| 64 | if interp_type == 'sph': |
| 65 | return alpha * cond + (1 - alpha**2)**0.5 * uncond |
| 66 | elif interp_type == 'lin': |
| 67 | return alpha * cond + (1 - alpha) * uncond |
| 68 | |
| 69 | def sample_t_r_v1(self, B, device): |
| 70 | samples = torch.rand(B, 2, device=device) |
| 71 | t_raw, r_raw = samples[:, 0], samples[:, 1] |
| 72 | t = torch.maximum(t_raw, r_raw) |
| 73 | g = _beta_ratio((B,), self.a, self.b, device) |
| 74 | r = t.detach() * (1. - g) |
| 75 | return t, r |
| 76 | |
| 77 | def _prepare_cfg_conditions(self, B, device, lq, z, cond_strength_aelq): |
| 78 | """Prepare lq_weak/lq_noised/lq_mixed and z_weak/z_noised/z_mixed for CFG training. |
| 79 | Fixed: use_aelq=True, use_venc=True, weak_cond_strength_venc=0, cond_strength_venc=1. |
| 80 | """ |
| 81 | cfg_mask = torch.rand(B, device=device) < self.cfg_ratio |
| 82 | cfg_indices = (cfg_mask > 0).nonzero(as_tuple=True)[0] |
| 83 | |
| 84 | weak_cond_strength_aelq = random.uniform(self.args.weak_cond_strength_aelq_list[0], self.args.weak_cond_strength_aelq_list[1]) |
| 85 | lq_weak = self.interpolate(lq, torch.zeros_like(lq), weak_cond_strength_aelq, self.interp_type) |
| 86 | lq_noised = self.interpolate(lq, torch.randn_like(lq), cond_strength_aelq, interp_type='sph') |
| 87 | lq_mixed = lq_noised.clone() |
| 88 | lq_mixed[cfg_indices] = lq_weak[cfg_indices] |
| 89 | |
| 90 | # weak_cond_strength_venc=0 => z_weak is zeros |