calculate vae-scale-factor and use everywhere

Signed-off-by: Vladimir Mandic <mandic00@live.com>

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**  

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():

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}')

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')

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']
 

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
 

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]

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
 
 

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):