(self, filename, dtype=torch.float16, *args, **kwargs)
| 333 | |
| 334 | class TransparentVAEDecoder(torch.nn.Module): |
| 335 | def __init__(self, filename, dtype=torch.float16, *args, **kwargs): |
| 336 | super().__init__(*args, **kwargs) |
| 337 | sd = sf.load_file(filename) |
| 338 | model = UNet1024(in_channels=3, out_channels=4) |
| 339 | model.load_state_dict(sd, strict=True) |
| 340 | model.to(dtype=dtype) |
| 341 | model.eval() |
| 342 | self.model = model |
| 343 | self.dtype = dtype |
| 344 | return |
| 345 | |
| 346 | @torch.no_grad() |
| 347 | def estimate_single_pass(self, pixel, latent): |