From b291c337a1c6d17c22811cfa3ae405dbaf87c3f4 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 1 Aug 2025 10:45:39 -0400 Subject: [PATCH] refactor internal post loop Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 6 +- extensions-builtin/sdnext-modernui | 2 +- modules/images.py | 7 +- modules/model_quant.py | 2 +- modules/model_te.py | 2 - modules/processing.py | 201 +++++++++++++------------- modules/sd_models.py | 1 + modules/sd_offload.py | 221 ++++++++++++++--------------- modules/shared.py | 198 +++++++++++++------------- pipelines/model_chroma.py | 4 - pipelines/model_cogview.py | 3 - pipelines/model_flux.py | 4 - pipelines/model_sana.py | 7 - 13 files changed, 321 insertions(+), 337 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e6e4babe7..38e2b6e34 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,14 +21,18 @@ - changed **default** values for offloading based on detected gpu memory see [offloading docs](https://vladmandic.github.io/sdnext-docs/Offload/) for details - new feature to specify which modules to offload always or never - in *settings -> models & loading -> offload always/never* + in *settings -> model offloading -> offload always/never* - new `highvram` profile provides significant performance boost on gpus with more than 24gb + - new `offload during pre-forward` option + in *settings -> model offloading* + switches from explicit offloading to implicit offloading on module execution change - **Features** - **Wan** select which stage to run: *first/second/both* with configurable *boundary ration* when running both stages in settings -> model options - prompt parser allow explict `BOS` and `EOS` tokens in prompt - **Nunchaku** support for *FLUX.1-Fill* and *FLUX.1-Depth* models - **Fixes** + - refactor legacy processing loop - fix Wan 2.2-5B I2V workflow - fix OpenVINO - fix inpaint image metadata diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index 783b26763..b679c1e45 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit 783b26763b4f8aca448c5c02e22abdeef513f5ce +Subproject commit b679c1e456a30aefa4d4946e439e53ca8efa7b25 diff --git a/modules/images.py b/modules/images.py index 6009e658c..c8d13f92c 100644 --- a/modules/images.py +++ b/modules/images.py @@ -45,7 +45,7 @@ def sanitize_filename_part(text, replace_spaces=True): def atomically_save_image(): Image.MAX_IMAGE_PIXELS = None # disable check in Pillow and rely on check below to allow large custom image sizes while True: - image, filename, extension, params, exifinfo, filename_txt = save_queue.get() + image, filename, extension, params, exifinfo, filename_txt, is_grid = save_queue.get() shared.state.image_history += 1 if len(exifinfo) > 2: with open(paths.params_path, "w", encoding="utf8") as file: @@ -111,7 +111,8 @@ def atomically_save_image(): shared.log.error(f'Save failed: file="{fn}" format={image_format} args={save_args} {e}') errors.display(e, 'Image save') size = os.path.getsize(fn) if os.path.exists(fn) else 0 - shared.log.info(f'Save: image="{fn}" type={image_format} width={image.width} height={image.height} size={size}') + what = 'grid' if is_grid else 'image' + shared.log.info(f'Save: {what}="{fn}" type={image_format} width={image.width} height={image.height} size={size}') if shared.opts.save_log_fn != '' and len(exifinfo) > 0: fn = os.path.join(paths.data_path, shared.opts.save_log_fn) @@ -206,7 +207,7 @@ def save_image(image, filename, extension = os.path.splitext(params.filename) filename_txt = f"{filename}.txt" if shared.opts.save_txt and len(exifinfo) > 0 else None shared.state.outputs(params.filename) - save_queue.put((params.image, filename, extension, params, exifinfo, filename_txt)) # actual save is executed in a thread that polls data from queue + save_queue.put((params.image, filename, extension, params, exifinfo, filename_txt, grid)) # actual save is executed in a thread that polls data from queue save_queue.join() if not hasattr(params.image, 'already_saved_as'): debug(f'Image marked: "{params.filename}"') diff --git a/modules/model_quant.py b/modules/model_quant.py index 5ce926114..d150cab2e 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -455,7 +455,7 @@ def sdnq_quantize_weights(sd_model): try: t0 = time.time() from modules import shared, devices, sd_models - log.debug(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights} dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32}") + log.debug(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights} dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} pre_forward={shared.opts.diffusers_offload_pre}") global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement sd_model = sd_models.apply_function_to_model(sd_model, sdnq_quantize_model, shared.opts.sdnq_quantize_weights, op="sdnq") diff --git a/modules/model_te.py b/modules/model_te.py index 049effd28..b47d73675 100644 --- a/modules/model_te.py +++ b/modules/model_te.py @@ -40,8 +40,6 @@ def load_t5(name=None, cache_dir=None): if torch.is_floating_point(param) and not is_param_float8_e4m3fn: param = param.to(devices.dtype) set_module_tensor_to_device(t5, param_name, device=0, value=param) - if shared.opts.diffusers_eval: - t5.eval() if t5.dtype != devices.dtype: try: t5 = t5.to(dtype=devices.dtype) diff --git a/modules/processing.py b/modules/processing.py index db138dce0..b160f52e7 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -250,6 +250,100 @@ def process_init(p: StableDiffusionProcessing): p.prompts, _ = extra_networks.parse_prompts(p.prompts) +def process_samples(p: StableDiffusionProcessing, samples): + out_images = [] + out_infotexts = [] + for i, sample in enumerate(samples): + debug(f'Processing result: index={i+1}/{len(samples)}') + p.batch_index = i + info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i) + if isinstance(sample, Image.Image) or (isinstance(sample, list) and isinstance(sample[0], Image.Image)): + image = sample + sample = np.array(sample) + else: + sample = validate_sample(sample) + image = Image.fromarray(sample) + + if p.restore_faces: + p.ops.append('restore') + if not p.do_not_save_samples and shared.opts.save_images_before_detailer: + images.save_image(Image.fromarray(sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-restore") + sample = face_restoration.restore_faces(sample, p) + if sample is not None: + image = Image.fromarray(sample) + + if p.detailer_enabled: + p.ops.append('detailer') + if not p.do_not_save_samples and shared.opts.save_images_before_detailer: + images.save_image(Image.fromarray(sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-detailer") + sample = detailer.detail(sample, p) + if sample is not None: + image = Image.fromarray(sample) + + if p.color_corrections is not None and i < len(p.color_corrections): + p.ops.append('color') + if not p.do_not_save_samples and shared.opts.save_images_before_color_correction: + orig = p.color_corrections + p.color_corrections = None + p.color_corrections = orig + image_without_cc = apply_overlay(image, p.paste_to, i, p.overlay_images) + images.save_image(image_without_cc, path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-color-correct") + image = apply_color_correction(p.color_corrections[i], image) + + if p.scripts is not None and isinstance(p.scripts, scripts_manager.ScriptRunner): + pp = scripts_manager.PostprocessImageArgs(image) + p.scripts.postprocess_image(p, pp) + if pp.image is not None: + image = pp.image + + if shared.opts.mask_apply_overlay: + image = apply_overlay(image, p.paste_to, i, p.overlay_images) + + if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([shared.opts.save_mask, shared.opts.save_mask_composite, shared.opts.return_mask, shared.opts.return_mask_composite]): + image_mask = p.mask_for_overlay.convert('RGB') + image1 = image.convert('RGBA').convert('RGBa') + image2 = Image.new('RGBa', image.size) + mask = images.resize_image(3, p.mask_for_overlay, image.width, image.height).convert('L') + image_mask_composite = Image.composite(image1, image2, mask).convert('RGBA') + if shared.opts.save_mask: + images.save_image(image_mask, p.outpath_samples, "", p.all_seeds[i], p.all_prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask") + if shared.opts.save_mask_composite: + images.save_image(image_mask_composite, p.outpath_samples, "", p.all_seeds[i], p.all_prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask-composite") + if shared.opts.return_mask: + out_infotexts.append(info) + out_images.append(image_mask) + if shared.opts.return_mask_composite: + out_infotexts.append(info) + out_images.append(image_mask_composite) + + if shared.opts.include_mask: + if shared.opts.mask_apply_overlay and p.overlay_images is not None and len(p.overlay_images) > 0: + p.image_mask = create_binary_mask(p.overlay_images[0]) + p.image_mask = ImageOps.invert(p.image_mask) + out_infotexts.append(info) + out_images.append(p.image_mask) + elif getattr(p, 'image_mask', None) is not None and isinstance(p.image_mask, Image.Image): + if getattr(p, 'mask_for_detailer', None) is not None: + out_infotexts.append(info) + out_images.append(p.mask_for_detailer) + else: + out_infotexts.append(info) + out_images.append(p.image_mask) + + if p.selected_scale_tab_after == 1: + p.width_after, p.height_after = int(image.width * p.scale_by_after), int(image.height * p.scale_by_after) + if p.resize_mode_after != 0 and p.resize_name_after != 'None': + image = images.resize_image(p.resize_mode_after, image, p.width_after, p.height_after, p.resize_name_after, context=p.resize_context_after) + + if shared.opts.samples_save and not p.do_not_save_samples and p.outpath_samples is not None: + images.save_image(image, p.outpath_samples, "", p.all_seeds[i], p.all_prompts[i], shared.opts.samples_format, info=info, p=p) # main save image + + image.info["parameters"] = info + out_infotexts.append(info) + out_images.append(image) + return out_images, out_infotexts + + def process_images_inner(p: StableDiffusionProcessing) -> Processed: """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch""" if type(p.prompt) == list: @@ -302,7 +396,9 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: processed = p.scripts.process_images(p) if processed is not None: samples = processed.images - infotexts += processed.infotexts + for script_image, script_infotext in zip(processed.images, processed.infotexts): + output_images.append(script_image) + infotexts.append(script_infotext) if samples is None: from modules.processing_diffusers import process_diffusers samples = process_diffusers(p) @@ -320,92 +416,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: p.scripts.postprocess_batch_list(p, batch_params, batch_number=n) samples = batch_params.images - for i, sample in enumerate(samples): - debug(f'Processing result: index={i+1}/{len(samples)} iteration={n+1}/{p.n_iter}') - p.batch_index = i - if isinstance(sample, Image.Image) or (isinstance(sample, list) and isinstance(sample[0], Image.Image)): - image = sample - sample = np.array(sample) - else: - sample = validate_sample(sample) - image = Image.fromarray(sample) - if p.restore_faces: - p.ops.append('restore') - if not p.do_not_save_samples and shared.opts.save_images_before_detailer: - info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i) - images.save_image(Image.fromarray(sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-restore") - sample = face_restoration.restore_faces(sample, p) - if sample is not None: - image = Image.fromarray(sample) - if p.detailer_enabled: - p.ops.append('detailer') - if not p.do_not_save_samples and shared.opts.save_images_before_detailer: - info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i) - images.save_image(Image.fromarray(sample), path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-detailer") - sample = detailer.detail(sample, p) - if sample is not None: - image = Image.fromarray(sample) - if p.color_corrections is not None and i < len(p.color_corrections): - p.ops.append('color') - if not p.do_not_save_samples and shared.opts.save_images_before_color_correction: - orig = p.color_corrections - p.color_corrections = None - p.color_corrections = orig - image_without_cc = apply_overlay(image, p.paste_to, i, p.overlay_images) - info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i) - images.save_image(image_without_cc, path=p.outpath_samples, basename="", seed=p.seeds[i], prompt=p.prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-color-correct") - image = apply_color_correction(p.color_corrections[i], image) - if p.scripts is not None and isinstance(p.scripts, scripts_manager.ScriptRunner): - pp = scripts_manager.PostprocessImageArgs(image) - p.scripts.postprocess_image(p, pp) - if pp.image is not None: - image = pp.image - if shared.opts.mask_apply_overlay: - image = apply_overlay(image, p.paste_to, i, p.overlay_images) - - info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i, all_negative_prompts=p.negative_prompts) - infotexts.append(info) - if isinstance(image, list): - for img in image: - img.info["parameters"] = info - output_images = image - else: - image.info["parameters"] = info - output_images.append(image) - devices.torch_gc() - del samples - - for i, image in enumerate(output_images): - is_grid = len(output_images) == p.batch_size * p.n_iter + 1 and i == 0 - # resize after - if p.selected_scale_tab_after == 1: - p.width_after, p.height_after = int(image.width * p.scale_by_after), int(image.height * p.scale_by_after) - if p.resize_mode_after != 0 and p.resize_name_after != 'None' and not is_grid: - image = images.resize_image(p.resize_mode_after, image, p.width_after, p.height_after, p.resize_name_after, context=p.resize_context_after) - - # save images - if shared.opts.samples_save and not p.do_not_save_samples and p.outpath_samples is not None: - info = create_infotext(p, p.prompts, p.seeds, p.subseeds, index=i) - if isinstance(image, list): - for img in image: - images.save_image(img, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p) # main save image - else: - images.save_image(image, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p) # main save image - - if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay and any([shared.opts.save_mask, shared.opts.save_mask_composite, shared.opts.return_mask, shared.opts.return_mask_composite]): - image_mask = p.mask_for_overlay.convert('RGB') - image1 = image.convert('RGBA').convert('RGBa') - image2 = Image.new('RGBa', image.size) - mask = images.resize_image(3, p.mask_for_overlay, image.width, image.height).convert('L') - image_mask_composite = Image.composite(image1, image2, mask).convert('RGBA') - if shared.opts.save_mask: - images.save_image(image_mask, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask") - if shared.opts.save_mask_composite: - images.save_image(image_mask_composite, p.outpath_samples, "", p.seeds[i], p.prompts[i], shared.opts.samples_format, info=info, p=p, suffix="-mask-composite") - if shared.opts.return_mask: - output_images.append(image_mask) - if shared.opts.return_mask_composite: - output_images.append(image_mask_composite) + batch_images, batch_infotexts = process_samples(p, samples) + for batch_image, batch_infotext in zip(batch_images, batch_infotexts): + output_images.append(batch_image) + infotexts.append(batch_infotext) if shared.cmd_opts.lowvram: devices.torch_gc(force=True, reason='lowvram') @@ -431,18 +445,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed: output_images.insert(0, grid) index_of_first_image = 1 if shared.opts.grid_save: - images.save_image(grid, p.outpath_grids, "", p.all_seeds[0], p.all_prompts[0], shared.opts.grid_format, info=grid_info, p=p, grid=True, suffix="-grid") # main save grid - - if shared.opts.include_mask: - if shared.opts.mask_apply_overlay and p.overlay_images is not None and len(p.overlay_images) > 0: - p.image_mask = create_binary_mask(p.overlay_images[0]) - p.image_mask = ImageOps.invert(p.image_mask) - output_images.append(p.image_mask) - elif getattr(p, 'image_mask', None) is not None and isinstance(p.image_mask, Image.Image): - if getattr(p, 'mask_for_detailer', None) is not None: - output_images.append(p.mask_for_detailer) - else: - output_images.append(p.image_mask) + images.save_image(grid, p.outpath_grids, "", p.all_seeds[0], p.all_prompts[0], shared.opts.grid_format, info=grid_info, p=p, grid=True) # main save grid processed = Processed( p, diff --git a/modules/sd_models.py b/modules/sd_models.py index 00b1ecb3e..1ba384732 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -146,6 +146,7 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True, except Exception as e: shared.log.error(f'Setting {op}: fused-qkv=True {e}') if shared.opts.diffusers_eval: + shared.log.debug(f'Setting {op}: eval=True') def eval_model(model, op=None, sd_model=None): # pylint: disable=unused-argument if hasattr(model, "requires_grad_"): model.requires_grad_(False) diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 69d0c6364..ff7997307 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -37,11 +37,7 @@ def get_signature(cls): def disable_offload(sd_model): if not getattr(sd_model, 'has_accelerate', False): return - if hasattr(sd_model, "_internal_dict"): - keys = sd_model._internal_dict.keys() # pylint: disable=protected-access - else: - keys = get_signature(sd_model).keys() - for module_name in keys: # pylint: disable=protected-access + for module_name in get_module_names(sd_model): module = getattr(sd_model, module_name, None) if isinstance(module, torch.nn.Module): network_layer_name = getattr(module, "network_layer_name", None) @@ -58,11 +54,10 @@ def set_accelerate(sd_model): def set_accelerate_to_module(model): if hasattr(model, "pipe"): set_accelerate_to_module(model.pipe) - if hasattr(model, "_internal_dict"): - for k in model._internal_dict.keys(): # pylint: disable=protected-access - component = getattr(model, k, None) - if isinstance(component, torch.nn.Module): - component.has_accelerate = True + for module_name in get_module_names(model): + component = getattr(model, module_name, None) + if isinstance(component, torch.nn.Module): + component.has_accelerate = True sd_model.has_accelerate = True set_accelerate_to_module(sd_model) @@ -179,6 +174,8 @@ class OffloadHook(accelerate.hooks.ModelHook): self.cpu = int(shared.cpu_memory * shared.opts.diffusers_offload_max_cpu_memory * 1024*1024*1024) self.offload_map = {} self.param_map = {} + self.last_pre = None + self.last_post = None gpu = f'{(shared.gpu_memory * shared.opts.diffusers_offload_min_gpu_memory):.2f}-{(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory):.2f}:{shared.gpu_memory:.2f}' shared.log.info(f'Offload: type=balanced op=init watermark={self.min_watermark}-{self.max_watermark} gpu={gpu} cpu={shared.cpu_memory:.3f} limit={shared.opts.cuda_mem_fraction:.2f} always={self.offload_always} never={self.offload_never}') self.validate() @@ -205,7 +202,50 @@ class OffloadHook(accelerate.hooks.ModelHook): def init_hook(self, module): return module + def offload_module(self, module, op='unk'): + try: + used_gpu, used_ram = devices.torch_gc(fast=True) + perc_gpu = used_gpu / shared.gpu_memory + prev_gpu = used_gpu + module_size = self.model_size() + module_cls = module.__class__.__name__ + op = f'{op}:skip' + if module_cls in self.offload_never: + op = f'{op}:never' + elif module_cls in self.offload_always: + op = f'{op}:always' + module = module.to(devices.cpu) + used_gpu -= module_size + elif perc_gpu > shared.opts.diffusers_offload_min_gpu_memory: + op = f'{op}:mem' + module = module.to(devices.cpu) + used_gpu -= module_size + if debug: + quant = getattr(module, "quantization_method", None) + debug_move(f'Offload: type=balanced op={op} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={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') + elif 'bitsandbytes' in str(e): + pass + else: + shared.log.error(f'Offload: type=balanced op=apply module={module.__name__} cls={module.__class__ if inspect.isclass(module) else None} {e}') + if os.environ.get('SD_MOVE_DEBUG', None): + errors.display(e, f'Offload: type=balanced op=apply module={module.__name__}') + def pre_forward(self, module, *args, **kwargs): + print('HERE', id(module), module.__class__.__name__) + if self.last_pre != module.__class__.__name__: # offload every other module first time when new module starts pre-forward + self.last_pre = module.__class__.__name__ + if shared.opts.diffusers_offload_pre: + debug_move(f'Offload: type=balanced op=pre module={self.last_pre}') + for pipe in get_pipe_variants(): + for module_name in get_module_names(pipe): + module_instance = getattr(pipe, module_name, None) + module_cls = module_instance.__class__.__name__ + if (module_cls != module.__class__.__name__) and (module_cls not in self.offload_never) and (not devices.same_device(module.device, devices.cpu)): + self.offload_module(module_instance, op='pre') + if not devices.same_device(module.device, devices.device): device_index = torch.device(devices.device).index if device_index is None: @@ -228,42 +268,62 @@ class OffloadHook(accelerate.hooks.ModelHook): return args, kwargs def post_forward(self, module, output): - 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() - module_cls = module.__class__.__name__ - prev_gpu = used_gpu - op = 'post:skip' - if module_cls in self.offload_never: - op = 'post:never' - elif module_cls in self.offload_always: - op = 'post:always' - module = module.to(devices.cpu) - used_gpu -= module_size - elif perc_gpu > shared.opts.diffusers_offload_min_gpu_memory: - op = 'post:mem' - module = module.to(devices.cpu) - used_gpu -= module_size - if debug: - quant = getattr(module, "quantization_method", None) - debug_move(f'Offload: type=balanced op={op} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={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') - elif 'bitsandbytes' in str(e): - pass - else: - shared.log.error(f'Offload: type=balanced op=apply module={module.__name__} cls={module.__class__ if inspect.isclass(module) else None} {e}') - if os.environ.get('SD_MOVE_DEBUG', None): - errors.display(e, f'Offload: type=balanced op=apply module={module.__name__}') + if self.last_post != module.__class__.__name__: + self.last_post = module.__class__.__name__ + if getattr(module, "offload_post", False) and (module.device != devices.cpu): + self.offload_module(module, op='post') return output def detach_hook(self, module): return module +def get_pipe_variants(pipe=None): + if pipe is None: + pipe = shared.sd_model + variants = [pipe] + if hasattr(pipe, "pipe"): + variants.append(pipe.pipe) + if hasattr(pipe, "prior_pipe"): + variants.append(pipe.prior_pipe) + if hasattr(pipe, "decoder_pipe"): + variants.append(pipe.decoder_pipe) + return variants + + +def get_module_names(pipe=None, exclude=[]): + if pipe is None: + pipe = shared.sd_model + if hasattr(pipe, "_internal_dict"): + modules_names = pipe._internal_dict.keys() # pylint: disable=protected-access + else: + modules_names = get_signature(pipe).keys() + modules_names = [m for m in modules_names if m not in exclude and not m.startswith('_')] + modules_names = [m for m in modules_names if isinstance(getattr(pipe, m, None), torch.nn.Module)] + return modules_names + + +def get_module_sizes(pipe=None, exclude=[]): + modules = {} + for module_name in get_module_names(pipe, exclude): + module_size = offload_hook_instance.offload_map.get(module_name, None) + if module_size is None: + module = getattr(pipe, module_name, None) + if not isinstance(module, torch.nn.Module): + continue + try: + module_size = sum(p.numel() * p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 + param_num = sum(p.numel() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 + except Exception as e: + shared.log.error(f'Offload: type=balanced op=calc module={module_name} {e}') + module_size = 0 + offload_hook_instance.offload_map[module_name] = module_size + offload_hook_instance.param_map[module_name] = param_num + modules[module_name] = module_size + modules = sorted(modules.items(), key=lambda x: x[1], reverse=True) + return modules + + def apply_balanced_offload(sd_model=None, exclude=[]): global offload_hook_instance # pylint: disable=global-statement if shared.opts.diffusers_offload_mode != "balanced": @@ -285,43 +345,12 @@ def apply_balanced_offload(sd_model=None, exclude=[]): cached = False offload_hook_instance = OffloadHook(checkpoint_name) - def get_pipe_modules(pipe): - if hasattr(pipe, "_internal_dict"): - modules_names = pipe._internal_dict.keys() # pylint: disable=protected-access - else: - modules_names = get_signature(pipe).keys() - modules_names = [m for m in modules_names if m not in exclude and not m.startswith('_')] - modules = {} - for module_name in modules_names: - module_size = offload_hook_instance.offload_map.get(module_name, None) - if module_size is None: - module = getattr(pipe, module_name, None) - if not isinstance(module, torch.nn.Module): - continue - try: - module_size = sum(p.numel() * p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 - param_num = sum(p.numel() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 - except Exception as e: - shared.log.error(f'Offload: type=balanced op=calc module={module_name} {e}') - module_size = 0 - offload_hook_instance.offload_map[module_name] = module_size - offload_hook_instance.param_map[module_name] = param_num - modules[module_name] = module_size - modules = sorted(modules.items(), key=lambda x: x[1], reverse=True) - return modules + if cached and shared.opts.diffusers_offload_pre: + debug_move(f'Offload: type=balanced op=apply skip') + return sd_model def apply_balanced_offload_to_module(pipe): - # shared.log.trace(f'Offload: type=balanced op=apply pipe={pipe.__class__.__name__}') - used_gpu, used_ram = devices.torch_gc(fast=True) - if hasattr(pipe, "_internal_dict"): - keys = pipe._internal_dict.keys() # pylint: disable=protected-access - else: - keys = get_signature(pipe).keys() - keys = [k for k in keys if k not in exclude and not k.startswith('_')] - offload_always = [m.strip() for m in re.split(';|,| ', shared.opts.diffusers_offload_always) if len(m.strip()) > 2] - offload_never = [m.strip() for m in re.split(';|,| ', shared.opts.diffusers_offload_never) if len(m.strip()) > 2] - for module_name, module_size in get_pipe_modules(pipe): # pylint: disable=protected-access - # shared.log.trace(f'Offload: type=balanced op=apply pipe={pipe.__class__.__name__} module={module_name} size={module_size:.3f}') + for module_name, _module_size in get_module_sizes(pipe, exclude): module = getattr(pipe, module_name, None) if module is None: continue @@ -332,38 +361,7 @@ def apply_balanced_offload(sd_model=None, exclude=[]): module = accelerate.hooks.remove_hook_from_module(module, recurse=True) except Exception as e: shared.log.warning(f'Offload remove hook: module={module_name} {e}') - perc_gpu = used_gpu / shared.gpu_memory - try: - prev_gpu = used_gpu - module_cls = module.__class__.__name__ - op = 'apply:skip' - if module_cls in offload_never: - op = 'apply:never' - elif module_cls in offload_always: - op = 'apply:always' - module = module.to(devices.cpu) - used_gpu -= module_size - elif perc_gpu > shared.opts.diffusers_offload_min_gpu_memory: - op = 'apply:mem' - module = module.to(devices.cpu) - used_gpu -= module_size - if debug: - quant = getattr(module, "quantization_method", None) - debug_move(f'Offload: type=balanced op={op} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={module_cls} size={module_size:.3f}') - quant = getattr(module, "quantization_method", None) - if not cached: - shared.log.debug(f'Model module={module_name} type={module_cls} dtype={module.dtype} quant={quant} params={offload_hook_instance.param_map[module_name]:.3f} size={offload_hook_instance.offload_map[module_name]:.3f}') - if debug: - debug_move(f'Offload: type=balanced op={op} gpu={prev_gpu:.3f}:{used_gpu:.3f} perc={perc_gpu:.2f} ram={used_ram:.3f} current={module.device} dtype={module.dtype} quant={quant} module={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') - elif 'bitsandbytes' in str(e): - pass - else: - shared.log.error(f'Offload: type=balanced op=apply module={module_name} {e}') - if os.environ.get('SD_MOVE_DEBUG', None): - errors.display(e, f'Offload: type=balanced op=apply module={module_name}') + offload_hook_instance.offload_module(module, op='apply') module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name) try: module = accelerate.hooks.add_hook_to_module(module, offload_hook_instance, append=True) @@ -378,13 +376,8 @@ def apply_balanced_offload(sd_model=None, exclude=[]): 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) - if hasattr(sd_model, "pipe"): - apply_balanced_offload_to_module(sd_model.pipe) - if hasattr(sd_model, "prior_pipe"): - apply_balanced_offload_to_module(sd_model.prior_pipe) - if hasattr(sd_model, "decoder_pipe"): - apply_balanced_offload_to_module(sd_model.decoder_pipe) + for pipe in get_pipe_variants(sd_model): + apply_balanced_offload_to_module(pipe) if shared.opts.layerwise_quantization or (hasattr(sd_model, "transformer") and getattr(sd_model.transformer, 'quantization_method', None) == 'LayerWise'): model_quant.apply_layerwise(sd_model, quiet=True) # need to reapply since hooks were removed/readded set_accelerate(sd_model) diff --git a/modules/shared.py b/modules/shared.py index 61be9d5d0..2036b2ffd 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -132,7 +132,7 @@ def list_samplers(): startup_offload_mode, startup_offload_min_gpu, startup_offload_max_gpu, startup_cross_attention, startup_sdp_options, startup_offload_always, startup_offload_never = get_default_modes(cmd_opts=cmd_opts, mem_stat=mem_stat) -options_templates.update(options_section(('sd', "Models & Loading"), { +options_templates.update(options_section(('sd', "Model Loading"), { "sd_backend": OptionInfo('diffusers', "Execution backend", gr.Radio, {"choices": ['diffusers', 'original'], "visible": False }), "diffusers_pipeline": OptionInfo('Autodetect', 'Model pipeline', gr.Dropdown, lambda: {"choices": list(shared_items.get_pipelines())}), "sd_model_checkpoint": OptionInfo(default_checkpoint, "Base model", DropdownEditable, lambda: {"choices": list_checkpoint_titles()}, refresh=refresh_checkpoints), @@ -140,20 +140,12 @@ options_templates.update(options_section(('sd', "Models & Loading"), { "sd_unet": OptionInfo("Default", "UNET model", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list), "latent_history": OptionInfo(16, "Latent history size", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}), - "offload_sep": OptionInfo("

Model Offloading

", "", gr.HTML), - "diffusers_offload_mode": OptionInfo(startup_offload_mode, "Model offload mode", gr.Radio, {"choices": ['none', 'balanced', 'group', 'model', 'sequential']}), - "diffusers_offload_min_gpu_memory": OptionInfo(startup_offload_min_gpu, "Balanced offload GPU low watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }), - "diffusers_offload_max_gpu_memory": OptionInfo(startup_offload_max_gpu, "Balanced offload GPU high watermark", gr.Slider, {"minimum": 0.1, "maximum": 1, "step": 0.01 }), - "diffusers_offload_max_cpu_memory": OptionInfo(0.90, "Balanced offload CPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": False }), - "diffusers_offload_always": OptionInfo(startup_offload_always, "Modules to always offload"), - "diffusers_offload_never": OptionInfo(startup_offload_never, "Modules to never offload"), - "advanced_sep": OptionInfo("

Advanced Options

", "", gr.HTML), "sd_checkpoint_autoload": OptionInfo(True, "Model auto-load on start"), "sd_checkpoint_autodownload": OptionInfo(True, "Model auto-download on demand"), "stream_load": OptionInfo(False, "Model load using streams", gr.Checkbox), - "diffusers_eval": OptionInfo(True, "Force model eval", gr.Checkbox, {"visible": False }), "diffusers_to_gpu": OptionInfo(False, "Model load model direct to GPU"), + "diffusers_eval": OptionInfo(True, "Force model eval", gr.Checkbox, {"visible": True }), "device_map": OptionInfo('default', "Model load device map", gr.Radio, {"choices": ['default', 'gpu', 'cpu'] }), "disable_accelerate": OptionInfo(False, "Disable accelerate", gr.Checkbox, {"visible": False }), "sd_checkpoint_cache": OptionInfo(0, "Cached models", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1, "visible": False }), @@ -169,6 +161,71 @@ options_templates.update(options_section(('model_options', "Models Options"), { "model_wan_boundary": OptionInfo(0.85, "Stage boundary ratio", gr.Slider, {"minimum": 0, "maximum": 1.0, "step": 0.05 }), })) +options_templates.update(options_section(('offload', "Model Offloading"), { + "offload_sep": OptionInfo("

Model Offloading

", "", gr.HTML), + "diffusers_offload_mode": OptionInfo(startup_offload_mode, "Model offload mode", gr.Radio, {"choices": ['none', 'balanced', 'group', 'model', 'sequential']}), + "diffusers_offload_pre": OptionInfo(False, "Offload during pre-forward"), + "diffusers_offload_min_gpu_memory": OptionInfo(startup_offload_min_gpu, "Balanced offload GPU low watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }), + "diffusers_offload_max_gpu_memory": OptionInfo(startup_offload_max_gpu, "Balanced offload GPU high watermark", gr.Slider, {"minimum": 0.1, "maximum": 1, "step": 0.01 }), + "diffusers_offload_max_cpu_memory": OptionInfo(0.90, "Balanced offload CPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": False }), + "diffusers_offload_always": OptionInfo(startup_offload_always, "Modules to always offload"), + "diffusers_offload_never": OptionInfo(startup_offload_never, "Modules to never offload"), +})) + +options_templates.update(options_section(("quantization", "Model Quantization"), { + "sdnq_quantize_sep": OptionInfo("

SDNQ: SD.Next Quantization

", "", gr.HTML), + "sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "Control", "VAE"]}), + "sdnq_quantize_mode": OptionInfo("auto", "Quantization mode", gr.Dropdown, {"choices": ["auto", "pre", "post"]}), + "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": sdnq_quant_modes}), + "sdnq_quantize_weights_mode_te": OptionInfo("Same as model", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ['Same as model'] + sdnq_quant_modes}), + "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1}), + "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize convolutional layers", gr.Checkbox), + "sdnq_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox), + "sdnq_use_quantized_matmul": OptionInfo(False, "Use quantized MatMul", gr.Checkbox), + "sdnq_use_quantized_matmul_conv": OptionInfo(False, "Use quantized MatMul with conv", gr.Checkbox), + "sdnq_quantize_with_gpu": OptionInfo(True, "Quantize using GPU", gr.Checkbox), + "sdnq_dequantize_fp32": OptionInfo(False, "Dequantize using full precision", gr.Checkbox), + "sdnq_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox), + + "bnb_quantization_sep": OptionInfo("

