refactor stable-cascade, fix taesd bf16, add skip-env cmd flag

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Change Log for SD.Next
 
-## Update for 2024-02-20
+## Update for 2024-02-21
 
 - **Improvements**:
   - **IP Adapter** major refactor  
@@ -78,6 +78,7 @@
     see *settings -> compute settings*  
 - **Fixes**:  
   - add variation seed to diffusers txt2img, thanks @AI-Casanova  
+  - add cmd param `--skip-env` to skip setting of environment parameters during sdnext load  
   - handle extensions that install conflicting versions of packages  
     `onnxruntime`, `opencv2-python`  
   - installer refresh package cache on any install  
@@ -100,6 +101,7 @@
   - fix controlnet inpaint mask  
   - fix theme list refresh  
   - fix extensions update information in ui  
+  - fix taesd with bfloat16
   - fix model merge manual merge settings, thanks @AI-Casanova  
   - fix gradio instant update issues for textboxes in quicksettings  
   - bind controlnet extension to last known working commit, thanks @Aptronymist  

TODO.md

@@ -11,6 +11,7 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 - diffusers public callbacks  
 - image2video: pia and vgen pipelines  
 - video2video
+- async lowvram: <https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14855>
 
 ## Control missing features
 

cli/run-benchmark.py

@@ -98,7 +98,7 @@ async def main():
             if i == 0:
                 log.info({ 'warmup': round(ts, 2) })
             else:
-                peak = gpu['session']['peak'] if 'session' in gpu else 0
+                peak = gpu['system']['used'] # gpu['session']['peak'] if 'session' in gpu else 0
                 log.info({ 'batch': batch[i], 'its': round(options.steps / (ts / batch[i]), 2), 'img': round(ts / batch[i], 2), 'wall': round(ts, 2), 'peak': gb(peak), 'oom': oom > 0 })
         else:
             await asyncio.sleep(10)

html/reference.json

@@ -73,6 +73,12 @@
     "desc": "(SVD) Image-to-Video is a latent diffusion model trained to generate short video clips from an image conditioning. This model was trained to generate 25 frames at resolution 576x1024 given a context frame of the same size, finetuned from SVD Image-to-Video [14 frames]. We also finetune the widely used f8-decoder for temporal consistency.",
     "preview": "stabilityai--stable-video-diffusion-img2vid-xt.jpg"
   },
