Foreign-Mirrors/Fooocus
1
0
Fork 0
mirror of https://github.com/lllyasviel/Fooocus.git synced 2026-08-16 13:13:16 +02:00
Code Issues Packages Projects Releases Wiki Activity

move expansion to managed device (#364)

Browse Source 
* 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
This commit is contained in:
lllyasviel
2023-09-13 12:48:27 -07:00
committed by GitHub
parent 53beede21d
commit e32f04da34
4 changed files with 17 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/expansion.py
+12 -1
View File
@@ -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)