From d16cc9a43891fd029183dfb73064052f03f20710 Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Thu, 15 Aug 2024 01:45:50 +0800
Subject: [PATCH] 'update'

---
 mz_fluxext_core.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mz_fluxext_core.py b/mz_fluxext_core.py
index 550d989..5021296 100644
--- a/mz_fluxext_core.py
+++ b/mz_fluxext_core.py
@@ -24,6 +24,8 @@ def Flux1PartialLoad_Patch(args={}):
         model.model.diffusion_model.txt_in.to("cpu")
         model.model.diffusion_model.pe_embedder.to("cpu")
 
+        torch.cuda.empty_cache()
+
     def other_to_cuda():
         model.model.diffusion_model.img_in.to("cuda")
         model.model.diffusion_model.time_in.to("cuda")