From cc03ebc58469a173fbab5ff5b3461a934f1b3b35 Mon Sep 17 00:00:00 2001 From: vladmandic Date: Fri, 30 Jan 2026 11:34:25 +0100 Subject: [PATCH] move vae to subfolder Signed-off-by: vladmandic --- CHANGELOG.md | 8 +- TODO.md | 1 + extensions-builtin/sd-extension-system-info | 2 +- modules/framepack/framepack_vae.py | 4 +- modules/options.py | 4 +- modules/processing_args.py | 4 +- modules/processing_correction.py | 3 +- modules/processing_vae.py | 27 +++-- modules/sd_samplers_common.py | 3 +- modules/shared.py | 4 +- modules/shared_items.py | 4 +- modules/{ => vae}/sd_vae_approx.py | 0 modules/vae/sd_vae_fal.py | 121 ++++++++++++++++++++ modules/{ => vae}/sd_vae_natten.py | 0 modules/{ => vae}/sd_vae_ostris.py | 0 modules/{ => vae}/sd_vae_remote.py | 0 modules/{ => vae}/sd_vae_repa.py | 0 modules/{ => vae}/sd_vae_stablecascade.py | 0 modules/{ => vae}/sd_vae_taesd.py | 0 modules/video_models/video_vae.py | 2 +- 20 files changed, 162 insertions(+), 25 deletions(-) rename modules/{ => vae}/sd_vae_approx.py (100%) create mode 100644 modules/vae/sd_vae_fal.py rename modules/{ => vae}/sd_vae_natten.py (100%) rename modules/{ => vae}/sd_vae_ostris.py (100%) rename modules/{ => vae}/sd_vae_remote.py (100%) rename modules/{ => vae}/sd_vae_repa.py (100%) rename modules/{ => vae}/sd_vae_stablecascade.py (100%) rename modules/{ => vae}/sd_vae_taesd.py (100%) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0a1cb9c60..ae7a2f8be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,6 @@ # Change Log for SD.Next -## Update for 2026-01-28 +## Update for 2026-01-30 - **Models** - [Tongyi-MAI Z-Image Base](https://tongyi-mai.github.io/Z-Image-blog/) @@ -10,7 +10,9 @@ - add SmilingWolf WD14/WaifuDiffusion tagger models, thanks @CalamitousFelicitousness - support comments in wildcard files, using `#` - support aliases in metadata skip params, thanks @CalamitousFelicitousness - - ui gallery add manual cache cleanup, thanks @awsr + - ui gallery improve cache cleanup and add manual option, thanks @awsr + - selectable options to add system info to metadata, thanks @Athari + see *settings -> image metadata* - **Schedulers** - schedulers documentation has new home: - add 13(!) new scheduler families @@ -31,9 +33,11 @@ **xpu**: update to `torch==2.10.0` **rocm**: update to `torch==2.10.0` **openvino**: update to `torch==2.10.0` and `openvino==2025.4.1` + - rocm: expand available gfx archs, thanks @crashingalexsan - rocm: set `MIOPEN_FIND_MODE=2` by default, thanks @crashingalexsan - relocate all json data files to `data/` folder existing data files are auto-migrated on startup + - refactor and improve connection monitor, thanks @awsr - further work on type consistency and type checking, thanks @awsr - log captured exceptions - improve temp folder handling and cleanup diff --git a/TODO.md b/TODO.md index 371c06764..1a6cedd90 100644 --- a/TODO.md +++ b/TODO.md @@ -6,6 +6,7 @@ ## Internal +- Feature: Flow-match `res4lyf` schedulers - Feature: Move `nunchaku` models to refernce instead of internal decision - Update: `transformers==5.0.0` - Feature: Unify *huggingface* and *diffusers* model folders diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info index bd33edfd2..ddf821483 160000 --- a/extensions-builtin/sd-extension-system-info +++ b/extensions-builtin/sd-extension-system-info @@ -1 +1 @@ -Subproject commit bd33edfd28f95e1366f3169f6faca532098866ad +Subproject commit ddf821483c8bcdac4868f15a9a838b45b4dd30ad diff --git a/modules/framepack/framepack_vae.py b/modules/framepack/framepack_vae.py index 908378a8b..77f20b415 100644 --- a/modules/framepack/framepack_vae.py +++ b/modules/framepack/framepack_vae.py @@ -43,7 +43,7 @@ def vae_decode_simple(latents): def vae_decode_tiny(latents): global taesd # pylint: disable=global-statement if taesd is None: - from modules import sd_vae_taesd + from modules.vae import sd_vae_taesd taesd, _variant = sd_vae_taesd.get_model(variant='TAE HunyuanVideo') shared.log.debug(f'Video VAE: type=Tiny cls={taesd.__class__.__name__} latents={latents.shape}') with devices.inference_context(): @@ -56,7 +56,7 @@ def vae_decode_tiny(latents): def vae_decode_remote(latents): - # from modules.sd_vae_remote import remote_decode + # from modules.vae.sd_vae_remote import remote_decode # images = remote_decode(latents, model_type='hunyuanvideo') from diffusers.utils.remote_utils import remote_decode images = remote_decode( diff --git a/modules/options.py b/modules/options.py index 6b551385b..83e9e4a11 100644 --- a/modules/options.py +++ b/modules/options.py @@ -11,8 +11,10 @@ if TYPE_CHECKING: from modules.ui_components import DropdownEditable -def options_section(section_identifier: tuple[str, str], options_dict: dict[str, OptionInfo | LegacyOption]): +def options_section(section_identifier: tuple[str, str], options_dict: dict[str, OptionInfo | LegacyOption]) -> dict[str, OptionInfo | LegacyOption]: """Set the `section` value for all OptionInfo/LegacyOption items""" + if len(section_identifier) > 2: + section_identifier = section_identifier[:2] for v in options_dict.values(): v.section = section_identifier return options_dict diff --git a/modules/processing_args.py b/modules/processing_args.py index ebd75d751..dc63ea84c 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -67,7 +67,7 @@ def task_specific_kwargs(p, model): if 'hires' not in p.ops: p.ops.append('img2img') if p.vae_type == 'Remote': - from modules.sd_vae_remote import remote_encode + from modules.vae.sd_vae_remote import remote_encode p.init_images = remote_encode(p.init_images) task_args = { 'image': p.init_images, @@ -117,7 +117,7 @@ def task_specific_kwargs(p, model): p.ops.append('inpaint') mask_image = p.task_args.get('image_mask', None) or getattr(p, 'image_mask', None) or getattr(p, 'mask', None) if p.vae_type == 'Remote': - from modules.sd_vae_remote import remote_encode + from modules.vae.sd_vae_remote import remote_encode p.init_images = remote_encode(p.init_images) # mask_image = remote_encode(mask_image) task_args = { diff --git a/modules/processing_correction.py b/modules/processing_correction.py index 0fecc4e7c..7069a8fa9 100644 --- a/modules/processing_correction.py +++ b/modules/processing_correction.py @@ -5,7 +5,8 @@ https://huggingface.co/blog/TimothyAlexisVass/explaining-the-sdxl-latent-space import os import torch -from modules import shared, sd_vae_taesd, devices +from modules import shared, devices +from modules.vae import sd_vae_taesd debug_enabled = os.environ.get('SD_HDR_DEBUG', None) is not None diff --git a/modules/processing_vae.py b/modules/processing_vae.py index ab2ea085e..72c385ac5 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -2,7 +2,8 @@ import os import time import numpy as np import torch -from modules import shared, devices, sd_models, sd_vae, sd_vae_taesd, errors +from modules import shared, devices, sd_models, sd_vae, errors +from modules.vae import sd_vae_taesd debug = os.environ.get('SD_VAE_DEBUG', None) is not None @@ -286,13 +287,13 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he if vae_type == 'Remote': jobid = shared.state.begin('Remote VAE') - from modules.sd_vae_remote import remote_decode + from modules.vae.sd_vae_remote import remote_decode tensors = remote_decode(latents=latents, width=width, height=height) shared.state.end(jobid) if tensors is not None and len(tensors) > 0: return vae_postprocess(tensors, model, output_type) if vae_type == 'Repa': - from modules.sd_vae_repa import repa_load + from modules.vae.sd_vae_repa import repa_load vae = repa_load(latents) vae_type = 'Full' if vae is not None: @@ -310,14 +311,17 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he latents = latents.unsqueeze(0) if latents.shape[-1] <= 4: # not a latent, likely an image decoded = latents.float().cpu().numpy() - elif vae_type == 'Full' and hasattr(model, "vae"): - decoded = full_vae_decode(latents=latents, model=model) - elif hasattr(model, "vqgan"): - decoded = full_vqgan_decode(latents=latents, model=model) - else: + elif vae_type == 'Tiny': decoded = taesd_vae_decode(latents=latents) if torch.is_tensor(decoded): decoded = 2.0 * decoded - 1.0 # typical normalized range + elif hasattr(model, "vqgan"): + decoded = full_vqgan_decode(latents=latents, model=model) + elif hasattr(model, "vae"): + decoded = full_vae_decode(latents=latents, model=model) + else: + shared.log.error('VAE not found in model') + decoded = [] images = vae_postprocess(decoded, model, output_type) if shared.cmd_opts.profile or debug: @@ -339,11 +343,14 @@ def vae_encode(image, model, vae_type='Full'): # pylint: disable=unused-variable shared.log.error('VAE not found in model') return [] tensor = f.to_tensor(image.convert("RGB")).unsqueeze(0).to(devices.device, devices.dtype_vae) - if vae_type == 'Full': + if vae_type == 'Tiny': + latents = taesd_vae_encode(image=tensor) + elif vae_type == 'Full' and hasattr(model, 'vae'): tensor = tensor * 2 - 1 latents = full_vae_encode(image=tensor, model=shared.sd_model) else: - latents = taesd_vae_encode(image=tensor) + shared.log.error('VAE not found in model') + latents = [] devices.torch_gc() shared.state.end(jobid) return latents diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index 3cdc91943..14ad69eb8 100644 --- a/modules/sd_samplers_common.py +++ b/modules/sd_samplers_common.py @@ -4,7 +4,8 @@ from collections import namedtuple import torch import torchvision.transforms as T from PIL import Image -from modules import shared, devices, processing, images, sd_vae_approx, sd_vae_taesd, sd_vae_stablecascade, sd_samplers, timer +from modules import shared, devices, processing, images, sd_samplers, timer +from modules.vae import sd_vae_approx, sd_vae_taesd, sd_vae_stablecascade SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options']) diff --git a/modules/shared.py b/modules/shared.py index 7e1676f49..b3dec7abf 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -582,9 +582,9 @@ options_templates.update(options_section(('saving-paths', "Image Paths"), { })) options_templates.update(options_section(('image-metadata', "Image Metadata"), { - "image_metadata": OptionInfo(True, "Include metadata in image"), + "image_metadata": OptionInfo(True, "Save metadata in image"), "save_txt": OptionInfo(False, "Save metadata to text file"), - "save_log_fn": OptionInfo("", "Append metadata to JSON file", component_args=hide_dirs), + "save_log_fn": OptionInfo("", "Save metadata to JSON file", component_args=hide_dirs), "disable_apply_params": OptionInfo('', "Restore from metadata: skip params", gr.Textbox), "disable_apply_metadata": OptionInfo(['sd_model_checkpoint', 'sd_vae', 'sd_unet', 'sd_text_encoder'], "Restore from metadata: skip settings", gr.Dropdown, lambda: {"multiselect":True, "choices": opts.list()}), })) diff --git a/modules/shared_items.py b/modules/shared_items.py index b973e4886..3177c80b8 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -93,8 +93,8 @@ def sd_vae_items(): def sd_taesd_items(): - import modules.sd_vae_taesd - return list(modules.sd_vae_taesd.TAESD_MODELS.keys()) + list(modules.sd_vae_taesd.CQYAN_MODELS.keys()) + import modules.vae.sd_vae_taesd + return list(modules.vae.sd_vae_taesd.TAESD_MODELS.keys()) + list(modules.vae.sd_vae_taesd.CQYAN_MODELS.keys()) def refresh_vae_list(): import modules.sd_vae diff --git a/modules/sd_vae_approx.py b/modules/vae/sd_vae_approx.py similarity index 100% rename from modules/sd_vae_approx.py rename to modules/vae/sd_vae_approx.py diff --git a/modules/vae/sd_vae_fal.py b/modules/vae/sd_vae_fal.py new file mode 100644 index 000000000..bd482a779 --- /dev/null +++ b/modules/vae/sd_vae_fal.py @@ -0,0 +1,121 @@ +import torch +import torch.nn as nn +from diffusers.models import AutoencoderTiny +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.autoencoders.vae import EncoderOutput, DecoderOutput +from diffusers.configuration_utils import ConfigMixin, register_to_config + +from modules import shared, devices + + +repo_id = "fal/FLUX.2-Tiny-AutoEncoder" +tiny_vae = None +prev_vae = None + + +def is_compatile(): + return shared.sd_model_type in ['f2'] + + +def load_fal_vae(): + if not hasattr(shared.sd_model, 'vae') or not is_compatile(): + return + global tiny_vae, prev_vae # pylint: disable=global-statement + if tiny_vae is None: + tiny_vae = Flux2TinyAutoEncoder.from_pretrained( + repo_id, + cache_dir=shared.opts.hfcache_dir, + ).to(device=devices.device, dtype=devices.dtype) + if prev_vae is None: + prev_vae = shared.sd_model.vae + shared.sd_model.vae = tiny_vae + shared.log.info(f'VAE load: cls={tiny_vae.__class__.__name__} repo_id={repo_id}') + + +def unload_fal_vae(): + global prev_vae # pylint: disable=global-statement + if not hasattr(shared.sd_model, 'vae'): + return + if prev_vae is not None: + shared.sd_model.vae = prev_vae + prev_vae = None + shared.log.info(f'VAE restore: cls={prev_vae.__class__.__name__}') + + +class Flux2TinyAutoEncoder(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + latent_channels: int = 128, + encoder_block_out_channels: list[int] = [64, 64, 64, 64], + decoder_block_out_channels: list[int] = [64, 64, 64, 64], + act_fn: str = "silu", + upsampling_scaling_factor: int = 2, + num_encoder_blocks: list[int] = [1, 3, 3, 3], + num_decoder_blocks: list[int] = [3, 3, 3, 1], + latent_magnitude: float = 3.0, + latent_shift: float = 0.5, + force_upcast: bool = False, + scaling_factor: float = 0.13025, + ) -> None: + super().__init__() + self.tiny_vae = AutoencoderTiny( + in_channels=in_channels, + out_channels=out_channels, + encoder_block_out_channels=encoder_block_out_channels, + decoder_block_out_channels=decoder_block_out_channels, + act_fn=act_fn, + latent_channels=latent_channels // 4, + upsampling_scaling_factor=upsampling_scaling_factor, + num_encoder_blocks=num_encoder_blocks, + num_decoder_blocks=num_decoder_blocks, + latent_magnitude=latent_magnitude, + latent_shift=latent_shift, + force_upcast=force_upcast, + scaling_factor=scaling_factor, + ) + self.extra_encoder = nn.Conv2d( + latent_channels // 4, latent_channels, + kernel_size=4, stride=2, padding=1 + ) + self.extra_decoder = nn.ConvTranspose2d( + latent_channels, latent_channels // 4, + kernel_size=4, stride=2, padding=1 + ) + self.residual_encoder = nn.Sequential( + nn.Conv2d(latent_channels, latent_channels, kernel_size=3, padding=1), + nn.GroupNorm(8, latent_channels), + nn.SiLU(), + nn.Conv2d(latent_channels, latent_channels, kernel_size=3, padding=1), + ) + self.residual_decoder = nn.Sequential( + nn.Conv2d(latent_channels // 4, latent_channels // 4, kernel_size=3, padding=1), + nn.GroupNorm(8, latent_channels // 4), + nn.SiLU(), + nn.Conv2d(latent_channels // 4, latent_channels // 4, kernel_size=3, padding=1), + ) + + def encode(self, x: torch.Tensor, return_dict: bool = True) -> EncoderOutput: + encoded = self.tiny_vae.encode(x, return_dict=False)[0] + compressed = self.extra_encoder(encoded) + enhanced = self.residual_encoder(compressed) + compressed + if return_dict: + return EncoderOutput(latent=enhanced) + return enhanced + + def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput: + decompressed = self.extra_decoder(z) + enhanced = self.residual_decoder(decompressed) + decompressed + decoded = self.tiny_vae.decode(enhanced, return_dict=False)[0] + if return_dict: + return DecoderOutput(sample=decoded) + return decoded + + def forward(self, sample: torch.Tensor, return_dict: bool = True) -> DecoderOutput: + encoded = self.encode(sample, return_dict=False)[0] + decoded = self.decode(encoded, return_dict=False)[0] + if return_dict: + return DecoderOutput(sample=decoded) + return decoded diff --git a/modules/sd_vae_natten.py b/modules/vae/sd_vae_natten.py similarity index 100% rename from modules/sd_vae_natten.py rename to modules/vae/sd_vae_natten.py diff --git a/modules/sd_vae_ostris.py b/modules/vae/sd_vae_ostris.py similarity index 100% rename from modules/sd_vae_ostris.py rename to modules/vae/sd_vae_ostris.py diff --git a/modules/sd_vae_remote.py b/modules/vae/sd_vae_remote.py similarity index 100% rename from modules/sd_vae_remote.py rename to modules/vae/sd_vae_remote.py diff --git a/modules/sd_vae_repa.py b/modules/vae/sd_vae_repa.py similarity index 100% rename from modules/sd_vae_repa.py rename to modules/vae/sd_vae_repa.py diff --git a/modules/sd_vae_stablecascade.py b/modules/vae/sd_vae_stablecascade.py similarity index 100% rename from modules/sd_vae_stablecascade.py rename to modules/vae/sd_vae_stablecascade.py diff --git a/modules/sd_vae_taesd.py b/modules/vae/sd_vae_taesd.py similarity index 100% rename from modules/sd_vae_taesd.py rename to modules/vae/sd_vae_taesd.py diff --git a/modules/video_models/video_vae.py b/modules/video_models/video_vae.py index e31108088..dd8ba233f 100644 --- a/modules/video_models/video_vae.py +++ b/modules/video_models/video_vae.py @@ -41,7 +41,7 @@ def vae_decode_tiny(latents): else: shared.log.warning(f'Decode: type=Tiny cls={shared.sd_model.__class__.__name__} not supported') return None - from modules import sd_vae_taesd + from modules.vae import sd_vae_taesd vae, variant = sd_vae_taesd.get_model(variant=variant) if vae is None: return None