+  "StabilityAI Stable Cascade": {
+    "path": "huggingface/stabilityai/stable-cascade",
+    "skip": true,
+    "desc": "Stable Cascade is a diffusion model built upon the Würstchen architecture and its main difference to other models like Stable Diffusion is that it is working at a much smaller latent space. Why is this important? The smaller the latent space, the faster you can run inference and the cheaper the training becomes. How small is the latent space? Stable Diffusion uses a compression factor of 8, resulting in a 1024x1024 image being encoded to 128x128. Stable Cascade achieves a compression factor of 42, meaning that it is possible to encode a 1024x1024 image to 24x24, while maintaining crisp reconstructions. The text-conditional model is then trained in the highly compressed latent space. Previous versions of this architecture, achieved a 16x cost reduction over Stable Diffusion 1.5",
+    "preview": "stabilityai--stable-cascade.jpg"
+  },
   "Segmind Vega": {
     "path": "segmind/Segmind-Vega",
     "desc": "The Segmind-Vega Model is a distilled version of the Stable Diffusion XL (SDXL), offering a remarkable 70% reduction in size and an impressive 100% speedup while retaining high-quality text-to-image generation capabilities. Trained on diverse datasets, including Grit and Midjourney scrape data, it excels at creating a wide range of visual content based on textual prompts. Employing a knowledge distillation strategy, Segmind-Vega leverages the teachings of several expert models, including SDXL, ZavyChromaXL, and JuggernautXL, to combine their strengths and produce compelling visual outputs.",

installer.py

@@ -975,6 +975,7 @@ def add_args(parser):
     group.add_argument('--skip-git', default = os.environ.get("SD_SKIPGIT",False), action='store_true', help = "Skips running all GIT operations, default: %(default)s")
     group.add_argument('--skip-torch', default = os.environ.get("SD_SKIPTORCH",False), action='store_true', help = "Skips running Torch checks, default: %(default)s")
     group.add_argument('--skip-all', default = os.environ.get("SD_SKIPALL",False), action='store_true', help = "Skips running all checks, default: %(default)s")
+    group.add_argument('--skip-env', default = os.environ.get("SD_SKIPENV",False), action='store_true', help = "Skips setting of env variables during startup, default: %(default)s")
     group.add_argument('--experimental', default = os.environ.get("SD_EXPERIMENTAL",False), action='store_true', help = "Allow unsupported versions of libraries, default: %(default)s")
     group.add_argument('--reinstall', default = os.environ.get("SD_REINSTALL",False), action='store_true', help = "Force reinstallation of all requirements, default: %(default)s")
     group.add_argument('--test', default = os.environ.get("SD_TEST",False), action='store_true', help = "Run test only and exit")

launch.py

@@ -202,7 +202,8 @@ def main():
         installer.log.info('Skipping GIT operations')
     installer.check_version()
     installer.log.info(f'Platform: {installer.print_dict(installer.get_platform())}')
-    installer.set_environment()
+    if not args.skip_env:
+        installer.set_environment()
     installer.check_torch()
     installer.check_onnx()
     installer.check_modified_files()

modules/modelloader.py

@@ -302,6 +302,8 @@ def load_reference(name: str):
         if v.get('path', '') == name:
             model_opts = v
             break
+    if model_opts.get('skip', False):
+        return True
     model_dir = download_diffusers_model(
         hub_id=name,
         cache_dir=shared.opts.diffusers_dir,

modules/processing_diffusers.py

@@ -165,7 +165,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         if hasattr(model, 'pipe'): # recurse
             model = model.pipe
         signature = inspect.signature(type(model).__call__, follow_wrapped=True)
-        possible = signature.parameters.keys()
+        possible = list(signature.parameters)
         debug(f'Diffusers pipeline possible: {possible}')
         if shared.opts.diffusers_generator_device == "Unset":
             generator_device = None
@@ -179,15 +179,16 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         prompts, negative_prompts, prompts_2, negative_prompts_2 = fix_prompts(prompts, negative_prompts, prompts_2, negative_prompts_2)
         parser = 'Fixed attention'
         clip_skip = kwargs.pop("clip_skip", 1)
+        steps = kwargs.get("num_inference_steps", 1)
         if shared.opts.prompt_attention != 'Fixed attention' and 'StableDiffusion' in model.__class__.__name__ and 'Onnx' not in model.__class__.__name__:
             try:
-                prompt_parser_diffusers.encode_prompts(model, p, prompts, negative_prompts, steps=kwargs.get("num_inference_steps", 1), clip_skip=clip_skip)
+                prompt_parser_diffusers.encode_prompts(model, p, prompts, negative_prompts, steps=steps, clip_skip=clip_skip)
                 parser = shared.opts.prompt_attention
             except Exception as e:
                 shared.log.error(f'Prompt parser encode: {e}')
                 if os.environ.get('SD_PROMPT_DEBUG', None) is not None:
                     errors.display(e, 'Prompt parser encode')
-        if parser == 'Fixed attention':
+        if 'clip_skip' in possible and parser == 'Fixed attention':
             if clip_skip == 1:
                 pass # clip_skip = None
             else:
@@ -211,17 +212,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             model.scheduler.noise_sampler_seed = p.seeds[0] # some schedulers have internal noise generator and do not use pipeline generator
         if 'noise_sampler_seed' in possible:
             args['noise_sampler_seed'] = p.seeds[0]
-        if hasattr(model, "decoder") and hasattr(model, "prior_prior") and 'prior_num_inference_steps' in possible:
-            steps = kwargs.pop("num_inference_steps", 20)
-            args["prior_num_inference_steps"] = steps
-            args["num_inference_steps"] = max(int(steps / 2), 1)
-        if hasattr(model, "decoder") and hasattr(model, "prior_prior") and 'prior_guidance_scale' in possible:
-            cfg_scale = kwargs.pop("guidance_scale", p.cfg_scale)
-            args["prior_guidance_scale"] = cfg_scale
-            # Using decoder_guidance_scale causes "Expected all tensors to be on the same device" errors right now
-            # Enabling model cpu offload fixes the error above
-            #args["decoder_guidance_scale"] = 0.0
-        elif 'guidance_scale' in possible:
+        if 'guidance_scale' in possible:
             args['guidance_scale'] = p.cfg_scale
         if 'generator' in possible and generator is not None:
             args['generator'] = generator
@@ -230,18 +221,40 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
         if 'output_type' in possible:
             if hasattr(model, 'vae'):
                 args['output_type'] = 'np' # only set latent if model has vae
+
+        # stable cascade
+        if 'StableCascade' in model.__class__.__name__:
+            kwargs.pop("guidance_scale") # remove
+            kwargs.pop("num_inference_steps") # remove
+            if 'prior_num_inference_steps' in possible:
+                args["prior_num_inference_steps"] = p.steps
+                args["num_inference_steps"] = p.refiner_steps
+            if 'prior_guidance_scale' in possible:
+                args["prior_guidance_scale"] = p.cfg_scale
+            if 'decoder_guidance_scale' in possible:
+                args["decoder_guidance_scale"] = p.image_cfg_scale
+            # TODO Stable Cascade callbacks are currently broken in combined pipeline so preview will not get triggered
+            if 'prior_callback_on_step_end' in possible:
+                possible.remove('callback_on_step_end')
+            if 'callback_on_step_end' in possible:
+                possible.remove('callback_on_step_end')
+            if 'callback' in possible:
+                possible.remove('callback')
+
+        # set callbacks
         if 'callback_steps' in possible:
             args['callback_steps'] = 1
-        if 'callback' in possible:
-            args['callback'] = diffusers_callback_legacy
-        elif 'callback_on_step_end_tensor_inputs' in possible:
-            if hasattr(model, "decoder") and hasattr(model, "prior_prior") and 'prior_guidance_scale' in possible:
-                args['prior_callback_on_step_end'] = diffusers_callback
+        if 'callback_on_step_end' in possible:
             args['callback_on_step_end'] = diffusers_callback
-            if 'prompt_embeds' in possible and 'negative_prompt_embeds' in possible and hasattr(model, '_callback_tensor_inputs'):
-                args['callback_on_step_end_tensor_inputs'] = model._callback_tensor_inputs # pylint: disable=protected-access
-            else:
-                args['callback_on_step_end_tensor_inputs'] = ['latents']
+            if 'callback_on_step_end_tensor_inputs' in possible:
+                if 'prompt_embeds' in possible and 'negative_prompt_embeds' in possible and hasattr(model, '_callback_tensor_inputs'):
+                    args['callback_on_step_end_tensor_inputs'] = model._callback_tensor_inputs # pylint: disable=protected-access
+                else:
+                    args['callback_on_step_end_tensor_inputs'] = ['latents']
+        elif 'callback' in possible:
+            args['callback'] = diffusers_callback_legacy
+
+        # handle remaining args
         for arg in kwargs:
             if arg in possible: # add kwargs
                 args[arg] = kwargs[arg]
@@ -261,6 +274,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 debug(f'Diffusers unknown task args: {k}={v}')
 
         sd_hijack_hypertile.hypertile_set(p, hr=len(getattr(p, 'init_images', [])) > 0)
+
+        # debug info
         clean = args.copy()
         clean.pop('callback', None)
         clean.pop('callback_steps', None)
@@ -284,8 +299,6 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
             txt += f' Clamp threshold={p.hdr_threshold} boundary={p.hdr_boundary}' if p.hdr_clamp else ' Clamp off'
             txt += f' Maximize boundary={p.hdr_max_boundry} center={p.hdr_max_center}' if p.hdr_maximize else ' Maximize off'
             shared.log.debug(txt)
-        # components = [{ k: getattr(v, 'device', None) } for k, v in model.components.items()]
-        # shared.log.debug(f'Diffuser pipeline components: {components}')
         if shared.cmd_opts.profile:
             t1 = time.time()
             shared.log.debug(f'Profile: pipeline args: {t1-t0:.2f}')

modules/sd_hijack_accelerate.py

@@ -58,3 +58,9 @@ def hijack_accelerate():
 
 def restore_accelerate():
     accelerate.utils.set_module_tensor_to_device = orig_method
+
+
+def hijack_hfhub():
+    import contextlib
+    import huggingface_hub.file_download
+    huggingface_hub.file_download.FileLock = contextlib.nullcontext

modules/sd_models.py

@@ -20,7 +20,7 @@ from omegaconf import OmegaConf
 import tomesd
 from transformers import logging as transformers_logging
 from ldm.util import instantiate_from_config
-from modules import paths, shared, shared_items, shared_state, modelloader, devices, script_callbacks, sd_vae, errors, hashes, sd_models_config, sd_models_compile
+from modules import paths, shared, shared_items, shared_state, modelloader, devices, script_callbacks, sd_vae, errors, hashes, sd_models_config, sd_models_compile, sd_hijack_accelerate
 from modules.timer import Timer
 from modules.memstats import memory_stats
 from modules.paths import models_path, script_path
@@ -123,6 +123,7 @@ class NoWatermark:
 
 def setup_model():
     list_models()
+    sd_hijack_accelerate.hijack_hfhub()
     if shared.backend == shared.Backend.ORIGINAL:
         enable_midas_autodownload()
 
@@ -189,6 +190,11 @@ def update_model_hashes():
 
 
 def get_closet_checkpoint_match(search_string):
+    if search_string.startswith('huggingface/'):
+        model_name = search_string.replace('huggingface/', '')
+        checkpoint_info = CheckpointInfo(model_name) # create a virutal model info
+        checkpoint_info.type = 'huggingface'
+        return checkpoint_info
     checkpoint_info = checkpoint_aliases.get(search_string, None)
     if checkpoint_info is not None:
         return checkpoint_info
@@ -813,23 +819,31 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
 
         shared.log.debug(f'Diffusers loading: path="{checkpoint_info.path}"')
         pipeline, model_type = detect_pipeline(checkpoint_info.path, op)
-        if os.path.isdir(checkpoint_info.path):
+        if os.path.isdir(checkpoint_info.path) or checkpoint_info.type == 'huggingface':
             files = shared.walk_files(checkpoint_info.path, ['.safetensors', '.bin', '.ckpt'])
             if 'variant' not in diffusers_load_config and any('diffusion_pytorch_model.fp16' in f for f in files): # deal with diffusers lack of variant fallback when loading
                 diffusers_load_config['variant'] = 'fp16'
             if model_type in ['Stable Cascade']: # forced pipeline
+                # TODO experimental stable cascade
                 try:
-                    # TODO: Add decoder and vqgan loader
-                    # vqgan can reuse the vae loader ui
-                    # decoder needs a new loader
-                    vae_file = diffusers_load_config.pop("vae", None)
-                    prior_model = "stabilityai/stable-cascade-prior" if "models--stabilityai--stable-cascade" in checkpoint_info.path and "models--stabilityai--stable-cascade-prior" not in checkpoint_info.path else checkpoint_info.path
-                    decoder_model = "stabilityai/stable-cascade" if vae_file is None else vae_file
-                    prior = diffusers.StableCascadePriorPipeline.from_pretrained(prior_model, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) # pylint: disable=no-member
-                    decoder = diffusers.StableCascadeDecoderPipeline.from_pretrained(decoder_model, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config) # pylint: disable=no-member
-                    sd_model = diffusers.StableCascadeCombinedPipeline(tokenizer=decoder.tokenizer, text_encoder=decoder.text_encoder, decoder=decoder.decoder, scheduler=decoder.scheduler, vqgan=decoder.vqgan, prior_prior=prior.prior, prior_scheduler=prior.scheduler, feature_extractor=prior.feature_extractor, image_encoder=prior.image_encoder) # pylint: disable=no-member
+                    diffusers_load_config.pop("vae", None)
+                    diffusers_load_config.pop("variant", None)
+                    decoder = diffusers.StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", cache_dir=shared.opts.diffusers_dir, revision="refs/pr/17", **diffusers_load_config)
+                    prior = diffusers.StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
+                    sd_model = diffusers.StableCascadeCombinedPipeline(
+                        tokenizer=decoder.tokenizer,
+                        text_encoder=decoder.text_encoder,
+                        decoder=decoder.decoder,
+                        scheduler=decoder.scheduler,
+                        vqgan=decoder.vqgan,
+                        prior_prior=prior.prior,
+                        prior_scheduler=prior.scheduler,
+                        feature_extractor=prior.feature_extractor,
+                        image_encoder=prior.image_encoder)
                 except Exception as e:
                     shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}')
+                    if debug_load:
+                        errors.display(e, 'Load')
                     return
             elif model_type in ['InstaFlow']: # forced pipeline
                 try:
@@ -837,6 +851,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                     sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
                 except Exception as e:
                     shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}')
