use torch.compile

bghira · Jul 7, 2023 · 94daee4 · 94daee4
1 parent 8451931
commit 94daee4
Show file tree

Hide file tree

Showing 3 changed files with 1,261 additions and 1,312 deletions.
diff --git a/discord_tron_client/classes/image_manipulation/diffusion.py b/discord_tron_client/classes/image_manipulation/diffusion.py
@@ -224,7 +224,7 @@ def get_pipe(
                     logging.error(f"Could not enable CPU offload on the model: {e}")
             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)
+            self.pipelines[model_id].unet = torch.compile(self.pipelines[model_id].unet, mode="reduce-overhead", fullgraph=True)
         else:
             logging.info(f"Keeping existing pipeline. Not creating any new ones.")
         self.last_pipe_type[model_id] = pipe_type