From b569d68b4f5eb197272e89b9198753d1c3c78721 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 26 Oct 2024 21:51:36 -0400 Subject: [PATCH] add ostris: experimental Signed-off-by: Vladimir Mandic --- extensions-builtin/Lora/networks.py | 2 ++ installer.py | 2 +- javascript/sdnext.css | 2 +- modules/sd_samplers_common.py | 4 +-- modules/sd_vae_ostris.py | 41 +++++++++++++++++++++++++++++ modules/sd_vae_taesd.py | 12 ++++++--- 6 files changed, 55 insertions(+), 8 deletions(-) create mode 100644 modules/sd_vae_ostris.py diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py index c0a8555e1..83aa6b40b 100644 --- a/extensions-builtin/Lora/networks.py +++ b/extensions-builtin/Lora/networks.py @@ -127,6 +127,8 @@ def load_diffusers(name, network_on_disk, lora_scale=shared.opts.extra_networks_ def load_network(name, network_on_disk) -> network.Network: + if not shared.sd_loaded: + return t0 = time.time() cached = lora_cache.get(name, None) if debug: diff --git a/installer.py b/installer.py index aca36056b..2cc1c74a7 100644 --- a/installer.py +++ b/installer.py @@ -731,7 +731,7 @@ def check_torch(): else: if args.use_zluda: log.warning("ZLUDA failed to initialize: no HIP SDK found") - log.info('Using CPU-only Torch') + log.warning('Torch: CPU-only version installed') torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision') if 'torch' in torch_command and not args.version: install(torch_command, 'torch torchvision', quiet=True) diff --git a/javascript/sdnext.css b/javascript/sdnext.css index 3c81e7d8e..cbd3cac83 100644 --- a/javascript/sdnext.css +++ b/javascript/sdnext.css @@ -38,7 +38,7 @@ td > div > span { overflow-y: auto; max-height: 3em; overflow-x: hidden; } .gradio-button.secondary-down, .gradio-button.secondary-down:hover { box-shadow: 1px 1px 1px rgba(0,0,0,0.25) inset, 0px 0px 3px rgba(0,0,0,0.15) inset; } .gradio-button.secondary-down:hover { background: var(--button-secondary-background-fill-hover); color: var(--button-secondary-text-color-hover); } .gradio-button.tool { max-width: min-content; min-width: min-content !important; font-size: 20px !important; color: var(--body-text-color) !important; align-self: end; margin-bottom: 4px; } -.gradio-checkbox { margin: 0.75em 1.5em 0 0; align-self: center; } +.gradio-checkbox { margin-right: 1em !important; align-self: center; } .gradio-column { min-width: min(160px, 100%) !important; } .gradio-container { max-width: unset !important; padding: var(--block-label-padding) !important; } .gradio-container .prose a, .gradio-container .prose a:visited{ color: unset; text-decoration: none; } diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index 1b1cd189a..a487fe9b7 100644 --- a/modules/sd_samplers_common.py +++ b/modules/sd_samplers_common.py @@ -44,7 +44,7 @@ def single_sample_to_image(sample, approximation=None): if sample.dtype == torch.bfloat16 and (approximation == 0 or approximation == 1): sample = sample.to(torch.float16) except Exception as e: - warn_once(f'live preview: {e}') + warn_once(f'Preview: {e}') if len(sample.shape) > 4: # likely unknown video latent (e.g. svd) return Image.new(mode="RGB", size=(512, 512)) @@ -82,7 +82,7 @@ def single_sample_to_image(sample, approximation=None): transform = T.ToPILImage() image = transform(x_sample) except Exception as e: - warn_once(f'live preview: {e}') + warn_once(f'Preview: {e}') image = Image.new(mode="RGB", size=(512, 512)) return image diff --git a/modules/sd_vae_ostris.py b/modules/sd_vae_ostris.py new file mode 100644 index 000000000..70542e9f5 --- /dev/null +++ b/modules/sd_vae_ostris.py @@ -0,0 +1,41 @@ +import time +import torch +import diffusers +from huggingface_hub import hf_hub_download +from safetensors.torch import load_file +from modules import shared, devices + + +decoder_id = "ostris/vae-kl-f8-d16" +adapter_id = "ostris/16ch-VAE-Adapters" + + +def load_vae(pipe): + if shared.sd_model_type == 'sd': + adapter_file = "16ch-VAE-Adapter-SD15-alpha.safetensors" + elif shared.sd_model_type == 'sdxl': + adapter_file = "16ch-VAE-Adapter-SDXL-alpha_v02.safetensors" + else: + shared.log.error('VAE: type=osiris unsupported model type') + return + t0 = time.time() + ckpt_file = hf_hub_download(adapter_id, adapter_file, cache_dir=shared.opts.hfcache_dir) + ckpt = load_file(ckpt_file) + lora_state_dict = {k: v for k, v in ckpt.items() if "lora" in k} + unet_state_dict = {k.replace("unet_", ""): v for k, v in ckpt.items() if "unet_" in k} + + pipe.unet.conv_in = torch.nn.Conv2d(16, 320, 3, 1, 1) + pipe.unet.conv_out = torch.nn.Conv2d(320, 16, 3, 1, 1) + pipe.unet.load_state_dict(unet_state_dict, strict=False) + pipe.unet.conv_in.to(devices.dtype) + pipe.unet.conv_out.to(devices.dtype) + pipe.unet.config.in_channels = 16 + pipe.unet.config.out_channels = 16 + + pipe.load_lora_weights(lora_state_dict, adapter_name=adapter_id) + # pipe.set_adapters(adapter_names=[adapter_id], adapter_weights=[0.8]) + pipe.fuse_lora(adapter_names=[adapter_id], lora_scale=0.8, fuse_unet=True) + + pipe.vae = diffusers.AutoencoderKL.from_pretrained(decoder_id, torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir) + t1 = time.time() + shared.log.info(f'VAE load: type=osiris decoder="{decoder_id}" adapter="{adapter_id}" time={t1-t0:.2f}s') diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index 5cd7fab7c..67886229f 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -171,20 +171,24 @@ def decode(latents): try: with devices.inference_context(): latents = latents.detach().clone().to(devices.device, dtype) - if len(latents.shape) == 3: + if len(latents.shape) == 3 and latents.shape[0] == 4: latents = latents.unsqueeze(0) image = vae.decoder(latents).clamp(0, 1).detach() image = 2.0 * image - 1.0 # typical normalized range except for preview which runs denormalization return image[0] - elif len(latents.shape) == 4: + elif len(latents.shape) == 4 and latents.shape[1] == 4: image = vae.decoder(latents).clamp(0, 1).detach() image = 2.0 * image - 1.0 # typical normalized range except for preview which runs denormalization return image else: - shared.log.error(f'TAESD decode unsupported latent type: {latents.shape}') + if not previous_warnings: + shared.log.error(f'TAESD decode unsupported latent type: {latents.shape}') + previous_warnings = True return latents except Exception as e: - shared.log.error(f'VAE decode taesd: {e}') + if not previous_warnings: + shared.log.error(f'VAE decode taesd: {e}') + previous_warnings = True return latents