Merge branch 'AUTOMATIC1111:master' into master

javascript/hints.js

@@ -8,8 +8,8 @@ titles = {
 	"DDIM": "Denoising Diffusion Implicit Models - best at inpainting",
 	"DPM adaptive": "Ignores step count - uses a number of steps determined by the CFG and resolution", 
 
-	"Batch count": "How many batches of images to create",
-	"Batch size": "How many image to create in a single batch",
+	"Batch count": "How many batches of images to create (has no impact on generation performance or VRAM usage)",
+	"Batch size": "How many image to create in a single batch (increases generation performance at cost of higher VRAM usage)",
     "CFG Scale": "Classifier Free Guidance Scale - how strongly the image should conform to prompt - lower values produce more creative results",
     "Seed": "A value that determines the output of random number generator - if you create an image with same parameters and seed as another image, you'll get the same result",
     "\u{1f3b2}\ufe0f": "Set seed to -1, which will cause a new random number to be used every time",

launch.py

@@ -242,7 +242,7 @@ def prepare_environment():
 
     sys.argv += shlex.split(commandline_args)
 
-    parser = argparse.ArgumentParser()
+    parser = argparse.ArgumentParser(add_help=False)
     parser.add_argument("--ui-settings-file", type=str, help="filename to use for ui settings", default='config.json')
     args, _ = parser.parse_known_args(sys.argv)
 

modules/api/api.py

@@ -498,7 +498,7 @@ class Api:
             if not apply_optimizations:
                 sd_hijack.undo_optimizations()
             try:
-                hypernetwork, filename = train_hypernetwork(*args)
+                hypernetwork, filename = train_hypernetwork(**args)
             except Exception as e:
                 error = e
             finally:

modules/esrgan_model_arch.py

@@ -1,5 +1,6 @@
 # this file is adapted from https://github.com/victorca25/iNNfer
 
+from collections import OrderedDict
 import math
 import functools
 import torch

modules/hypernetworks/hypernetwork.py

@@ -380,8 +380,8 @@ def apply_single_hypernetwork(hypernetwork, context_k, context_v, layer=None):
         layer.hyper_k = hypernetwork_layers[0]
         layer.hyper_v = hypernetwork_layers[1]
 
-    context_k = hypernetwork_layers[0](context_k)
-    context_v = hypernetwork_layers[1](context_v)
+    context_k = devices.cond_cast_unet(hypernetwork_layers[0](devices.cond_cast_float(context_k)))
+    context_v = devices.cond_cast_unet(hypernetwork_layers[1](devices.cond_cast_float(context_v)))
     return context_k, context_v
 
 
@@ -496,7 +496,7 @@ def create_hypernetwork(name, enable_sizes, overwrite_old, layer_structure=None,
     shared.reload_hypernetworks()
 
 
-def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, varsize, steps, clip_grad_mode, clip_grad_value, shuffle_tags, tag_drop_out, latent_sampling_method, create_image_every, save_hypernetwork_every, template_filename, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height):
+def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, varsize, steps, clip_grad_mode, clip_grad_value, shuffle_tags, tag_drop_out, latent_sampling_method, use_weight, create_image_every, save_hypernetwork_every, template_filename, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height):
     # images allows training previews to have infotext. Importing it at the top causes a circular import problem.
     from modules import images
 
@@ -554,7 +554,7 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
 
     pin_memory = shared.opts.pin_memory
 
-    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=hypernetwork_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize)
+    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=hypernetwork_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, include_cond=True, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight)
 
     if shared.opts.save_training_settings_to_txt:
         saved_params = dict(
@@ -640,13 +640,19 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi
                 
                 with devices.autocast():
                     x = batch.latent_sample.to(devices.device, non_blocking=pin_memory)
+                    if use_weight:
+                        w = batch.weight.to(devices.device, non_blocking=pin_memory)
                     if tag_drop_out != 0 or shuffle_tags:
                         shared.sd_model.cond_stage_model.to(devices.device)
                         c = shared.sd_model.cond_stage_model(batch.cond_text).to(devices.device, non_blocking=pin_memory)
                         shared.sd_model.cond_stage_model.to(devices.cpu)
                     else:
                         c = stack_conds(batch.cond).to(devices.device, non_blocking=pin_memory)
-                    loss = shared.sd_model(x, c)[0] / gradient_step
+                    if use_weight:
+                        loss = shared.sd_model.weighted_forward(x, c, w)[0] / gradient_step
+                        del w
+                    else:
+                        loss = shared.sd_model.forward(x, c)[0] / gradient_step
                     del x
                     del c
 

modules/img2img.py

@@ -75,6 +75,8 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args):
 
             if not save_normally:
                 os.makedirs(output_dir, exist_ok=True)
