move expansion to managed device (#364)

* move expansion to managed device

* move expansion to managed device

* move expansion to managed device

* move expansion to managed device

* move expansion to managed device

* move expansion to managed device

modules/expansion.py

@@ -1,6 +1,9 @@
 import torch
+import comfy.model_management as model_management
+
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, set_seed
 from modules.path import fooocus_expansion_path
+from comfy.sd import ModelPatcher
 
 
 fooocus_magic_split = [
@@ -27,14 +30,22 @@ class FooocusExpansion:
     def __init__(self):
         self.tokenizer = AutoTokenizer.from_pretrained(fooocus_expansion_path)
         self.model = AutoModelForCausalLM.from_pretrained(fooocus_expansion_path)
+
+        load_device = model_management.text_encoder_device()
+        offload_device = model_management.text_encoder_offload_device()
+        self.patcher = ModelPatcher(self.model, load_device=load_device, offload_device=offload_device)
+
         self.pipe = pipeline('text-generation',
                              model=self.model,
                              tokenizer=self.tokenizer,
                              device='cpu',
                              torch_dtype=torch.float32)
-        print('Fooocus Expansion engine loaded.')
+
+        print(f'Fooocus Expansion engine loaded.')
 
     def __call__(self, prompt, seed):
+        model_management.load_model_gpu(self.patcher)
+        self.pipe.device = self.patcher.load_device
         seed = int(seed)
         set_seed(seed)
         origin = safe_str(prompt)

modules/patch.py

@@ -76,5 +76,7 @@ def text_encoder_device_patched():
 
 def patch_all():
     comfy.model_management.text_encoder_device = text_encoder_device_patched
+    print(f'Fooocus Text Processing Pipelines are retargeted to {str(comfy.model_management.text_encoder_device())}')
+
     comfy.k_diffusion.external.DiscreteEpsDDPMDenoiser.forward = patched_discrete_eps_ddpm_denoiser_forward
     comfy.model_base.SDXL.encode_adm = sdxl_encode_adm_patched