diff --git a/modules/ltx/ltx_process.py b/modules/ltx/ltx_process.py index c684dcee2..aa99ad6b1 100644 --- a/modules/ltx/ltx_process.py +++ b/modules/ltx/ltx_process.py @@ -683,6 +683,8 @@ def run_ltx(task_id, h, w = pixels.shape[-2], pixels.shape[-1] else: w, h = p.width, p.height + del pixels + resolution = f'{w}x{h}' if num_frames > 0 else None summary = timer.process.summary(min_time=0.25, total=False).replace('=', ' ') memory = shared.mem_mon.summary() diff --git a/modules/minimax/minimax_video.py b/modules/minimax/minimax_video.py index aa1a41136..ed63ca194 100644 --- a/modules/minimax/minimax_video.py +++ b/modules/minimax/minimax_video.py @@ -3,7 +3,7 @@ import time from PIL import Image import numpy as np from modules.logger import log -from modules import shared, processing, timer, progress, paths, sd_models, scripts_manager, call_queue, memstats, processing_video +from modules import shared, devices, processing, timer, progress, paths, sd_models, scripts_manager, call_queue, memstats, processing_video from modules.video_models import models_def, video_save, video_utils @@ -130,7 +130,8 @@ def generate(task_id, _ui_state, _processed: processing.Processed = scripts_manager.scripts_video.run(p, *args) processed = processing.process_images(p) - sd_models.offload_ondemand(shared.sd_model, reason='finish') + sd_models.offload_ondemand(shared.sd_model, reason='finish', force=True) # force offload all loaded modules to cpu + devices.torch_gc(force=True) # free gpu memory before saving video # init vars pixels = None @@ -172,6 +173,7 @@ def generate(task_id, _ui_state, metadata={}, ) _n, _c, _t, h, w = pixels.shape + del pixels t1 = time.time() progress.finish_task(task_id) diff --git a/modules/modular_load.py b/modules/modular_load.py index 2a251712f..b8e8b8cc0 100644 --- a/modules/modular_load.py +++ b/modules/modular_load.py @@ -42,7 +42,7 @@ def install_state_hook(pipe): def _pre_transformer_hook(module, args): # pylint: disable=unused-argument new_phase = set_phase('Generate', module) if new_phase: - sd_offload.offload_ondemand(pipe, exclude=['transformer', 'transformer_ref'], reason='generate') + sd_offload.offload_ondemand(pipe, exclude=['transformer', 'transformer_ref'], reason='generate', force=hasattr(pipe, 'sdnext_force_offload')) if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0: shared.state.sampling_steps = pipe.num_timesteps if shared.state.paused: @@ -58,21 +58,21 @@ def install_state_hook(pipe): def _pre_text_encode_hook(module, args): # pylint: disable=unused-argument new_phase = set_phase('Text Encode', module) if new_phase: - sd_offload.offload_ondemand(pipe, exclude=['text_encoder'], reason='text encode') + sd_offload.offload_ondemand(pipe, exclude=['text_encoder'], reason='text encode', force=hasattr(pipe, 'sdnext_force_offload')) if shared.state.interrupted or shared.state.skipped: raise AssertionError('Interrupted...') def _pre_vae_decode_hook(module, args): # pylint: disable=unused-argument new_phase = set_phase('Decode', module) if new_phase: - sd_offload.offload_ondemand(pipe, exclude=['vae', 'audio_vae'], reason='vae decode') + sd_offload.offload_ondemand(pipe, exclude=['vae', 'audio_vae'], reason='vae decode', force=hasattr(pipe, 'sdnext_force_offload')) if shared.state.interrupted or shared.state.skipped: # fires per tile, so tiled decodes abort promptly raise AssertionError('Interrupted...') def _pre_vae_encode_hook(module, args): # pylint: disable=unused-argument new_phase = set_phase('Encode', module) if new_phase: - sd_offload.offload_ondemand(pipe, exclude=['vae', 'audio_vae'], reason='vae encode') + sd_offload.offload_ondemand(pipe, exclude=['vae', 'audio_vae'], reason='vae encode', force=hasattr(pipe, 'sdnext_force_offload')) if shared.state.interrupted or shared.state.skipped: # fires per tile, so tiled encodes abort promptly raise AssertionError('Interrupted...') diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index d7361fc87..ddd4e0171 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -617,6 +617,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing): if hasattr(p, 'dummy'): images = [Image.new(mode='RGB', size=(p.width, p.height))] return images + if 'base' not in p.skip: output = process_base(p) else: diff --git a/modules/processing_prompt.py b/modules/processing_prompt.py index 699e4fc94..cab96ce23 100644 --- a/modules/processing_prompt.py +++ b/modules/processing_prompt.py @@ -116,6 +116,9 @@ def set_prompt(p: StableDiffusionProcessing, prompt_parser_diffusers.embedder = None prompt_attention = 'fixed' + if not hasattr(shared.sd_model, 'orig_encode_prompt'): + log.debug(f'Encode: prompt="{prompts}" negative="{negative_prompts}" embedder={prompt_parser_diffusers.embedder is not None} attention={prompt_attention}') + prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompt_batch(p, prompts, negative_prompts, prompts_2, negative_prompts_2) prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompt_model(cls, prompts, negative_prompts, prompts_2, negative_prompts_2) diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 8bdb6fd48..f04b785e0 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -19,7 +19,6 @@ debug_move = log.trace if debug else lambda *args, **kwargs: None offload_allow_none = ['sd', 'sdxl'] offload_post = ['h1'] offload_hook_instance = None -group_offload_vae_limit = 1.0 # GB; vae-class components above this rest on cpu and onload whole at encode/decode balanced_offload_exclude = ['CogView4Pipeline', 'MeissonicPipeline'] no_split_module_classes = [ "Linear", "Conv1d", "Conv2d", "Conv3d", "ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d", "Embedding", @@ -180,11 +179,10 @@ def set_group_resident(module): module.to(devices.device) -def group_offload_role(module_name: str, module) -> str: - cls = module.__class__.__name__ - if 'vae' in module_name.lower() or cls.startswith(('Autoencoder', 'VQModel', 'AsymmetricAutoencoder', 'ConsistencyDecoder')): - return 'resident' - if module_name.startswith(('text_encoder', 'image_encoder', 'safety_checker')): +def group_offload_role(module_name) -> str: + if any(m in module_name for m in ['vae']): + return 'vae' + if any(m in module_name for m in ['text_encoder', 'image_encoder', 'safety_checker']): return 'aux' return 'main' @@ -208,24 +206,22 @@ class OnDemandHook(accelerate.hooks.ModelHook): if param is not None and not devices.same_device(param.device, devices.device): t0 = time.time() module.to(devices.device, non_blocking=shared.opts.diffusers_offload_nonblocking) - dt = time.time() - t0 - process_timer.add('onload', dt) - log.debug(f'Offload: type=ondemand op=onload module={module.__class__.__name__} nonblocking={shared.opts.diffusers_offload_nonblocking} time={dt:.3f}') + t1 = time.time() + process_timer.add('onload', t1 - t0) + debug_move(f'Offload: type=ondemand op=onload module={module.__class__.__name__} nonblocking={shared.opts.diffusers_offload_nonblocking} time={t1 - t0:.3f}') # working so no need to log return args, kwargs -def set_group_vae(sd_model, module, module_name: str) -> str: - """Placement policy for vae-class components, which never take group hooks. Small - components stay resident; components above group_offload_vae_limit rest on cpu and - onload whole when their decode or encode entry point fires.""" - size_gb, _params = get_module_size(module) - if size_gb < group_offload_vae_limit or not has_entry_bridge(module): - set_group_resident(module) +def apply_group_offload_vae(sd_model, module, module_name: str) -> str: + """Placement policy for vae-class components, which never take group hooks. Components onload whole when their decode or encode entry point fires.""" + if not has_entry_bridge(module): + log.warning(f'Offload: type=group module={module_name} class={module.__class__.__name__} no entry bridge') + set_group_resident(module) # TODO group offload: this will fail as vae will end up on cpu and there will be nothing to pull it back on gpu module.sdnext_ondemand = False # a lingering stamp would let the seams offload a component with no onload hook names = getattr(sd_model, 'sdnext_ondemand_modules', None) or [] if module_name in names: sd_model.sdnext_ondemand_modules = [n for n in names if n != module_name] - return 'resident' + return False if not getattr(module, 'sdnext_ondemand', False) or not hasattr(module, '_hf_hook'): if hasattr(module, '_hf_hook'): module = accelerate.hooks.remove_hook_from_module(module, recurse=True) @@ -237,14 +233,17 @@ def set_group_vae(sd_model, module, module_name: str) -> str: names = getattr(sd_model, 'sdnext_ondemand_modules', None) or [] if module_name not in names: sd_model.sdnext_ondemand_modules = names + [module_name] - return 'ondemand' + return True -def offload_ondemand(sd_model, include=[], exclude=[], reason=''): +def offload_ondemand(sd_model, include=[], exclude=[], reason='', force=False): """Return on-demand components to cpu once their outputs are materialized.""" if sd_model is None: return - names = getattr(sd_model, 'sdnext_ondemand_modules', None) + if force: # if force, all loaded modules are candidates + names = [name for name, component in sd_model.components.items() if isinstance(component, torch.nn.Module)] + else: + names = getattr(sd_model, 'sdnext_ondemand_modules', None) if not names and hasattr(sd_model, 'pipe'): sd_model = sd_model.pipe names = getattr(sd_model, 'sdnext_ondemand_modules', None) @@ -260,7 +259,8 @@ def offload_ondemand(sd_model, include=[], exclude=[], reason=''): module.to(devices.cpu, non_blocking=shared.opts.diffusers_offload_nonblocking) dt = time.time() - t0 process_timer.add('offload', dt) - log.debug(f'Offload: type=ondemand op=offload module={module_name} nonblocking={shared.opts.diffusers_offload_nonblocking} reason="{reason}" time={dt:.3f}') + debug_move(f'Offload: type=ondemand op=offload module={module_name} nonblocking={shared.opts.diffusers_offload_nonblocking} reason="{reason}" time={dt:.3f}') + devices.torch_gc() def report_group_stats(sd_model, module_names): @@ -286,25 +286,20 @@ def apply_modular_group_offload(sd_model): sd_model.sdnext_modular_offload_warned = True log.warning(f'Offload: desired={shared.opts.diffusers_offload_mode} override=group reason="modular pipeline"') applied = [] - for name in ('transformer', 'transformer_ref'): - transformer = getattr(sd_model, name, None) - if transformer is not None and apply_group_offload_component(transformer, name, main=True): - applied.append(name) - text_encoder = getattr(sd_model, 'text_encoder', None) - if text_encoder is not None: - # offload targets the inner model when present: conditioning may call it directly, - # and hooks on the wrapper forward would never fire - if apply_group_offload_component(getattr(text_encoder, 'model', text_encoder), 'text_encoder', main=False): - applied.append('text_encoder') - for name in ('vae', 'audio_vae'): - component = getattr(sd_model, name, None) - if component is not None: - placement = set_group_vae(sd_model, component, name) - applied.append(f'{name}:{placement}') - # has_accelerate stays unset: group hooks are not accelerate hooks, and the modular - # pipeline's own to() skips group-offloaded components when move_model runs - if any(':' not in name for name in applied): - log.info(f'Offload: type=group type={shared.opts.group_offload_type} modules={applied}') + loaded = [name for name, component in sd_model.components.items() if isinstance(component, torch.nn.Module)] + for name in loaded: + module = getattr(sd_model, name, None) + if 'text_encoder' in name: + if apply_group_offload_component(getattr(module, 'model', module), name, main=False): + applied.append(name) + if 'vae' in name: + if apply_group_offload_vae(sd_model, module, name): + applied.append(name) + else: + if apply_group_offload_component(module, name, main=True): + applied.append(name) + # has_accelerate stays unset: group hooks are not accelerate hooks, and the modular pipeline's own to() skips group-offloaded components when move_model runs + log.info(f'Offload: type=group type={shared.opts.group_offload_type} modules={applied}') report_group_stats(sd_model, ('transformer', 'transformer_ref', 'text_encoder', 'vae', 'audio_vae')) @@ -315,9 +310,9 @@ def apply_group_offload(sd_model): if not isinstance(module, torch.nn.Module): continue try: - role = group_offload_role(module_name, module) - if role == 'resident': - if set_group_vae(sd_model, module, module_name) == 'ondemand': + role = group_offload_role(module_name) + if role == 'vae': + if apply_group_offload_vae(sd_model, module, module_name) == 'ondemand': ondemand.append(module_name) else: resident.append(module_name) diff --git a/modules/styles.py b/modules/styles.py index 05837c4cd..6fdbb672d 100644 --- a/modules/styles.py +++ b/modules/styles.py @@ -143,7 +143,8 @@ def apply_file_wildcards(prompt, replaced = None, not_found = None, recursion=0, if '|' in choice: choice = random.choice(choice.split('|')).strip(' []{}\n') prompt = prompt.replace(f"__{wildcard}__", choice, 1) - log.debug(f'Apply wildcard: select="{wildcard}" choice="{choice}" file="{file}" choices={len(lines)}') + if debug_enabled: + log.debug(f'Apply wildcard: select="{wildcard}" choice="{choice}" file="{file}" choices={len(lines)}') replaced.append(wildcard) return prompt, True except Exception as e: diff --git a/modules/video_models/video_run.py b/modules/video_models/video_run.py index 91e683e91..861b743e4 100644 --- a/modules/video_models/video_run.py +++ b/modules/video_models/video_run.py @@ -339,6 +339,7 @@ def run(selected: models_def.Model, *, mp4_interpolate=mp4_interpolate, metadata={}, ) + del pixels return VideoResult(images=processed.images, video_path=video_file, thumb_path=thumb_file, num_frames=num_frames, fps=float(save_fps), has_audio=waveform is not None, still=False, processed=processed) diff --git a/pipelines/model_minimax.py b/pipelines/model_minimax.py index a29bc3907..7bb4ca3ce 100644 --- a/pipelines/model_minimax.py +++ b/pipelines/model_minimax.py @@ -27,6 +27,7 @@ def load_minimax(checkpoint_info, diffusers_load_config = None, workflow: str | if pipe is None: return None pipe.sd_checkpoint_info = checkpoint_info + pipe.sdnext_force_offload = True # very large model so this triggers force-offload on each stage if hasattr(pipe, 'min_duration') and hasattr(pipe, 'fps'): pipe.sdnext_supported_min_frames = int(pipe.min_duration * pipe.fps) # fresh pipes report the true floor; still mode gates per instance