Merge pull request #26 from AUTOMATIC1111/master

sync branches

README.md

@@ -155,6 +155,8 @@ Licenses for borrowed code can be found in `Settings -> Licenses` screen, and al
 - Idea for Composable Diffusion - https://github.com/energy-based-model/Compositional-Visual-Generation-with-Composable-Diffusion-Models-PyTorch
 - xformers - https://github.com/facebookresearch/xformers
 - DeepDanbooru - interrogator for anime diffusers https://github.com/KichangKim/DeepDanbooru
+- Sampling in float32 precision from a float16 UNet - marunine for the idea, Birch-san for the example Diffusers implementation (https://github.com/Birch-san/diffusers-play/tree/92feee6)
+- Instruct pix2pix - Tim Brooks (star), Aleksander Holynski (star), Alexei A. Efros (no star) - https://github.com/timothybrooks/instruct-pix2pix
 - Security advice - RyotaK
 - Initial Gradio script - posted on 4chan by an Anonymous user. Thank you Anonymous user.
 - (You)

extensions-builtin/Lora/lora.py

@@ -166,7 +166,10 @@ def lora_forward(module, input, res):
     for lora in loaded_loras:
         module = lora.modules.get(lora_layer_name, None)
         if module is not None:
-            res = res + module.up(module.down(input)) * lora.multiplier * (module.alpha / module.up.weight.shape[1] if module.alpha else 1.0)
+            if shared.opts.lora_apply_to_outputs and res.shape == input.shape:
+                res = res + module.up(module.down(res)) * lora.multiplier * (module.alpha / module.up.weight.shape[1] if module.alpha else 1.0)
+            else:
+                res = res + module.up(module.down(input)) * lora.multiplier * (module.alpha / module.up.weight.shape[1] if module.alpha else 1.0)
 
     return res
 

extensions-builtin/Lora/scripts/lora_script.py

@@ -3,7 +3,7 @@ import torch
 import lora
 import extra_networks_lora
 import ui_extra_networks_lora
-from modules import script_callbacks, ui_extra_networks, extra_networks
+from modules import script_callbacks, ui_extra_networks, extra_networks, shared
 
 
 def unload():
@@ -28,3 +28,8 @@ torch.nn.Conv2d.forward = lora.lora_Conv2d_forward
 script_callbacks.on_model_loaded(lora.assign_lora_names_to_compvis_modules)
 script_callbacks.on_script_unloaded(unload)
 script_callbacks.on_before_ui(before_ui)
+
+
+shared.options_templates.update(shared.options_section(('extra_networks', "Extra Networks"), {
+    "lora_apply_to_outputs": shared.OptionInfo(False, "Apply Lora to outputs rather than inputs when possible (experimental)"),
+}))

launch.py

@@ -48,10 +48,19 @@ def extract_opt(args, name):
     return args, is_present, opt
 
 
-def run(command, desc=None, errdesc=None, custom_env=None):
+def run(command, desc=None, errdesc=None, custom_env=None, live=False):
     if desc is not None:
         print(desc)
 
+    if live:
+        result = subprocess.run(command, shell=True, env=os.environ if custom_env is None else custom_env)
+        if result.returncode != 0:
+            raise RuntimeError(f"""{errdesc or 'Error running command'}.
+Command: {command}
+Error code: {result.returncode}""")
+
+        return ""
+
     result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, env=os.environ if custom_env is None else custom_env)
 
     if result.returncode != 0:
@@ -108,18 +117,18 @@ def git_clone(url, dir, name, commithash=None):
         if commithash is None:
             return
 
-        current_hash = run(f'"{git}" -C {dir} rev-parse HEAD', None, f"Couldn't determine {name}'s hash: {commithash}").strip()
+        current_hash = run(f'"{git}" -C "{dir}" rev-parse HEAD', None, f"Couldn't determine {name}'s hash: {commithash}").strip()
         if current_hash == commithash:
             return
 
-        run(f'"{git}" -C {dir} fetch', f"Fetching updates for {name}...", f"Couldn't fetch {name}")
-        run(f'"{git}" -C {dir} checkout {commithash}', f"Checking out commit for {name} with hash: {commithash}...", f"Couldn't checkout commit {commithash} for {name}")
+        run(f'"{git}" -C "{dir}" fetch', f"Fetching updates for {name}...", f"Couldn't fetch {name}")
+        run(f'"{git}" -C "{dir}" checkout {commithash}', f"Checking out commit for {name} with hash: {commithash}...", f"Couldn't checkout commit {commithash} for {name}")
         return
 
     run(f'"{git}" clone "{url}" "{dir}"', f"Cloning {name} into {dir}...", f"Couldn't clone {name}")
 
     if commithash is not None:
-        run(f'"{git}" -C {dir} checkout {commithash}', None, "Couldn't checkout {name}'s hash: {commithash}")
+        run(f'"{git}" -C "{dir}" checkout {commithash}', None, "Couldn't checkout {name}'s hash: {commithash}")
 
         
 def version_check(commit):
@@ -219,9 +228,9 @@ def prepare_environment():
 
     print(f"Python {sys.version}")
     print(f"Commit hash: {commit}")
-    
+
     if reinstall_torch or not is_installed("torch") or not is_installed("torchvision"):
-        run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch")
+        run(f'"{python}" -m {torch_command}', "Installing torch and torchvision", "Couldn't install torch", live=True)
 
     if not skip_torch_cuda_test:
         run_python("import torch; assert torch.cuda.is_available(), 'Torch is not able to use GPU; add --skip-torch-cuda-test to COMMANDLINE_ARGS variable to disable this check'")

modules/deepbooru_model.py

@@ -2,6 +2,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from modules import devices
+
 # see https://github.com/AUTOMATIC1111/TorchDeepDanbooru for more
 
 
@@ -196,7 +198,7 @@ class DeepDanbooruModel(nn.Module):
         t_358, = inputs
         t_359 = t_358.permute(*[0, 3, 1, 2])
         t_359_padded = F.pad(t_359, [2, 3, 2, 3], value=0)
-        t_360 = self.n_Conv_0(t_359_padded)
+        t_360 = self.n_Conv_0(t_359_padded.to(self.n_Conv_0.bias.dtype) if devices.unet_needs_upcast else t_359_padded)
         t_361 = F.relu(t_360)
         t_361 = F.pad(t_361, [0, 1, 0, 1], value=float('-inf'))
         t_362 = self.n_MaxPool_0(t_361)

modules/devices.py

@@ -79,6 +79,8 @@ cpu = torch.device("cpu")
 device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = None
 dtype = torch.float16
 dtype_vae = torch.float16
+dtype_unet = torch.float16
+unet_needs_upcast = False
 
 
 def randn(seed, shape):
@@ -106,6 +108,10 @@ def autocast(disable=False):
     return torch.autocast("cuda")
 
 
+def without_autocast(disable=False):
+    return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
+
+
 class NansException(Exception):
     pass
 
@@ -123,7 +129,7 @@ def test_for_nans(x, where):
         message = "A tensor with all NaNs was produced in Unet."
 
         if not shared.cmd_opts.no_half:
-            message += " This could be either because there's not enough precision to represent the picture, or because your video card does not support half type. Try using --no-half commandline argument to fix this."
+            message += " This could be either because there's not enough precision to represent the picture, or because your video card does not support half type. Try setting the \"Upcast cross attention layer to float32\" option in Settings > Stable Diffusion or using the --no-half commandline argument to fix this."
 
     elif where == "vae":
         message = "A tensor with all NaNs was produced in VAE."

modules/images.py

@@ -195,7 +195,7 @@ def draw_grid_annotations(im, width, height, hor_texts, ver_texts):
     ver_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing * len(lines) for lines in
                         ver_texts]
 
-    pad_top = max(hor_text_heights) + line_spacing * 2
+    pad_top = 0 if sum(hor_text_heights) == 0 else max(hor_text_heights) + line_spacing * 2
 
     result = Image.new("RGB", (im.width + pad_left, im.height + pad_top), "white")
     result.paste(im, (pad_left, pad_top))

modules/interrogate.py

@@ -82,9 +82,16 @@ class InterrogateModels:
 
         return self.loaded_categories
 
+    def create_fake_fairscale(self):
+        class FakeFairscale:
+            def checkpoint_wrapper(self):
+                pass
+
+        sys.modules["fairscale.nn.checkpoint.checkpoint_activations"] = FakeFairscale
+
     def load_blip_model(self):
-        with paths.Prioritize("BLIP"):
-            import models.blip
+        self.create_fake_fairscale()
+        import models.blip
 
         files = modelloader.load_models(
             model_path=os.path.join(paths.models_path, "BLIP"),

modules/processing.py

@@ -13,7 +13,7 @@ from skimage import exposure
 from typing import Any, Dict, List, Optional
 
 import modules.sd_hijack
-from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, script_callbacks, extra_networks
+from modules import devices, prompt_parser, masking, sd_samplers, lowvram, generation_parameters_copypaste, script_callbacks, extra_networks, sd_vae_approx
 from modules.sd_hijack import model_hijack
 from modules.shared import opts, cmd_opts, state
 import modules.shared as shared
@@ -172,7 +172,8 @@ class StableDiffusionProcessing:
         midas_in = torch.from_numpy(transformed["midas_in"][None, ...]).to(device=shared.device)
         midas_in = repeat(midas_in, "1 ... -> n ...", n=self.batch_size)
 
-        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(source_image))
+        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(source_image.to(devices.dtype_vae) if devices.unet_needs_upcast else source_image))
+        conditioning_image = conditioning_image.float() if devices.unet_needs_upcast else conditioning_image
         conditioning = torch.nn.functional.interpolate(
             self.sd_model.depth_model(midas_in),
             size=conditioning_image.shape[2:],
@@ -184,7 +185,12 @@ class StableDiffusionProcessing:
         conditioning = 2. * (conditioning - depth_min) / (depth_max - depth_min) - 1.
         return conditioning
 
-    def inpainting_image_conditioning(self, source_image, latent_image, image_mask = None):
+    def edit_image_conditioning(self, source_image):
+        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(source_image))
+
+        return conditioning_image
+
+    def inpainting_image_conditioning(self, source_image, latent_image, image_mask=None):
         self.is_using_inpainting_conditioning = True
 
         # Handle the different mask inputs
@@ -203,7 +209,7 @@ class StableDiffusionProcessing:
 
         # Create another latent image, this time with a masked version of the original input.
         # Smoothly interpolate between the masked and unmasked latent conditioning image using a parameter.
-        conditioning_mask = conditioning_mask.to(source_image.device).to(source_image.dtype)
+        conditioning_mask = conditioning_mask.to(device=source_image.device, dtype=source_image.dtype)
         conditioning_image = torch.lerp(
             source_image,
             source_image * (1.0 - conditioning_mask),
@@ -211,7 +217,7 @@ class StableDiffusionProcessing:
         )
 
         # Encode the new masked image using first stage of network.
-        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(conditioning_image))
+        conditioning_image = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(conditioning_image.to(devices.dtype_vae) if devices.unet_needs_upcast else conditioning_image))
 
         # Create the concatenated conditioning tensor to be fed to `c_concat`
         conditioning_mask = torch.nn.functional.interpolate(conditioning_mask, size=latent_image.shape[-2:])
@@ -225,10 +231,13 @@ class StableDiffusionProcessing:
         # HACK: Using introspection as the Depth2Image model doesn't appear to uniquely
         # identify itself with a field common to all models. The conditioning_key is also hybrid.
         if isinstance(self.sd_model, LatentDepth2ImageDiffusion):
-            return self.depth2img_image_conditioning(source_image)
+            return self.depth2img_image_conditioning(source_image.float() if devices.unet_needs_upcast else source_image)
+
+        if self.sd_model.cond_stage_key == "edit":
+            return self.edit_image_conditioning(source_image)
 
         if self.sampler.conditioning_key in {'hybrid', 'concat'}:
-            return self.inpainting_image_conditioning(source_image, latent_image, image_mask=image_mask)
+            return self.inpainting_image_conditioning(source_image.float() if devices.unet_needs_upcast else source_image, latent_image, image_mask=image_mask)
 
         # Dummy zero conditioning if we're not using inpainting or depth model.
         return latent_image.new_zeros(latent_image.shape[0], 5, 1, 1)
@@ -408,7 +417,7 @@ def create_random_tensors(shape, seeds, subseeds=None, subseed_strength=0.0, see
 
 def decode_first_stage(model, x):
     with devices.autocast(disable=x.dtype == devices.dtype_vae):
-        x = model.decode_first_stage(x)
+        x = model.decode_first_stage(x.to(devices.dtype_vae) if devices.unet_needs_upcast else x)
 
     return x
 
@@ -568,6 +577,10 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         with devices.autocast():
             p.init(p.all_prompts, p.all_seeds, p.all_subseeds)
 
+            # for OSX, loading the model during sampling changes the generated picture, so it is loaded here
+            if shared.opts.live_previews_enable and opts.show_progress_type == "Approx NN":
+                sd_vae_approx.model()
+
             if not p.disable_extra_networks:
                 extra_networks.activate(p, extra_network_data)
 
@@ -610,7 +623,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             if p.n_iter > 1:
                 shared.state.job = f"Batch {n+1} out of {p.n_iter}"
 
-            with devices.autocast():
+            with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast():
                 samples_ddim = p.sample(conditioning=c, unconditional_conditioning=uc, seeds=seeds, subseeds=subseeds, subseed_strength=p.subseed_strength, prompts=prompts)
 
             x_samples_ddim = [decode_first_stage(p.sd_model, samples_ddim[i:i+1].to(dtype=devices.dtype_vae))[0].cpu() for i in range(samples_ddim.size(0))]
@@ -988,7 +1001,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
 
         image = torch.from_numpy(batch_images)
         image = 2. * image - 1.
-        image = image.to(shared.device)
+        image = image.to(device=shared.device, dtype=devices.dtype_vae if devices.unet_needs_upcast else None)
 
         self.init_latent = self.sd_model.get_first_stage_encoding(self.sd_model.encode_first_stage(image))
 

modules/sd_hijack_ip2p.py

@@ -0,0 +1,13 @@
+import collections
+import os.path
+import sys
+import gc
+import time
+
+def should_hijack_ip2p(checkpoint_info):
+    from modules import sd_models
+
+    ckpt_basename = os.path.basename(checkpoint_info.filename).lower()
+    cfg_basename = os.path.basename(sd_models.find_checkpoint_config(checkpoint_info)).lower()
+
+    return "pix2pix" in ckpt_basename and not "pix2pix" in cfg_basename

modules/sd_hijack_optimizations.py

@@ -9,7 +9,7 @@ from torch import einsum
 from ldm.util import default
 from einops import rearrange
 
-from modules import shared, errors
+from modules import shared, errors, devices
 from modules.hypernetworks import hypernetwork
 
 from .sub_quadratic_attention import efficient_dot_product_attention
@@ -52,18 +52,25 @@ def split_cross_attention_forward_v1(self, x, context=None, mask=None):
     q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
     del q_in, k_in, v_in
 
-    r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)
-    for i in range(0, q.shape[0], 2):
-        end = i + 2
-        s1 = einsum('b i d, b j d -> b i j', q[i:end], k[i:end])
-        s1 *= self.scale
+    dtype = q.dtype
+    if shared.opts.upcast_attn:
+        q, k, v = q.float(), k.float(), v.float()
 