+                if processed_image.mode == 'RGBA':
+                    processed_image = processed_image.convert("RGB")
                 processed_image.save(os.path.join(output_dir, filename))
 
 

modules/script_callbacks.py

@@ -46,6 +46,18 @@ class CFGDenoiserParams:
         """Total number of sampling steps planned"""
 
 
+class CFGDenoisedParams:
+    def __init__(self, x, sampling_step, total_sampling_steps):
+        self.x = x
+        """Latent image representation in the process of being denoised"""
+
+        self.sampling_step = sampling_step
+        """Current Sampling step number"""
+
+        self.total_sampling_steps = total_sampling_steps
+        """Total number of sampling steps planned"""
+
+
 class UiTrainTabParams:
     def __init__(self, txt2img_preview_params):
         self.txt2img_preview_params = txt2img_preview_params
@@ -68,6 +80,7 @@ callback_map = dict(
     callbacks_before_image_saved=[],
     callbacks_image_saved=[],
     callbacks_cfg_denoiser=[],
+    callbacks_cfg_denoised=[],
     callbacks_before_component=[],
     callbacks_after_component=[],
     callbacks_image_grid=[],
@@ -150,6 +163,14 @@ def cfg_denoiser_callback(params: CFGDenoiserParams):
             report_exception(c, 'cfg_denoiser_callback')
 
 
+def cfg_denoised_callback(params: CFGDenoisedParams):
+    for c in callback_map['callbacks_cfg_denoised']:
+        try:
+            c.callback(params)
+        except Exception:
+            report_exception(c, 'cfg_denoised_callback')
+
+
 def before_component_callback(component, **kwargs):
     for c in callback_map['callbacks_before_component']:
         try:
@@ -283,6 +304,14 @@ def on_cfg_denoiser(callback):
     add_callback(callback_map['callbacks_cfg_denoiser'], callback)
 
 
+def on_cfg_denoised(callback):
+    """register a function to be called in the kdiffussion cfg_denoiser method after building the inner model inputs.
+    The callback is called with one argument:
+        - params: CFGDenoisedParams - parameters to be passed to the inner model and sampling state details.
+    """
+    add_callback(callback_map['callbacks_cfg_denoised'], callback)
+
+
 def on_before_component(callback):
     """register a function to be called before a component is created.
     The callback is called with arguments:

modules/sd_hijack.py

@@ -1,5 +1,6 @@
 import torch
 from torch.nn.functional import silu
+from types import MethodType
 
 import modules.textual_inversion.textual_inversion
 from modules import devices, sd_hijack_optimizations, shared, sd_hijack_checkpoint
@@ -76,6 +77,54 @@ def fix_checkpoint():
     pass
 
 
+def weighted_loss(sd_model, pred, target, mean=True):
+    #Calculate the weight normally, but ignore the mean
+    loss = sd_model._old_get_loss(pred, target, mean=False)
+    
+    #Check if we have weights available
+    weight = getattr(sd_model, '_custom_loss_weight', None)
+    if weight is not None:
+        loss *= weight
+    
+    #Return the loss, as mean if specified
+    return loss.mean() if mean else loss
+
+def weighted_forward(sd_model, x, c, w, *args, **kwargs):
+    try:
+        #Temporarily append weights to a place accessible during loss calc
+        sd_model._custom_loss_weight = w
+        
+        #Replace 'get_loss' with a weight-aware one. Otherwise we need to reimplement 'forward' completely
+        #Keep 'get_loss', but don't overwrite the previous old_get_loss if it's already set
+        if not hasattr(sd_model, '_old_get_loss'):
+            sd_model._old_get_loss = sd_model.get_loss
+        sd_model.get_loss = MethodType(weighted_loss, sd_model)
+
+        #Run the standard forward function, but with the patched 'get_loss'
+        return sd_model.forward(x, c, *args, **kwargs)
+    finally:
+        try:
+            #Delete temporary weights if appended
+            del sd_model._custom_loss_weight
+        except AttributeError as e:
+            pass
+            
+        #If we have an old loss function, reset the loss function to the original one
+        if hasattr(sd_model, '_old_get_loss'):
+            sd_model.get_loss = sd_model._old_get_loss
+            del sd_model._old_get_loss
+
+def apply_weighted_forward(sd_model):
+    #Add new function 'weighted_forward' that can be called to calc weighted loss
+    sd_model.weighted_forward = MethodType(weighted_forward, sd_model)
+
+def undo_weighted_forward(sd_model):
+    try:
+        del sd_model.weighted_forward
+    except AttributeError as e:
+        pass
+
+
 class StableDiffusionModelHijack:
     fixes = None
     comments = []
@@ -104,6 +153,10 @@ class StableDiffusionModelHijack:
             m.cond_stage_model.model.token_embedding = EmbeddingsWithFixes(m.cond_stage_model.model.token_embedding, self)
             m.cond_stage_model = sd_hijack_open_clip.FrozenOpenCLIPEmbedderWithCustomWords(m.cond_stage_model, self)
 
+        apply_weighted_forward(m)
+        if m.cond_stage_key == "edit":
+            sd_hijack_unet.hijack_ddpm_edit()
+
         self.optimization_method = apply_optimizations()
 
         self.clip = m.cond_stage_model
@@ -132,6 +185,7 @@ class StableDiffusionModelHijack:
             m.cond_stage_model = m.cond_stage_model.wrapped
 
         undo_optimizations()
+        undo_weighted_forward(m)
 
         self.apply_circular(False)
         self.layers = None

modules/sd_hijack_inpainting.py

@@ -11,6 +11,7 @@ import ldm.models.diffusion.plms
 from ldm.models.diffusion.ddpm import LatentDiffusion
 from ldm.models.diffusion.plms import PLMSSampler
 from ldm.models.diffusion.ddim import DDIMSampler, noise_like
+from ldm.models.diffusion.sampling_util import norm_thresholding
 
 
 @torch.no_grad()

modules/sd_hijack_unet.py

@@ -44,6 +44,7 @@ def apply_model(orig_func, self, x_noisy, t, cond, **kwargs):
     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)
@@ -53,6 +54,16 @@ class GELUHijack(torch.nn.GELU, torch.nn.Module):
         else:
             return torch.nn.GELU.forward(self, x)
 
+
+ddpm_edit_hijack = None
+def hijack_ddpm_edit():
+    global ddpm_edit_hijack
+    if not ddpm_edit_hijack:
+        CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.decode_first_stage', first_stage_sub, first_stage_cond)
+        CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.encode_first_stage', first_stage_sub, first_stage_cond)
+        ddpm_edit_hijack = CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.apply_model', apply_model, unet_needs_upcast)
+
+
 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, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast)

modules/sd_samplers_kdiffusion.py

@@ -8,6 +8,7 @@ from modules import prompt_parser, devices, sd_samplers_common
 from modules.shared import opts, state
 import modules.shared as shared
 from modules.script_callbacks import CFGDenoiserParams, cfg_denoiser_callback
+from modules.script_callbacks import CFGDenoisedParams, cfg_denoised_callback
 
 samplers_k_diffusion = [
     ('Euler a', 'sample_euler_ancestral', ['k_euler_a', 'k_euler_ancestral'], {}),
@@ -136,6 +137,9 @@ class CFGDenoiser(torch.nn.Module):
 
             x_out[-uncond.shape[0]:] = self.inner_model(x_in[-uncond.shape[0]:], sigma_in[-uncond.shape[0]:], cond={"c_crossattn": [uncond], "c_concat": [image_cond_in[-uncond.shape[0]:]]})
 
+        denoised_params = CFGDenoisedParams(x_out, state.sampling_step, state.sampling_steps)
+        cfg_denoised_callback(denoised_params)
+
         devices.test_for_nans(x_out, "unet")
 
         if opts.live_preview_content == "Prompt":

modules/shared.py

@@ -81,6 +81,7 @@ parser.add_argument("--freeze-settings", action='store_true', help="disable edit
 parser.add_argument("--ui-settings-file", type=str, help="filename to use for ui settings", default=os.path.join(data_path, 'config.json'))
 parser.add_argument("--gradio-debug",  action='store_true', help="launch gradio with --debug option")
 parser.add_argument("--gradio-auth", type=str, help='set gradio authentication like "username:password"; or comma-delimit multiple like "u1:p1,u2:p2,u3:p3"', default=None)
+parser.add_argument("--gradio-auth-path", type=str, help='set gradio authentication file path ex. "/path/to/auth/file" same auth format as --gradio-auth', default=None)
 parser.add_argument("--gradio-img2img-tool", type=str, help='does not do anything')
 parser.add_argument("--gradio-inpaint-tool", type=str, help="does not do anything")
 parser.add_argument("--opt-channelslast", action='store_true', help="change memory type for stable diffusion to channels last")

modules/shared_items.py

@@ -20,4 +20,4 @@ def sd_vae_items():
 def refresh_vae_list():
     import modules.sd_vae
 
-    return modules.sd_vae.refresh_vae_list
+    modules.sd_vae.refresh_vae_list()

modules/textual_inversion/dataset.py

@@ -19,9 +19,10 @@ re_numbers_at_start = re.compile(r"^[-\d]+\s*")
 
 
 class DatasetEntry:
-    def __init__(self, filename=None, filename_text=None, latent_dist=None, latent_sample=None, cond=None, cond_text=None, pixel_values=None):
+    def __init__(self, filename=None, filename_text=None, latent_dist=None, latent_sample=None, cond=None, cond_text=None, pixel_values=None, weight=None):
         self.filename = filename
         self.filename_text = filename_text
+        self.weight = weight
         self.latent_dist = latent_dist
         self.latent_sample = latent_sample
         self.cond = cond
@@ -30,7 +31,7 @@ class DatasetEntry:
 
 
 class PersonalizedBase(Dataset):
-    def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, cond_model=None, device=None, template_file=None, include_cond=False, batch_size=1, gradient_step=1, shuffle_tags=False, tag_drop_out=0, latent_sampling_method='once', varsize=False):
+    def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, cond_model=None, device=None, template_file=None, include_cond=False, batch_size=1, gradient_step=1, shuffle_tags=False, tag_drop_out=0, latent_sampling_method='once', varsize=False, use_weight=False):
         re_word = re.compile(shared.opts.dataset_filename_word_regex) if len(shared.opts.dataset_filename_word_regex) > 0 else None
 
         self.placeholder_token = placeholder_token
@@ -56,10 +57,16 @@ class PersonalizedBase(Dataset):
 
         print("Preparing dataset...")
         for path in tqdm.tqdm(self.image_paths):
+            alpha_channel = None
             if shared.state.interrupted:
                 raise Exception("interrupted")
             try:
-                image = Image.open(path).convert('RGB')
+                image = Image.open(path)
+                #Currently does not work for single color transparency
+                #We would need to read image.info['transparency'] for that
+                if use_weight and 'A' in image.getbands():
+                    alpha_channel = image.getchannel('A')
+                image = image.convert('RGB')
                 if not varsize:
                     image = image.resize((width, height), PIL.Image.BICUBIC)
             except Exception:
@@ -87,17 +94,35 @@ class PersonalizedBase(Dataset):
             with devices.autocast():
                 latent_dist = model.encode_first_stage(torchdata.unsqueeze(dim=0))
 
-            if latent_sampling_method == "once" or (latent_sampling_method == "deterministic" and not isinstance(latent_dist, DiagonalGaussianDistribution)):
-                latent_sample = model.get_first_stage_encoding(latent_dist).squeeze().to(devices.cpu)
-                latent_sampling_method = "once"
-                entry = DatasetEntry(filename=path, filename_text=filename_text, latent_sample=latent_sample)
-            elif latent_sampling_method == "deterministic":
-                # Works only for DiagonalGaussianDistribution
-                latent_dist.std = 0
-                latent_sample = model.get_first_stage_encoding(latent_dist).squeeze().to(devices.cpu)
-                entry = DatasetEntry(filename=path, filename_text=filename_text, latent_sample=latent_sample)
-            elif latent_sampling_method == "random":
-                entry = DatasetEntry(filename=path, filename_text=filename_text, latent_dist=latent_dist)
+            #Perform latent sampling, even for random sampling.
+            #We need the sample dimensions for the weights
+            if latent_sampling_method == "deterministic":
+                if isinstance(latent_dist, DiagonalGaussianDistribution):
+                    # Works only for DiagonalGaussianDistribution
+                    latent_dist.std = 0
+                else:
+                    latent_sampling_method = "once"
+            latent_sample = model.get_first_stage_encoding(latent_dist).squeeze().to(devices.cpu)
+
+            if use_weight and alpha_channel is not None:
+                channels, *latent_size = latent_sample.shape
+                weight_img = alpha_channel.resize(latent_size)
+                npweight = np.array(weight_img).astype(np.float32)
+                #Repeat for every channel in the latent sample
+                weight = torch.tensor([npweight] * channels).reshape([channels] + latent_size)
+                #Normalize the weight to a minimum of 0 and a mean of 1, that way the loss will be comparable to default.
+                weight -= weight.min()
+                weight /= weight.mean()
+            elif use_weight:
+                #If an image does not have a alpha channel, add a ones weight map anyway so we can stack it later
+                weight = torch.ones(latent_sample.shape)
+            else:
+                weight = None
+            
+            if latent_sampling_method == "random":
+                entry = DatasetEntry(filename=path, filename_text=filename_text, latent_dist=latent_dist, weight=weight)
+            else:
+                entry = DatasetEntry(filename=path, filename_text=filename_text, latent_sample=latent_sample, weight=weight)
 
             if not (self.tag_drop_out != 0 or self.shuffle_tags):
                 entry.cond_text = self.create_text(filename_text)
@@ -110,6 +135,7 @@ class PersonalizedBase(Dataset):
             del torchdata
             del latent_dist
             del latent_sample
+            del weight
 
         self.length = len(self.dataset)
         self.groups = list(groups.values())
@@ -195,6 +221,10 @@ class BatchLoader:
         self.cond_text = [entry.cond_text for entry in data]
         self.cond = [entry.cond for entry in data]
         self.latent_sample = torch.stack([entry.latent_sample for entry in data]).squeeze(1)
+        if all(entry.weight is not None for entry in data):
+            self.weight = torch.stack([entry.weight for entry in data]).squeeze(1)
+        else:
+            self.weight = None
         #self.emb_index = [entry.emb_index for entry in data]
         #print(self.latent_sample.device)
 

modules/textual_inversion/textual_inversion.py

@@ -351,7 +351,7 @@ def validate_train_inputs(model_name, learn_rate, batch_size, gradient_step, dat
         assert log_directory, "Log directory is empty"
 
 
-def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, varsize, steps, clip_grad_mode, clip_grad_value, shuffle_tags, tag_drop_out, latent_sampling_method, create_image_every, save_embedding_every, template_filename, save_image_with_stored_embedding, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height):
+def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_step, data_root, log_directory, training_width, training_height, varsize, steps, clip_grad_mode, clip_grad_value, shuffle_tags, tag_drop_out, latent_sampling_method, use_weight, create_image_every, save_embedding_every, template_filename, save_image_with_stored_embedding, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height):
     save_embedding_every = save_embedding_every or 0
     create_image_every = create_image_every or 0
     template_file = textual_inversion_templates.get(template_filename, None)
@@ -410,7 +410,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
 
     pin_memory = shared.opts.pin_memory
 
-    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=embedding_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize)
+    ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=embedding_name, model=shared.sd_model, cond_model=shared.sd_model.cond_stage_model, device=devices.device, template_file=template_file, batch_size=batch_size, gradient_step=gradient_step, shuffle_tags=shuffle_tags, tag_drop_out=tag_drop_out, latent_sampling_method=latent_sampling_method, varsize=varsize, use_weight=use_weight)
 
     if shared.opts.save_training_settings_to_txt:
         save_settings_to_file(log_directory, {**dict(model_name=checkpoint.model_name, model_hash=checkpoint.shorthash, num_of_dataset_images=len(ds), num_vectors_per_token=len(embedding.vec)), **locals()})
@@ -480,6 +480,8 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
             
                 with devices.autocast():
                     x = batch.latent_sample.to(devices.device, non_blocking=pin_memory)
+                    if use_weight:
+                        w = batch.weight.to(devices.device, non_blocking=pin_memory)
                     c = shared.sd_model.cond_stage_model(batch.cond_text)
 
                     if is_training_inpainting_model:
@@ -490,7 +492,11 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
                     else:
                         cond = c
 
-                    loss = shared.sd_model(x, cond)[0] / gradient_step
+                    if use_weight:
+                        loss = shared.sd_model.weighted_forward(x, cond, w)[0] / gradient_step
+                        del w
+                    else:
+                        loss = shared.sd_model.forward(x, cond)[0] / gradient_step
                     del x
 
                     _loss_step += loss.item()

modules/ui.py

@@ -1280,6 +1280,8 @@ def create_ui():
                         create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable', value=500, precision=0, elem_id="train_create_image_every")
                         save_embedding_every = gr.Number(label='Save a copy of embedding to log directory every N steps, 0 to disable', value=500, precision=0, elem_id="train_save_embedding_every")
 
+                    use_weight = gr.Checkbox(label="Use PNG alpha channel as loss weight", value=False, elem_id="use_weight")
+
                     save_image_with_stored_embedding = gr.Checkbox(label='Save images with embedding in PNG chunks', value=True, elem_id="train_save_image_with_stored_embedding")
                     preview_from_txt2img = gr.Checkbox(label='Read parameters (prompt, etc...) from txt2img tab when making previews', value=False, elem_id="train_preview_from_txt2img")
 
@@ -1393,6 +1395,7 @@ def create_ui():
                 shuffle_tags,
                 tag_drop_out,
                 latent_sampling_method,
+                use_weight,
                 create_image_every,
                 save_embedding_every,
                 template_file,
@@ -1426,6 +1429,7 @@ def create_ui():
                 shuffle_tags,
                 tag_drop_out,
                 latent_sampling_method,
+                use_weight,
                 create_image_every,
                 save_embedding_every,
                 template_file,

modules/ui_extra_networks.py

@@ -76,6 +76,10 @@ class ExtraNetworksPage:
                 while subdir.startswith("/"):
                     subdir = subdir[1:]
 
+                is_empty = len(os.listdir(x)) == 0
+                if not is_empty and not subdir.endswith("/"):
+                    subdir = subdir + "/"
+
                 subdirs[subdir] = 1
 
         if subdirs:
@@ -94,11 +98,13 @@ class ExtraNetworksPage:
             dirs = "".join([f"<li>{x}</li>" for x in self.allowed_directories_for_previews()])
             items_html = shared.html("extra-networks-no-cards.html").format(dirs=dirs)
 
+        self_name_id = self.name.replace(" ", "_")
+
         res = f"""
-<div id='{tabname}_{self.name}_subdirs' class='extra-network-subdirs extra-network-subdirs-{view}'>
+<div id='{tabname}_{self_name_id}_subdirs' class='extra-network-subdirs extra-network-subdirs-{view}'>
 {subdirs_html}
 </div>
-<div id='{tabname}_{self.name}_cards' class='extra-network-{view}'>
+<div id='{tabname}_{self_name_id}_cards' class='extra-network-{view}'>
 {items_html}
 </div>
 """

scripts/prompt_matrix.py

@@ -54,7 +54,7 @@ class Script(scripts.Script):
                 prompt_type = gr.Radio(["positive", "negative"], label="Select prompt", elem_id=self.elem_id("prompt_type"), value="positive")
                 variations_delimiter = gr.Radio(["comma", "space"], label="Select joining char", elem_id=self.elem_id("variations_delimiter"), value="comma")
             with gr.Column():
-                margin_size = gr.Slider(label="Grid margins (px)", min=0, max=500, value=0, step=2, elem_id=self.elem_id("margin_size"))
+                margin_size = gr.Slider(label="Grid margins (px)", minimum=0, maximum=500, value=0, step=2, elem_id=self.elem_id("margin_size"))
 
         return [put_at_start, different_seeds, prompt_type, variations_delimiter, margin_size]
 
@@ -99,8 +99,8 @@ class Script(scripts.Script):
         p.prompt_for_display = positive_prompt
         processed = process_images(p)
 
-        grid = images.image_grid(processed.images, p.batch_size, rows=1 << ((len(prompt_matrix_parts) - 1) // 2))
-        grid = images.draw_prompt_matrix(grid, p.width, p.height, prompt_matrix_parts, margin_size)
+        grid = images.image_grid(processed.images, p.batch_size, rows=1 << ((len(prompt_matrix_parts) - 1) // 2)) 
+        grid = images.draw_prompt_matrix(grid, processed.images[0].width, processed.images[1].height, prompt_matrix_parts, margin_size)
         processed.images.insert(0, grid)
         processed.index_of_first_image = 1
         processed.infotexts.insert(0, processed.infotexts[0])

scripts/xyz_grid.py

@@ -25,6 +25,8 @@ from modules.ui_components import ToolButton
 
 fill_values_symbol = "\U0001f4d2"  # 📒
 
+AxisInfo = namedtuple('AxisInfo', ['axis', 'values'])
+
 
 def apply_field(field):
     def fun(p, x, xs):
@@ -186,6 +188,7 @@ axis_options = [
     AxisOption("Steps", int, apply_field("steps")),
     AxisOptionTxt2Img("Hires steps", int, apply_field("hr_second_pass_steps")),
     AxisOption("CFG Scale", float, apply_field("cfg_scale")),
+    AxisOptionImg2Img("Image CFG Scale", float, apply_field("image_cfg_scale")),
     AxisOption("Prompt S/R", str, apply_prompt, format_value=format_value),
     AxisOption("Prompt order", str_permutations, apply_order, format_value=format_value_join_list),
     AxisOptionTxt2Img("Sampler", str, apply_sampler, format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
@@ -358,7 +361,7 @@ class Script(scripts.Script):
                 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"))
             with gr.Column():
-                margin_size = gr.Slider(label="Grid margins (px)", min=0, max=500, value=0, step=2, elem_id=self.elem_id("margin_size"))
+                margin_size = gr.Slider(label="Grid margins (px)", minimum=0, maximum=500, value=0, step=2, elem_id=self.elem_id("margin_size"))
         
         with gr.Row(variant="compact", elem_id="swap_axes"):
             swap_xy_axes_button = gr.Button(value="Swap X/Y axes", elem_id="xy_grid_swap_axes_button")
@@ -520,6 +523,10 @@ class Script(scripts.Script):
 
         grid_infotext = [None]
 
+        state.xyz_plot_x = AxisInfo(x_opt, xs)
+        state.xyz_plot_y = AxisInfo(y_opt, ys)
+        state.xyz_plot_z = AxisInfo(z_opt, zs)
+
         # 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.

webui.py

@@ -207,6 +207,14 @@ def webui():
         if cmd_opts.gradio_queue:
             shared.demo.queue(64)
 
+        gradio_auth_creds = []
+        if cmd_opts.gradio_auth:
+            gradio_auth_creds += cmd_opts.gradio_auth.strip('"').replace('\n', '').split(',')
+        if cmd_opts.gradio_auth_path:
+            with open(cmd_opts.gradio_auth_path, 'r', encoding="utf8") as file:
+                for line in file.readlines():
+                    gradio_auth_creds += [x.strip() for x in line.split(',')]
+
         app, local_url, share_url = shared.demo.launch(
             share=cmd_opts.share,
             server_name=server_name,
@@ -214,7 +222,7 @@ def webui():
             ssl_keyfile=cmd_opts.tls_keyfile,
             ssl_certfile=cmd_opts.tls_certfile,
             debug=cmd_opts.gradio_debug,
-            auth=[tuple(cred.split(':')) for cred in cmd_opts.gradio_auth.strip('"').split(',')] if cmd_opts.gradio_auth else None,
+            auth=[tuple(cred.split(':')) for cred in gradio_auth_creds] if gradio_auth_creds else None,
             inbrowser=cmd_opts.autolaunch,
             prevent_thread_lock=True
         )