diff --git a/CHANGELOG.md b/CHANGELOG.md index a7f366c29..ea39339bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -88,6 +88,7 @@ And (*as always*) many bugfixes and improvements to existing features! - prompt parser allow explict `BOS` and `EOS` tokens in prompt - **Nunchaku** support for *FLUX.1-Fill* and *FLUX.1-Depth* models - update requirements/packages + - use model vae scale-factor for image width/heigt calculations - **Other** - **prompt enhance** add `allura-org/Gemma-3-Glitter-4B` model support - remove **LDSR** diff --git a/modules/control/run.py b/modules/control/run.py index 57cd402f0..bc25d27bb 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -13,7 +13,7 @@ from modules.control.units import lite # Kohya ControlLLLite from modules.control.units import t2iadapter # TencentARC T2I-Adapter from modules.control.units import reference # ControlNet-Reference from modules.control.processor import preprocess_image -from modules import devices, shared, errors, processing, images, sd_models, scripts_manager, masking +from modules import devices, shared, errors, processing, images, sd_models, sd_vae, scripts_manager, masking from modules.processing_class import StableDiffusionProcessingControl from modules.ui_common import infotext_to_html from modules.api import script @@ -384,10 +384,11 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg p.selected_scale_tab_mask = 1 # hires/refine defined outside of main init + vae_scale_factor = sd_vae.get_vae_scale_factor() if p.enable_hr and (p.hr_resize_x == 0 or p.hr_resize_y == 0): - p.hr_upscale_to_x, p.hr_upscale_to_y = 8 * int(p.width_before * p.hr_scale / 8), 8 * int(p.height_before * p.hr_scale / 8) + p.hr_upscale_to_x, p.hr_upscale_to_y = vae_scale_factor * int(p.width_before * p.hr_scale / vae_scale_factor), vae_scale_factor * int(p.height_before * p.hr_scale / vae_scale_factor) elif p.enable_hr and (p.hr_upscale_to_x == 0 or p.hr_upscale_to_y == 0): - p.hr_upscale_to_x, p.hr_upscale_to_y = 8 * int(p.hr_resize_x / 8), 8 * int(p.hr_resize_y / 8) + p.hr_upscale_to_x, p.hr_upscale_to_y = 8 * int(p.hr_resize_x / vae_scale_factor), vae_scale_factor * int(p.hr_resize_y / vae_scale_factor) global p_extra_args # pylint: disable=global-statement for k, v in p_extra_args.items(): diff --git a/modules/control/tile.py b/modules/control/tile.py index de9df1131..1d6478edc 100644 --- a/modules/control/tile.py +++ b/modules/control/tile.py @@ -1,6 +1,6 @@ import time from PIL import Image -from modules import shared, processing, images, sd_models +from modules import shared, processing, images, sd_models, sd_vae def get_tile(image: Image.Image, x: int, y: int, sx: int, sy: int) -> Image.Image: @@ -23,17 +23,18 @@ def run_tiling(p: processing.StableDiffusionProcessing, input_image: Image.Image sx, sy = p.control_tile.split('x') sx = int(sx) sy = int(sy) + vae_scale_factor = sd_vae.get_vae_scale_factor() if sx <= 0 or sy <= 0: raise ValueError('Control Tile: invalid tile size') control_image = p.task_args.get('control_image', None) or p.task_args.get('image', None) control_upscaled = None if isinstance(control_image, list) and len(control_image) > 0: - w, h = 8 * int(sx * control_image[0].width) // 8, 8 * int(sy * control_image[0].height) // 8 + w, h = vae_scale_factor * int(sx * control_image[0].width) // vae_scale_factor, vae_scale_factor * int(sy * control_image[0].height) // vae_scale_factor control_upscaled = images.resize_image(resize_mode=1 if sx==sy else 5, im=control_image[0], width=w, height=h, context='add with forward') init_image = p.override or input_image init_upscaled = None if init_image is not None: - w, h = 8 * int(sx * init_image.width) // 8, 8 * int(sy * init_image.height) // 8 + w, h = vae_scale_factor * int(sx * init_image.width) // vae_scale_factor, vae_scale_factor * int(sy * init_image.height) // vae_scale_factor init_upscaled = images.resize_image(resize_mode=1 if sx==sy else 5, im=init_image, width=w, height=h, context='add with forward') t1 = time.time() shared.log.debug(f'Control Tile: scale={sx}x{sy} resize={"fixed" if sx==sy else "context"} control={control_upscaled} init={init_upscaled} time={t1-t0:.3f}') diff --git a/modules/processing.py b/modules/processing.py index 268661647..675ac9160 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -158,11 +158,12 @@ def process_images(p: StableDiffusionProcessing) -> Processed: shared.prompt_styles.apply_styles_to_extra(p) shared.prompt_styles.extract_comments(p) + vae_scale_factor = sd_vae.get_vae_scale_factor() if p.width is not None: - p.width = 8 * int(p.width / 8) + p.width = vae_scale_factor * int(p.width / vae_scale_factor) if p.height is not None: - p.height = 8 * int(p.height / 8) + p.height = vae_scale_factor * int(p.height / vae_scale_factor) script_callbacks.before_process_callback(p) timer.process.record('pre') diff --git a/modules/processing_args.py b/modules/processing_args.py index 4cee55eb4..c34cdc2b9 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -7,7 +7,7 @@ import inspect import torch import numpy as np from PIL import Image -from modules import shared, errors, sd_models, processing, processing_vae, processing_helpers, sd_hijack_hypertile, prompt_parser_diffusers, timer, extra_networks +from modules import shared, errors, sd_models, processing, processing_vae, processing_helpers, sd_hijack_hypertile, prompt_parser_diffusers, timer, extra_networks, sd_vae from modules.processing_callbacks import diffusers_callback_legacy, diffusers_callback, set_callbacks_p from modules.processing_helpers import resize_hires, fix_prompts, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, get_generator, set_latents, apply_circular # pylint: disable=unused-import from modules.api import helpers @@ -19,6 +19,7 @@ disable_pbar = os.environ.get('SD_DISABLE_PBAR', None) is not None def task_specific_kwargs(p, model): + vae_scale_factor = sd_vae.get_vae_scale_factor(model) task_args = {} is_img2img_model = bool('Zero123' in shared.sd_model.__class__.__name__) if len(getattr(p, 'init_images', [])) > 0: @@ -30,8 +31,8 @@ def task_specific_kwargs(p, model): p.ops.append('txt2img') if hasattr(p, 'width') and hasattr(p, 'height'): task_args = { - 'width': 8 * math.ceil(p.width / 8), - 'height': 8 * math.ceil(p.height / 8), + 'width': vae_scale_factor * math.ceil(p.width / vae_scale_factor), + 'height': vae_scale_factor * math.ceil(p.height / vae_scale_factor), } elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0: if shared.sd_model_type == 'sdxl' and hasattr(model, 'register_to_config'): @@ -50,19 +51,18 @@ def task_specific_kwargs(p, model): } if model.__class__.__name__ == 'FluxImg2ImgPipeline' or model.__class__.__name__ == 'FluxKontextPipeline': # needs explicit width/height if torch.is_tensor(p.init_images[0]): - p.width, p.height = p.init_images[0].shape[-1] * 16, p.init_images[0].shape[-2] * 16 + p.width, p.height = p.init_images[0].shape[-1] * vae_scale_factor, p.init_images[0].shape[-2] * vae_scale_factor else: - p.width, p.height = 8 * math.ceil(p.init_images[0].width / 8), 8 * math.ceil(p.init_images[0].height / 8) + p.width, p.height = 8 * math.ceil(p.init_images[0].width / vae_scale_factor), 8 * math.ceil(p.init_images[0].height / vae_scale_factor) if model.__class__.__name__ == 'FluxKontextPipeline': aspect_ratio = p.width / p.height - vae_scale_factor = 16 max_area = max(p.width, p.height)**2 p.width, p.height = round((max_area * aspect_ratio) ** 0.5), round((max_area / aspect_ratio) ** 0.5) p.width, p.height = p.width // vae_scale_factor * vae_scale_factor, p.height // vae_scale_factor * vae_scale_factor task_args['max_area'] = max_area task_args['width'], task_args['height'] = p.width, p.height elif model.__class__.__name__ == 'OmniGenPipeline' or model.__class__.__name__ == 'OmniGen2Pipeline': - p.width, p.height = 16 * math.ceil(p.init_images[0].width / 16), 16 * math.ceil(p.init_images[0].height / 16) + p.width, p.height = vae_scale_factor * math.ceil(p.init_images[0].width / vae_scale_factor), vae_scale_factor * math.ceil(p.init_images[0].height / vae_scale_factor) task_args = { 'width': p.width, 'height': p.height, @@ -71,8 +71,8 @@ def task_specific_kwargs(p, model): elif sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INSTRUCT and len(getattr(p, 'init_images', [])) > 0: p.ops.append('instruct') task_args = { - 'width': 8 * math.ceil(p.width / 8) if hasattr(p, 'width') else None, - 'height': 8 * math.ceil(p.height / 8) if hasattr(p, 'height') else None, + 'width': vae_scale_factor * math.ceil(p.width / vae_scale_factor) if hasattr(p, 'width') else None, + 'height': vae_scale_factor * math.ceil(p.height / vae_scale_factor) if hasattr(p, 'height') else None, 'image': p.init_images, 'strength': p.denoising_strength, } @@ -122,11 +122,6 @@ def task_specific_kwargs(p, model): 'target_subject_category': getattr(p, 'prompt', '').split()[-1], 'output_type': 'pil', } - if model.__class__.__name__ in ['StableDiffusion3Pipeline', 'WanPipeline']: - p.width = 16 * (p.width // 16) - p.height = 16 * (p.height // 16) - task_args['width'] = p.width - task_args['height'] = p.height if debug_enabled: debug_log(f'Process task specific args: {task_args}') return task_args @@ -383,18 +378,19 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t # handle missing resolution if args.get('image', None) is not None and ('width' not in args or 'height' not in args): if 'width' in possible and 'height' in possible: + vae_scale_factor = sd_vae.get_vae_scale_factor(model) if isinstance(args['image'], torch.Tensor) or isinstance(args['image'], np.ndarray): - args['width'] = 8 * args['image'].shape[-1] - args['height'] = 8 * args['image'].shape[-2] + args['width'] = vae_scale_factor * args['image'].shape[-1] + args['height'] = vae_scale_factor * args['image'].shape[-2] elif isinstance(args['image'], Image.Image): args['width'] = args['image'].width args['height'] = args['image'].height elif isinstance(args['image'][0], torch.Tensor) or isinstance(args['image'][0], np.ndarray): - args['width'] = 8 * args['image'][0].shape[-1] - args['height'] = 8 * args['image'][0].shape[-2] + args['width'] = vae_scale_factor * args['image'][0].shape[-1] + args['height'] = vae_scale_factor * args['image'][0].shape[-2] else: - args['width'] = 8 * math.ceil(args['image'][0].width / 8) - args['height'] = 8 * math.ceil(args['image'][0].height / 8) + args['width'] = vae_scale_factor * math.ceil(args['image'][0].width / vae_scale_factor) + args['height'] = vae_scale_factor * math.ceil(args['image'][0].height / vae_scale_factor) if 'max_area' in possible and 'width' in args and 'height' in args and 'max_area' not in args: args['max_area'] = args['width'] * args['height'] diff --git a/modules/processing_class.py b/modules/processing_class.py index 4a9d93b90..daff8ccea 100644 --- a/modules/processing_class.py +++ b/modules/processing_class.py @@ -6,7 +6,7 @@ from typing import Any, Dict, List from dataclasses import dataclass, field import numpy as np from PIL import Image, ImageOps -from modules import shared, images, scripts_manager, masking, sd_models, processing_helpers +from modules import shared, images, scripts_manager, masking, sd_models, sd_vae, processing_helpers debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -449,10 +449,11 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): def init(self, all_prompts=None, all_seeds=None, all_subseeds=None): if self.init_images is not None and len(self.init_images) > 0: + vae_scale_factor = sd_vae.get_vae_scale_factor() if self.width is None or self.width == 0: - self.width = int(8 * (self.init_images[0].width * self.scale_by // 8)) + self.width = int(vae_scale_factor * (self.init_images[0].width * self.scale_by // vae_scale_factor)) if self.height is None or self.height == 0: - self.height = int(8 * (self.init_images[0].height * self.scale_by // 8)) + self.height = int(vae_scale_factor * (self.init_images[0].height * self.scale_by // vae_scale_factor)) if getattr(self, 'image_mask', None) is not None: shared.sd_model = sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.INPAINTING) elif getattr(self, 'init_images', None) is not None: @@ -554,7 +555,8 @@ class StableDiffusionProcessingControl(StableDiffusionProcessingImg2Img): self.hr_force = force self.hr_upscaler = upscaler if use_scale: - self.hr_upscale_to_x, self.hr_upscale_to_y = 8 * int(self.width * scale / 8), 8 * int(self.height * scale / 8) + vae_scale_factor = sd_vae.get_vae_scale_factor() + self.hr_upscale_to_x, self.hr_upscale_to_y = vae_scale_factor * int(self.width * scale / vae_scale_factor), vae_scale_factor * int(self.height * scale / vae_scale_factor) else: self.hr_upscale_to_x, self.hr_upscale_to_y = self.hr_resize_x, self.hr_resize_y diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 4777463ba..8729d2465 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -8,7 +8,7 @@ import numpy as np import cv2 from PIL import Image from blendmodes.blend import blendLayers, BlendType -from modules import shared, devices, images, sd_models, sd_samplers, sd_hijack_hypertile, processing_vae, timer +from modules import shared, devices, images, sd_models, sd_samplers, sd_vae, sd_hijack_hypertile, processing_vae, timer debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -282,7 +282,8 @@ def resize_init_images(p): if getattr(p, 'image', None) is not None and getattr(p, 'init_images', None) is None: p.init_images = [p.image] if getattr(p, 'init_images', None) is not None and len(p.init_images) > 0: - tgt_width, tgt_height = 8 * math.ceil(p.init_images[0].width / 8), 8 * math.ceil(p.init_images[0].height / 8) + vae_scale_factor = sd_vae.get_vae_scale_factor() + tgt_width, tgt_height = vae_scale_factor * math.ceil(p.init_images[0].width / vae_scale_factor), vae_scale_factor * math.ceil(p.init_images[0].height / vae_scale_factor) if p.init_images[0].size != (tgt_width, tgt_height): shared.log.debug(f'Resizing init images: original={p.init_images[0].width}x{p.init_images[0].height} target={tgt_width}x{tgt_height}') p.init_images = [images.resize_image(1, image, tgt_width, tgt_height, upscaler_name=None) for image in p.init_images] diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 6637e0e5b..62273045b 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -168,7 +168,8 @@ def full_vae_decode(latents, model): if debug: log_debug(f'VAE memory: {shared.mem_mon.read()}') vae_name = os.path.splitext(os.path.basename(sd_vae.loaded_vae_file))[0] if sd_vae.loaded_vae_file is not None else "default" - shared.log.debug(f'Decode: vae="{vae_name}" upcast={upcast} slicing={getattr(model.vae, "use_slicing", None)} tiling={getattr(model.vae, "use_tiling", None)} latents={list(latents.shape)}:{latents.device}:{latents.dtype} time={t1-t0:.3f}') + vae_scale_factor = sd_vae.get_vae_scale_factor(model) + shared.log.debug(f'Decode: vae="{vae_name}" scale={vae_scale_factor} upcast={upcast} slicing={getattr(model.vae, "use_slicing", None)} tiling={getattr(model.vae, "use_tiling", None)} latents={list(latents.shape)}:{latents.device}:{latents.dtype} time={t1-t0:.3f}') return decoded diff --git a/modules/sd_vae.py b/modules/sd_vae.py index b6c3982e9..2578ff196 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -12,6 +12,35 @@ checkpoint_info = None vae_path = os.path.abspath(os.path.join(paths.models_path, 'VAE')) debug = os.environ.get('SD_LOAD_DEBUG', None) is not None unspecified = object() +vae_scale_override = { + 'WanPipeline': 16, +} + + +def get_vae_scale_factor(model=None): + patch_size = 1 + if model is None: + model = shared.sd_model + if model is None: + vae_scale_factor = 8 + elif model.__class__.__name__ in vae_scale_override: + vae_scale_factor = vae_scale_override[model.__class__.__name__] + elif hasattr(model, 'vae_scale_factor_spatial'): + vae_scale_factor = model.vae_scale_factor_spatial + elif hasattr(model, 'vae_scale_factor'): + vae_scale_factor = model.vae_scale_factor + elif hasattr(model, 'pipe') and hasattr(model.pipe, 'vae_scale_factor'): + vae_scale_factor = model.pipe.vae_scale_factor + elif hasattr(model, 'config') and hasattr(model.config, 'vae_scale_factor'): + vae_scale_factor = model.config.vae_scale_factor + else: + shared.log.warning(f'VAE: cls={model.__class__.__name__ if model else "None"} scale=unknown') + vae_scale_factor = 8 + if hasattr(model, 'patch_size'): + patch_size = model.patch_size + if debug: + shared.log.trace(f'VAE: cls={model.__class__.__name__ if model else "None"} scale={vae_scale_factor} patch={patch_size}') + return vae_scale_factor * patch_size def load_vae_dict(filename):