+                    if debug_load:
+                        errors.display(e, 'Load')
                     return
             elif model_type in ['SegMoE']: # forced pipeline
                 try:
@@ -845,6 +861,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                     sd_model = sd_model.pipe # segmoe pipe does its stuff in __init__ and __call__ is the original pipeline
                 except Exception as e:
                     shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}')
+                    if debug_load:
+                        errors.display(e, 'Load')
                     return
             elif 'ONNX' in model_type: # forced pipeline
                 sd_model = pipeline.from_pretrained(checkpoint_info.path)
@@ -913,11 +931,10 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                     if shared.opts.disable_accelerate:
                         from diffusers.utils import import_utils
                         import_utils._accelerate_available = False # pylint: disable=protected-access
-                    import modules.sd_hijack_accelerate
                     if shared.opts.diffusers_to_gpu:
-                        modules.sd_hijack_accelerate.hijack_accelerate()
+                        sd_hijack_accelerate.hijack_accelerate()
                     else:
-                        modules.sd_hijack_accelerate.restore_accelerate()
+                        sd_hijack_accelerate.restore_accelerate()
                     sd_model = pipeline.from_single_file(checkpoint_info.path, **diffusers_load_config)
                     if sd_model is not None and hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config') and 'inpainting' in checkpoint_info.path.lower():
                         shared.log.debug('Model patch: type=inpaint')

