mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-06-09 13:02:08 +08:00
Make loras work with --async-offload (#7824)
This commit is contained in:
parent
0dcc75ca54
commit
ac10a0d69e
17
comfy/ops.py
17
comfy/ops.py
@ -22,6 +22,7 @@ import comfy.model_management
|
|||||||
from comfy.cli_args import args, PerformanceFeature
|
from comfy.cli_args import args, PerformanceFeature
|
||||||
import comfy.float
|
import comfy.float
|
||||||
import comfy.rmsnorm
|
import comfy.rmsnorm
|
||||||
|
import contextlib
|
||||||
|
|
||||||
cast_to = comfy.model_management.cast_to #TODO: remove once no more references
|
cast_to = comfy.model_management.cast_to #TODO: remove once no more references
|
||||||
|
|
||||||
@ -38,20 +39,28 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
|
|||||||
device = input.device
|
device = input.device
|
||||||
|
|
||||||
offload_stream = comfy.model_management.get_offload_stream(device)
|
offload_stream = comfy.model_management.get_offload_stream(device)
|
||||||
|
if offload_stream is not None:
|
||||||
|
wf_context = offload_stream
|
||||||
|
else:
|
||||||
|
wf_context = contextlib.nullcontext()
|
||||||
|
|
||||||
bias = None
|
bias = None
|
||||||
non_blocking = comfy.model_management.device_supports_non_blocking(device)
|
non_blocking = comfy.model_management.device_supports_non_blocking(device)
|
||||||
if s.bias is not None:
|
if s.bias is not None:
|
||||||
has_function = len(s.bias_function) > 0
|
has_function = len(s.bias_function) > 0
|
||||||
bias = comfy.model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=has_function, stream=offload_stream)
|
bias = comfy.model_management.cast_to(s.bias, bias_dtype, device, non_blocking=non_blocking, copy=has_function, stream=offload_stream)
|
||||||
|
|
||||||
if has_function:
|
if has_function:
|
||||||
for f in s.bias_function:
|
with wf_context:
|
||||||
bias = f(bias)
|
for f in s.bias_function:
|
||||||
|
bias = f(bias)
|
||||||
|
|
||||||
has_function = len(s.weight_function) > 0
|
has_function = len(s.weight_function) > 0
|
||||||
weight = comfy.model_management.cast_to(s.weight, dtype, device, non_blocking=non_blocking, copy=has_function, stream=offload_stream)
|
weight = comfy.model_management.cast_to(s.weight, dtype, device, non_blocking=non_blocking, copy=has_function, stream=offload_stream)
|
||||||
if has_function:
|
if has_function:
|
||||||
for f in s.weight_function:
|
with wf_context:
|
||||||
weight = f(weight)
|
for f in s.weight_function:
|
||||||
|
weight = f(weight)
|
||||||
|
|
||||||
comfy.model_management.sync_stream(device, offload_stream)
|
comfy.model_management.sync_stream(device, offload_stream)
|
||||||
return weight, bias
|
return weight, bias
|
||||||
|
Loading…
x
Reference in New Issue
Block a user