(self, w=512, h=512)
| 261 | |
| 262 | @torch.no_grad() |
| 263 | def sample(self, w=512, h=512): |
| 264 | x = torch.randn((1, self.channels, h//self.spatial_compression, w//self.spatial_compression), device='cuda') |
| 265 | if self.is_video_vae: |
| 266 | x = x.unsqueeze(2) |
| 267 | timesteps = self.scheduler.timesteps |
| 268 | for i, step in enumerate(tqdm(timesteps, desc='Sampling')): |
| 269 | t = step / 1000 |
| 270 | t = t.float().view(1) |
| 271 | inputs = (x, t, *self.conds) |
| 272 | v = self.pipeline_model(inputs).float() |
| 273 | if self.sample_cfg > 1: |
| 274 | inputs_uncond = (x, t, *self.unconds) |
| 275 | v_uncond = self.pipeline_model(inputs_uncond).float() |
| 276 | v = v_uncond + self.sample_cfg*(v - v_uncond) |
| 277 | x = self.scheduler.step(v, step, x, return_dict=False)[0] |
| 278 | vae = self.get_vae() |
| 279 | if isinstance(vae, nn.Module): |
| 280 | vae = vae.to('cuda') |
| 281 | else: |
| 282 | vae.load_model_if_needed() |
| 283 | img = self.vae_decode(x) |
| 284 | if isinstance(vae, nn.Module): |
| 285 | vae = vae.to('cpu') |
| 286 | else: |
| 287 | model_management.unload_all_models() |
| 288 | if img.ndim == 5: |
| 289 | # in case of video VAE |
| 290 | img = img.squeeze(1) |
| 291 | return img |
| 292 | |
| 293 | |
| 294 | class BasePipeline(CommonPipeline): |
nothing calls this directly
no test coverage detected