BitsAndBytes

", "", gr.HTML), + "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "VAE"]}), + "bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ["nf4", "fp8", "fp4"]}), + "bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Dropdown, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"]}), + + "quanto_quantization_sep": OptionInfo("

Optimum Quanto

", "", gr.HTML), + "quanto_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM"]}), + "quanto_quantization_type": OptionInfo("int8", "Quantization weights type", gr.Dropdown, {"choices": ["float8", "int8", "int4", "int2"]}), + + "optimum_quanto_sep": OptionInfo("

Optimum Quanto: post-load

", "", gr.HTML), + "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "Control", "VAE"]}), + "optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ["qint8", "qfloat8_e4m3fn", "qfloat8_e5m2", "qint4", "qint2"]}), + "optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ["none", "qint8", "qfloat8_e4m3fn", "qfloat8_e5m2"]}), + "optimum_quanto_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox), + + "torchao_sep": OptionInfo("

TorchAO

", "", gr.HTML), + "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "Control", "VAE"]}), + "torchao_quantization_mode": OptionInfo("auto", "Quantization mode", gr.Dropdown, {"choices": ["auto", "pre", "post"]}), + "torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ["int4_weight_only", "int8_dynamic_activation_int4_weight", "int8_weight_only", "int8_dynamic_activation_int8_weight", "float8_weight_only", "float8_dynamic_activation_float8_weight", "float8_static_activation_float8_weight"]}), + + "layerwise_quantization_sep": OptionInfo("

