Pin VAE on GPU for async overlap when ``cpu_offload`` would block overlap.
(self, vae_decoder: Any)
| 61 | return decoder |
| 62 | |
| 63 | def prepare_vae(self, vae_decoder: Any) -> None: |
| 64 | """Pin VAE on GPU for async overlap when ``cpu_offload`` would block overlap.""" |
| 65 | self._vae_decoder = vae_decoder |
| 66 | if self.enabled: |
| 67 | logger.info("[AsyncVAEChunkDecoder] async VAE decode enabled") |
| 68 | if not self.enabled or not getattr(vae_decoder, "cpu_offload", False): |
| 69 | return |
| 70 | logger.info("[AsyncVAEChunkDecoder] async mode: VAE cpu_offload was True, forcing False and pinning VAE on GPU") |
| 71 | self._saved_vae_cpu_offload = True |
| 72 | vae_decoder.cpu_offload = False |
| 73 | if hasattr(vae_decoder, "to_cuda"): |
| 74 | vae_decoder.to_cuda() |
| 75 | |
| 76 | def _restore_vae(self) -> None: |
| 77 | vae_decoder = self._vae_decoder |
no test coverage detected