Add some option to model_options for the text encoder.

load_device, offload_device and the initial_device can now be set.
thomastraum · Sep 17, 2024 · d514bb3 · d514bb3
1 parent 0849c80
commit d514bb3
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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))