Layerwise Casting

", "", gr.HTML), + "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "TE"]}), + "layerwise_quantization_storage": OptionInfo("float8_e4m3fn", "Layerwise casting storage", gr.Dropdown, {"choices": ["float8_e4m3fn", "float8_e5m2"]}), + "layerwise_quantization_nonblocking": OptionInfo(False, "Layerwise non-blocking operations", gr.Checkbox), + + "nunchaku_sep": OptionInfo("

Nunchaku Engine

", "", gr.HTML), + "nunchaku_quantization": OptionInfo([], "SVDQuant enabled", gr.CheckboxGroup, {"choices": ["Model", "TE"]}), + "nunchaku_attention": OptionInfo(False, "Nunchaku attention", gr.Checkbox), + "nunchaku_offload": OptionInfo(False, "Nunchaku offloading", gr.Checkbox), + + "nncf_compress_sep": OptionInfo("

NNCF: Neural Network Compression Framework

", "", gr.HTML, {"visible": cmd_opts.use_openvino}), + "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE"], "visible": cmd_opts.use_openvino}), + "nncf_compress_weights_mode": OptionInfo("INT8_SYM", "Quantization type", gr.Dropdown, {"choices": ["INT8", "INT4_ASYM", "INT8_SYM", "INT4_SYM", "NF4"], "visible": cmd_opts.use_openvino}), + "nncf_compress_weights_raito": OptionInfo(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}), + "nncf_compress_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": cmd_opts.use_openvino}), + "nncf_quantize": OptionInfo([], "Static Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE"], "visible": cmd_opts.use_openvino}), + "nncf_quantize_mode": OptionInfo("INT8", "OpenVINO activations mode", gr.Dropdown, {"choices": ["INT8", "FP8_E4M3", "FP8_E5M2"], "visible": cmd_opts.use_openvino}), +})) + options_templates.update(options_section(('vae_encoder', "Variational Auto Encoder"), { "sd_vae": OptionInfo("Automatic", "VAE model", gr.Dropdown, lambda: {"choices": shared_items.sd_vae_items()}, refresh=shared_items.refresh_vae_list), "diffusers_vae_upcast": OptionInfo("default", "VAE upcasting", gr.Radio, {"choices": ['default', 'true', 'false']}), @@ -254,60 +311,6 @@ options_templates.update(options_section(('backends', "Backend Settings"), { "directml_catch_nan": OptionInfo(False, "DirectML retry ops for NaN", gr.Checkbox, {"visible": devices.backend == "directml"}), })) -options_templates.update(options_section(("quantization", "Quantization Settings"), { - "sdnq_quantize_sep": OptionInfo("

SDNQ: SD.Next Quantization

", "", gr.HTML), - "sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "Control", "VAE"]}), - "sdnq_quantize_mode": OptionInfo("auto", "Quantization mode", gr.Dropdown, {"choices": ["auto", "pre", "post"]}), - "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": sdnq_quant_modes}), - "sdnq_quantize_weights_mode_te": OptionInfo("Same as model", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ['Same as model'] + sdnq_quant_modes}), - "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1}), - "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize convolutional layers", gr.Checkbox), - "sdnq_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox), - "sdnq_use_quantized_matmul": OptionInfo(False, "Use quantized MatMul", gr.Checkbox), - "sdnq_use_quantized_matmul_conv": OptionInfo(False, "Use quantized MatMul with conv", gr.Checkbox), - "sdnq_quantize_with_gpu": OptionInfo(True, "Quantize using GPU", gr.Checkbox), - "sdnq_dequantize_fp32": OptionInfo(False, "Dequantize using full precision", gr.Checkbox), - "sdnq_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox), - - "bnb_quantization_sep": OptionInfo("