-        s2 = s1.softmax(dim=-1)
-        del s1
+    with devices.without_autocast(disable=not shared.opts.upcast_attn):
+        r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
+        for i in range(0, q.shape[0], 2):
+            end = i + 2
+            s1 = einsum('b i d, b j d -> b i j', q[i:end], k[i:end])
+            s1 *= self.scale
+    
+            s2 = s1.softmax(dim=-1)
+            del s1
+    
+            r1[i:end] = einsum('b i j, b j d -> b i d', s2, v[i:end])
+            del s2
+        del q, k, v
 
-        r1[i:end] = einsum('b i j, b j d -> b i d', s2, v[i:end])
-        del s2
-    del q, k, v
+    r1 = r1.to(dtype)
 
     r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
     del r1
@@ -82,45 +89,52 @@ def split_cross_attention_forward(self, x, context=None, mask=None):
     k_in = self.to_k(context_k)
     v_in = self.to_v(context_v)
 
-    k_in *= self.scale
+    dtype = q_in.dtype
+    if shared.opts.upcast_attn:
+        q_in, k_in, v_in = q_in.float(), k_in.float(), v_in if v_in.device.type == 'mps' else v_in.float()
 
-    del context, x
+    with devices.without_autocast(disable=not shared.opts.upcast_attn):
+        k_in = k_in * self.scale
+    
+        del context, x
+    
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
+        del q_in, k_in, v_in
+    
+        r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
+    
+        mem_free_total = get_available_vram()
+    
+        gb = 1024 ** 3
+        tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()
+        modifier = 3 if q.element_size() == 2 else 2.5
+        mem_required = tensor_size * modifier
+        steps = 1
+    
+        if mem_required > mem_free_total:
+            steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2)))
+            # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "
+            #       f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")
+    
+        if steps > 64:
+            max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
+            raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
+                               f'Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free')
+    
+        slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
+        for i in range(0, q.shape[1], slice_size):
+            end = i + slice_size
+            s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k)
+    
+            s2 = s1.softmax(dim=-1, dtype=q.dtype)
+            del s1
+    
+            r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
+            del s2
+    
+        del q, k, v
 
-    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
-    del q_in, k_in, v_in
-
-    r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
-
-    mem_free_total = get_available_vram()
-
-    gb = 1024 ** 3
-    tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()
-    modifier = 3 if q.element_size() == 2 else 2.5
-    mem_required = tensor_size * modifier
-    steps = 1
-
-    if mem_required > mem_free_total:
-        steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2)))
-        # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "
-        #       f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")
-
-    if steps > 64:
-        max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
-        raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
-                           f'Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free')
-
-    slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
-    for i in range(0, q.shape[1], slice_size):
-        end = i + slice_size
-        s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k)
-
-        s2 = s1.softmax(dim=-1, dtype=q.dtype)
-        del s1
-
-        r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
-        del s2
-
-    del q, k, v
+    r1 = r1.to(dtype)
 
     r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
     del r1
@@ -204,12 +218,20 @@ def split_cross_attention_forward_invokeAI(self, x, context=None, mask=None):
     context = default(context, x)
 
     context_k, context_v = hypernetwork.apply_hypernetworks(shared.loaded_hypernetworks, context)
-    k = self.to_k(context_k) * self.scale
+    k = self.to_k(context_k)
     v = self.to_v(context_v)
     del context, context_k, context_v, x
 
-    q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
-    r = einsum_op(q, k, v)
+    dtype = q.dtype
+    if shared.opts.upcast_attn:
+        q, k, v = q.float(), k.float(), v if v.device.type == 'mps' else v.float()
+
+    with devices.without_autocast(disable=not shared.opts.upcast_attn):
+        k = k * self.scale
+    
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+        r = einsum_op(q, k, v)
+    r = r.to(dtype)
     return self.to_out(rearrange(r, '(b h) n d -> b n (h d)', h=h))
 
 # -- End of code from https://github.com/invoke-ai/InvokeAI --
@@ -234,8 +256,14 @@ def sub_quad_attention_forward(self, x, context=None, mask=None):
     k = k.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)
     v = v.unflatten(-1, (h, -1)).transpose(1,2).flatten(end_dim=1)
 
+    dtype = q.dtype
+    if shared.opts.upcast_attn:
+        q, k = q.float(), k.float()
+
     x = sub_quad_attention(q, k, v, q_chunk_size=shared.cmd_opts.sub_quad_q_chunk_size, kv_chunk_size=shared.cmd_opts.sub_quad_kv_chunk_size, chunk_threshold=shared.cmd_opts.sub_quad_chunk_threshold, use_checkpoint=self.training)
 
+    x = x.to(dtype)
+
     x = x.unflatten(0, (-1, h)).transpose(1,2).flatten(start_dim=2)
 
     out_proj, dropout = self.to_out
@@ -268,15 +296,16 @@ def sub_quad_attention(q, k, v, q_chunk_size=1024, kv_chunk_size=None, kv_chunk_
         query_chunk_size = q_tokens
         kv_chunk_size = k_tokens
 
-    return efficient_dot_product_attention(
-        q,
-        k,
-        v,
-        query_chunk_size=q_chunk_size,
-        kv_chunk_size=kv_chunk_size,
-        kv_chunk_size_min = kv_chunk_size_min,
-        use_checkpoint=use_checkpoint,
-    )
+    with devices.without_autocast(disable=q.dtype == v.dtype):
+        return efficient_dot_product_attention(
+            q,
+            k,
+            v,
+            query_chunk_size=q_chunk_size,
+            kv_chunk_size=kv_chunk_size,
+            kv_chunk_size_min = kv_chunk_size_min,
+            use_checkpoint=use_checkpoint,
+        )
 
 
 def get_xformers_flash_attention_op(q, k, v):
@@ -306,8 +335,14 @@ def xformers_attention_forward(self, x, context=None, mask=None):
     q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b n h d', h=h), (q_in, k_in, v_in))
     del q_in, k_in, v_in
 
+    dtype = q.dtype
+    if shared.opts.upcast_attn:
+        q, k = q.float(), k.float()
+
     out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=get_xformers_flash_attention_op(q, k, v))
 