modules/sd_samplers_common.py

@@ -3,7 +3,7 @@ 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_cascade_previewer, sd_samplers
+from modules import shared, devices, processing, images, sd_vae_approx, sd_vae_taesd, sd_vae_stablecascade, sd_samplers
 
 
 SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
@@ -48,6 +48,7 @@ def single_sample_to_image(sample, approximation=None):
             sd_cascade = True
         if len(sample.shape) == 4 and sample.shape[0]: # likely animatediff latent
             sample = sample.permute(1, 0, 2, 3)[0]
+
         if shared.backend == shared.Backend.DIFFUSERS: # [-x,x] to [-5,5]
             sample_max = torch.max(sample)
             if sample_max > 5:
@@ -56,7 +57,7 @@ def single_sample_to_image(sample, approximation=None):
             if sample_min < -5:
                 sample = sample * (5 / abs(sample_min))
         if sd_cascade:
-            x_sample = sd_cascade_previewer.decode(sample)
+            x_sample = sd_vae_stablecascade.decode(sample)
         elif approximation == 0: # Simple
             x_sample = sd_vae_approx.cheap_approximation(sample) * 0.5 + 0.5
         elif approximation == 1: # Approximate

modules/sd_vae_stablecascade.py

@@ -1,5 +1,4 @@
 import os
