(
self, model: Callable[..., torch.Tensor], scale: float
)
| 190 | ) |
| 191 | |
| 192 | def _uncond_guide_model( |
| 193 | self, model: Callable[..., torch.Tensor], scale: float |
| 194 | ) -> Callable[..., torch.Tensor]: |
| 195 | def model_fn(x_t, ts, **kwargs): |
| 196 | half = x_t[: len(x_t) // 2] |
| 197 | combined = torch.cat([half, half], dim=0) |
| 198 | model_out = model(combined, ts, **kwargs) |
| 199 | eps, rest = model_out[:, :3], model_out[:, 3:] |
| 200 | cond_eps, uncond_eps = torch.chunk(eps, 2, dim=0) |
| 201 | half_eps = uncond_eps + scale * (cond_eps - uncond_eps) |
| 202 | eps = torch.cat([half_eps, half_eps], dim=0) |
| 203 | return torch.cat([eps, rest], dim=1) |
| 204 | |
| 205 | return model_fn |
| 206 | |
| 207 | def split_model_output( |
| 208 | self, |
no outgoing calls
no test coverage detected