+    out = out.to(dtype)
+
     out = rearrange(out, 'b n h d -> b n (h d)', h=h)
     return self.to_out(out)
 
@@ -378,10 +413,14 @@ def xformers_attnblock_forward(self, x):
         v = self.v(h_)
         b, c, h, w = q.shape
         q, k, v = map(lambda t: rearrange(t, 'b c h w -> b (h w) c'), (q, k, v))
+        dtype = q.dtype
+        if shared.opts.upcast_attn:
+            q, k = q.float(), k.float()
         q = q.contiguous()
         k = k.contiguous()
         v = v.contiguous()
         out = xformers.ops.memory_efficient_attention(q, k, v, op=get_xformers_flash_attention_op(q, k, v))
+        out = out.to(dtype)
         out = rearrange(out, 'b (h w) c -> b c h w', h=h)
         out = self.proj_out(out)
         return x + out

modules/sd_hijack_unet.py

@@ -1,4 +1,8 @@
 import torch
+from packaging import version
+
+from modules import devices
+from modules.sd_hijack_utils import CondFunc
 
 
 class TorchHijackForUnet:
@@ -28,3 +32,31 @@ class TorchHijackForUnet:
 
 
 th = TorchHijackForUnet()
+
+
+# Below are monkey patches to enable upcasting a float16 UNet for float32 sampling
+def apply_model(orig_func, self, x_noisy, t, cond, **kwargs):
+
+    if isinstance(cond, dict):
+        for y in cond.keys():
+            cond[y] = [x.to(devices.dtype_unet) if isinstance(x, torch.Tensor) else x for x in cond[y]]
+
+    with devices.autocast():
+        return orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs).float()
+
+class GELUHijack(torch.nn.GELU, torch.nn.Module):
+    def __init__(self, *args, **kwargs):
+        torch.nn.GELU.__init__(self, *args, **kwargs)
+    def forward(self, x):
+        if devices.unet_needs_upcast:
+            return torch.nn.GELU.forward(self.float(), x.float()).to(devices.dtype_unet)
+        else:
+            return torch.nn.GELU.forward(self, x)
+
+unet_needs_upcast = lambda *args, **kwargs: devices.unet_needs_upcast
+CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model, unet_needs_upcast)
+CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).to(devices.dtype_unet), unet_needs_upcast)
+if version.parse(torch.__version__) <= version.parse("1.13.1"):
+    CondFunc('ldm.modules.diffusionmodules.util.GroupNorm32.forward', lambda orig_func, self, *args, **kwargs: orig_func(self.float(), *args, **kwargs), unet_needs_upcast)
+    CondFunc('ldm.modules.attention.GEGLU.forward', lambda orig_func, self, x: orig_func(self.float(), x.float()).to(devices.dtype_unet), unet_needs_upcast)
+    CondFunc('open_clip.transformer.ResidualAttentionBlock.__init__', lambda orig_func, *args, **kwargs: kwargs.update({'act_layer': GELUHijack}) and False or orig_func(*args, **kwargs), lambda _, *args, **kwargs: kwargs.get('act_layer') is None or kwargs['act_layer'] == torch.nn.GELU)

modules/sd_hijack_utils.py

@@ -0,0 +1,28 @@
+import importlib
+
+class CondFunc:
+    def __new__(cls, orig_func, sub_func, cond_func):
+        self = super(CondFunc, cls).__new__(cls)
+        if isinstance(orig_func, str):
+            func_path = orig_func.split('.')
+            for i in range(len(func_path)-1, -1, -1):
+                try:
+                    resolved_obj = importlib.import_module('.'.join(func_path[:i]))
+                    break
+                except ImportError:
+                    pass
+            for attr_name in func_path[i:-1]:
+                resolved_obj = getattr(resolved_obj, attr_name)
+            orig_func = getattr(resolved_obj, func_path[-1])
+            setattr(resolved_obj, func_path[-1], lambda *args, **kwargs: self(*args, **kwargs))
+        self.__init__(orig_func, sub_func, cond_func)
+        return lambda *args, **kwargs: self(*args, **kwargs)
+    def __init__(self, orig_func, sub_func, cond_func):
+        self.__orig_func = orig_func
+        self.__sub_func = sub_func
+        self.__cond_func = cond_func
+    def __call__(self, *args, **kwargs):
+        if not self.__cond_func or self.__cond_func(self.__orig_func, *args, **kwargs):
+            return self.__sub_func(self.__orig_func, *args, **kwargs)
+        else:
+            return self.__orig_func(*args, **kwargs)

modules/sd_models.py

@@ -17,6 +17,7 @@ from ldm.util import instantiate_from_config
 from modules import shared, modelloader, devices, script_callbacks, sd_vae, sd_disable_initialization, errors, hashes
 from modules.paths import models_path
 from modules.sd_hijack_inpainting import do_inpainting_hijack, should_hijack_inpainting
+from modules.sd_hijack_ip2p import should_hijack_ip2p
 
 model_dir = "Stable-diffusion"
 model_path = os.path.abspath(os.path.join(models_path, model_dir))
@@ -257,16 +258,24 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo):
 
         if not shared.cmd_opts.no_half:
             vae = model.first_stage_model
+            depth_model = getattr(model, 'depth_model', None)
 
             # with --no-half-vae, remove VAE from model when doing half() to prevent its weights from being converted to float16
             if shared.cmd_opts.no_half_vae:
                 model.first_stage_model = None
+            # with --upcast-sampling, don't convert the depth model weights to float16
+            if shared.cmd_opts.upcast_sampling and depth_model:
+                model.depth_model = None
 
             model.half()
             model.first_stage_model = vae
+            if depth_model:
+                model.depth_model = depth_model
 
         devices.dtype = torch.float32 if shared.cmd_opts.no_half else torch.float16
         devices.dtype_vae = torch.float32 if shared.cmd_opts.no_half or shared.cmd_opts.no_half_vae else torch.float16
+        devices.dtype_unet = model.model.diffusion_model.dtype
+        devices.unet_needs_upcast = shared.cmd_opts.upcast_sampling and devices.dtype == torch.float16 and devices.dtype_unet == torch.float16
 
         model.first_stage_model.to(devices.dtype_vae)
 
@@ -365,6 +374,15 @@ def load_model(checkpoint_info=None):
         sd_config.model.params.unet_config.params.in_channels = 9
         sd_config.model.params.finetune_keys = None
 