-import torch
 from torch import nn
 import safetensors
 from modules import devices, paths
@@ -69,7 +68,7 @@ def load_model(model_path):
 
 def decode(latents):
     from modules import shared
-    global preview_model
+    global preview_model # pylint: disable=global-statement
     if preview_model is None:
         model_path = os.path.join(paths.models_path, "VAE-approx", "sd_cascade_previewer.safetensors")
         download_model(model_path)

modules/sd_vae_taesd.py

@@ -77,7 +77,6 @@ def download_model(model_path):
     model_name = os.path.basename(model_path)
     model_url = f'https://github.com/madebyollin/taesd/raw/main/{model_name}'
     if not os.path.exists(model_path):
-        import torch
         from modules.shared import log
         os.makedirs(os.path.dirname(model_path), exist_ok=True)
         log.info(f'Downloading TAESD decoder: {model_path}')
@@ -109,6 +108,7 @@ def decode(latents):
     model_class = shared.sd_model_type
     if model_class == 'ldm':
         model_class = 'sd'
+    dtype = devices.dtype_vae if devices.dtype_vae != torch.bfloat16 else torch.float16 # taesd does not support bf16
     if 'sd' not in model_class:
         shared.log.warning(f'TAESD unsupported model type: {model_class}')
         return Image.new('RGB', (8, 8), color = (0, 0, 0))
@@ -120,10 +120,10 @@ def decode(latents):
             taesd_models[f'{model_class}-decoder'] = TAESD(decoder_path=model_path, encoder_path=None)
             shared.log.debug(f'VAE load: type=taesd model={model_path}')
             vae = taesd_models[f'{model_class}-decoder']
-            vae.decoder.to(devices.device, devices.dtype_vae)
+            vae.decoder.to(devices.device, dtype)
     try:
         with devices.inference_context():
-            latents = latents.detach().clone().to(devices.device, devices.dtype_vae)
+            latents = latents.detach().clone().to(devices.device, dtype)
             if len(latents.shape) == 3:
                 latents = latents.unsqueeze(0)
                 image = vae.decoder(latents).clamp(0, 1).detach()

modules/ui_sections.py

@@ -206,7 +206,7 @@ def create_hires_inputs(tab):
         with gr.Group(visible=shared.backend == shared.Backend.DIFFUSERS):
             with gr.Row(elem_id=f"{tab}_refiner_row1", variant="compact"):
                 refiner_start = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Refiner start', value=0.8, elem_id=f"{tab}_refiner_start")
-                refiner_steps = gr.Slider(minimum=0, maximum=99, step=1, label="Refiner steps", elem_id=f"{tab}_refiner_steps", value=5)
+                refiner_steps = gr.Slider(minimum=0, maximum=99, step=1, label="Refiner steps", elem_id=f"{tab}_refiner_steps", value=10)
             with gr.Row(elem_id=f"{tab}_refiner_row3", variant="compact"):
                 refiner_prompt = gr.Textbox(value='', label='Secondary prompt', elem_id=f"{tab}_refiner_prompt")
             with gr.Row(elem_id="txt2img_refiner_row4", variant="compact"):