diff --git a/modules/processing.py b/modules/processing.py index ebbaf7272..095eba54c 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -483,4 +483,5 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: for stat in stats[:20]: frame = stat.traceback[0] shared.log.debug(f' file="{frame.filename}":{frame.lineno} size={stat.size}') + devices.torch_gc(force=True) return processed diff --git a/modules/processing_args.py b/modules/processing_args.py index a716b685e..4ce552825 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -12,7 +12,8 @@ from modules.processing_helpers import resize_hires, fix_prompts, calculate_base from modules.api import helpers -debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None +debug_enabled = os.environ.get('SD_DIFFUSERS_DEBUG', None) +debug_log = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None def task_specific_kwargs(p, model): @@ -93,7 +94,8 @@ def task_specific_kwargs(p, model): 'target_subject_category': getattr(p, 'prompt', '').split()[-1], 'output_type': 'pil', } - debug(f'Diffusers task specific args: {task_args}') + if debug_enabled: + debug_log(f'Diffusers task specific args: {task_args}') return task_args @@ -108,7 +110,8 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2 signature = inspect.signature(type(model).__call__, follow_wrapped=True) possible = list(signature.parameters) - debug(f'Diffusers pipeline possible: {possible}') + if debug_enabled: + debug_log(f'Diffusers pipeline possible: {possible}') prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2) steps = kwargs.get("num_inference_steps", None) or len(getattr(p, 'timesteps', ['1'])) clip_skip = kwargs.pop("clip_skip", 1) @@ -159,6 +162,8 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2 args['negative_prompt'] = negative_prompts[0] else: args['negative_prompt'] = negative_prompts + if prompt_parser_diffusers.embedder is not None and not prompt_parser_diffusers.embedder.scheduled_prompt: # not scheduled so we dont need it anymore + prompt_parser_diffusers.embedder = None if 'clip_skip' in possible and parser == 'fixed': if clip_skip == 1: @@ -248,14 +253,16 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2 if arg in possible: args[arg] = task_kwargs[arg] task_args = getattr(p, 'task_args', {}) - debug(f'Diffusers task args: {task_args}') + if debug_enabled: + debug_log(f'Diffusers task args: {task_args}') for k, v in task_args.items(): if k in possible: args[k] = v else: - debug(f'Diffusers unknown task args: {k}={v}') + debug_log(f'Diffusers unknown task args: {k}={v}') cross_attention_args = getattr(p, 'cross_attention_kwargs', {}) - debug(f'Diffusers cross-attention args: {cross_attention_args}') + if debug_enabled: + debug_log(f'Diffusers cross-attention args: {cross_attention_args}') for k, v in cross_attention_args.items(): if args.get('cross_attention_kwargs', None) is None: args['cross_attention_kwargs'] = {} @@ -273,7 +280,7 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2 # handle implicit controlnet if 'control_image' in possible and 'control_image' not in args and 'image' in args: - debug('Diffusers: set control image') + debug_log('Diffusers: set control image') args['control_image'] = args['image'] sd_hijack_hypertile.hypertile_set(p, hr=len(getattr(p, 'init_images', [])) > 0) @@ -309,5 +316,6 @@ def set_pipeline_args(p, model, prompts: list, negative_prompts: list, prompts_2 if shared.cmd_opts.profile: t1 = time.time() shared.log.debug(f'Profile: pipeline args: {t1-t0:.2f}') - debug(f'Diffusers pipeline args: {args}') + if debug_enabled: + debug_log(f'Diffusers pipeline args: {args}') return args diff --git a/modules/sd_models.py b/modules/sd_models.py index 83bf6f994..37567962c 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -443,23 +443,24 @@ def apply_balanced_offload(sd_model): max_memory = getattr(module, "balanced_offload_max_memory", None) module = accelerate.hooks.remove_hook_from_module(module, recurse=True) try: - if used_gpu > 100 * shared.opts.diffusers_offload_min_gpu_memory: - debug_move(f'Balanced offload: gpu={used_gpu} ram={used_ram} current={module.device} target={devices.cpu} component={module.__class__.__name__}') - module = module.to(devices.cpu, non_blocking=True) - used_gpu, used_ram = devices.torch_gc(fast=True) - else: - debug_move(f'Balanced offload: gpu={used_gpu} ram={used_ram} current={module.device} target={devices.cpu} component={module.__class__.__name__}') - module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name) - module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True) - module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access - if network_layer_name: - module.network_layer_name = network_layer_name - if device_map and max_memory: - module.balanced_offload_device_map = device_map - module.balanced_offload_max_memory = max_memory + do_offload = used_gpu > 100 * shared.opts.diffusers_offload_min_gpu_memory + debug_move(f'Balanced offload: gpu={used_gpu} ram={used_ram} current={module.device} dtype={module.dtype} op={"move" if do_offload else "skip"} component={module.__class__.__name__}') + if do_offload: + module = module.to(devices.cpu) + used_gpu, used_ram = devices.torch_gc(fast=True, force=True) except Exception as e: if 'bitsandbytes' not in str(e): shared.log.error(f'Balanced offload: module={module_name} {e}') + if os.environ.get('SD_MOVE_DEBUG', None): + errors.display(e, f'Balanced offload: module={module_name}') + module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name) + module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True) + module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access + if network_layer_name: + module.network_layer_name = network_layer_name + if device_map and max_memory: + module.balanced_offload_device_map = device_map + module.balanced_offload_max_memory = max_memory apply_balanced_offload_to_module(sd_model) if hasattr(sd_model, "pipe"):