(self, vqgan_checkpoint, replicate=False)
| 13 | |
| 14 | class VQGAN: |
| 15 | def __init__(self, vqgan_checkpoint, replicate=False): |
| 16 | assert vqgan_checkpoint != '' |
| 17 | self.replicate = replicate |
| 18 | self.config = VQGANConfig.get_default_config() |
| 19 | self.params = pickle.load(open_file(vqgan_checkpoint, 'rb')) |
| 20 | if replicate: |
| 21 | self.params = jax_utils.replicate(self.params) |
| 22 | else: |
| 23 | self.params = jax.jit(lambda x: x)(self.params) |
| 24 | self.model = VQGANModel(self.config) |
| 25 | |
| 26 | def _wrap_fn(self, fn): |
| 27 | if self.replicate: |
nothing calls this directly
no test coverage detected