From ba58dbc8bcda4f16e9b264ef5715da5256319001 Mon Sep 17 00:00:00 2001
From: lllyasviel <lyuminzhang@outlook.com>
Date: Tue, 3 Oct 2023 17:14:38 -0700
Subject: [PATCH] speed up again (#527)

---
 fooocus_version.py          |  2 +-
 modules/async_worker.py     |  3 +++
 modules/default_pipeline.py | 14 ++++++++++++++
 modules/expansion.py        |  1 -
 modules/patch.py            |  6 ++++++
 update_log.md               |  4 ++++
 6 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/fooocus_version.py b/fooocus_version.py
index 37a051b0..a499f681 100644
--- a/fooocus_version.py
+++ b/fooocus_version.py
@@ -1 +1 @@
-version = '2.0.84'
+version = '2.0.85'
diff --git a/modules/async_worker.py b/modules/async_worker.py
index 75b8c2f1..b9f0d5e6 100644
--- a/modules/async_worker.py
+++ b/modules/async_worker.py
@@ -271,6 +271,7 @@ def worker():
             refiner_model_name=refiner_model_name,
             base_model_name=base_model_name,
             loras=loras)
+        pipeline.prepare_text_encoder(async_call=False)
 
         progressbar(3, 'Processing prompts ...')
 
@@ -401,6 +402,8 @@ def worker():
             print(f'Generating and saving time: {execution_time:.2f} seconds')
 
         outputs.append(['results', results])
+
+        pipeline.prepare_text_encoder(async_call=True)
         return
 
     while True:
diff --git a/modules/default_pipeline.py b/modules/default_pipeline.py
index 17fe3091..78d0f3c9 100644
--- a/modules/default_pipeline.py
+++ b/modules/default_pipeline.py
@@ -196,6 +196,20 @@ refresh_everything(
 expansion = FooocusExpansion()
 
 
+@torch.no_grad()
+@torch.inference_mode()
+def prepare_text_encoder(async_call=True):
+    if async_call:
+        # TODO: make sure that this is always called in an async way so that users cannot feel it.
+        pass
+    assert_model_integrity()
+    comfy.model_management.load_models_gpu([xl_base_patched.clip.patcher, expansion.patcher])
+    return
+
+
+prepare_text_encoder(async_call=True)
+
+
 @torch.no_grad()
 @torch.inference_mode()
 def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback, sampler_name, scheduler_name, latent=None, denoise=1.0, tiled=False, cfg_scale=7.0):
diff --git a/modules/expansion.py b/modules/expansion.py
index a07f0cfa..e3658897 100644
--- a/modules/expansion.py
+++ b/modules/expansion.py
@@ -47,7 +47,6 @@ class FooocusExpansion:
         print(f'Fooocus Expansion engine loaded for {load_device}.')
 
     def __call__(self, prompt, seed):
-        model_management.load_model_gpu(self.patcher)
         seed = int(seed)
         set_seed(seed)
         origin = safe_str(prompt)
diff --git a/modules/patch.py b/modules/patch.py
index 55badb43..7a41c9e3 100644
--- a/modules/patch.py
+++ b/modules/patch.py
@@ -402,7 +402,13 @@ def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=
         return self.out(h)
 
 
+def text_encoder_device_patched():
+    # Fooocus's style system uses text encoder much more times than comfy so this makes things much faster.
+    return comfy.model_management.get_torch_device()
+
+
 def patch_all():
+    comfy.model_management.text_encoder_device = text_encoder_device_patched
     comfy.model_patcher.ModelPatcher.calculate_weight = calculate_weight_patched
     comfy.ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = patched_unet_forward
     comfy.k_diffusion.sampling.sample_dpmpp_fooocus_2m_sde_inpaint_seamless = sample_dpmpp_fooocus_2m_sde_inpaint_seamless
diff --git a/update_log.md b/update_log.md
index 8df10931..535876e5 100644
--- a/update_log.md
+++ b/update_log.md
@@ -1,3 +1,7 @@
+# 2.0.85
+
+* Speed Up Again
+
 # 2.0.80
 
 * Improved the scheduling of ADM guidance and CFG mimicking for better visual quality in high frequency domain and small objects.