Use different torch compile opts

bghira · Jul 1, 2023 · 3b0a4a7 · 3b0a4a7
1 parent a210f2e
commit 3b0a4a7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/discord_tron_client/classes/image_manipulation/diffusion.py b/discord_tron_client/classes/image_manipulation/diffusion.py
@@ -241,7 +241,7 @@ def get_pipe(
                     self.pipelines[model_id].enable_model_cpu_offload()
                 except Exception as e:
                     logging.error(f"Could not enable CPU offload on the model: {e}")
-            self.pipelines[model_id].unet = torch.compile(self.pipelines[model_id].unet)
+            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