+    if should_hijack_ip2p(checkpoint_info):
+        sd_config.model.target = "modules.models.diffusion.ddpm_edit.LatentDiffusion"
+        sd_config.model.params.conditioning_key = "hybrid"
+        sd_config.model.params.first_stage_key = "edited"
+        sd_config.model.params.cond_stage_key = "edit"
+        sd_config.model.params.image_size = 16
+        sd_config.model.params.unet_config.params.in_channels = 8
+        sd_config.model.params.unet_config.params.out_channels = 4
+
     if not hasattr(sd_config.model.params, "use_ema"):
         sd_config.model.params.use_ema = False
 
@@ -372,6 +390,8 @@ def load_model(checkpoint_info=None):
 
     if shared.cmd_opts.no_half:
         sd_config.model.params.unet_config.params.use_fp16 = False
+    elif shared.cmd_opts.upcast_sampling:
+        sd_config.model.params.unet_config.params.use_fp16 = True
 
     timer = Timer()
 
@@ -447,7 +467,7 @@ def reload_model_weights(sd_model=None, info=None):
 
     checkpoint_config = find_checkpoint_config(current_checkpoint_info)
 
-    if current_checkpoint_info is None or checkpoint_config != find_checkpoint_config(checkpoint_info) or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info):
+    if current_checkpoint_info is None or checkpoint_config != find_checkpoint_config(checkpoint_info) or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info) or should_hijack_ip2p(checkpoint_info) != should_hijack_ip2p(sd_model.sd_checkpoint_info):
         del sd_model
         checkpoints_loaded.clear()
         load_model(checkpoint_info)

modules/sd_samplers.py

@@ -454,7 +454,7 @@ class KDiffusionSampler:
     def initialize(self, p):
         self.model_wrap_cfg.mask = p.mask if hasattr(p, 'mask') else None
         self.model_wrap_cfg.nmask = p.nmask if hasattr(p, 'nmask') else None
-        self.model_wrap.step = 0
+        self.model_wrap_cfg.step = 0
         self.eta = p.eta or opts.eta_ancestral
 
         k_diffusion.sampling.torch = TorchHijack(self.sampler_noises if self.sampler_noises is not None else [])

modules/shared.py

