VRAM: Fallback should move to CUDA

Signed-off-by: bghira <[email protected]>
bghira · Jul 8, 2023 · a3a8ba3 · a3a8ba3
1 parent f115230
commit a3a8ba3
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/discord_tron_client/classes/image_manipulation/diffusion.py b/discord_tron_client/classes/image_manipulation/diffusion.py
@@ -233,6 +233,9 @@ def get_pipe(
                     logging.error(
                         f"Could not enable sequential CPU offload on the model: {e}"
                     )
+            else:
+                logging.info(f'Moving pipe to CUDA early, because no offloading is being used.')
+                self.pipelines[model_id].to(self.device)
             torch._dynamo.config.suppress_errors = True
             torch._dynamo.config.log_level = logging.WARNING
             self.pipelines[model_id].unet = torch.compile(self.pipelines[model_id].unet, mode="reduce-overhead", fullgraph=True)