From d514bb38ee11477298ba2fb361c7d1479d1b6a82 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 17 Sep 2024 03:49:54 -0400 Subject: [PATCH] Add some option to model_options for the text encoder. load_device, offload_device and the initial_device can now be set. --- comfy/sd.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/comfy/sd.py b/comfy/sd.py index 07310b9d4e2..8c5b058ceae 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -70,14 +70,14 @@ def __init__(self, target=None, embedding_directory=None, no_init=False, tokeniz clip = target.clip tokenizer = target.tokenizer - load_device = model_management.text_encoder_device() - offload_device = model_management.text_encoder_offload_device() + load_device = model_options.get("load_device", model_management.text_encoder_device()) + offload_device = model_options.get("offload_device", model_management.text_encoder_offload_device()) dtype = model_options.get("dtype", None) if dtype is None: dtype = model_management.text_encoder_dtype(load_device) params['dtype'] = dtype - params['device'] = model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype)) + params['device'] = model_options.get("initial_device", model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype))) params['model_options'] = model_options self.cond_stage_model = clip(**(params))