@@ -45,6 +45,7 @@ parser.add_argument("--lowram", action='store_true', help="load stable diffusion
 parser.add_argument("--always-batch-cond-uncond", action='store_true', help="disables cond/uncond batching that is enabled to save memory with --medvram or --lowvram")
 parser.add_argument("--unload-gfpgan", action='store_true', help="does not do anything.")
 parser.add_argument("--precision", type=str, help="evaluate at this precision", choices=["full", "autocast"], default="autocast")
+parser.add_argument("--upcast-sampling", action='store_true', help="upcast sampling. No effect with --no-half. Usually produces similar results to --no-half with better performance while using less memory.")
 parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site")
 parser.add_argument("--ngrok", type=str, help="ngrok authtoken, alternative to gradio --share", default=None)
 parser.add_argument("--ngrok-region", type=str, help="The region in which ngrok should start.", default="us")
@@ -410,6 +411,7 @@ options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "comma_padding_backtrack": OptionInfo(20, "Increase coherency by padding from the last comma within n tokens when using more than 75 tokens", gr.Slider, {"minimum": 0, "maximum": 74, "step": 1 }),
     "CLIP_stop_at_last_layers": OptionInfo(1, "Clip skip", gr.Slider, {"minimum": 1, "maximum": 12, "step": 1}),
     "extra_networks_default_multiplier": OptionInfo(1.0, "Multiplier for extra networks", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+    "upcast_attn": OptionInfo(False, "Upcast cross attention layer to float32"),
 }))
 
 options_templates.update(options_section(('compatibility', "Compatibility"), {

modules/sub_quadratic_attention.py

@@ -67,7 +67,7 @@ def _summarize_chunk(
     max_score, _ = torch.max(attn_weights, -1, keepdim=True)
     max_score = max_score.detach()
     exp_weights = torch.exp(attn_weights - max_score)
-    exp_values = torch.bmm(exp_weights, value)
+    exp_values = torch.bmm(exp_weights, value) if query.device.type == 'mps' else torch.bmm(exp_weights, value.to(exp_weights.dtype)).to(value.dtype)
     max_score = max_score.squeeze(-1)
     return AttnChunk(exp_values, exp_weights.sum(dim=-1), max_score)
 
@@ -129,7 +129,7 @@ def _get_attention_scores_no_kv_chunking(
     )
     attn_probs = attn_scores.softmax(dim=-1)
     del attn_scores
-    hidden_states_slice = torch.bmm(attn_probs, value)
+    hidden_states_slice = torch.bmm(attn_probs, value) if query.device.type == 'mps' else torch.bmm(attn_probs, value.to(attn_probs.dtype)).to(value.dtype)
     return hidden_states_slice
 
 

modules/textual_inversion/textual_inversion.py

@@ -194,7 +194,7 @@ class EmbeddingDatabase:
         if not os.path.isdir(embdir.path):
             return
 
-        for root, dirs, fns in os.walk(embdir.path):
+        for root, dirs, fns in os.walk(embdir.path, followlinks=True):
             for fn in fns:
                 try:
                     fullfn = os.path.join(root, fn)

modules/ui_extra_networks.py

@@ -117,8 +117,13 @@ def create_ui(container, button, tabname):
     ui.button_save_preview = gr.Button('Save preview', elem_id=tabname+"_save_preview", visible=False)
     ui.preview_target_filename = gr.Textbox('Preview save filename', elem_id=tabname+"_preview_filename", visible=False)
 
-    button.click(fn=lambda: gr.update(visible=True), inputs=[], outputs=[container])
-    button_close.click(fn=lambda: gr.update(visible=False), inputs=[], outputs=[container])
+    def toggle_visibility(is_visible):
+        is_visible = not is_visible
+        return is_visible, gr.update(visible=is_visible)
+
+    state_visible = gr.State(value=False)
+    button.click(fn=toggle_visibility, inputs=[state_visible], outputs=[state_visible, container])
+    button_close.click(fn=toggle_visibility, inputs=[state_visible], outputs=[state_visible, container])
 
     def refresh():
         res = []

requirements.txt

@@ -1,7 +1,6 @@
 blendmodes
 accelerate
 basicsr
-fairscale==0.4.4
 fonts
 font-roboto
 gfpgan
@@ -17,7 +16,7 @@ pytorch_lightning==1.7.7
 realesrgan
 scikit-image>=0.19
 timm==0.4.12
-transformers==4.19.2
+transformers==4.25.1
 torch
 einops
 jsonmerge

requirements_versions.txt

@@ -1,5 +1,5 @@
 blendmodes==2022
-transformers==4.19.2
+transformers==4.25.1
 accelerate==0.12.0
 basicsr==1.4.2
 gfpgan==1.3.8
@@ -14,7 +14,6 @@ scikit-image==0.19.2
 fonts
 font-roboto
 timm==0.6.7
-fairscale==0.4.9
 piexif==1.1.3
 einops==0.4.1
 jsonmerge==1.8.0

scripts/xyz_grid.py

@@ -205,26 +205,30 @@ axis_options = [
 ]
 
 
-def draw_xy_grid(p, xs, ys, x_labels, y_labels, cell, draw_legend, include_lone_images, swap_axes_processing_order):
-    ver_texts = [[images.GridAnnotation(y)] for y in y_labels]
+def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend, include_lone_images, include_sub_grids, first_axes_processed, second_axes_processed):
     hor_texts = [[images.GridAnnotation(x)] for x in x_labels]
+    ver_texts = [[images.GridAnnotation(y)] for y in y_labels]
+    title_texts = [[images.GridAnnotation(z)] for z in z_labels]
 
     # Temporary list of all the images that are generated to be populated into the grid.
     # Will be filled with empty images for any individual step that fails to process properly
-    image_cache = [None] * (len(xs) * len(ys))
+    image_cache = [None] * (len(xs) * len(ys) * len(zs))
 
     processed_result = None
     cell_mode = "P"
     cell_size = (1, 1)
 
-    state.job_count = len(xs) * len(ys) * p.n_iter
+    state.job_count = len(xs) * len(ys) * len(zs) * p.n_iter
 
-    def process_cell(x, y, ix, iy):
+    def process_cell(x, y, z, ix, iy, iz):
         nonlocal image_cache, processed_result, cell_mode, cell_size
 
-        state.job = f"{ix + iy * len(xs) + 1} out of {len(xs) * len(ys)}"
+        def index(ix, iy, iz):
+            return ix + iy * len(xs) + iz * len(xs) * len(ys)
 
-        processed: Processed = cell(x, y)
+        state.job = f"{index(ix, iy, iz) + 1} out of {len(xs) * len(ys) * len(zs)}"
+
+        processed: Processed = cell(x, y, z)
 
         try:
             # this dereference will throw an exception if the image was not processed
@@ -238,33 +242,65 @@ def draw_xy_grid(p, xs, ys, x_labels, y_labels, cell, draw_legend, include_lone_
                 cell_size = processed_image.size
                 processed_result.images = [Image.new(cell_mode, cell_size)]
 
-            image_cache[ix + iy * len(xs)] = processed_image
+            image_cache[index(ix, iy, iz)] = processed_image
             if include_lone_images:
                 processed_result.images.append(processed_image)
                 processed_result.all_prompts.append(processed.prompt)
                 processed_result.all_seeds.append(processed.seed)
                 processed_result.infotexts.append(processed.infotexts[0])
         except:
-            image_cache[ix + iy * len(xs)] = Image.new(cell_mode, cell_size)
+            image_cache[index(ix, iy, iz)] = Image.new(cell_mode, cell_size)
 
-    if swap_axes_processing_order:
+    if first_axes_processed == 'x':
         for ix, x in enumerate(xs):
-            for iy, y in enumerate(ys):
-                process_cell(x, y, ix, iy)
-    else:
+            if second_axes_processed == 'y':
+                for iy, y in enumerate(ys):
+                    for iz, z in enumerate(zs):
+                        process_cell(x, y, z, ix, iy, iz)
+            else:
+                for iz, z in enumerate(zs):
+                    for iy, y in enumerate(ys):
+                        process_cell(x, y, z, ix, iy, iz)
+    elif first_axes_processed == 'y':
         for iy, y in enumerate(ys):
-            for ix, x in enumerate(xs):
-                process_cell(x, y, ix, iy)
+            if second_axes_processed == 'x':
+                for ix, x in enumerate(xs):
+                    for iz, z in enumerate(zs):
+                        process_cell(x, y, z, ix, iy, iz)
+            else:
+                for iz, z in enumerate(zs):
+                    for ix, x in enumerate(xs):
+                        process_cell(x, y, z, ix, iy, iz)
+    elif first_axes_processed == 'z':
+        for iz, z in enumerate(zs):
+            if second_axes_processed == 'x':
+                for ix, x in enumerate(xs):
+                    for iy, y in enumerate(ys):
+                        process_cell(x, y, z, ix, iy, iz)
+            else:
+                for iy, y in enumerate(ys):
+                    for ix, x in enumerate(xs):
+                        process_cell(x, y, z, ix, iy, iz)
 
     if not processed_result:
-        print("Unexpected error: draw_xy_grid failed to return even a single processed image")
+        print("Unexpected error: draw_xyz_grid failed to return even a single processed image")
         return Processed(p, [])
 
-    grid = images.image_grid(image_cache, rows=len(ys))
-    if draw_legend:
-        grid = images.draw_grid_annotations(grid, cell_size[0], cell_size[1], hor_texts, ver_texts)
+    grids = [None] * len(zs)
+    for i in range(len(zs)):
+        start_index = i * len(xs) * len(ys)
+        end_index = start_index + len(xs) * len(ys)
+        grid = images.image_grid(image_cache[start_index:end_index], rows=len(ys))
+        if draw_legend:
+            grid = images.draw_grid_annotations(grid, cell_size[0], cell_size[1], hor_texts, ver_texts)
+        
+        grids[i] = grid        
+        if include_sub_grids and len(zs) > 1:
+            processed_result.images.insert(i+1, grid)
 
-    processed_result.images[0] = grid
+    original_grid_size = grids[0].size
+    grids = images.image_grid(grids, rows=1)
+    processed_result.images[0] = images.draw_grid_annotations(grids, original_grid_size[0], original_grid_size[1], title_texts, [[images.GridAnnotation()]])
 
     return processed_result
 
@@ -291,7 +327,7 @@ re_range_count_float = re.compile(r"\s*([+-]?\s*\d+(?:.\d*)?)\s*-\s*([+-]?\s*\d+
 
 class Script(scripts.Script):
     def title(self):
-        return "X/Y plot"
+        return "X/Y/Z plot"
 
     def ui(self, is_img2img):
         self.current_axis_options = [x for x in axis_options if type(x) == AxisOption or x.is_img2img == is_img2img]
@@ -301,24 +337,36 @@ class Script(scripts.Script):
                 with gr.Row():
                     x_type = gr.Dropdown(label="X type", choices=[x.label for x in self.current_axis_options], value=self.current_axis_options[1].label, type="index", elem_id=self.elem_id("x_type"))
                     x_values = gr.Textbox(label="X values", lines=1, elem_id=self.elem_id("x_values"))
-                    fill_x_button = ToolButton(value=fill_values_symbol, elem_id="xy_grid_fill_x_tool_button", visible=False)
+                    fill_x_button = ToolButton(value=fill_values_symbol, elem_id="xyz_grid_fill_x_tool_button", visible=False)
 
                 with gr.Row():
                     y_type = gr.Dropdown(label="Y type", choices=[x.label for x in self.current_axis_options], value=self.current_axis_options[0].label, type="index", elem_id=self.elem_id("y_type"))
                     y_values = gr.Textbox(label="Y values", lines=1, elem_id=self.elem_id("y_values"))
-                    fill_y_button = ToolButton(value=fill_values_symbol, elem_id="xy_grid_fill_y_tool_button", visible=False)
+                    fill_y_button = ToolButton(value=fill_values_symbol, elem_id="xyz_grid_fill_y_tool_button", visible=False)
+
+                with gr.Row():
+                    z_type = gr.Dropdown(label="Z type", choices=[x.label for x in self.current_axis_options], value=self.current_axis_options[0].label, type="index", elem_id=self.elem_id("z_type"))
+                    z_values = gr.Textbox(label="Z values", lines=1, elem_id=self.elem_id("z_values"))
+                    fill_z_button = ToolButton(value=fill_values_symbol, elem_id="xyz_grid_fill_z_tool_button", visible=False)
 
         with gr.Row(variant="compact", elem_id="axis_options"):
             draw_legend = gr.Checkbox(label='Draw legend', value=True, elem_id=self.elem_id("draw_legend"))
-            include_lone_images = gr.Checkbox(label='Include Separate Images', value=False, elem_id=self.elem_id("include_lone_images"))
+            include_lone_images = gr.Checkbox(label='Include Sub Images', value=False, elem_id=self.elem_id("include_lone_images"))
+            include_sub_grids = gr.Checkbox(label='Include Sub Grids', value=False, elem_id=self.elem_id("include_sub_grids"))
             no_fixed_seeds = gr.Checkbox(label='Keep -1 for seeds', value=False, elem_id=self.elem_id("no_fixed_seeds"))
-            swap_axes_button = gr.Button(value="Swap axes", elem_id="xy_grid_swap_axes_button")
+            swap_xy_axes_button = gr.Button(value="Swap X/Y axes", elem_id="xy_grid_swap_axes_button")
+            swap_yz_axes_button = gr.Button(value="Swap Y/Z axes", elem_id="yz_grid_swap_axes_button")
+            swap_xz_axes_button = gr.Button(value="Swap X/Z axes", elem_id="xz_grid_swap_axes_button")
 
-        def swap_axes(x_type, x_values, y_type, y_values):
-            return self.current_axis_options[y_type].label, y_values, self.current_axis_options[x_type].label, x_values
+        def swap_axes(axis1_type, axis1_values, axis2_type, axis2_values):
+            return self.current_axis_options[axis2_type].label, axis2_values, self.current_axis_options[axis1_type].label, axis1_values
 
-        swap_args = [x_type, x_values, y_type, y_values]
-        swap_axes_button.click(swap_axes, inputs=swap_args, outputs=swap_args)
+        xy_swap_args = [x_type, x_values, y_type, y_values]
+        swap_xy_axes_button.click(swap_axes, inputs=xy_swap_args, outputs=xy_swap_args)
+        yz_swap_args = [y_type, y_values, z_type, z_values]
+        swap_yz_axes_button.click(swap_axes, inputs=yz_swap_args, outputs=yz_swap_args)
+        xz_swap_args = [x_type, x_values, z_type, z_values]
+        swap_xz_axes_button.click(swap_axes, inputs=xz_swap_args, outputs=xz_swap_args)
 
         def fill(x_type):
             axis = self.current_axis_options[x_type]
@@ -326,16 +374,18 @@ class Script(scripts.Script):
 
         fill_x_button.click(fn=fill, inputs=[x_type], outputs=[x_values])
         fill_y_button.click(fn=fill, inputs=[y_type], outputs=[y_values])
+        fill_z_button.click(fn=fill, inputs=[z_type], outputs=[z_values])
 
         def select_axis(x_type):
             return gr.Button.update(visible=self.current_axis_options[x_type].choices is not None)
 
         x_type.change(fn=select_axis, inputs=[x_type], outputs=[fill_x_button])
         y_type.change(fn=select_axis, inputs=[y_type], outputs=[fill_y_button])
+        z_type.change(fn=select_axis, inputs=[z_type], outputs=[fill_z_button])
 
-        return [x_type, x_values, y_type, y_values, draw_legend, include_lone_images, no_fixed_seeds]
+        return [x_type, x_values, y_type, y_values, z_type, z_values, draw_legend, include_lone_images, include_sub_grids, no_fixed_seeds]
 
-    def run(self, p, x_type, x_values, y_type, y_values, draw_legend, include_lone_images, no_fixed_seeds):
+    def run(self, p, x_type, x_values, y_type, y_values, z_type, z_values, draw_legend, include_lone_images, include_sub_grids, no_fixed_seeds):
         if not no_fixed_seeds:
             modules.processing.fix_seed(p)
 
@@ -409,6 +459,9 @@ class Script(scripts.Script):
         y_opt = self.current_axis_options[y_type]
         ys = process_axis(y_opt, y_values)
 
+        z_opt = self.current_axis_options[z_type]
+        zs = process_axis(z_opt, z_values)
+
         def fix_axis_seeds(axis_opt, axis_list):
             if axis_opt.label in ['Seed', 'Var. seed']:
                 return [int(random.randrange(4294967294)) if val is None or val == '' or val == -1 else val for val in axis_list]
@@ -418,21 +471,26 @@ class Script(scripts.Script):
         if not no_fixed_seeds:
             xs = fix_axis_seeds(x_opt, xs)
             ys = fix_axis_seeds(y_opt, ys)
+            zs = fix_axis_seeds(z_opt, zs)
 
         if x_opt.label == 'Steps':
-            total_steps = sum(xs) * len(ys)
+            total_steps = sum(xs) * len(ys) * len(zs)
         elif y_opt.label == 'Steps':
-            total_steps = sum(ys) * len(xs)
+            total_steps = sum(ys) * len(xs) * len(zs)
+        elif z_opt.label == 'Steps':
+            total_steps = sum(zs) * len(xs) * len(ys)
         else:
-            total_steps = p.steps * len(xs) * len(ys)
+            total_steps = p.steps * len(xs) * len(ys) * len(zs)
 
         if isinstance(p, StableDiffusionProcessingTxt2Img) and p.enable_hr:
             if x_opt.label == "Hires steps":
-                total_steps += sum(xs) * len(ys)
+                total_steps += sum(xs) * len(ys) * len(zs)
             elif y_opt.label == "Hires steps":
-                total_steps += sum(ys) * len(xs)
+                total_steps += sum(ys) * len(xs) * len(zs)
+            elif z_opt.label == "Hires steps":
+                total_steps += sum(zs) * len(xs) * len(ys)
             elif p.hr_second_pass_steps:
-                total_steps += p.hr_second_pass_steps * len(xs) * len(ys)
+                total_steps += p.hr_second_pass_steps * len(xs) * len(ys) * len(zs)
             else:
                 total_steps *= 2
 
@@ -440,7 +498,8 @@ class Script(scripts.Script):
 
         image_cell_count = p.n_iter * p.batch_size
         cell_console_text = f"; {image_cell_count} images per cell" if image_cell_count > 1 else ""
-        print(f"X/Y plot will create {len(xs) * len(ys) * image_cell_count} images on a {len(xs)}x{len(ys)} grid{cell_console_text}. (Total steps to process: {total_steps})")
+        plural_s = 's' if len(zs) > 1 else ''
+        print(f"X/Y plot will create {len(xs) * len(ys) * len(zs) * image_cell_count} images on {len(zs)} {len(xs)}x{len(ys)} grid{plural_s}{cell_console_text}. (Total steps to process: {total_steps})")
         shared.total_tqdm.updateTotal(total_steps)
 
         grid_infotext = [None]
@@ -448,15 +507,35 @@ class Script(scripts.Script):
         # If one of the axes is very slow to change between (like SD model
         # checkpoint), then make sure it is in the outer iteration of the nested
         # `for` loop.
-        swap_axes_processing_order = x_opt.cost > y_opt.cost
+        first_axes_processed = 'x'
+        second_axes_processed = 'y'
+        if x_opt.cost > y_opt.cost and x_opt.cost > z_opt.cost:
+            first_axes_processed = 'x'
+            if y_opt.cost > z_opt.cost:
+                second_axes_processed = 'y'
+            else:
+                second_axes_processed = 'z'
+        elif y_opt.cost > x_opt.cost and y_opt.cost > z_opt.cost:
+            first_axes_processed = 'y'
+            if x_opt.cost > z_opt.cost:
+                second_axes_processed = 'x'
+            else:
+                second_axes_processed = 'z'
+        elif z_opt.cost > x_opt.cost and z_opt.cost > y_opt.cost:
+            first_axes_processed = 'z'
+            if x_opt.cost > y_opt.cost:
+                second_axes_processed = 'x'
+            else:
+                second_axes_processed = 'y'
 
-        def cell(x, y):
+        def cell(x, y, z):
             if shared.state.interrupted:
                 return Processed(p, [], p.seed, "")
 
             pc = copy(p)
             x_opt.apply(pc, x, xs)
             y_opt.apply(pc, y, ys)
+            z_opt.apply(pc, z, zs)
 
             res = process_images(pc)
 
@@ -475,24 +554,34 @@ class Script(scripts.Script):
                     if y_opt.label in ["Seed", "Var. seed"] and not no_fixed_seeds:
                         pc.extra_generation_params["Fixed Y Values"] = ", ".join([str(y) for y in ys])
 
+                if z_opt.label != 'Nothing':
+                    pc.extra_generation_params["Z Type"] = z_opt.label
+                    pc.extra_generation_params["Z Values"] = z_values
+                    if z_opt.label in ["Seed", "Var. seed"] and not no_fixed_seeds:
+                        pc.extra_generation_params["Fixed Z Values"] = ", ".join([str(z) for z in zs])
+
                 grid_infotext[0] = processing.create_infotext(pc, pc.all_prompts, pc.all_seeds, pc.all_subseeds)
 
             return res
 
         with SharedSettingsStackHelper():
-            processed = draw_xy_grid(
+            processed = draw_xyz_grid(
                 p,
                 xs=xs,
                 ys=ys,
+                zs=zs,
                 x_labels=[x_opt.format_value(p, x_opt, x) for x in xs],
                 y_labels=[y_opt.format_value(p, y_opt, y) for y in ys],
+                z_labels=[z_opt.format_value(p, z_opt, z) for z in zs],
                 cell=cell,
                 draw_legend=draw_legend,
                 include_lone_images=include_lone_images,
-                swap_axes_processing_order=swap_axes_processing_order
+                include_sub_grids=include_sub_grids,
+                first_axes_processed=first_axes_processed,
+                second_axes_processed=second_axes_processed
             )
 
         if opts.grid_save:
-            images.save_image(processed.images[0], p.outpath_grids, "xy_grid", info=grid_infotext[0], extension=opts.grid_format, prompt=p.prompt, seed=processed.seed, grid=True, p=p)
+            images.save_image(processed.images[0], p.outpath_grids, "xyz_grid", info=grid_infotext[0], extension=opts.grid_format, prompt=p.prompt, seed=processed.seed, grid=True, p=p)
 
         return processed

style.css

@@ -857,6 +857,7 @@ footer {
     white-space: nowrap;
     text-overflow: ellipsis;
     background: rgba(0,0,0,.5);
+    color: white;
 }
 
 .extra-network-thumbs .card:hover .actions .name {

webui-macos-env.sh

@@ -10,7 +10,7 @@ then
 fi
 
 export install_dir="$HOME"
-export COMMANDLINE_ARGS="--skip-torch-cuda-test --no-half --use-cpu interrogate"
+export COMMANDLINE_ARGS="--skip-torch-cuda-test --upcast-sampling --use-cpu interrogate"
 export TORCH_COMMAND="pip install torch==1.12.1 torchvision==0.13.1"
 export K_DIFFUSION_REPO="https://github.com/brkirch/k-diffusion.git"
 export K_DIFFUSION_COMMIT_HASH="51c9778f269cedb55a4d88c79c0246d35bdadb71"

webui.bat

@@ -3,17 +3,28 @@
 if not defined PYTHON (set PYTHON=python)
 if not defined VENV_DIR (set "VENV_DIR=%~dp0%venv")
 
+
 set ERROR_REPORTING=FALSE
 
 mkdir tmp 2>NUL
 
 %PYTHON% -c "" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :start_venv
+if %ERRORLEVEL% == 0 goto :check_pip
 echo Couldn't launch python
 goto :show_stdout_stderr
 
+:check_pip
+%PYTHON% -mpip --help >tmp/stdout.txt 2>tmp/stderr.txt
+if %ERRORLEVEL% == 0 goto :start_venv
+if "%PIP_INSTALLER_LOCATION%" == "" goto :show_stdout_stderr
+%PYTHON% "%PIP_INSTALLER_LOCATION%" >tmp/stdout.txt 2>tmp/stderr.txt
+if %ERRORLEVEL% == 0 goto :start_venv
+echo Couldn't install pip
+goto :show_stdout_stderr
+
 :start_venv
 if ["%VENV_DIR%"] == ["-"] goto :skip_venv
+if ["%SKIP_VENV%"] == ["1"] goto :skip_venv
 
 dir "%VENV_DIR%\Scripts\Python.exe" >tmp/stdout.txt 2>tmp/stderr.txt
 if %ERRORLEVEL% == 0 goto :activate_venv
@@ -28,13 +39,13 @@ goto :show_stdout_stderr
 :activate_venv
 set PYTHON="%VENV_DIR%\Scripts\Python.exe"
 echo venv %PYTHON%
+
+:skip_venv
 if [%ACCELERATE%] == ["True"] goto :accelerate
 goto :launch
 
-:skip_venv
-
 :accelerate
-echo "Checking for accelerate"
+echo Checking for accelerate
 set ACCELERATE="%VENV_DIR%\Scripts\accelerate.exe"
 if EXIST %ACCELERATE% goto :accelerate_launch
 
@@ -44,7 +55,7 @@ pause
 exit /b
 
 :accelerate_launch
-echo "Accelerating"
+echo Accelerating
 %ACCELERATE% launch --num_cpu_threads_per_process=6 launch.py
 pause
 exit /b