BitsAndBytes

", "", gr.HTML), - "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "VAE"]}), - "bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ["nf4", "fp8", "fp4"]}), - "bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Dropdown, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"]}), - - "quanto_quantization_sep": OptionInfo("

Optimum Quanto

", "", gr.HTML), - "quanto_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM"]}), - "quanto_quantization_type": OptionInfo("int8", "Quantization weights type", gr.Dropdown, {"choices": ["float8", "int8", "int4", "int2"]}), - - "optimum_quanto_sep": OptionInfo("

Optimum Quanto: post-load

", "", gr.HTML), - "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "Control", "VAE"]}), - "optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ["qint8", "qfloat8_e4m3fn", "qfloat8_e5m2", "qint4", "qint2"]}), - "optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ["none", "qint8", "qfloat8_e4m3fn", "qfloat8_e5m2"]}), - "optimum_quanto_shuffle_weights": OptionInfo(False, "Shuffle weights in post mode", gr.Checkbox), - - "torchao_sep": OptionInfo("

TorchAO

", "", gr.HTML), - "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "LLM", "Control", "VAE"]}), - "torchao_quantization_mode": OptionInfo("auto", "Quantization mode", gr.Dropdown, {"choices": ["auto", "pre", "post"]}), - "torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ["int4_weight_only", "int8_dynamic_activation_int4_weight", "int8_weight_only", "int8_dynamic_activation_int8_weight", "float8_weight_only", "float8_dynamic_activation_float8_weight", "float8_static_activation_float8_weight"]}), - - "layerwise_quantization_sep": OptionInfo("

