From 7700276b5013ff6a4ce9776bfe6b6919326654b0 Mon Sep 17 00:00:00 2001
From: lllyasviel <lyuminzhang@outlook.com>
Date: Wed, 13 Sep 2023 03:43:15 -0700
Subject: [PATCH] move text encoder to GPU (#360)

---
 fooocus_version.py | 2 +-
 modules/patch.py   | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/fooocus_version.py b/fooocus_version.py
index e83c613b..024361b7 100644
--- a/fooocus_version.py
+++ b/fooocus_version.py
@@ -1 +1 @@
-version = '2.0.3'
+version = '2.0.4'
diff --git a/modules/patch.py b/modules/patch.py
index 08474956..09ec6fc2 100644
--- a/modules/patch.py
+++ b/modules/patch.py
@@ -70,7 +70,8 @@ def sdxl_encode_adm_patched(self, **kwargs):
 
 
 def text_encoder_device_patched():
-    return torch.device("cpu")
+    # 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():