Added Model rescale and prepared a release upgrade

# Conflicts:
#	jobs/process/BaseSDTrainProcess.py
#	jobs/process/TrainSliderProcess.py

README.md

@@ -40,6 +40,8 @@ pip3 install -r requirements.txt
 I have so many hodge podge scripts I am going to be moving over to this that I use in my ML work. But this is what is
 here so far.
 
+---
+
 ### LoRA (lierla), LoCON (LyCORIS) extractor
 
 It is based on the extractor in the [LyCORIS](https://github.com/KohakuBlueleaf/LyCORIS) tool, but adding some QOL features
@@ -64,6 +66,31 @@ Most people used fixed, which is traditional fixed dimension extraction.
 
 `process` is an array of different processes to run. You can add a few and mix and match. One LoRA, one LyCON, etc.
 
+---
+
+### LoRA Rescale
+
+Change `<lora:my_lora:4.6>` to `<lora:my_lora:1.0>` or whatever you want with the same effect. 
+A tool for rescaling a LoRA's weights. Should would with LoCON as well, but I have not tested it.
+It all runs off a config file, which you can find an example of in  `config/examples/mod_lora_scale.yml`.
+Just copy that file, into the `config` folder, and rename it to `whatever_you_want.yml`.
+Then you can edit the file to your liking. and call it like so:
+
+```bash
+python3 run.py config/whatever_you_want.yml
+```
+
+You can also put a full path to a config file, if you want to keep it somewhere else.
+
+```bash
+python3 run.py "/home/user/whatever_you_want.yml"
+```
+
+More notes on how it works are available in the example config file itself. This is useful when making 
+all LoRAs, as the ideal weight is rarely 1.0, but now you can fix that. For sliders, they can have weird scales form -2 to 2
+or even -15 to 15. This will allow you to dile it in so they all have your desired scale
+
+---
 
 ### LoRA Slider Trainer
 
@@ -108,13 +135,32 @@ Just went in and out. It is much worse on smaller faces than shown here.
 
 ## TODO
 - [X] Add proper regs on sliders
-- [ ] Add SDXL support (base model only for now)
+- [X] Add SDXL support (base model only for now)
 - [ ] Add plain erasing
 - [ ] Make Textual inversion network trainer (network that spits out TI embeddings)
 
 ---
 
 ## Change Log
+
+#### 2021-08-01
+Major changes and update. New LoRA rescale tool, look above for details. Added better metadata so
+Automatic1111 knows what the base model is. Added some experiments and a ton of updates. This thing is still unstable
+at the moment, so hopefully there are not breaking changes. 
+
+Unfortunately, I am too lazy to write a proper changelog with all the changes.
+
+I added SDXL training to sliders... but.. it does not work properly. 
+The slider training relies on a model's ability to understand that an unconditional (negative prompt)
+means you do not want that concept in the output. SDXL does not understand this for whatever reason, 
+which makes separating out
+concepts within the model hard. I am sure the community will find a way to fix this 
+over time, but for now, it is not 
+going to work properly. And if any of you are thinking "Could we maybe fix it by adding 1 or 2 more text
+encoders to the model as well as a few more entirely separate diffusion networks?" No. God no. It just needs a little
+training without every experimental new paper added to it. The KISS principal. 
+
+
 #### 2021-07-30
 Added "anchors" to the slider trainer. This allows you to set a prompt that will be used as a 
 regularizer. You can set the network multiplier to force spread consistency at high weights

config/examples/mod_lora_scale.yaml

@@ -0,0 +1,48 @@
+---
+job: mod
+config:
+  name: name_of_your_model_v1
+  process:
+    - type: rescale_lora
+      # path to your current lora model
+      input_path: "/path/to/lora/lora.safetensors"
+      # output path for your new lora model, can be the same as input_path to replace
+      output_path: "/path/to/lora/output_lora_v1.safetensors"
+      # replaces meta with the meta below (plus minimum meta fields)
+      # if false, we will leave the meta alone except for updating hashes (sd-script hashes)
+      replace_meta: true
+      # how to adjust, we can scale the up_down weights or the alpha
+      # up_down is the default and probably the best, they will both net the same outputs
+      # would only affect rare NaN cases and maybe merging with old merge tools
+      scale_target: 'up_down'
+      # precision to save, fp16 is the default and standard
+      save_dtype: fp16
+      # current_weight is the ideal weight you use as a multiplier when using the lora
+      # IE in automatic1111 <lora:my_lora:6.0> the 6.0 is the current_weight
+      # you can do negatives here too if you want to flip the lora
+      current_weight: 6.0
+      # target_weight is the ideal weight you use as a multiplier when using the lora
+      # instead of the one above. IE in automatic1111 instead of using <lora:my_lora:6.0>
+      # we want to use <lora:my_lora:1.0> so 1.0 is the target_weight
+      target_weight: 1.0
+
+      # base model for the lora
+      # this is just used to add meta so automatic111 knows which model it is for
+      # assume v1.5 if these are not set
+      is_xl: false
+      is_v2: false
+meta:
+  # this is only used if you set replace_meta to true above
+  name: "[name]"  # [name] gets replaced with the name above
+  description: A short description of your lora
+  trigger_words:
+  - put
+  - trigger
+  - words
+  - here
+  version: '0.1'
+  creator:
+    name: Your Name
+    email: your@email.com
+    website: https://yourwebsite.com
+  any: All meta data above is arbitrary, it can be whatever you want.

info.py

@@ -3,6 +3,6 @@ from collections import OrderedDict
 v = OrderedDict()
 v["name"] = "ai-toolkit"
 v["repo"] = "https://github.com/ostris/ai-toolkit"
-v["version"] = "0.0.1"
+v["version"] = "0.0.2"
 
 software_meta = v

jobs/ModJob.py

@@ -0,0 +1,28 @@
+import os
+from collections import OrderedDict
+from jobs import BaseJob
+from toolkit.metadata import get_meta_for_safetensors
+from toolkit.train_tools import get_torch_dtype
+
+process_dict = {
+    'rescale_lora': 'ModRescaleLoraProcess',
+}
+
+
+class ModJob(BaseJob):
+
+    def __init__(self, config: OrderedDict):
+        super().__init__(config)
+        self.device = self.get_conf('device', 'cpu')
+
+        # loads the processes from the config
+        self.load_processes(process_dict)
+
+    def run(self):
+        super().run()
+
+        print("")
+        print(f"Running  {len(self.process)} process{'' if len(self.process) == 1 else 'es'}")
+
+        for process in self.process:
+            process.run()

jobs/TrainJob.py

@@ -17,6 +17,7 @@ sys.path.append(REPOS_ROOT)
 process_dict = {
     'vae': 'TrainVAEProcess',
     'slider': 'TrainSliderProcess',
+    'slider_old': 'TrainSliderProcessOld',
     'lora_hack': 'TrainLoRAHack',
     'rescale_sd': 'TrainSDRescaleProcess',
 }

jobs/__init__.py

@@ -2,3 +2,4 @@ from .BaseJob import BaseJob
 from .ExtractJob import ExtractJob
 from .TrainJob import TrainJob
 from .MergeJob import MergeJob
+from .ModJob import ModJob

jobs/process/BaseSDTrainProcess.py

@@ -3,23 +3,23 @@ import time
 from collections import OrderedDict
 import os
 
-import diffusers
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import rescale_noise_cfg
-from safetensors import safe_open
 
-from library import sdxl_train_util, sdxl_model_util
-from toolkit.kohya_model_util import load_vae
 from toolkit.lora_special import LoRASpecialNetwork
 from toolkit.optimizer import get_optimizer
 from toolkit.paths import REPOS_ROOT
 import sys
 
+from toolkit.pipelines import CustomStableDiffusionXLPipeline, CustomStableDiffusionPipeline
+
 sys.path.append(REPOS_ROOT)
 sys.path.append(os.path.join(REPOS_ROOT, 'leco'))
 
-from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, DDPMScheduler
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, KDPM2DiscreteScheduler, PNDMScheduler, \
+    DDIMScheduler, DDPMScheduler
 
 from jobs.process import BaseTrainProcess
-from toolkit.metadata import get_meta_for_safetensors, load_metadata_from_safetensors
+from toolkit.metadata import get_meta_for_safetensors, load_metadata_from_safetensors, add_base_model_info_to_meta
 from toolkit.train_tools import get_torch_dtype, apply_noise_offset
 import gc
 
@@ -41,8 +41,9 @@ VAE_SCALE_FACTOR = 8  # 2 ** (len(vae.config.block_out_channels) - 1) = 8
 
 
 class BaseSDTrainProcess(BaseTrainProcess):
-    def __init__(self, process_id: int, job, config: OrderedDict):
+    def __init__(self, process_id: int, job, config: OrderedDict, custom_pipeline=None):
         super().__init__(process_id, job, config)
+        self.custom_pipeline = custom_pipeline
         self.step_num = 0
         self.start_step = 0
         self.device = self.get_conf('device', self.job.device)
@@ -53,8 +54,13 @@ class BaseSDTrainProcess(BaseTrainProcess):
         self.model_config = ModelConfig(**self.get_conf('model', {}))
         self.save_config = SaveConfig(**self.get_conf('save', {}))
         self.sample_config = SampleConfig(**self.get_conf('sample', {}))
-        self.first_sample_config = SampleConfig(
+        first_sample_config = self.get_conf('first_sample', None)
-            **self.get_conf('first_sample', {})) if 'first_sample' in self.config else self.sample_config
+        if first_sample_config is not None:
+            self.has_first_sample_requested = True
+            self.first_sample_config = SampleConfig(**first_sample_config)
+        else:
+            self.has_first_sample_requested = False
+            self.first_sample_config = self.sample_config
         self.logging_config = LogingConfig(**self.get_conf('logging', {}))
         self.optimizer = None
         self.lr_scheduler = None
@@ -186,6 +192,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
                             num_inference_steps=sample_config.sample_steps,
                             guidance_scale=sample_config.guidance_scale,
                             negative_prompt=neg,
+                            guidance_rescale=0.7,
                         ).images[0]
                     else:
                         img = pipeline(
@@ -230,18 +237,26 @@ class BaseSDTrainProcess(BaseTrainProcess):
         # self.sd.tokenizer.to(original_device_dict['tokenizer'])
 
     def update_training_metadata(self):
-        dict = OrderedDict({
+        o_dict = OrderedDict({
             "training_info": self.get_training_info()
         })
         if self.model_config.is_v2:
-            dict['ss_v2'] = True
+            o_dict['ss_v2'] = True
+            o_dict['ss_base_model_version'] = 'sd_2.1'
 
-        if self.model_config.is_xl:
+        elif self.model_config.is_xl:
-            dict['ss_base_model_version'] = 'sdxl_1.0'
+            o_dict['ss_base_model_version'] = 'sdxl_1.0'
+        else:
+            o_dict['ss_base_model_version'] = 'sd_1.5'
 
-        dict['ss_output_name'] = self.job.name
+        o_dict = add_base_model_info_to_meta(
+            o_dict,
+            is_v2=self.model_config.is_v2,
+            is_xl=self.model_config.is_xl,
+        )
+        o_dict['ss_output_name'] = self.job.name
 
-        self.add_meta(dict)
+        self.add_meta(o_dict)
 
     def get_training_info(self):
         info = OrderedDict({
@@ -299,6 +314,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
             )
 
         self.print(f"Saved to {file_path}")
+        self.clean_up_saves()
 
     # Called before the model is loaded
     def hook_before_model_load(self):
@@ -321,7 +337,6 @@ class BaseSDTrainProcess(BaseTrainProcess):
     ):
         if height is None and pixel_height is None:
             raise ValueError("height or pixel_height must be specified")
-            raise ValueError("height or pixel_height must be specified")
         if width is None and pixel_width is None:
             raise ValueError("width or pixel_width must be specified")
         if height is None:
@@ -372,25 +387,15 @@ class BaseSDTrainProcess(BaseTrainProcess):
             text_embeddings: PromptEmbeds,
             timestep: int,
             guidance_scale=7.5,
-            guidance_rescale=0.7,
+            guidance_rescale=0, # 0.7
             add_time_ids=None,
             **kwargs,
     ):
+
         if self.sd.is_xl:
             if add_time_ids is None:
                 add_time_ids = self.get_time_ids_from_latents(latents)
-            # todo LECOs code looks like it is omitting noise_pred
+
-            # noise_pred = train_util.predict_noise_xl(
-            #     self.sd.unet,
-            #     self.sd.noise_scheduler,
-            #     timestep,
-            #     latents,
-            #     text_embeddings.text_embeds,
-            #     text_embeddings.pooled_embeds,
-            #     add_time_ids,
-            #     guidance_scale=guidance_scale,
-            #     guidance_rescale=guidance_rescale
-            # )
             latent_model_input = torch.cat([latents] * 2)
 
             latent_model_input = self.sd.noise_scheduler.scale_model_input(latent_model_input, timestep)
@@ -410,25 +415,31 @@ class BaseSDTrainProcess(BaseTrainProcess):
 
             # perform guidance
             noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-            guided_target = noise_pred_uncond + guidance_scale * (
+            noise_pred = noise_pred_uncond + guidance_scale * (
                     noise_pred_text - noise_pred_uncond
             )
 
             # https://github.com/huggingface/diffusers/blob/7a91ea6c2b53f94da930a61ed571364022b21044/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L775
-            # noise_pred = rescale_noise_cfg(
+            if guidance_rescale > 0.0:
-            #     noise_pred, noise_pred_text, guidance_rescale=guidance_rescale
+                # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
-            # )
+                noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
-
-            noise_pred = guided_target
 
         else:
-            noise_pred = train_util.predict_noise(
+            # if we are doing classifier free guidance, need to double up
-                self.sd.unet,
+            latent_model_input = torch.cat([latents] * 2)
-                self.sd.noise_scheduler,
+
+            latent_model_input = self.sd.noise_scheduler.scale_model_input(latent_model_input, timestep)
+
+            # predict the noise residual
+            noise_pred = self.sd.unet(
+                latent_model_input,
                 timestep,
-                latents,
+                encoder_hidden_states=text_embeddings.text_embeds,
-                text_embeddings.text_embeds if hasattr(text_embeddings, 'text_embeds') else text_embeddings,
+            ).sample
-                guidance_scale=guidance_scale
+
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
             )
 
         return noise_pred
@@ -493,56 +504,80 @@ class BaseSDTrainProcess(BaseTrainProcess):
 
         dtype = get_torch_dtype(self.train_config.dtype)
 
+        # TODO handle other schedulers
+        # sch = KDPM2DiscreteScheduler
+        sch = DDPMScheduler
+        # do our own scheduler
+        prediction_type = "v_prediction" if self.model_config.is_v_pred else "epsilon"
+        scheduler = sch(
+            num_train_timesteps=1000,
+            beta_start=0.00085,
+            beta_end=0.0120,
+            beta_schedule="scaled_linear",
+            clip_sample=False,
+            prediction_type=prediction_type,
+        )
         if self.model_config.is_xl:
-
+            if self.custom_pipeline is not None:
-            pipe = StableDiffusionXLPipeline.from_single_file(
+                pipln = self.custom_pipeline
+            else:
+                pipln = CustomStableDiffusionXLPipeline
+            pipe = pipln.from_single_file(
                 self.model_config.name_or_path,
                 dtype=dtype,
-                scheduler_type='pndm',
+                scheduler_type='ddpm',
-                device=self.device_torch
+                device=self.device_torch,
-            )
+            ).to(self.device_torch)
+
             text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
             tokenizer = [pipe.tokenizer, pipe.tokenizer_2]
-            unet = pipe.unet
-            noise_scheduler = pipe.scheduler
-            vae = pipe.vae.to('cpu', dtype=dtype)
-            vae.eval()
-            vae.set_use_memory_efficient_attention_xformers(True)
-
             for text_encoder in text_encoders:
                 text_encoder.to(self.device_torch, dtype=dtype)
                 text_encoder.requires_grad_(False)
                 text_encoder.eval()
-
             text_encoder = text_encoders
-            tokenizer = tokenizer
-            del pipe
-            flush()
-
-
         else:
-            tokenizer, text_encoder, unet, noise_scheduler = model_util.load_models(
+            if self.custom_pipeline is not None:
+                pipln = self.custom_pipeline
+            else:
+                pipln = CustomStableDiffusionPipeline
+            pipe = pipln.from_single_file(
                 self.model_config.name_or_path,
-                scheduler_name=self.train_config.noise_scheduler,
+                dtype=dtype,
-                v2=self.model_config.is_v2,
+                scheduler_type='dpm',
-                v_pred=self.model_config.is_v_pred,
+                device=self.device_torch,
-            )
+                load_safety_checker=False,
-
+            ).to(self.device_torch)
+            pipe.register_to_config(requires_safety_checker=False)
+            text_encoder = pipe.text_encoder
             text_encoder.to(self.device_torch, dtype=dtype)
+            text_encoder.requires_grad_(False)
             text_encoder.eval()
-            vae = load_vae(self.model_config.name_or_path, dtype=dtype).to('cpu', dtype=dtype)
+            tokenizer = pipe.tokenizer
-            vae.eval()
+
+        # scheduler doesn't get set sometimes, so we set it here
+        pipe.scheduler = scheduler
+
+        unet = pipe.unet
+        noise_scheduler = pipe.scheduler
+        vae = pipe.vae.to('cpu', dtype=dtype)
+        vae.eval()
+        vae.requires_grad_(False)
         flush()
 
-
+        self.sd = StableDiffusion(
-        # just for now or of we want to load a custom one
+            vae,
-        # put on cpu for now, we only need it when sampling
+            tokenizer,
-        # vae = load_vae(self.model_config.name_or_path, dtype=dtype).to('cpu', dtype=dtype)
+            text_encoder,
-        # vae.eval()
+            unet,
-        self.sd = StableDiffusion(vae, tokenizer, text_encoder, unet, noise_scheduler, is_xl=self.model_config.is_xl)
+            noise_scheduler,
+            is_xl=self.model_config.is_xl,
+            pipeline=pipe
+        )
 
         unet.to(self.device_torch, dtype=dtype)
         if self.train_config.xformers:
+            vae.set_use_memory_efficient_attention_xformers(True)
             unet.enable_xformers_memory_efficient_attention()
         if self.train_config.gradient_checkpointing:
             unet.enable_gradient_checkpointing()
@@ -588,19 +623,26 @@ class BaseSDTrainProcess(BaseTrainProcess):
                 self.network.multiplier = 1.0
 
 
-
         else:
             params = []
             # assume dreambooth/finetune
             if self.train_config.train_text_encoder:
-                text_encoder.requires_grad_(True)
+                if self.sd.is_xl:
-                text_encoder.train()
+                    for te in text_encoder:
-                params += text_encoder.parameters()
+                        te.requires_grad_(True)
+                        te.train()
+                        params += te.parameters()
+                else:
+                    text_encoder.requires_grad_(True)
+                    text_encoder.train()
+                    params += text_encoder.parameters()
             if self.train_config.train_unet:
                 unet.requires_grad_(True)
                 unet.train()
                 params += unet.parameters()
 
+            # TODO recover save if training network. Maybe load from beginning
+
         ### HOOK ###
         params = self.hook_add_extra_train_params(params)
 
@@ -613,7 +655,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
             self.train_config.lr_scheduler,
             optimizer,
             max_iterations=self.train_config.steps,
-            lr_min=self.train_config.lr / 100,  # not sure why leco did this, but ill do it to
+            lr_min=self.train_config.lr / 100,
         )
 
         self.lr_scheduler = lr_scheduler
@@ -621,20 +663,25 @@ class BaseSDTrainProcess(BaseTrainProcess):
         ### HOOK ###
         self.hook_before_train_loop()
 
+        if self.has_first_sample_requested:
+            self.print("Generating first sample from first sample config")
+            self.sample(0, is_first=True)
+
         # sample first
         if self.train_config.skip_first_sample:
             self.print("Skipping first sample due to config setting")
         else:
             self.print("Generating baseline samples before training")
-            self.sample(0, is_first=True)
+            self.sample(0)
 
         self.progress_bar = tqdm(
             total=self.train_config.steps,
             desc=self.job.name,
-            leave=True
+            leave=True,
+            initial=self.step_num,
+            iterable=range(0, self.train_config.steps),
         )
-        # set it to our current step in case it was updated from a load
+
-        self.progress_bar.update(self.step_num)
         # self.step_num = 0
         for step in range(self.step_num, self.train_config.steps):
             # todo handle dataloader here maybe, not sure
@@ -642,7 +689,8 @@ class BaseSDTrainProcess(BaseTrainProcess):
             ### HOOK ###
             loss_dict = self.hook_train_loop()
 
-            if self.train_config.optimizer.lower().startswith('dadaptation'):
+            if self.train_config.optimizer.lower().startswith('dadaptation') or \
+                    self.train_config.optimizer.lower().startswith('prodigy'):
                 learning_rate = (
                         optimizer.param_groups[0]["d"] *
                         optimizer.param_groups[0]["lr"]

jobs/process/ModRescaleLoraProcess.py

@@ -0,0 +1,100 @@
+import gc
+import os
+from collections import OrderedDict
+from typing import ForwardRef
+
+import torch
+from safetensors.torch import save_file, load_file
+
+from jobs.process.BaseProcess import BaseProcess
+from toolkit.metadata import get_meta_for_safetensors, load_metadata_from_safetensors, add_model_hash_to_meta, \
+    add_base_model_info_to_meta
+from toolkit.train_tools import get_torch_dtype
+
+
+class ModRescaleLoraProcess(BaseProcess):
+    process_id: int
+    config: OrderedDict
+    progress_bar: ForwardRef('tqdm') = None
+
+    def __init__(
+            self,
+            process_id: int,
+            job,
+            config: OrderedDict
+    ):
+        super().__init__(process_id, job, config)
+        self.input_path = self.get_conf('input_path', required=True)
+        self.output_path = self.get_conf('output_path', required=True)
+        self.replace_meta = self.get_conf('replace_meta', default=False)
+        self.save_dtype = self.get_conf('save_dtype', default='fp16', as_type=get_torch_dtype)
+        self.current_weight = self.get_conf('current_weight', required=True, as_type=float)
+        self.target_weight = self.get_conf('target_weight', required=True, as_type=float)
+        self.scale_target = self.get_conf('scale_target', default='up_down')  # alpha or up_down
+        self.is_xl = self.get_conf('is_xl', default=False, as_type=bool)
+        self.is_v2 = self.get_conf('is_v2', default=False, as_type=bool)
+
+        self.progress_bar = None
+
+    def run(self):
+        super().run()
+        source_state_dict = load_file(self.input_path)
+        source_meta = load_metadata_from_safetensors(self.input_path)
+
+        if self.replace_meta:
+            self.meta.update(
+                add_base_model_info_to_meta(
+                    self.meta,
+                    is_xl=self.is_xl,
+                    is_v2=self.is_v2,
+                )
+            )
+            save_meta = get_meta_for_safetensors(self.meta, self.job.name)
+        else:
+            save_meta = get_meta_for_safetensors(source_meta, self.job.name, add_software_info=False)
+
+        # save
+        os.makedirs(os.path.dirname(self.output_path), exist_ok=True)
+
+        new_state_dict = OrderedDict()
+
+        for key in list(source_state_dict.keys()):
+            v = source_state_dict[key]
+            v = v.detach().clone().to("cpu").to(get_torch_dtype('fp32'))
+
+            # all loras have an alpha, up weight and down weight
+            #  - "lora_te_text_model_encoder_layers_0_mlp_fc1.alpha",
+            #  - "lora_te_text_model_encoder_layers_0_mlp_fc1.lora_down.weight",
+            #  - "lora_te_text_model_encoder_layers_0_mlp_fc1.lora_up.weight",
+            # we can rescale by adjusting the alpha or the up weights, or the up and down weights
+            # I assume doing both up and down would be best all around, but I'm not sure
+            # some locons also have mid weights, we will leave those alone for now, will work without them
+
+            # when adjusting alpha, it is used to calculate the multiplier in a lora module
+            #  - scale = alpha / lora_dim
+            #  - output = layer_out + lora_up_out * multiplier * scale
+            total_module_scale = torch.tensor(self.current_weight / self.target_weight) \
+                .to("cpu", dtype=get_torch_dtype('fp32'))
+            num_modules_layers = 2  # up and down
+            up_down_scale = torch.pow(total_module_scale, 1.0 / num_modules_layers) \
+                .to("cpu", dtype=get_torch_dtype('fp32'))
+            # only update alpha
+            if self.scale_target == 'alpha' and key.endswith('.alpha'):
+                v = v * total_module_scale
+            if self.scale_target == 'up_down' and key.endswith('.lora_up.weight') or key.endswith('.lora_down.weight'):
+                # would it be better to adjust the up weights for fp16 precision? Doing both should reduce chance of NaN
+                v = v * up_down_scale
+            new_state_dict[key] = v.to(get_torch_dtype(self.save_dtype))
+
+        save_meta = add_model_hash_to_meta(new_state_dict, save_meta)
+        save_file(new_state_dict, self.output_path, save_meta)
+
+        # cleanup incase there are other jobs
+        del new_state_dict
+        del source_state_dict
+        del source_meta
+
+        torch.cuda.empty_cache()
+        gc.collect()
+
+        print(f"Saved to {self.output_path}")

jobs/process/TrainSDRescaleProcess.py

@@ -5,6 +5,7 @@ from collections import OrderedDict
 import os
 from typing import Optional
 
+import numpy as np
 from safetensors.torch import load_file, save_file
 from tqdm import tqdm
 
@@ -14,6 +15,7 @@ from toolkit.paths import REPOS_ROOT
 import sys
 
 from toolkit.stable_diffusion_model import PromptEmbeds
+from toolkit.train_pipelines import TransferStableDiffusionXLPipeline
 
 sys.path.append(REPOS_ROOT)
 sys.path.append(os.path.join(REPOS_ROOT, 'leco'))
@@ -61,7 +63,8 @@ class PromptEmbedsCache:
 
 class TrainSDRescaleProcess(BaseSDTrainProcess):
     def __init__(self, process_id: int, job, config: OrderedDict):
-        super().__init__(process_id, job, config)
+        # pass our custom pipeline to super so it sets it up
+        super().__init__(process_id, job, config, custom_pipeline=TransferStableDiffusionXLPipeline)
         self.step_num = 0
         self.start_step = 0
         self.device = self.get_conf('device', self.job.device)
@@ -129,9 +132,12 @@ class TrainSDRescaleProcess(BaseSDTrainProcess):
                     print(f"Saving prompt tensors to {self.rescale_config.prompt_tensors}")
                     state_dict = {}
                     for prompt_txt, prompt_embeds in cache.prompts.items():
-                        state_dict[f"te:{prompt_txt}"] = prompt_embeds.text_embeds.to("cpu", dtype=get_torch_dtype('fp16'))
+                        state_dict[f"te:{prompt_txt}"] = prompt_embeds.text_embeds.to("cpu",
+                                                                                      dtype=get_torch_dtype('fp16'))
                         if prompt_embeds.pooled_embeds is not None:
-                            state_dict[f"pe:{prompt_txt}"] = prompt_embeds.pooled_embeds.to("cpu", dtype=get_torch_dtype('fp16'))
+                            state_dict[f"pe:{prompt_txt}"] = prompt_embeds.pooled_embeds.to("cpu",
+                                                                                            dtype=get_torch_dtype(
+                                                                                                'fp16'))
                     save_file(state_dict, self.rescale_config.prompt_tensors)
 
             self.print("Encoding complete.")
@@ -157,37 +163,25 @@ class TrainSDRescaleProcess(BaseSDTrainProcess):
             torch.randint(0, len(self.prompt_txt_list), (1,)).item()
         ]
         prompt = self.prompt_cache[prompt_txt].to(device=self.device_torch, dtype=dtype)
+        prompt.text_embeds.to(device=self.device_torch, dtype=dtype)
         neutral = self.prompt_cache[""].to(device=self.device_torch, dtype=dtype)
+        neutral.text_embeds.to(device=self.device_torch, dtype=dtype)
+        if hasattr(prompt, 'pooled_embeds') \
+                and hasattr(neutral, 'pooled_embeds') \
+                and prompt.pooled_embeds is not None \
+                and neutral.pooled_embeds is not None:
+            prompt.pooled_embeds.to(device=self.device_torch, dtype=dtype)
+            neutral.pooled_embeds.to(device=self.device_torch, dtype=dtype)
+
         if prompt is None:
             raise ValueError(f"Prompt {prompt_txt} is not in cache")
 
-        prompt_batch = train_tools.concat_prompt_embeddings(
-            prompt,
-            neutral,
-            self.train_config.batch_size,
-        )
-
-        noise_scheduler = self.sd.noise_scheduler
-        optimizer = self.optimizer
-        lr_scheduler = self.lr_scheduler
         loss_function = torch.nn.MSELoss()
 
-        def get_noise_pred(p, n, gs, cts, dn):
-            return self.predict_noise(
-                latents=dn,
-                text_embeddings=train_tools.concat_prompt_embeddings(
-                    p,  # unconditional
-                    n,  # positive
-                    self.train_config.batch_size,
-                ),
-                timestep=cts,
-                guidance_scale=gs,
-            )
-
         with torch.no_grad():
-            self.sd.noise_scheduler.set_timesteps(
+            # self.sd.noise_scheduler.set_timesteps(
-                self.train_config.max_denoising_steps, device=self.device_torch
+            #     self.train_config.max_denoising_steps, device=self.device_torch
-            )
+            # )
 
             self.optimizer.zero_grad()
 
@@ -197,81 +191,76 @@ class TrainSDRescaleProcess(BaseSDTrainProcess):
             ).item()
 
             # get noise
-            noise = self.get_latent_noise(
+            latents = self.get_latent_noise(
                 pixel_height=self.rescale_config.from_resolution,
                 pixel_width=self.rescale_config.from_resolution,
             ).to(self.device_torch, dtype=dtype)
 
-            # get latents
+            self.sd.pipeline.to(self.device_torch)
-            latents = noise * self.sd.noise_scheduler.init_noise_sigma
+            torch.set_default_device(self.device_torch)
-            latents = latents.to(self.device_torch, dtype=dtype)
-            #
-            # # predict without network
-            # assert self.network.is_active is False
-            # denoised_latents = self.diffuse_some_steps(
-            #     latents,  # pass simple noise latents
-            #     prompt_batch,
-            #     start_timesteps=0,
-            #     total_timesteps=timesteps_to,
-            #     guidance_scale=3,
-            # )
-            # noise_scheduler.set_timesteps(1000)
-            #
-            # current_timestep = noise_scheduler.timesteps[
-            #     int(timesteps_to * 1000 / self.train_config.max_denoising_steps)
-            # ]
 
-            current_timestep = 0
+            # turn off progress bar
-            denoised_latents = latents
+            self.sd.pipeline.set_progress_bar_config(disable=True)
-            # get noise prediction at full scale
+
-            from_prediction = get_noise_pred(
+            # get random guidance scale from 1.0 to 10.0
-                prompt, neutral, 1, current_timestep, denoised_latents
+            guidance_scale = torch.rand(1).item() * 9.0 + 1.0
+
+            loss_arr = []
+
+
+            max_len_timestep_str = len(str(self.train_config.max_denoising_steps))
+            # pad with spaces
+            timestep_str = str(timesteps_to).rjust(max_len_timestep_str, " ")
+            new_description = f"{self.job.name} ts: {timestep_str}"
+            self.progress_bar.set_description(new_description)
+
+            def pre_condition_callback(target_pred, input_latents):
+                # handle any manipulations before feeding to our network
+                reduced_pred = self.reduce_size_fn(target_pred)
+                reduced_latents = self.reduce_size_fn(input_latents)
+                self.optimizer.zero_grad()
+                return reduced_pred, reduced_latents
+
+            def each_step_callback(noise_target, noise_train_pred):
+                noise_target.requires_grad = False
+                loss = loss_function(noise_target, noise_train_pred)
+                loss_arr.append(loss.item())
+                loss.backward()
+                self.optimizer.step()
+                self.lr_scheduler.step()
+                self.optimizer.zero_grad()
+
+            # run the pipeline
+            self.sd.pipeline.transfer_diffuse(
+                num_inference_steps=timesteps_to,
+                latents=latents,
+                prompt_embeds=prompt.text_embeds,
+                negative_prompt_embeds=neutral.text_embeds,
+                pooled_prompt_embeds=prompt.pooled_embeds,
+                negative_pooled_prompt_embeds=neutral.pooled_embeds,
+                output_type="latent",
+                num_images_per_prompt=self.train_config.batch_size,
+                guidance_scale=guidance_scale,
+                network=self.network,
+                target_unet=self.sd.unet,
+                pre_condition_callback=pre_condition_callback,
+                each_step_callback=each_step_callback,
             )
 
-            reduced_from_prediction = self.reduce_size_fn(from_prediction).to("cpu", dtype=torch.float32)
-
-            # get noise prediction at reduced scale
-            to_denoised_latents = self.reduce_size_fn(denoised_latents)
-
-        # start gradient
-        optimizer.zero_grad()
-        self.network.multiplier = 1.0
-        with self.network:
-            assert self.network.is_active is True
-            to_prediction = get_noise_pred(
-                prompt, neutral, 1, current_timestep, to_denoised_latents
-            ).to("cpu", dtype=torch.float32)
-
-        reduced_from_prediction.requires_grad = False
-        from_prediction.requires_grad = False
-
-        loss = loss_function(
-            reduced_from_prediction,
-            to_prediction,
-        )
-
-        loss_float = loss.item()
-
-        loss = loss.to(self.device_torch)
-
-        loss.backward()
-        optimizer.step()
-        lr_scheduler.step()
-
-        del (
-            reduced_from_prediction,
-            from_prediction,
-            to_denoised_latents,
-            to_prediction,
-            latents,
-        )
         flush()
 
         # reset network
         self.network.multiplier = 1.0
 
+        # average losses
+        s = 0
+        for num in loss_arr:
+            s += num
+
+        avg_loss = s / len(loss_arr)
+
         loss_dict = OrderedDict(
-            {'loss': loss_float},
+            {'loss': avg_loss},
         )
 
         return loss_dict

jobs/process/TrainSliderProcess.py

@@ -1,10 +1,14 @@
 # ref:
 # - https://github.com/p1atdev/LECO/blob/main/train_lora.py
+import random
 import time
 from collections import OrderedDict
 import os
 from typing import Optional
 
+from safetensors.torch import save_file, load_file
+from tqdm import tqdm
+
 from toolkit.config_modules import SliderConfig
 from toolkit.paths import REPOS_ROOT
 import sys
@@ -36,27 +40,45 @@ class EncodedPromptPair:
     def __init__(
             self,
             target_class,
-            positive,
+            target_class_with_neutral,
-            negative,
+            positive_target,
+            positive_target_with_neutral,
+            negative_target,
+            negative_target_with_neutral,
             neutral,
-            width=512,
+            empty_prompt,
-            height=512,
+            both_targets,
             action=ACTION_TYPES_SLIDER.ERASE_NEGATIVE,
             multiplier=1.0,
             weight=1.0
     ):
         self.target_class = target_class
-        self.positive = positive
+        self.target_class_with_neutral = target_class_with_neutral
-        self.negative = negative
+        self.positive_target = positive_target
+        self.positive_target_with_neutral = positive_target_with_neutral
+        self.negative_target = negative_target
+        self.negative_target_with_neutral = negative_target_with_neutral
         self.neutral = neutral
-        self.width = width
+        self.empty_prompt = empty_prompt
-        self.height = height
+        self.both_targets = both_targets
-        self.action: int = action
         self.multiplier = multiplier
+        self.action: int = action
         self.weight = weight
 
+    # simulate torch to for tensors
+    def to(self, *args, **kwargs):
+        self.target_class = self.target_class.to(*args, **kwargs)
+        self.positive_target = self.positive_target.to(*args, **kwargs)
+        self.positive_target_with_neutral = self.positive_target_with_neutral.to(*args, **kwargs)
+        self.negative_target = self.negative_target.to(*args, **kwargs)
+        self.negative_target_with_neutral = self.negative_target_with_neutral.to(*args, **kwargs)
+        self.neutral = self.neutral.to(*args, **kwargs)
+        self.empty_prompt = self.empty_prompt.to(*args, **kwargs)
+        self.both_targets = self.both_targets.to(*args, **kwargs)
+        return self
 
-class PromptEmbedsCache:  # 使いまわしたいので
+
+class PromptEmbedsCache:
     prompts: dict[str, PromptEmbeds] = {}
 
     def __setitem__(self, __name: str, __value: PromptEmbeds) -> None:
@@ -84,6 +106,7 @@ class EncodedAnchor:
 class TrainSliderProcess(BaseSDTrainProcess):
     def __init__(self, process_id: int, job, config: OrderedDict):
         super().__init__(process_id, job, config)
+        self.prompt_txt_list = None
         self.step_num = 0
         self.start_step = 0
         self.device = self.get_conf('device', self.job.device)
@@ -97,24 +120,76 @@ class TrainSliderProcess(BaseSDTrainProcess):
         pass
 
     def hook_before_train_loop(self):
+        self.print(f"Loading prompt file from {self.slider_config.prompt_file}")
+
+        # read line by line from file
+        if self.slider_config.prompt_file:
+            with open(self.slider_config.prompt_file, 'r') as f:
+                self.prompt_txt_list = f.readlines()
+                # clean empty lines
+                self.prompt_txt_list = [line.strip() for line in self.prompt_txt_list if len(line.strip()) > 0]
+
+            self.print(f"Loaded {len(self.prompt_txt_list)} prompts. Encoding them..")
+
+
+            if not self.slider_config.prompt_tensors:
+                # shuffle
+                random.shuffle(self.prompt_txt_list)
+                # trim to max steps
+                self.prompt_txt_list = self.prompt_txt_list[:self.train_config.steps]
+                # trim list to our max steps
+
         cache = PromptEmbedsCache()
-        prompt_pairs: list[EncodedPromptPair] = []
 
         # get encoded latents for our prompts
         with torch.no_grad():
-            neutral = ""
+            if self.slider_config.prompt_tensors is not None:
-            for target in self.slider_config.targets:
+                # check to see if it exists
-                # build the cache
+                if os.path.exists(self.slider_config.prompt_tensors):
-                for prompt in [
+                    # load it.
-                    target.target_class,
+                    self.print(f"Loading prompt tensors from {self.slider_config.prompt_tensors}")
-                    target.positive,
+                    prompt_tensors = load_file(self.slider_config.prompt_tensors, device='cpu')
-                    target.negative,
+                    # add them to the cache
-                    neutral  # empty neutral
+                    for prompt_txt, prompt_tensor in tqdm(prompt_tensors.items(), desc="Loading prompts", leave=False):
-                ]:
+                        if prompt_txt.startswith("te:"):
-                    if cache[prompt] is None:
+                            prompt = prompt_txt[3:]
-                        cache[prompt] = self.sd.encode_prompt(prompt)
+                            # text_embeds
-                for resolution in self.slider_config.resolutions:
+                            text_embeds = prompt_tensor
-                    width, height = resolution
+                            pooled_embeds = None
+                            # find pool embeds
+                            if f"pe:{prompt}" in prompt_tensors:
+                                pooled_embeds = prompt_tensors[f"pe:{prompt}"]
+
+                            # make it
+                            prompt_embeds = PromptEmbeds([text_embeds, pooled_embeds])
+                            cache[prompt] = prompt_embeds.to(device='cpu', dtype=torch.float32)
+
+            if len(cache.prompts) == 0:
+                print("Prompt tensors not found. Encoding prompts..")
+                empty_prompt = ""
+                # encode empty_prompt
+                cache[empty_prompt] = self.sd.encode_prompt(empty_prompt)
+
+                neutral_list = self.prompt_txt_list if self.prompt_txt_list is not None else [""]
+
+                for neutral in tqdm(neutral_list, desc="Encoding prompts", leave=False):
+                    for target in self.slider_config.targets:
+                        prompt_list = [
+                            f"{target.target_class}", # target_class
+                            f"{target.target_class} {neutral}", # target_class with neutral
+                            f"{target.positive}",  # positive_target
+                            f"{target.positive} {neutral}",  # positive_target with neutral
+                            f"{target.negative}",  # negative_target
+                            f"{target.negative} {neutral}",  # negative_target with neutral
+                            f"{neutral}",  # neutral
+                            f"{target.positive} {target.negative}",  # both targets
+                            f"{target.negative} {target.positive}",  # both targets
+                        ]
+                        for p in prompt_list:
+                            # build the cache
+                            if cache[p] is None:
+                                cache[p] = self.sd.encode_prompt(p).to(device="cpu", dtype=torch.float32)
+
                     erase_negative = len(target.positive.strip()) == 0
                     enhance_positive = len(target.negative.strip()) == 0
 
@@ -127,63 +202,100 @@ class TrainSliderProcess(BaseSDTrainProcess):
                     # if we don't do this, we will get different contrast and focus.
                     # we only perform actions of enhancing and erasing on the negative
                     # todo work on way to do all of this in one shot
+                if self.slider_config.prompt_tensors:
+                    print(f"Saving prompt tensors to {self.slider_config.prompt_tensors}")
+                    state_dict = {}
+                    for prompt_txt, prompt_embeds in cache.prompts.items():
+                        state_dict[f"te:{prompt_txt}"] = prompt_embeds.text_embeds.to("cpu",
+                                                                                      dtype=get_torch_dtype('fp16'))
+                        if prompt_embeds.pooled_embeds is not None:
+                            state_dict[f"pe:{prompt_txt}"] = prompt_embeds.pooled_embeds.to("cpu",
+                                                                                            dtype=get_torch_dtype(
+                                                                                                'fp16'))
+                    save_file(state_dict, self.slider_config.prompt_tensors)
+
+            prompt_pairs = []
+            for neutral in tqdm(neutral_list, desc="Encoding prompts", leave=False):
+                for target in self.slider_config.targets:
+                    erase_negative = len(target.positive.strip()) == 0
+                    enhance_positive = len(target.negative.strip()) == 0
+
+                    both = not erase_negative and not enhance_positive
 
                     if both or erase_negative:
+                        print("Encoding erase negative")
                         prompt_pairs += [
                             # erase standard
                             EncodedPromptPair(
                                 target_class=cache[target.target_class],
-                                positive=cache[target.positive],
+                                target_class_with_neutral=cache[f"{target.target_class} {neutral}"],
-                                negative=cache[target.negative],
+                                positive_target=cache[f"{target.positive}"],
+                                positive_target_with_neutral=cache[f"{target.positive} {neutral}"],
+                                negative_target=cache[f"{target.negative}"],
+                                negative_target_with_neutral=cache[f"{target.negative} {neutral}"],
                                 neutral=cache[neutral],
-                                width=width,
-                                height=height,
                                 action=ACTION_TYPES_SLIDER.ERASE_NEGATIVE,
                                 multiplier=target.multiplier,
+                                both_targets=cache[f"{target.positive} {target.negative}"],
+                                empty_prompt=cache[""],
                                 weight=target.weight
                             ),
                         ]
                     if both or enhance_positive:
+                        print("Encoding enhance positive")
                         prompt_pairs += [
                             # enhance standard, swap pos neg
                             EncodedPromptPair(
                                 target_class=cache[target.target_class],
-                                positive=cache[target.negative],
+                                target_class_with_neutral=cache[f"{target.target_class} {neutral}"],
-                                negative=cache[target.positive],
+                                positive_target=cache[f"{target.negative}"],
+                                positive_target_with_neutral=cache[f"{target.negative} {neutral}"],
+                                negative_target=cache[f"{target.positive}"],
+                                negative_target_with_neutral=cache[f"{target.positive} {neutral}"],
                                 neutral=cache[neutral],
-                                width=width,
-                                height=height,
                                 action=ACTION_TYPES_SLIDER.ENHANCE_NEGATIVE,
                                 multiplier=target.multiplier,
+                                both_targets=cache[f"{target.positive} {target.negative}"],
+                                empty_prompt=cache[""],
                                 weight=target.weight
                             ),
                         ]
-                    if both or enhance_positive:
+                    # if both or enhance_positive:
+                    if both:
+                        print("Encoding erase positive (inverse)")
                         prompt_pairs += [
                             # erase inverted
                             EncodedPromptPair(
                                 target_class=cache[target.target_class],
-                                positive=cache[target.negative],
+                                target_class_with_neutral=cache[f"{target.target_class} {neutral}"],
-                                negative=cache[target.positive],
+                                positive_target=cache[f"{target.negative}"],
+                                positive_target_with_neutral=cache[f"{target.negative} {neutral}"],
+                                negative_target=cache[f"{target.positive}"],
+                                negative_target_with_neutral=cache[f"{target.positive} {neutral}"],
                                 neutral=cache[neutral],
-                                width=width,
-                                height=height,
                                 action=ACTION_TYPES_SLIDER.ERASE_NEGATIVE,
+                                both_targets=cache[f"{target.positive} {target.negative}"],
+                                empty_prompt=cache[""],
                                 multiplier=target.multiplier * -1.0,
                                 weight=target.weight
                             ),
                         ]
-                    if both or erase_negative:
+                    # if both or erase_negative:
+                    if both:
+                        print("Encoding enhance negative (inverse)")
                         prompt_pairs += [
                             # enhance inverted
                             EncodedPromptPair(
                                 target_class=cache[target.target_class],
-                                positive=cache[target.positive],
+                                target_class_with_neutral=cache[f"{target.target_class} {neutral}"],
-                                negative=cache[target.negative],
+                                positive_target=cache[f"{target.positive}"],
+                                positive_target_with_neutral=cache[f"{target.positive} {neutral}"],
+                                negative_target=cache[f"{target.negative}"],
+                                negative_target_with_neutral=cache[f"{target.negative} {neutral}"],
+                                both_targets=cache[f"{target.positive} {target.negative}"],
                                 neutral=cache[neutral],
-                                width=width,
-                                height=height,
                                 action=ACTION_TYPES_SLIDER.ENHANCE_NEGATIVE,
+                                empty_prompt=cache[""],
                                 multiplier=target.multiplier * -1.0,
                                 weight=target.weight
                             ),
@@ -208,6 +320,7 @@ class TrainSliderProcess(BaseSDTrainProcess):
                     )
                 ]
 
+
         # move to cpu to save vram
         # We don't need text encoder anymore, but keep it on cpu for sampling
         # if text encoder is list
@@ -225,17 +338,22 @@ class TrainSliderProcess(BaseSDTrainProcess):
     def hook_train_loop(self):
         dtype = get_torch_dtype(self.train_config.dtype)
 
+        # get random multiplier between 1 and 3
+        rand_weight = 1
+        # rand_weight = torch.rand((1,)).item() * 2 + 1
+
         # get a random pair
         prompt_pair: EncodedPromptPair = self.prompt_pairs[
             torch.randint(0, len(self.prompt_pairs), (1,)).item()
         ]
+        # move to device and dtype
+        prompt_pair.to(self.device_torch, dtype=dtype)
+
+        # get a random resolution
+        height, width = self.slider_config.resolutions[
+            torch.randint(0, len(self.slider_config.resolutions), (1,)).item()
+        ]
 
-        height = prompt_pair.height
-        width = prompt_pair.width
-        target_class = prompt_pair.target_class
-        neutral = prompt_pair.neutral
-        negative = prompt_pair.negative
-        positive = prompt_pair.positive
         weight = prompt_pair.weight
         multiplier = prompt_pair.multiplier
 
@@ -245,12 +363,12 @@ class TrainSliderProcess(BaseSDTrainProcess):
         lr_scheduler = self.lr_scheduler
         loss_function = torch.nn.MSELoss()
 
-        def get_noise_pred(p, n, gs, cts, dn):
+        def get_noise_pred(neg, pos, gs, cts, dn):
             return self.predict_noise(
                 latents=dn,
                 text_embeddings=train_tools.concat_prompt_embeddings(
-                    p,  # unconditional
+                    neg,  # negative prompt
-                    n,  # positive
+                    pos,  # positive prompt
                     self.train_config.batch_size,
                 ),
                 timestep=cts,
@@ -258,7 +376,7 @@ class TrainSliderProcess(BaseSDTrainProcess):
             )
 
         # set network multiplier
-        self.network.multiplier = multiplier
+        self.network.multiplier = multiplier * rand_weight
 
         with torch.no_grad():
             self.sd.noise_scheduler.set_timesteps(
@@ -284,12 +402,12 @@ class TrainSliderProcess(BaseSDTrainProcess):
 
             with self.network:
                 assert self.network.is_active
-                self.network.multiplier = multiplier
+                self.network.multiplier = multiplier * rand_weight
                 denoised_latents = self.diffuse_some_steps(
                     latents,  # pass simple noise latents
                     train_tools.concat_prompt_embeddings(
-                        positive,  # unconditional
+                        prompt_pair.positive_target,  # unconditional
-                        target_class,  # target
+                        prompt_pair.target_class,  # target
                         self.train_config.batch_size,
                     ),
                     start_timesteps=0,
@@ -304,15 +422,27 @@ class TrainSliderProcess(BaseSDTrainProcess):
             ]
 
             positive_latents = get_noise_pred(
-                positive, negative, 1, current_timestep, denoised_latents
+                prompt_pair.positive_target,  # negative prompt
+                prompt_pair.negative_target,  # positive prompt
+                1,
+                current_timestep,
+                denoised_latents
             ).to("cpu", dtype=torch.float32)
 
             neutral_latents = get_noise_pred(
-                positive, neutral, 1, current_timestep, denoised_latents
+                prompt_pair.positive_target,  # negative prompt
+                prompt_pair.empty_prompt,  # positive prompt (normally neutral
+                1,
+                current_timestep,
+                denoised_latents
             ).to("cpu", dtype=torch.float32)
 
             unconditional_latents = get_noise_pred(
-                positive, positive, 1, current_timestep, denoised_latents
+                prompt_pair.positive_target,  # negative prompt
+                prompt_pair.positive_target,  # positive prompt
+                1,
+                current_timestep,
+                denoised_latents
             ).to("cpu", dtype=torch.float32)
 
         anchor_loss = None
@@ -328,18 +458,22 @@ class TrainSliderProcess(BaseSDTrainProcess):
             with self.network:
                 # anchor whatever weight  prompt pair is using
                 pos_nem_mult = 1.0 if prompt_pair.multiplier > 0 else -1.0
-                self.network.multiplier = anchor.multiplier * pos_nem_mult
+                self.network.multiplier = anchor.multiplier * pos_nem_mult * rand_weight
 
                 anchor_pred_noise = get_noise_pred(
                     anchor.prompt, anchor.neg_prompt, 1, current_timestep, denoised_latents
                 ).to("cpu", dtype=torch.float32)
 
-                self.network.multiplier = prompt_pair.multiplier
+                self.network.multiplier = prompt_pair.multiplier * rand_weight
 
         with self.network:
-            self.network.multiplier = prompt_pair.multiplier
+            self.network.multiplier = prompt_pair.multiplier * rand_weight
             target_latents = get_noise_pred(
-                positive, target_class, 1, current_timestep, denoised_latents
+                prompt_pair.positive_target,
+                prompt_pair.target_class,
+                1,
+                current_timestep,
+                denoised_latents
             ).to("cpu", dtype=torch.float32)
 
             # if self.logging_config.verbose:
@@ -391,6 +525,8 @@ class TrainSliderProcess(BaseSDTrainProcess):
             target_latents,
             latents,
         )
+        # move back to cpu
+        prompt_pair.to("cpu")
         flush()
 
         # reset network

jobs/process/TrainSliderProcessOld.py

@@ -0,0 +1,406 @@
+# ref:
+# - https://github.com/p1atdev/LECO/blob/main/train_lora.py
+import time
+from collections import OrderedDict
+import os
+from typing import Optional
+
+from toolkit.config_modules import SliderConfig
+from toolkit.paths import REPOS_ROOT
+import sys
+
+from toolkit.stable_diffusion_model import PromptEmbeds
+
+sys.path.append(REPOS_ROOT)
+sys.path.append(os.path.join(REPOS_ROOT, 'leco'))
+from toolkit.train_tools import get_torch_dtype, apply_noise_offset
+import gc
+from toolkit import train_tools
+
+import torch
+from leco import train_util, model_util
+from .BaseSDTrainProcess import BaseSDTrainProcess, StableDiffusion
+
+
+class ACTION_TYPES_SLIDER:
+    ERASE_NEGATIVE = 0
+    ENHANCE_NEGATIVE = 1
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+class EncodedPromptPair:
+    def __init__(
+            self,
+            target_class,
+            positive,
+            negative,
+            neutral,
+            width=512,
+            height=512,
+            action=ACTION_TYPES_SLIDER.ERASE_NEGATIVE,
+            multiplier=1.0,
+            weight=1.0
+    ):
+        self.target_class = target_class
+        self.positive = positive
+        self.negative = negative
+        self.neutral = neutral
+        self.width = width
+        self.height = height
+        self.action: int = action
+        self.multiplier = multiplier
+        self.weight = weight
+
+
+class PromptEmbedsCache:  # 使いまわしたいので
+    prompts: dict[str, PromptEmbeds] = {}
+
+    def __setitem__(self, __name: str, __value: PromptEmbeds) -> None:
+        self.prompts[__name] = __value
+
+    def __getitem__(self, __name: str) -> Optional[PromptEmbeds]:
+        if __name in self.prompts:
+            return self.prompts[__name]
+        else:
+            return None
+
+
+class EncodedAnchor:
+    def __init__(
+            self,
+            prompt,
+            neg_prompt,
+            multiplier=1.0
+    ):
+        self.prompt = prompt
+        self.neg_prompt = neg_prompt
+        self.multiplier = multiplier
+
+
+class TrainSliderProcessOld(BaseSDTrainProcess):
+    def __init__(self, process_id: int, job, config: OrderedDict):
+        super().__init__(process_id, job, config)
+        self.step_num = 0
+        self.start_step = 0
+        self.device = self.get_conf('device', self.job.device)
+        self.device_torch = torch.device(self.device)
+        self.slider_config = SliderConfig(**self.get_conf('slider', {}))
+        self.prompt_cache = PromptEmbedsCache()
+        self.prompt_pairs: list[EncodedPromptPair] = []
+        self.anchor_pairs: list[EncodedAnchor] = []
+
+    def before_model_load(self):
+        pass
+
+    def hook_before_train_loop(self):
+        cache = PromptEmbedsCache()
+        prompt_pairs: list[EncodedPromptPair] = []
+
+        # get encoded latents for our prompts
+        with torch.no_grad():
+            neutral = ""
+            for target in self.slider_config.targets:
+                # build the cache
+                for prompt in [
+                    target.target_class,
+                    target.positive,
+                    target.negative,
+                    neutral  # empty neutral
+                ]:
+                    if cache[prompt] is None:
+                        cache[prompt] = self.sd.encode_prompt(prompt)
+                for resolution in self.slider_config.resolutions:
+                    width, height = resolution
+                    only_erase = len(target.positive.strip()) == 0
+                    only_enhance = len(target.negative.strip()) == 0
+
+                    both = not only_erase and not only_enhance
+
+                    if only_erase and only_enhance:
+                        raise ValueError("target must have at least one of positive or negative or both")
+                    # for slider we need to have an enhancer, an eraser, and then
+                    # an inverse with negative weights to balance the network
+                    # if we don't do this, we will get different contrast and focus.
+                    # we only perform actions of enhancing and erasing on the negative
+                    # todo work on way to do all of this in one shot
+
+                    if both or only_erase:
+                        prompt_pairs += [
+                            # erase standard
+                            EncodedPromptPair(
+                                target_class=cache[target.target_class],
+                                positive=cache[target.positive],
+                                negative=cache[target.negative],
+                                neutral=cache[neutral],
+                                width=width,
+                                height=height,
+                                action=ACTION_TYPES_SLIDER.ERASE_NEGATIVE,
+                                multiplier=target.multiplier,
+                                weight=target.weight
+                            ),
+                        ]
+                    if both or only_enhance:
+                        prompt_pairs += [
+                            # enhance standard, swap pos neg
+                            EncodedPromptPair(
+                                target_class=cache[target.target_class],
+                                positive=cache[target.negative],
+                                negative=cache[target.positive],
+                                neutral=cache[neutral],
+                                width=width,
+                                height=height,
+                                action=ACTION_TYPES_SLIDER.ENHANCE_NEGATIVE,
+                                multiplier=target.multiplier,
+                                weight=target.weight
+                            ),
+                        ]
+                    if both:
+                        prompt_pairs += [
+                            # erase inverted
+                            EncodedPromptPair(
+                                target_class=cache[target.target_class],
+                                positive=cache[target.negative],
+                                negative=cache[target.positive],
+                                neutral=cache[neutral],
+                                width=width,
+                                height=height,
+                                action=ACTION_TYPES_SLIDER.ERASE_NEGATIVE,
+                                multiplier=target.multiplier * -1.0,
+                                weight=target.weight
+                            ),
+                        ]
+                        prompt_pairs += [
+                            # enhance inverted
+                            EncodedPromptPair(
+                                target_class=cache[target.target_class],
+                                positive=cache[target.positive],
+                                negative=cache[target.negative],
+                                neutral=cache[neutral],
+                                width=width,
+                                height=height,
+                                action=ACTION_TYPES_SLIDER.ENHANCE_NEGATIVE,
+                                multiplier=target.multiplier * -1.0,
+                                weight=target.weight
+                            ),
+                        ]
+
+            # setup anchors
+            anchor_pairs = []
+            for anchor in self.slider_config.anchors:
+                # build the cache
+                for prompt in [
+                    anchor.prompt,
+                    anchor.neg_prompt  # empty neutral
+                ]:
+                    if cache[prompt] == None:
+                        cache[prompt] = self.sd.encode_prompt(prompt)
+
+                anchor_pairs += [
+                    EncodedAnchor(
+                        prompt=cache[anchor.prompt],
+                        neg_prompt=cache[anchor.neg_prompt],
+                        multiplier=anchor.multiplier
+                    )
+                ]
+
+        # move to cpu to save vram
+        # We don't need text encoder anymore, but keep it on cpu for sampling
+        # if text encoder is list
+        if isinstance(self.sd.text_encoder, list):
+            for encoder in self.sd.text_encoder:
+                encoder.to("cpu")
+        else:
+            self.sd.text_encoder.to("cpu")
+        self.prompt_cache = cache
+        self.prompt_pairs = prompt_pairs
+        self.anchor_pairs = anchor_pairs
+        flush()
+        # end hook_before_train_loop
+
+    def hook_train_loop(self):
+        dtype = get_torch_dtype(self.train_config.dtype)
+
+        # get a random pair
+        prompt_pair: EncodedPromptPair = self.prompt_pairs[
+            torch.randint(0, len(self.prompt_pairs), (1,)).item()
+        ]
+
+        height = prompt_pair.height
+        width = prompt_pair.width
+        target_class = prompt_pair.target_class
+        neutral = prompt_pair.neutral
+        negative = prompt_pair.negative
+        positive = prompt_pair.positive
+        weight = prompt_pair.weight
+        multiplier = prompt_pair.multiplier
+
+        unet = self.sd.unet
+        noise_scheduler = self.sd.noise_scheduler
+        optimizer = self.optimizer
+        lr_scheduler = self.lr_scheduler
+        loss_function = torch.nn.MSELoss()
+
+        def get_noise_pred(p, n, gs, cts, dn):
+            return self.predict_noise(
+                latents=dn,
+                text_embeddings=train_tools.concat_prompt_embeddings(
+                    p,  # unconditional
+                    n,  # positive
+                    self.train_config.batch_size,
+                ),
+                timestep=cts,
+                guidance_scale=gs,
+            )
+
+        # set network multiplier
+        self.network.multiplier = multiplier
+
+        with torch.no_grad():
+            self.sd.noise_scheduler.set_timesteps(
+                self.train_config.max_denoising_steps, device=self.device_torch
+            )
+
+            self.optimizer.zero_grad()
+
+            # ger a random number of steps
+            timesteps_to = torch.randint(
+                1, self.train_config.max_denoising_steps, (1,)
+            ).item()
+
+            # get noise
+            noise = self.get_latent_noise(
+                pixel_height=height,
+                pixel_width=width,
+            ).to(self.device_torch, dtype=dtype)
+
+            # get latents
+            latents = noise * self.sd.noise_scheduler.init_noise_sigma
+            latents = latents.to(self.device_torch, dtype=dtype)
+
+            with self.network:
+                assert self.network.is_active
+                self.network.multiplier = multiplier
+                denoised_latents = self.diffuse_some_steps(
+                    latents,  # pass simple noise latents
+                    train_tools.concat_prompt_embeddings(
+                        positive,  # unconditional
+                        target_class,  # target
+                        self.train_config.batch_size,
+                    ),
+                    start_timesteps=0,
+                    total_timesteps=timesteps_to,
+                    guidance_scale=3,
+                )
+
+            noise_scheduler.set_timesteps(1000)
+
+            current_timestep = noise_scheduler.timesteps[
+                int(timesteps_to * 1000 / self.train_config.max_denoising_steps)
+            ]
+
+            positive_latents = get_noise_pred(
+                positive, negative, 1, current_timestep, denoised_latents
+            ).to("cpu", dtype=torch.float32)
+
+            neutral_latents = get_noise_pred(
+                positive, neutral, 1, current_timestep, denoised_latents
+            ).to("cpu", dtype=torch.float32)
+
+            unconditional_latents = get_noise_pred(
+                positive, positive, 1, current_timestep, denoised_latents
+            ).to("cpu", dtype=torch.float32)
+
+        anchor_loss = None
+        if len(self.anchor_pairs) > 0:
+            # get a random anchor pair
+            anchor: EncodedAnchor = self.anchor_pairs[
+                torch.randint(0, len(self.anchor_pairs), (1,)).item()
+            ]
+            with torch.no_grad():
+                anchor_target_noise = get_noise_pred(
+                    anchor.prompt, anchor.neg_prompt, 1, current_timestep, denoised_latents
+                ).to("cpu", dtype=torch.float32)
+            with self.network:
+                # anchor whatever weight  prompt pair is using
+                pos_nem_mult = 1.0 if prompt_pair.multiplier > 0 else -1.0
+                self.network.multiplier = anchor.multiplier * pos_nem_mult
+
+                anchor_pred_noise = get_noise_pred(
+                    anchor.prompt, anchor.neg_prompt, 1, current_timestep, denoised_latents
+                ).to("cpu", dtype=torch.float32)
+
+                self.network.multiplier = prompt_pair.multiplier
+
+        with self.network:
+            self.network.multiplier = prompt_pair.multiplier
+            target_latents = get_noise_pred(
+                positive, target_class, 1, current_timestep, denoised_latents
+            ).to("cpu", dtype=torch.float32)
+
+            # if self.logging_config.verbose:
+            #     self.print("target_latents:", target_latents[0, 0, :5, :5])
+
+        positive_latents.requires_grad = False
+        neutral_latents.requires_grad = False
+        unconditional_latents.requires_grad = False
+        if len(self.anchor_pairs) > 0:
+            anchor_target_noise.requires_grad = False
+            anchor_loss = loss_function(
+                anchor_target_noise,
+                anchor_pred_noise,
+            )
+        erase = prompt_pair.action == ACTION_TYPES_SLIDER.ERASE_NEGATIVE
+        guidance_scale = 1.0
+
+        offset = guidance_scale * (positive_latents - unconditional_latents)
+
+        offset_neutral = neutral_latents
+        if erase:
+            offset_neutral -= offset
+        else:
+            # enhance
+            offset_neutral += offset
+
+        loss = loss_function(
+            target_latents,
+            offset_neutral,
+        ) * weight
+
+        loss_slide = loss.item()
+
+        if anchor_loss is not None:
+            loss += anchor_loss
+
+        loss_float = loss.item()
+
+        loss = loss.to(self.device_torch)
+
+        loss.backward()
+        optimizer.step()
+        lr_scheduler.step()
+
+        del (
+            positive_latents,
+            neutral_latents,
+            unconditional_latents,
+            target_latents,
+            latents,
+        )
+        flush()
+
+        # reset network
+        self.network.multiplier = 1.0
+
+        loss_dict = OrderedDict(
+            {'loss': loss_float},
+        )
+        if anchor_loss is not None:
+            loss_dict['sl_l'] = loss_slide
+            loss_dict['an_l'] = anchor_loss.item()
+
+        return loss_dict
+        # end hook_train_loop

jobs/process/__init__.py

@@ -6,5 +6,7 @@ from .BaseTrainProcess import BaseTrainProcess
 from .TrainVAEProcess import TrainVAEProcess
 from .BaseMergeProcess import BaseMergeProcess
 from .TrainSliderProcess import TrainSliderProcess
+from .TrainSliderProcessOld import TrainSliderProcessOld
 from .TrainLoRAHack import TrainLoRAHack
-from .TrainSDRescaleProcess import TrainSDRescaleProcess
+from .TrainSDRescaleProcess import TrainSDRescaleProcess
+from .ModRescaleLoraProcess import ModRescaleLoraProcess

run.py

@@ -36,6 +36,14 @@ def main():
         action='store_true',
         help='Continue running additional jobs even if a job fails'
     )
+
+    # flag to continue if failed job
+    parser.add_argument(
+        '-n', '--name',
+        type=str,
+        default=None,
+        help='Name to replace [name] tag in config file, useful for shared config file'
+    )
     args = parser.parse_args()
 
     config_file_list = args.config_file_list
@@ -49,7 +57,7 @@ def main():
 
     for config_file in config_file_list:
         try:
-            job = get_job(config_file)
+            job = get_job(config_file, args.name)
             job.run()
             job.cleanup()
             jobs_completed += 1

testing/compare_keys.py

@@ -0,0 +1,89 @@
+import argparse
+import os
+
+import torch
+from safetensors.torch import load_file
+from collections import OrderedDict
+import json
+# this was just used to match the vae keys to the diffusers keys
+# you probably wont need this. Unless they change them.... again... again
+# on second thought, you probably will
+
+device = torch.device('cpu')
+dtype = torch.float32
+
+parser = argparse.ArgumentParser()
+
+# require at lease one config file
+parser.add_argument(
+    'file_1',
+    nargs='+',
+    type=str,
+    help='Path to first safe tensor file'
+)
+
+parser.add_argument(
+    'file_2',
+    nargs='+',
+    type=str,
+    help='Path to second safe tensor file'
+)
+
+args = parser.parse_args()
+
+find_matches = False
+
+state_dict_file_1 = load_file(args.file_1[0])
+state_dict_1_keys = list(state_dict_file_1.keys())
+
+state_dict_file_2 = load_file(args.file_2[0])
+state_dict_2_keys = list(state_dict_file_2.keys())
+keys_in_both = []
+
+keys_not_in_state_dict_2 = []
+for key in state_dict_1_keys:
+    if key not in state_dict_2_keys:
+        keys_not_in_state_dict_2.append(key)
+
+keys_not_in_state_dict_1 = []
+for key in state_dict_2_keys:
+    if key not in state_dict_1_keys:
+        keys_not_in_state_dict_1.append(key)
+
+keys_in_both = []
+for key in state_dict_1_keys:
+    if key in state_dict_2_keys:
+        keys_in_both.append(key)
+
+# sort them
+keys_not_in_state_dict_2.sort()
+keys_not_in_state_dict_1.sort()
+keys_in_both.sort()
+
+
+json_data = {
+    "both": keys_in_both,
+    "state_dict_2": keys_not_in_state_dict_2,
+    "state_dict_1": keys_not_in_state_dict_1
+}
+json_data = json.dumps(json_data, indent=4)
+
+remaining_diffusers_values = OrderedDict()
+for key in keys_not_in_state_dict_1:
+    remaining_diffusers_values[key] = state_dict_file_2[key]
+
+# print(remaining_diffusers_values.keys())
+
+remaining_ldm_values = OrderedDict()
+for key in keys_not_in_state_dict_2:
+    remaining_ldm_values[key] = state_dict_file_1[key]
+
+# print(json_data)
+
+project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+json_save_path = os.path.join(project_root, 'config', 'keys.json')
+json_matched_save_path = os.path.join(project_root, 'config', 'matched.json')
+json_duped_save_path = os.path.join(project_root, 'config', 'duped.json')
+
+with open(json_save_path, 'w') as f:
+    f.write(json_data)

toolkit/config.py

@@ -15,22 +15,24 @@ def get_cwd_abs_path(path):
     return path
 
 
-def preprocess_config(config: OrderedDict):
+def preprocess_config(config: OrderedDict, name: str = None):
     if "job" not in config:
         raise ValueError("config file must have a job key")
     if "config" not in config:
         raise ValueError("config file must have a config section")
-    if "name" not in config["config"]:
+    if "name" not in config["config"] and name is None:
         raise ValueError("config file must have a config.name key")
     # we need to replace tags. For now just [name]
-    name = config["config"]["name"]
+    if name is not None:
+        config["config"]["name"] = name
+    else:
+        name = config["config"]["name"]
     config_string = json.dumps(config)
     config_string = config_string.replace("[name]", name)
     config = json.loads(config_string, object_pairs_hook=OrderedDict)
     return config
 
 
-
 # Fixes issue where yaml doesnt load exponents correctly
 fixed_loader = yaml.SafeLoader
 fixed_loader.add_implicit_resolver(
@@ -44,7 +46,8 @@ fixed_loader.add_implicit_resolver(
     |\\.(?:nan|NaN|NAN))$''', re.X),
     list(u'-+0123456789.'))
 
-def get_config(config_file_path):
+
+def get_config(config_file_path, name=None):
     # first check if it is in the config folder
     config_path = os.path.join(TOOLKIT_ROOT, 'config', config_file_path)
     # see if it is in the config folder with any of the possible extensions if it doesnt have one
@@ -75,4 +78,4 @@ def get_config(config_file_path):
     else:
         raise ValueError(f"Config file {config_file_path} must be a json or yaml file")
 
-    return preprocess_config(config)
+    return preprocess_config(config, name)

toolkit/config_modules.py

@@ -99,3 +99,5 @@ class SliderConfig:
         anchors = [SliderConfigAnchors(**anchor) for anchor in anchors]
         self.anchors: List[SliderConfigAnchors] = anchors
         self.resolutions: List[List[int]] = kwargs.get('resolutions', [[512, 512]])
+        self.prompt_file: str = kwargs.get('prompt_file', None)
+        self.prompt_tensors: str = kwargs.get('prompt_tensors', None)

toolkit/job.py

@@ -1,8 +1,8 @@
 from toolkit.config import get_config
 
 
-def get_job(config_path):
+def get_job(config_path, name=None):
-    config = get_config(config_path)
+    config = get_config(config_path, name)
     if not config['job']:
         raise ValueError('config file is invalid. Missing "job" key')
 
@@ -13,6 +13,9 @@ def get_job(config_path):
     if job == 'train':
         from jobs import TrainJob
         return TrainJob(config)
+    if job == 'mod':
+        from jobs import ModJob
+        return ModJob(config)
 
     # elif job == 'train':
     #     from jobs import TrainJob

toolkit/lora.py

@@ -6,12 +6,14 @@
 import os
 import math
 from typing import Optional, List, Type, Set, Literal
+from collections import OrderedDict
 
 import torch
 import torch.nn as nn
 from diffusers import UNet2DConditionModel
 from safetensors.torch import save_file
 
+from toolkit.metadata import add_model_hash_to_meta
 
 UNET_TARGET_REPLACE_MODULE_TRANSFORMER = [
     "Transformer2DModel",  # どうやらこっちの方らしい？ # attn1, 2
@@ -31,7 +33,7 @@ TRAINING_METHODS = Literal[
     "innoxattn",  # train all layers except self attention layers
     "selfattn",  # ESD-u, train only self attention layers
     "xattn",  # ESD-x, train only x attention layers
-    "full",  #  train all layers
+    "full",  # train all layers
     # "notime",
     # "xlayer",
     # "outxattn",
@@ -48,12 +50,12 @@ class LoRAModule(nn.Module):
     """
 
     def __init__(
-        self,
+            self,
-        lora_name,
+            lora_name,
-        org_module: nn.Module,
+            org_module: nn.Module,
-        multiplier=1.0,
+            multiplier=1.0,
-        lora_dim=4,
+            lora_dim=4,
-        alpha=1,
+            alpha=1,
     ):
         """if alpha == 0 or None, alpha is rank (no scaling)."""
         super().__init__()
@@ -102,19 +104,19 @@ class LoRAModule(nn.Module):
 
     def forward(self, x):
         return (
-            self.org_forward(x)
+                self.org_forward(x)
-            + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+                + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
         )
 
 
 class LoRANetwork(nn.Module):
     def __init__(
-        self,
+            self,
-        unet: UNet2DConditionModel,
+            unet: UNet2DConditionModel,
-        rank: int = 4,
+            rank: int = 4,
-        multiplier: float = 1.0,
+            multiplier: float = 1.0,
-        alpha: float = 1.0,
+            alpha: float = 1.0,
-        train_method: TRAINING_METHODS = "full",
+            train_method: TRAINING_METHODS = "full",
     ) -> None:
         super().__init__()
 
@@ -140,7 +142,7 @@ class LoRANetwork(nn.Module):
         lora_names = set()
         for lora in self.unet_loras:
             assert (
-                lora.lora_name not in lora_names
+                    lora.lora_name not in lora_names
             ), f"duplicated lora name: {lora.lora_name}. {lora_names}"
             lora_names.add(lora.lora_name)
 
@@ -157,13 +159,13 @@ class LoRANetwork(nn.Module):
         torch.cuda.empty_cache()
 
     def create_modules(
-        self,
+            self,
-        prefix: str,
+            prefix: str,
-        root_module: nn.Module,
+            root_module: nn.Module,
-        target_replace_modules: List[str],
+            target_replace_modules: List[str],
-        rank: int,
+            rank: int,
-        multiplier: float,
+            multiplier: float,
-        train_method: TRAINING_METHODS,
+            train_method: TRAINING_METHODS,
     ) -> list:
         loras = []
 
@@ -212,6 +214,8 @@ class LoRANetwork(nn.Module):
 
     def save_weights(self, file, dtype=None, metadata: Optional[dict] = None):
         state_dict = self.state_dict()
+        if metadata is None:
+            metadata = OrderedDict()
 
         if dtype is not None:
             for key in list(state_dict.keys()):
@@ -221,9 +225,10 @@ class LoRANetwork(nn.Module):
 
         for key in list(state_dict.keys()):
             if not key.startswith("lora"):
-                # lora以外除外
+                # remove any not lora
                 del state_dict[key]
 
+        metadata = add_model_hash_to_meta(state_dict, metadata)
         if os.path.splitext(file)[1] == ".safetensors":
             save_file(state_dict, file, metadata)
         else:

toolkit/metadata.py

@@ -1,18 +1,23 @@
 import json
 from collections import OrderedDict
+from io import BytesIO
 
+import safetensors
 from safetensors import safe_open
 
 from info import software_meta
+from toolkit.train_tools import addnet_hash_legacy
+from toolkit.train_tools import addnet_hash_safetensors
 
 
-def get_meta_for_safetensors(meta: OrderedDict, name=None) -> OrderedDict:
+def get_meta_for_safetensors(meta: OrderedDict, name=None, add_software_info=True) -> OrderedDict:
     # stringify the meta and reparse OrderedDict to replace [name] with name
     meta_string = json.dumps(meta)
     if name is not None:
         meta_string = meta_string.replace("[name]", name)
     save_meta = json.loads(meta_string, object_pairs_hook=OrderedDict)
-    save_meta["software"] = software_meta
+    if add_software_info:
+        save_meta["software"] = software_meta
     # safetensors can only be one level deep
     for key, value in save_meta.items():
         # if not float, int, bool, or str, convert to json string
@@ -21,6 +26,46 @@ def get_meta_for_safetensors(meta: OrderedDict, name=None) -> OrderedDict:
     return save_meta
 
 
+def add_model_hash_to_meta(state_dict, meta: OrderedDict) -> OrderedDict:
+    """Precalculate the model hashes needed by sd-webui-additional-networks to
+    save time on indexing the model later."""
+
+    # Because writing user metadata to the file can change the result of
+    # sd_models.model_hash(), only retain the training metadata for purposes of
+    # calculating the hash, as they are meant to be immutable
+    metadata = {k: v for k, v in meta.items() if k.startswith("ss_")}
+
+    bytes = safetensors.torch.save(state_dict, metadata)
+    b = BytesIO(bytes)
+
+    model_hash = addnet_hash_safetensors(b)
+    legacy_hash = addnet_hash_legacy(b)
+    meta["sshs_model_hash"] = model_hash
+    meta["sshs_legacy_hash"] = legacy_hash
+    return meta
+
+
+def add_base_model_info_to_meta(
+        meta: OrderedDict,
+        base_model: str = None,
+        is_v1: bool = False,
+        is_v2: bool = False,
+        is_xl: bool = False,
+) -> OrderedDict:
+    if base_model is not None:
+        meta['ss_base_model'] = base_model
+    elif is_v2:
+        meta['ss_v2'] = True
+        meta['ss_base_model_version'] = 'sd_2.1'
+
+    elif is_xl:
+        meta['ss_base_model_version'] = 'sdxl_1.0'
+    else:
+        # default to v1.5
+        meta['ss_base_model_version'] = 'sd_1.5'
+    return meta
+
+
 def parse_metadata_from_safetensors(meta: OrderedDict) -> OrderedDict:
     parsed_meta = OrderedDict()
     for key, value in meta.items():

toolkit/optimizer.py

@@ -27,6 +27,17 @@ def get_optimizer(
             optimizer = dadaptation.DAdaptAdam(params, lr=use_lr, **optimizer_params)
             # warn user that dadaptation is deprecated
             print("WARNING: Dadaptation optimizer type has been changed to DadaptationAdam. Please update your config.")
+    elif lower_type.startswith("prodigy"):
+        from prodigyopt import Prodigy
+
+        print("Using Prodigy optimizer")
+        use_lr = learning_rate
+        if use_lr < 0.1:
+            # dadaptation uses different lr that is values of 0.1 to 1.0. default to 1.0
+            use_lr = 1.0
+        # let net be the neural network you want to train
+        # you can choose weight decay value based on your problem, 0 by default
+        optimizer = Prodigy(params, lr=use_lr, **optimizer_params)
     elif lower_type.endswith("8bit"):
         import bitsandbytes
 
@@ -43,6 +54,8 @@ def get_optimizer(
     elif lower_type == 'lion':
         from lion_pytorch import Lion
         return Lion(params, lr=learning_rate, **optimizer_params)
+    elif lower_type == 'adagrad':
+        optimizer = torch.optim.Adagrad(params, lr=float(learning_rate), **optimizer_params)
     else:
         raise ValueError(f'Unknown optimizer type {optimizer_type}')
     return optimizer

toolkit/pipelines.py

@@ -0,0 +1,534 @@
+from typing import Union, List, Optional, Dict, Any, Tuple, Callable
+
+import torch
+from diffusers import StableDiffusionXLPipeline, StableDiffusionPipeline
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import rescale_noise_cfg
+
+
+class CustomStableDiffusionXLPipeline(StableDiffusionXLPipeline):
+    # def __init__(self, *args, **kwargs):
+    #     super().__init__(*args, **kwargs)
+
+    def predict_noise(
+            self,
+            prompt: Union[str, List[str]] = None,
+            prompt_2: Optional[Union[str, List[str]]] = None,
+            num_inference_steps: int = 50,
+            guidance_scale: float = 5.0,
+            negative_prompt: Optional[Union[str, List[str]]] = None,
+            negative_prompt_2: Optional[Union[str, List[str]]] = None,
+            num_images_per_prompt: Optional[int] = 1,
+            eta: float = 0.0,
+            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+            latents: Optional[torch.FloatTensor] = None,
+            prompt_embeds: Optional[torch.FloatTensor] = None,
+            negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+            pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            guidance_rescale: float = 0.0,
+            crops_coords_top_left: Tuple[int, int] = (0, 0),
+            timestep: Optional[int] = None,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in both text-encoders
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            denoising_end (`float`, *optional*):
+                When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
+                completed before it is intentionally prematurely terminated. As a result, the returned sample will
+                still retain a substantial amount of noise as determined by the discrete timesteps selected by the
+                scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
+                "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
+                Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
+                of a plain tuple.
+            callback (`Callable`, *optional*):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function will be called. If not specified, the callback will be
+                called at every step.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
+            guidance_rescale (`float`, *optional*, defaults to 0.7):
+                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
+                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
+                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+                Guidance rescale factor should fix overexposure when using zero terminal SNR.
+            original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
+                `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
+                explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+                `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
+                `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
+                `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                For most cases, `target_size` should be set to the desired height and width of the generated image. If
+                not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
+                section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+        # if not predict_noise:
+        #     # call parent
+        #     return super().__call__(
+        #         prompt=prompt,
+        #         prompt_2=prompt_2,
+        #         height=height,
+        #         width=width,
+        #         num_inference_steps=num_inference_steps,
+        #         denoising_end=denoising_end,
+        #         guidance_scale=guidance_scale,
+        #         negative_prompt=negative_prompt,
+        #         negative_prompt_2=negative_prompt_2,
+        #         num_images_per_prompt=num_images_per_prompt,
+        #         eta=eta,
+        #         generator=generator,
+        #         latents=latents,
+        #         prompt_embeds=prompt_embeds,
+        #         negative_prompt_embeds=negative_prompt_embeds,
+        #         pooled_prompt_embeds=pooled_prompt_embeds,
+        #         negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+        #         output_type=output_type,
+        #         return_dict=return_dict,
+        #         callback=callback,
+        #         callback_steps=callback_steps,
+        #         cross_attention_kwargs=cross_attention_kwargs,
+        #         guidance_rescale=guidance_rescale,
+        #         original_size=original_size,
+        #         crops_coords_top_left=crops_coords_top_left,
+        #         target_size=target_size,
+        #     )
+
+        # 0. Default height and width to unet
+        height = self.default_sample_size * self.vae_scale_factor
+        width = self.default_sample_size * self.vae_scale_factor
+
+        original_size = (height, width)
+        target_size = (height, width)
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+        )
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+        )
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 7. Prepare added time ids & embeddings
+        add_text_embeds = pooled_prompt_embeds
+        add_time_ids = self._get_add_time_ids(
+            original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype
+        ).to(device)  # TODO DOES NOT CAST ORIGINALLY
+
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+            add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
+
+        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+
+        latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)
+
+        # predict the noise residual
+        added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+        noise_pred = self.unet(
+            latent_model_input,
+            timestep,
+            encoder_hidden_states=prompt_embeds,
+            cross_attention_kwargs=cross_attention_kwargs,
+            added_cond_kwargs=added_cond_kwargs,
+            return_dict=False,
+        )[0]
+
+        # perform guidance
+        if do_classifier_free_guidance:
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+        if do_classifier_free_guidance and guidance_rescale > 0.0:
+            # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+            noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+        return noise_pred
+
+    def enable_model_cpu_offload(self, gpu_id=0):
+        print('Called cpu offload', gpu_id)
+        # fuck off
+        pass
+
+
+class CustomStableDiffusionPipeline(StableDiffusionPipeline):
+
+    # replace the call so it matches SDXL call so we can use the same code and also stop early
+    def __call__(
+            self,
+            prompt: Union[str, List[str]] = None,
+            prompt_2: Optional[Union[str, List[str]]] = None,
+            height: Optional[int] = None,
+            width: Optional[int] = None,
+            num_inference_steps: int = 50,
+            denoising_end: Optional[float] = None,
+            guidance_scale: float = 5.0,
+            negative_prompt: Optional[Union[str, List[str]]] = None,
+            negative_prompt_2: Optional[Union[str, List[str]]] = None,
+            num_images_per_prompt: Optional[int] = 1,
+            eta: float = 0.0,
+            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+            latents: Optional[torch.FloatTensor] = None,
+            prompt_embeds: Optional[torch.FloatTensor] = None,
+            negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+            pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            output_type: Optional[str] = "pil",
+            return_dict: bool = True,
+            callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+            callback_steps: int = 1,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            guidance_rescale: float = 0.0,
+            original_size: Optional[Tuple[int, int]] = None,
+            crops_coords_top_left: Tuple[int, int] = (0, 0),
+            target_size: Optional[Tuple[int, int]] = None,
+    ):
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+        )
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+        )
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        # 7.1 Apply denoising_end
+        if denoising_end is not None and type(denoising_end) == float and denoising_end > 0 and denoising_end < 1:
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (denoising_end * self.scheduler.config.num_train_timesteps)
+                )
+            )
+            num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
+            timesteps = timesteps[:num_inference_steps]
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                # predict the noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                if do_classifier_free_guidance and guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+        if not output_type == "latent":
+            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
+        else:
+            image = latents
+            has_nsfw_concept = None
+
+        if has_nsfw_concept is None:
+            do_denormalize = [True] * image.shape[0]
+        else:
+            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
+
+        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+
+        # Offload last model to CPU
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.final_offload_hook.offload()
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
+
+    # some of the inputs are to keep it compatible with sdx
+    def predict_noise(
+            self,
+            prompt: Union[str, List[str]] = None,
+            prompt_2: Optional[Union[str, List[str]]] = None,
+            num_inference_steps: int = 50,
+            guidance_scale: float = 5.0,
+            negative_prompt: Optional[Union[str, List[str]]] = None,
+            negative_prompt_2: Optional[Union[str, List[str]]] = None,
+            num_images_per_prompt: Optional[int] = 1,
+            eta: float = 0.0,
+            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+            latents: Optional[torch.FloatTensor] = None,
+            prompt_embeds: Optional[torch.FloatTensor] = None,
+            negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+            pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            guidance_rescale: float = 0.0,
+            crops_coords_top_left: Tuple[int, int] = (0, 0),
+            timestep: Optional[int] = None,
+    ):
+
+        # 0. Default height and width to unet
+        height = self.unet.config.sample_size * self.vae_scale_factor
+        width = self.unet.config.sample_size * self.vae_scale_factor
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+        )
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+        )
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # expand the latents if we are doing classifier free guidance
+        latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+        latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep)
+
+        # predict the noise residual
+        noise_pred = self.unet(
+            latent_model_input,
+            timestep,
+            encoder_hidden_states=prompt_embeds,
+            cross_attention_kwargs=cross_attention_kwargs,
+            return_dict=False,
+        )[0]
+
+        # perform guidance
+        if do_classifier_free_guidance:
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+        if do_classifier_free_guidance and guidance_rescale > 0.0:
+            # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+            noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+        return noise_pred

toolkit/stable_diffusion_model.py

@@ -1,3 +1,4 @@
+import typing
 from typing import Union, OrderedDict
 import sys
 import os
@@ -29,14 +30,22 @@ class PromptEmbeds:
             self.text_embeds = args
             self.pooled_embeds = None
 
-    def to(self, **kwargs):
+    def to(self, *args, **kwargs):
-        self.text_embeds = self.text_embeds.to(**kwargs)
+        self.text_embeds = self.text_embeds.to(*args, **kwargs)
         if self.pooled_embeds is not None:
-            self.pooled_embeds = self.pooled_embeds.to(**kwargs)
+            self.pooled_embeds = self.pooled_embeds.to(*args, **kwargs)
         return self
 
 
+# if is type checking
+if typing.TYPE_CHECKING:
+    from diffusers import StableDiffusionPipeline
+    from toolkit.pipelines import CustomStableDiffusionXLPipeline
+
+
 class StableDiffusion:
+    pipeline: Union[None, 'StableDiffusionPipeline', 'CustomStableDiffusionXLPipeline']
+
     def __init__(
             self,
             vae,
@@ -44,7 +53,8 @@ class StableDiffusion:
             text_encoder,
             unet,
             noise_scheduler,
-            is_xl=False
+            is_xl=False,
+            pipeline=None,
     ):
         # text encoder has a list of 2 for xl
         self.vae = vae
@@ -53,6 +63,7 @@ class StableDiffusion:
         self.unet = unet
         self.noise_scheduler = noise_scheduler
         self.is_xl = is_xl
+        self.pipeline = pipeline
 
     def encode_prompt(self, prompt, num_images_per_prompt=1) -> PromptEmbeds:
         prompt = prompt

toolkit/train_pipelines.py

@@ -0,0 +1,316 @@
+from typing import Optional, Tuple, Callable, Dict, Any, Union, List
+
+import torch
+from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipelineOutput
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import rescale_noise_cfg
+
+from toolkit.lora_special import LoRASpecialNetwork
+from toolkit.pipelines import CustomStableDiffusionXLPipeline
+
+
+class TransferStableDiffusionXLPipeline(CustomStableDiffusionXLPipeline):
+    def transfer_diffuse(
+            self,
+            prompt: Union[str, List[str]] = None,
+            prompt_2: Optional[Union[str, List[str]]] = None,
+            height: Optional[int] = None,
+            width: Optional[int] = None,
+            num_inference_steps: int = 50,
+            denoising_end: Optional[float] = None,
+            guidance_scale: float = 5.0,
+            negative_prompt: Optional[Union[str, List[str]]] = None,
+            negative_prompt_2: Optional[Union[str, List[str]]] = None,
+            num_images_per_prompt: Optional[int] = 1,
+            eta: float = 0.0,
+            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+            latents: Optional[torch.FloatTensor] = None,
+            prompt_embeds: Optional[torch.FloatTensor] = None,
+            negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+            pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+            output_type: Optional[str] = "pil",
+            return_dict: bool = True,
+            callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+            callback_steps: int = 1,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            guidance_rescale: float = 0.0,
+            original_size: Optional[Tuple[int, int]] = None,
+            crops_coords_top_left: Tuple[int, int] = (0, 0),
+            target_size: Optional[Tuple[int, int]] = None,
+            target_unet: Optional[torch.nn.Module] = None,
+            pre_condition_callback = None,
+            each_step_callback = None,
+            network: Optional[LoRASpecialNetwork] = None,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                used in both text-encoders
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            denoising_end (`float`, *optional*):
+                When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
+                completed before it is intentionally prematurely terminated. As a result, the returned sample will
+                still retain a substantial amount of noise as determined by the discrete timesteps selected by the
+                scheduler. The denoising_end parameter should ideally be utilized when this pipeline forms a part of a
+                "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
+                Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output)
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+                `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+                input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
+                of a plain tuple.
+            callback (`Callable`, *optional*):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function will be called. If not specified, the callback will be
+                called at every step.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
+            guidance_rescale (`float`, *optional*, defaults to 0.7):
+                Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
+                Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
+                [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+                Guidance rescale factor should fix overexposure when using zero terminal SNR.
+            original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
+                `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
+                explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+                `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
+                `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
+                `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
+                [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+            target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+                For most cases, `target_size` should be set to the desired height and width of the generated image. If
+                not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
+                section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+        # 0. Default height and width to unet
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            callback_steps,
+            negative_prompt,
+            negative_prompt_2,
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+        )
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+        )
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Prepare added time ids & embeddings
+        add_text_embeds = pooled_prompt_embeds
+        add_time_ids = self._get_add_time_ids(
+            original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype
+        )
+
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+            add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
+
+        # 8. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        # 7.1 Apply denoising_end
+        if denoising_end is not None and type(denoising_end) == float and denoising_end > 0 and denoising_end < 1:
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (denoising_end * self.scheduler.config.num_train_timesteps)
+                )
+            )
+            num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
+            timesteps = timesteps[:num_inference_steps]
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                # predict the noise residual
+                added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    added_cond_kwargs=added_cond_kwargs,
+                    return_dict=False,
+                )[0]
+
+                conditioned_noise_pred, conditioned_latent_model_input = pre_condition_callback(
+                    noise_pred.clone().detach(),
+                    latent_model_input.clone().detach(),
+                )
+
+                # start grad
+                with torch.enable_grad():
+                    with network:
+                        assert network.is_active
+                        noise_train_pred = target_unet(
+                            conditioned_latent_model_input,
+                            t,
+                            encoder_hidden_states=prompt_embeds,
+                            cross_attention_kwargs=cross_attention_kwargs,
+                            added_cond_kwargs=added_cond_kwargs,
+                            return_dict=False,
+                        )[0]
+                    each_step_callback(conditioned_noise_pred, noise_train_pred)
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                if do_classifier_free_guidance and guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+

toolkit/train_tools.py

@@ -1,4 +1,5 @@
 import argparse
+import hashlib
 import json
 import os
 import time
@@ -399,3 +400,29 @@ def concat_prompt_embeddings(
             [unconditional.pooled_embeds, conditional.pooled_embeds]
         ).repeat_interleave(n_imgs, dim=0)
     return PromptEmbeds([text_embeds, pooled_embeds])
+
+
+def addnet_hash_safetensors(b):
+    """New model hash used by sd-webui-additional-networks for .safetensors format files"""
+    hash_sha256 = hashlib.sha256()
+    blksize = 1024 * 1024
+
+    b.seek(0)
+    header = b.read(8)
+    n = int.from_bytes(header, "little")
+
+    offset = n + 8
+    b.seek(offset)
+    for chunk in iter(lambda: b.read(blksize), b""):
+        hash_sha256.update(chunk)
+
+    return hash_sha256.hexdigest()
+
+
+def addnet_hash_legacy(b):
+    """Old model hash used by sd-webui-additional-networks for .safetensors format files"""
+    m = hashlib.sha256()
+
+    b.seek(0x100000)
+    m.update(b.read(0x10000))
+    return m.hexdigest()[0:8]