Layerwise Casting

", "", gr.HTML), - "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "TE"]}), - "layerwise_quantization_storage": OptionInfo("float8_e4m3fn", "Layerwise casting storage", gr.Dropdown, {"choices": ["float8_e4m3fn", "float8_e5m2"]}), - "layerwise_quantization_nonblocking": OptionInfo(False, "Layerwise non-blocking operations", gr.Checkbox), - - "nunchaku_sep": OptionInfo("

Nunchaku Engine

", "", gr.HTML), - "nunchaku_quantization": OptionInfo([], "SVDQuant enabled", gr.CheckboxGroup, {"choices": ["Model", "TE"]}), - "nunchaku_attention": OptionInfo(False, "Nunchaku attention", gr.Checkbox), - "nunchaku_offload": OptionInfo(False, "Nunchaku offloading", gr.Checkbox), - - "nncf_compress_sep": OptionInfo("

NNCF: Neural Network Compression Framework

", "", gr.HTML, {"visible": cmd_opts.use_openvino}), - "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE"], "visible": cmd_opts.use_openvino}), - "nncf_compress_weights_mode": OptionInfo("INT8_SYM", "Quantization type", gr.Dropdown, {"choices": ["INT8", "INT4_ASYM", "INT8_SYM", "INT4_SYM", "NF4"], "visible": cmd_opts.use_openvino}), - "nncf_compress_weights_raito": OptionInfo(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}), - "nncf_compress_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": cmd_opts.use_openvino}), - "nncf_quantize": OptionInfo([], "Static Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "TE", "VAE"], "visible": cmd_opts.use_openvino}), - "nncf_quantize_mode": OptionInfo("INT8", "OpenVINO activations mode", gr.Dropdown, {"choices": ["INT8", "FP8_E4M3", "FP8_E5M2"], "visible": cmd_opts.use_openvino}), -})) - options_templates.update(options_section(('advanced', "Pipeline Modifiers"), { "clip_skip_sep": OptionInfo("

CLiP Skip

", "", gr.HTML), "clip_skip_enabled": OptionInfo(False, "CLiP skip enabled"), @@ -433,7 +436,7 @@ options_templates.update(options_section(('system-paths', "System Paths"), { options_templates.update(options_section(('saving-images', "Image Options"), { "samples_save": OptionInfo(True, "Save all generated images"), - "keep_incomplete": OptionInfo(False, "Keep incomplete images"), + "keep_incomplete": OptionInfo(True, "Keep incomplete images"), "samples_format": OptionInfo('jpg', 'File format', gr.Dropdown, {"choices": ["jpg", "png", "webp", "tiff", "jp2", "jxl"]}), "jpeg_quality": OptionInfo(90, "Image quality", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}), "img_max_size_mp": OptionInfo(1000, "Maximum image size (MP)", gr.Slider, {"minimum": 100, "maximum": 2000, "step": 1}), @@ -544,31 +547,6 @@ options_templates.update(options_section(('live-preview', "Live Previews"), { "notification_audio_path": OptionInfo("html/notification.mp3","Path to notification sound", component_args=hide_dirs, folder=True), })) -options_templates.update(options_section(('sampler-params', "Sampler Settings"), { - "show_samplers": OptionInfo([], "Show samplers in user interface", gr.CheckboxGroup, lambda: {"choices": [x.name for x in list_samplers()]}), - 'eta_noise_seed_delta': OptionInfo(0, "Noise seed delta (eta)", gr.Number, {"precision": 0}), - "scheduler_eta": OptionInfo(1.0, "Noise multiplier (eta)", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), - "schedulers_solver_order": OptionInfo(0, "Solver order (where", gr.Slider, {"minimum": 0, "maximum": 5, "step": 1, "visible": False}), - "schedulers_use_loworder": OptionInfo(True, "Use simplified solvers in final steps", gr.Checkbox, {"visible": False}), - "schedulers_prediction_type": OptionInfo("default", "Override model prediction type", gr.Radio, {"choices": ['default', 'epsilon', 'sample', 'v_prediction'], "visible": False}), - "schedulers_sigma": OptionInfo("default", "Sigma algorithm", gr.Radio, {"choices": ['default', 'karras', 'exponential', 'polyexponential'], "visible": False}), - "schedulers_beta_schedule": OptionInfo("default", "Beta schedule", gr.Dropdown, {"choices": ['default', 'linear', 'scaled_linear', 'squaredcos_cap_v2', 'sigmoid'], "visible": False}), - "schedulers_use_thresholding": OptionInfo(False, "Use dynamic thresholding", gr.Checkbox, {"visible": False}), - "schedulers_timestep_spacing": OptionInfo("default", "Timestep spacing", gr.Dropdown, {"choices": ['default', 'linspace', 'leading', 'trailing'], "visible": False}), - 'schedulers_timesteps': OptionInfo('', "Timesteps", gr.Textbox, {"visible": False}), - "schedulers_rescale_betas": OptionInfo(False, "Rescale betas with zero terminal SNR", gr.Checkbox, {"visible": False}), - 'schedulers_beta_start': OptionInfo(0, "Beta start", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.00001}), - 'schedulers_beta_end': OptionInfo(0, "Beta end", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.00001}), - 'schedulers_timesteps_range': OptionInfo(1000, "Timesteps range", gr.Slider, {"minimum": 250, "maximum": 4000, "step": 1}), - 'schedulers_shift': OptionInfo(3, "Sampler shift", gr.Slider, {"minimum": 0.1, "maximum": 10, "step": 0.1, "visible": False}), - 'schedulers_dynamic_shift': OptionInfo(False, "Sampler dynamic shift", gr.Checkbox, {"visible": False}), - 'schedulers_sigma_adjust': OptionInfo(1.0, "Sigma adjust", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01, "visible": False}), - 'schedulers_sigma_adjust_min': OptionInfo(0.2, "Sigma adjust start", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}), - 'schedulers_sigma_adjust_max': OptionInfo(0.8, "Sigma adjust end", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}), - 'uni_pc_variant': OptionInfo("bh2", "UniPC variant", gr.Radio, {"choices": ["bh1", "bh2", "vary_coeff"], "visible": False}), - 'uni_pc_skip_type': OptionInfo("time_uniform", "UniPC skip type", gr.Radio, {"choices": ["time_uniform", "time_quadratic", "logSNR"], "visible": False}), -})) - options_templates.update(options_section(('postprocessing', "Postprocessing"), { 'postprocessing_enable_in_main_ui': OptionInfo([], "Additional postprocessing operations", gr.Dropdown, lambda: {"multiselect":True, "choices": [x.name for x in shared_items.postprocessing_scripts()]}), 'postprocessing_operation_order': OptionInfo([], "Postprocessing operation order", gr.Dropdown, lambda: {"multiselect":True, "choices": [x.name for x in shared_items.postprocessing_scripts()], "visible": False }), @@ -609,14 +587,6 @@ options_templates.update(options_section(('postprocessing', "Postprocessing"), { "upscaler_tile_overlap": OptionInfo(8, "Upscaler tile overlap", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}), })) -options_templates.update(options_section(('control', "Control Options"), { - "control_hires": OptionInfo(False, "Use control during hires", gr.Checkbox, {"visible": False}), - "control_max_units": OptionInfo(4, "Maximum number of units", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1, "visible": False}), - "control_tiles": OptionInfo("1x1, 1x2, 1x3, 1x4, 2x1, 2x1, 2x2, 2x3, 2x4, 3x1, 3x2, 3x3, 3x4, 4x1, 4x2, 4x3, 4x4", "Tiling options", gr.Textbox, {"visible": False}), - "control_move_processor": OptionInfo(False, "Processor move to CPU after use", gr.Checkbox, {"visible": False}), - "control_unload_processor": OptionInfo(False, "Processor unload after use", gr.Checkbox, {"visible": False}), -})) - options_templates.update(options_section(('interrogate', "Interrogate"), { "interrogate_default_type": OptionInfo("OpenCLiP", "Default type", gr.Radio, {"choices": ["OpenCLiP", "VLM", "DeepBooru"]}), "interrogate_offload": OptionInfo(True, "Offload models "), @@ -710,10 +680,42 @@ options_templates.update(options_section(('extra_networks', "Networks"), { })) options_templates.update(options_section((None, "Hidden options"), { + # internal options "diffusers_version": OptionInfo("", "Diffusers version", gr.Textbox, {"visible": False}), "disabled_extensions": OptionInfo([], "Disable these extensions", gr.Textbox, {"visible": False}), "sd_checkpoint_hash": OptionInfo("", "SHA256 hash of the current checkpoint", gr.Textbox, {"visible": False}), "tooltips": OptionInfo("UI Tooltips", "UI tooltips", gr.Radio, {"choices": ["None", "Browser default", "UI tooltips"], "visible": False}), + + # control settings are handled separately + "control_hires": OptionInfo(False, "Use control during hires", gr.Checkbox, {"visible": False}), + "control_max_units": OptionInfo(4, "Maximum number of units", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1, "visible": False}), + "control_tiles": OptionInfo("1x1, 1x2, 1x3, 1x4, 2x1, 2x1, 2x2, 2x3, 2x4, 3x1, 3x2, 3x3, 3x4, 4x1, 4x2, 4x3, 4x4", "Tiling options", gr.Textbox, {"visible": False}), + "control_move_processor": OptionInfo(False, "Processor move to CPU after use", gr.Checkbox, {"visible": False}), + "control_unload_processor": OptionInfo(False, "Processor unload after use", gr.Checkbox, {"visible": False}), + + # sampler settings are handled separately + "show_samplers": OptionInfo([], "Show samplers in user interface", gr.CheckboxGroup, lambda: {"choices": [x.name for x in list_samplers()], "visible": False}), + 'eta_noise_seed_delta': OptionInfo(0, "Noise seed delta (eta)", gr.Number, {"precision": 0, "visible": False}), + "scheduler_eta": OptionInfo(1.0, "Noise multiplier (eta)", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}), + "schedulers_solver_order": OptionInfo(0, "Solver order (where", gr.Slider, {"minimum": 0, "maximum": 5, "step": 1, "visible": False}), + "schedulers_use_loworder": OptionInfo(True, "Use simplified solvers in final steps", gr.Checkbox, {"visible": False}), + "schedulers_prediction_type": OptionInfo("default", "Override model prediction type", gr.Radio, {"choices": ['default', 'epsilon', 'sample', 'v_prediction'], "visible": False}), + "schedulers_sigma": OptionInfo("default", "Sigma algorithm", gr.Radio, {"choices": ['default', 'karras', 'exponential', 'polyexponential'], "visible": False}), + "schedulers_beta_schedule": OptionInfo("default", "Beta schedule", gr.Dropdown, {"choices": ['default', 'linear', 'scaled_linear', 'squaredcos_cap_v2', 'sigmoid'], "visible": False}), + "schedulers_use_thresholding": OptionInfo(False, "Use dynamic thresholding", gr.Checkbox, {"visible": False}), + "schedulers_timestep_spacing": OptionInfo("default", "Timestep spacing", gr.Dropdown, {"choices": ['default', 'linspace', 'leading', 'trailing'], "visible": False}), + 'schedulers_timesteps': OptionInfo('', "Timesteps", gr.Textbox, {"visible": False}), + "schedulers_rescale_betas": OptionInfo(False, "Rescale betas with zero terminal SNR", gr.Checkbox, {"visible": False}), + 'schedulers_beta_start': OptionInfo(0, "Beta start", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.00001}), + 'schedulers_beta_end': OptionInfo(0, "Beta end", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.00001}), + 'schedulers_timesteps_range': OptionInfo(1000, "Timesteps range", gr.Slider, {"minimum": 250, "maximum": 4000, "step": 1}), + 'schedulers_shift': OptionInfo(3, "Sampler shift", gr.Slider, {"minimum": 0.1, "maximum": 10, "step": 0.1, "visible": False}), + 'schedulers_dynamic_shift': OptionInfo(False, "Sampler dynamic shift", gr.Checkbox, {"visible": False}), + 'schedulers_sigma_adjust': OptionInfo(1.0, "Sigma adjust", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01, "visible": False}), + 'schedulers_sigma_adjust_min': OptionInfo(0.2, "Sigma adjust start", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}), + 'schedulers_sigma_adjust_max': OptionInfo(0.8, "Sigma adjust end", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}), + 'uni_pc_variant': OptionInfo("bh2", "UniPC variant", gr.Radio, {"choices": ["bh1", "bh2", "vary_coeff"], "visible": False}), + 'uni_pc_skip_type': OptionInfo("time_uniform", "UniPC skip type", gr.Radio, {"choices": ["time_uniform", "time_quadratic", "logSNR"], "visible": False}), })) diff --git a/pipelines/model_chroma.py b/pipelines/model_chroma.py index 4964543ac..e976bfa56 100644 --- a/pipelines/model_chroma.py +++ b/pipelines/model_chroma.py @@ -33,8 +33,6 @@ def load_chroma_quanto(checkpoint_info): with torch.device("meta"): transformer = diffusers.ChromaTransformer2DModel.from_config(os.path.join(repo_path, "transformer", "config.json")).to(dtype=dtype) quanto.requantize(transformer, state_dict, quantization_map, device=torch.device("cpu")) - if shared.opts.diffusers_eval: - transformer.eval() transformer_dtype = transformer.dtype if transformer_dtype != devices.dtype: try: @@ -61,8 +59,6 @@ def load_chroma_quanto(checkpoint_info): with torch.device("meta"): text_encoder = transformers.T5EncoderModel(t5_config).to(dtype=dtype) quanto.requantize(text_encoder, state_dict, quantization_map, device=torch.device("cpu")) - if shared.opts.diffusers_eval: - text_encoder.eval() text_encoder_dtype = text_encoder.dtype if text_encoder_dtype != devices.dtype: try: diff --git a/pipelines/model_cogview.py b/pipelines/model_cogview.py index 7d59293ec..400038dc3 100644 --- a/pipelines/model_cogview.py +++ b/pipelines/model_cogview.py @@ -73,9 +73,6 @@ def load_cogview4(checkpoint_info, diffusers_load_config={}): cache_dir=shared.opts.diffusers_dir, **load_args, ) - if shared.opts.diffusers_eval: - pipe.text_encoder.eval() - pipe.transformer.eval() pipe.enable_model_cpu_offload() devices.torch_gc() return pipe diff --git a/pipelines/model_flux.py b/pipelines/model_flux.py index dd113a0d4..5c1ba745b 100644 --- a/pipelines/model_flux.py +++ b/pipelines/model_flux.py @@ -33,8 +33,6 @@ def load_flux_quanto(checkpoint_info): with torch.device("meta"): transformer = diffusers.FluxTransformer2DModel.from_config(os.path.join(repo_path, "transformer", "config.json")).to(dtype=dtype) quanto.requantize(transformer, state_dict, quantization_map, device=torch.device("cpu")) - if shared.opts.diffusers_eval: - transformer.eval() transformer_dtype = transformer.dtype if transformer_dtype != devices.dtype: try: @@ -61,8 +59,6 @@ def load_flux_quanto(checkpoint_info): with torch.device("meta"): text_encoder_2 = transformers.T5EncoderModel(t5_config).to(dtype=dtype) quanto.requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cpu")) - if shared.opts.diffusers_eval: - text_encoder_2.eval() text_encoder_2_dtype = text_encoder_2.dtype if text_encoder_2_dtype != devices.dtype: try: diff --git a/pipelines/model_sana.py b/pipelines/model_sana.py index 1c04ff1e3..7400a1d97 100644 --- a/pipelines/model_sana.py +++ b/pipelines/model_sana.py @@ -78,13 +78,6 @@ def load_sana(checkpoint_info, kwargs={}): except Exception as e: shared.log.error(f'Load model: type=Sana {e}') - try: - if shared.opts.diffusers_eval: - pipe.text_encoder.eval() - pipe.transformer.eval() - except Exception: - pass - sd_hijack_te.init_hijack(pipe) t1 = time.time() shared.log.debug(f'Load model: type=Sana target={devices.dtype} te={pipe.text_encoder.dtype} transformer={pipe.transformer.dtype} vae={pipe.vae.dtype} time={t1-t0:.2f}')