Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

offload cleanup

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-04-15 10:06:34 -04:00
parent ecdd37747c
commit ee1a4c607d
5 changed files with 31 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+4
View File
@@ -1,5 +1,9 @@
# Change Log for SD.Next
## Update for 2025-04-15
- **HiDream** optimized offloading, now works in 12GB VRAM / 26GB RAM
## Update for 2025-04-14
- [CFG-Zero](https://github.com/WeichenFan/CFG-Zero-star) new guidance method optimized for flow-matching models
installer.py
+1 -1
View File
@@ -538,7 +538,7 @@ def check_diffusers():
t_start = time.time()
if args.skip_all or args.skip_git or args.experimental:
return
sha = 'a8f5134c113da402a93580ef7a021557e816c98d' # diffusers commit hash
sha = 'b6156aafe998eb57902efd3b8cce9a6fde35c1ea' # diffusers commit hash
pkg = pkg_resources.working_set.by_key.get('diffusers', None)
minor = int(pkg.version.split('.')[1] if pkg is not None else 0)
cur = opts.get('diffusers_version', '') if minor > 0 else ''
modules/model_hidream.py
+7 -3
View File
@@ -24,7 +24,7 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
**quant_args,
)
if shared.opts.diffusers_offload_mode != 'none':
transformer = transformer.to(devices.cpu)
sd_models.move_model(transformer, devices.cpu)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
shared.log.debug(f'Load model: type=HiDream te3="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
@@ -36,7 +36,7 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
**quant_args,
)
if shared.opts.diffusers_offload_mode != 'none':
text_encoder_3 = text_encoder_3.to(devices.cpu)
sd_models.move_model(text_encoder_3, devices.cpu)
load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='LLM', device_map=True)
shared.log.debug(f'Load model: type=HiDream te4="{shared.opts.model_h1_llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
@@ -55,7 +55,7 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
**load_args,
)
if shared.opts.diffusers_offload_mode != 'none':
text_encoder_4 = text_encoder_4.to(devices.cpu)
sd_models.move_model(text_encoder_4, devices.cpu)
load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
shared.log.debug(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
@@ -69,6 +69,10 @@ def load_hidream(checkpoint_info, diffusers_load_config={}):
**load_args,
)
sd_hijack_te.init_hijack(pipe)
del text_encoder_3
del text_encoder_4
del tokenizer_4
del transformer
devices.torch_gc()
return pipe
modules/sd_models.py
+2
View File
@@ -225,6 +225,8 @@ def move_model(model, device=None, force=False):
pass # ignore model move if sequential offload is enabled
elif 'Params4bit' in str(e0) or 'Params8bit' in str(e0):
pass # ignore model move if quantization is enabled
elif 'already been set to the correct devices' in str(e0):
pass # ignore errors on pre-quant models
else:
raise e0
t1 = time.time()
modules/sd_offload.py
+17 -13
View File
@@ -9,8 +9,10 @@ from modules import shared, devices, errors, model_quant
from modules.timer import process as process_timer
debug_move = log.trace if os.environ.get('SD_MOVE_DEBUG', None) is not None else lambda *args, **kwargs: None
should_offload = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4']
debug = os.environ.get('SD_MOVE_DEBUG', None) is not None
debug_move = log.trace if debug else lambda *args, **kwargs: None
offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4']
offload_post = ['h1']
offload_hook_instance = None
balanced_offload_exclude = ['OmniGenPipeline', 'CogView4Pipeline']
@@ -66,7 +68,7 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False):
if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate):
sd_model.has_accelerate = False
if shared.opts.diffusers_offload_mode == "none":
if shared.sd_model_type in should_offload or 'video' in shared.sd_model_type:
if shared.sd_model_type in offload_warn or 'video' in shared.sd_model_type:
shared.log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model')
else:
shared.log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}')
@@ -175,19 +177,20 @@ class OffloadHook(accelerate.hooks.ModelHook):
return args, kwargs
def post_forward(self, module, output):
if getattr(module, "do_offload", False) and shared.opts.te_hijack and module.device != devices.cpu:
if getattr(module, "offload_post", False) and module.device != devices.cpu:
used_gpu, used_ram = devices.torch_gc(fast=True)
perc_gpu = used_gpu / shared.gpu_memory
try:
module_size = self.model_size()
prev_gpu = used_gpu
do_offload = (perc_gpu > shared.opts.diffusers_offload_min_gpu_memory)
if do_offload:
offload_now = perc_gpu > shared.opts.diffusers_offload_min_gpu_memory
if offload_now:
module = module.to(devices.cpu)
used_gpu -= module_size
cls = module.__class__.__name__
quant = getattr(module, "quantization_method", None)
debug_move(f'Offload: type=balanced op={"move post forward" if do_offload else "skip post forward"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}')
if debug:
cls = module.__class__.__name__
quant = getattr(module, "quantization_method", None)
debug_move(f'Offload: type=balanced op={"post" if offload_now else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}')
except Exception as e:
if 'out of memory' in str(e):
devices.torch_gc(fast=True, force=True, reason='oom')
@@ -269,15 +272,16 @@ def apply_balanced_offload(sd_model=None, exclude=[]):
perc_gpu = used_gpu / shared.gpu_memory
try:
prev_gpu = used_gpu
do_offload = (perc_gpu > shared.opts.diffusers_offload_min_gpu_memory) and (module.device != devices.cpu)
if do_offload:
offload_now = (perc_gpu > shared.opts.diffusers_offload_min_gpu_memory) and (module.device != devices.cpu)
if offload_now:
module = module.to(devices.cpu)
used_gpu -= module_size
cls = module.__class__.__name__
quant = getattr(module, "quantization_method", None)
if not cached:
shared.log.debug(f'Model module={module_name} type={cls} dtype={module.dtype} quant={quant} params={offload_hook_instance.param_map[module_name]:.3f} size={offload_hook_instance.offload_map[module_name]:.3f}')
debug_move(f'Offload: type=balanced op={"move" if do_offload else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}')
if debug:
debug_move(f'Offload: type=balanced op={"move" if offload_now else "skip"} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={cls} size={module_size:.3f}')
except Exception as e:
if 'out of memory' in str(e):
devices.torch_gc(fast=True, force=True, reason='oom')
@@ -295,7 +299,7 @@ def apply_balanced_offload(sd_model=None, exclude=[]):
if device_map and max_memory:
module.balanced_offload_device_map = device_map
module.balanced_offload_max_memory = max_memory
module.do_offload = bool("HiDreamImage" in sd_model.__class__.__name__ and module_name.startswith("text_encoder"))
module.offload_post = shared.sd_model_type in [offload_post] and shared.opts.te_hijack and module_name.startswith("text_encoder")
devices.torch_gc(fast=True, force=True, reason='offload')
apply_balanced_offload_to_module(sd_model)