Upload 5 files

SDXL-LoRA-RNPD.ipynb

@@ -0,0 +1,317 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b01332d1-1384-4405-8af6-335c768da6e2",
+   "metadata": {},
+   "source": [
+    "## SDXL LoRA Trainer by TheLastBen https://github.com/TheLastBen/fast-stable-diffusion, if you encounter any issues, feel free to discuss them."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8f82bb3b-76de-4e2c-9251-df918f8f2cbe",
+   "metadata": {},
+   "source": [
+    "# Dependencies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "3d144e06-1f7a-467b-9cf1-452bf773f0ab",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d1e84d74d92c46f8aa78c03f50a0d0d8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Button(button_style='success', description='Done!', disabled=True, icon='check', style=ButtonStyle(), tooltip=…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# Install the dependencies\n",
+    "\n",
+    "force_reinstall= False\n",
+    "\n",
+    "# Set to true only if you want to install the dependencies again.\n",
+    "\n",
+    "#--------------------\n",
+    "with open('/dev/null', 'w') as devnull:import requests, os, time, importlib;open('/workspace/sdxllorarunpod.py', 'wb').write(requests.get('https://huggingface.co/datasets/TheLastBen/RNPD/raw/main/Scripts/sdxllorarunpod.py').content);os.chdir('/workspace');import sdxllorarunpod;importlib.reload(sdxllorarunpod);from sdxllorarunpod import *;restored=False;restoreda=False;Deps(force_reinstall)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "461b7686-e4aa-4fa8-ab6f-5a6acbf4c601",
+   "metadata": {},
+   "source": [
+    "# Download the model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "2f705bd1-35c9-49bd-84fd-03a1348cbe83",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[1;32mUsing SDXL model\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Run the cell to download the model\n",
+    "\n",
+    "#-------------\n",
+    "MODEL_NAMExl=dls_xlf(\"\", \"\", \"\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8e22327b-e0c3-424c-82e1-fb7f8a815c0b",
+   "metadata": {},
+   "source": [
+    "# Create/Load a Session"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "ac69c221-205a-40d2-b42e-6c8d515a43cc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[1;32mCreating session...\n",
+      "\u001b[1;32mSession created, proceed to uploading instance images\n"
+     ]
+    }
+   ],
+   "source": [
+    "Session_Name = \"aether_skin_230808_SDXL_LoRA_128_dim_50_epochs\"\n",
+    "\n",
+    "# Enter the session name, it if it exists, it will load it, otherwise it'll create an new session.\n",
+    "\n",
+    "#-----------------\n",
+    "[WORKSPACE, Session_Name, INSTANCE_NAME, OUTPUT_DIR, SESSION_DIR, INSTANCE_DIR, CAPTIONS_DIR, MDLPTH, MODEL_NAMExl]=sess_xl(Session_Name, MODEL_NAMExl if 'MODEL_NAMExl' in locals() else \"\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5d239e77-f7fd-404b-8006-081f15326412",
+   "metadata": {},
+   "source": [
+    "# Train LoRA"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c54a7335-8402-42f2-9a71-9da99f6ea604",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m'########:'########:::::'###::::'####:'##::: ##:'####:'##::: ##::'######:::\n",
+      "... ##..:: ##.... ##:::'## ##:::. ##:: ###:: ##:. ##:: ###:: ##:'##... ##::\n",
+      "::: ##:::: ##:::: ##::'##:. ##::: ##:: ####: ##:: ##:: ####: ##: ##:::..:::\n",
+      "::: ##:::: ########::'##:::. ##:: ##:: ## ## ##:: ##:: ## ## ##: ##::'####:\n",
+      "::: ##:::: ##.. ##::: #########:: ##:: ##. ####:: ##:: ##. ####: ##::: ##::\n",
+      "::: ##:::: ##::. ##:: ##.... ##:: ##:: ##:. ###:: ##:: ##:. ###: ##::: ##::\n",
+      "::: ##:::: ##:::. ##: ##:::: ##:'####: ##::. ##:'####: ##::. ##:. ######:::\n",
+      ":::..:::::..:::::..::..:::::..::....::..::::..::....::..::::..:::......::::\n",
+      "\u001b[0m\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Progress:  71%|███████   | 676/950 [06:22<02:23,  1.91it/s, loss=0.245, lr=5.75e-7]  "
+     ]
+    }
+   ],
+   "source": [
+    "Resume_Training= False\n",
+    "\n",
+    "# If you're not satisfied with the result, Set to True, run again the cell and it will continue training the current model.\n",
+    "\n",
+    "\n",
+    "Training_Epochs= 50\n",
+    "\n",
+    "# Epoch = Number of steps/images.\n",
+    "\n",
+    "Learning_Rate= \"3e-6\"\n",
+    "\n",
+    "# keep it between 1e-6 and 6e-6\n",
+    "\n",
+    "\n",
+    "External_Captions= True\n",
+    "\n",
+    "# Load the captions from a text file for each instance image.\n",
+    "\n",
+    "\n",
+    "LoRA_Dim = 128\n",
+    "\n",
+    "# Dimension of the LoRa model, between 64 and 128 is good enough.\n",
+    "\n",
+    "\n",
+    "Resolution= 1024\n",
+    "\n",
+    "# 1024 is the native resolution.\n",
+    "\n",
+    "\n",
+    "Save_VRAM = False\n",
+    "\n",
+    "# Use as low as 9.7GB VRAM with Dim = 64, but slightly slower training.\n",
+    "\n",
+    "#-----------------\n",
+    "dbtrainxl(Resume_Training, Training_Epochs, Learning_Rate, LoRA_Dim, False, Resolution, MODEL_NAMExl, SESSION_DIR, INSTANCE_DIR, CAPTIONS_DIR, External_Captions, INSTANCE_NAME, Session_Name, OUTPUT_DIR, 0.03, Save_VRAM)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e2751798-508e-47ad-8e54-95188bdab051",
+   "metadata": {
+    "jp-MarkdownHeadingCollapsed": true,
+    "tags": []
+   },
+   "source": [
+    "# Test the Trained Model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d1bc48d6-1526-44c6-ab7c-cc1538c7f61c",
+   "metadata": {},
+   "source": [
+    "# ComfyUI"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "26272665-16de-4042-a7a4-6b9205ff3309",
+   "metadata": {
+    "scrolled": true,
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "Args=\"--listen --port 3000\"\n",
+    "\n",
+    "\n",
+    "Download_SDXL_Model= True\n",
+    "\n",
+    "\n",
+    "Huggingface_token_optional= \"\"\n",
+    "\n",
+    "# Restore your backed-up Comfy folder by entering your huggingface token, leave it empty to start fresh or continue with the existing sd folder (if any).\n",
+    "\n",
+    "#--------------------\n",
+    "restored=sdcmff(Huggingface_token_optional, MDLPTH, Download_SDXL_Model, restored)\n",
+    "!python /workspace/ComfyUI/main.py $Args"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "410520ca-7352-4fc4-907b-cb53f661074e",
+   "metadata": {},
+   "source": [
+    "# A1111"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "351f18d5-f723-4d25-b1ae-1296a22c6d8c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "User = \"\"\n",
+    "\n",
+    "Password= \"\"\n",
+    "\n",
+    "# Add credentials to your Gradio interface (optional).\n",
+    "\n",
+    "Download_SDXL_Model= True\n",
+    "\n",
+    "\n",
+    "Huggingface_token_optional= \"\"\n",
+    "\n",
+    "# Restore your backed-up SD folder by entering your huggingface token, leave it empty to start fresh or continue with the existing sd folder (if any).\n",
+    "\n",
+    "#-----------------\n",
+    "configf, restoreda=test(MDLPTH, User, Password, Huggingface_token_optional, Download_SDXL_Model, restoreda)\n",
+    "!python /workspace/sd/stable-diffusion-webui/webui.py $configf"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "093d64a7-3d4e-4197-8075-4ed11c7f0ae8",
+   "metadata": {},
+   "source": [
+    "# Free up space"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "370ba58a-d58d-4a80-9575-8c6e094e2626",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Display a list of sessions from which you can remove any session you don't need anymore\n",
+    "\n",
+    "#-------------------------\n",
+    "clean()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

lora_sdxl.py

@@ -0,0 +1,1128 @@
+# LoRA network module
+# reference:
+# https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
+# https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
+
+import math
+import os
+from typing import Dict, List, Optional, Tuple, Type, Union
+from diffusers import AutoencoderKL
+from transformers import CLIPTextModel
+import numpy as np
+import torch
+import re
+
+
+RE_UPDOWN = re.compile(r"(up|down)_blocks_(\d+)_(resnets|upsamplers|downsamplers|attentions)_(\d+)_")
+
+
+class LoRAModule(torch.nn.Module):
+    """
+    replaces forward method of the original Linear, instead of replacing the original Linear module.
+    """
+
+    def __init__(
+        self,
+        lora_name,
+        org_module: torch.nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        dropout=None,
+        rank_dropout=None,
+        module_dropout=None,
+    ):
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
+        super().__init__()
+        self.lora_name = lora_name
+
+        if org_module.__class__.__name__ == "Conv2d":
+            in_dim = org_module.in_channels
+            out_dim = org_module.out_channels
+        else:
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+
+        # if limit_rank:
+        #   self.lora_dim = min(lora_dim, in_dim, out_dim)
+        #   if self.lora_dim != lora_dim:
+        #     print(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
+        # else:
+        self.lora_dim = lora_dim
+
+        if org_module.__class__.__name__ == "Conv2d":
+            kernel_size = org_module.kernel_size
+            stride = org_module.stride
+            padding = org_module.padding
+            self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False)
+            self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
+        else:
+            self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
+            self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)
+
+        if type(alpha) == torch.Tensor:
+            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+        alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
+        self.scale = alpha / self.lora_dim
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
+
+        # same as microsoft's
+        torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+        torch.nn.init.zeros_(self.lora_up.weight)
+
+        self.multiplier = multiplier
+        self.org_module = org_module  # remove in applying
+        self.dropout = dropout
+        self.rank_dropout = rank_dropout
+        self.module_dropout = module_dropout
+
+    def apply_to(self):
+        self.org_forward = self.org_module.forward
+        self.org_module.forward = self.forward
+        del self.org_module
+
+    def forward(self, x):
+        org_forwarded = self.org_forward(x)
+
+        # module dropout
+        if self.module_dropout is not None and self.training:
+            if torch.rand(1) < self.module_dropout:
+                return org_forwarded
+
+        lx = self.lora_down(x)
+
+        # normal dropout
+        if self.dropout is not None and self.training:
+            lx = torch.nn.functional.dropout(lx, p=self.dropout)
+
+        # rank dropout
+        if self.rank_dropout is not None and self.training:
+            mask = torch.rand((lx.size(0), self.lora_dim), device=lx.device) > self.rank_dropout
+            if len(lx.size()) == 3:
+                mask = mask.unsqueeze(1)  # for Text Encoder
+            elif len(lx.size()) == 4:
+                mask = mask.unsqueeze(-1).unsqueeze(-1)  # for Conv2d
+            lx = lx * mask
+
+            # scaling for rank dropout: treat as if the rank is changed
+            # maskから計算することも考えられるが、augmentation的な効果を期待してrank_dropoutを用いる
+            scale = self.scale * (1.0 / (1.0 - self.rank_dropout))  # redundant for readability
+        else:
+            scale = self.scale
+
+        lx = self.lora_up(lx)
+
+        return org_forwarded + lx * self.multiplier * scale
+
+
+class LoRAInfModule(LoRAModule):
+    def __init__(
+        self,
+        lora_name,
+        org_module: torch.nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        **kwargs,
+    ):
+        # no dropout for inference
+        super().__init__(lora_name, org_module, multiplier, lora_dim, alpha)
+
+        self.org_module_ref = [org_module]  # 後から参照できるように
+        self.enabled = True
+
+        # check regional or not by lora_name
+        self.text_encoder = False
+        if lora_name.startswith("lora_te_"):
+            self.regional = False
+            self.use_sub_prompt = True
+            self.text_encoder = True
+        elif "attn2_to_k" in lora_name or "attn2_to_v" in lora_name:
+            self.regional = False
+            self.use_sub_prompt = True
+        elif "time_emb" in lora_name:
+            self.regional = False
+            self.use_sub_prompt = False
+        else:
+            self.regional = True
+            self.use_sub_prompt = False
+
+        self.network: LoRANetwork = None
+
+    def set_network(self, network):
+        self.network = network
+
+    # freezeしてマージする
+    def merge_to(self, sd, dtype, device):
+        # get up/down weight
+        up_weight = sd["lora_up.weight"].to(torch.float).to(device)
+        down_weight = sd["lora_down.weight"].to(torch.float).to(device)
+
+        # extract weight from org_module
+        org_sd = self.org_module.state_dict()
+        weight = org_sd["weight"].to(torch.float)
+
+        # merge weight
+        if len(weight.size()) == 2:
+            # linear
+            weight = weight + self.multiplier * (up_weight @ down_weight) * self.scale
+        elif down_weight.size()[2:4] == (1, 1):
+            # conv2d 1x1
+            weight = (
+                weight
+                + self.multiplier
+                * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
+                * self.scale
+            )
+        else:
+            # conv2d 3x3
+            conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
+            # print(conved.size(), weight.size(), module.stride, module.padding)
+            weight = weight + self.multiplier * conved * self.scale
+
+        # set weight to org_module
+        org_sd["weight"] = weight.to(dtype)
+        self.org_module.load_state_dict(org_sd)
+
+    # 復元できるマージのため、このモジュールのweightを返す
+    def get_weight(self, multiplier=None):
+        if multiplier is None:
+            multiplier = self.multiplier
+
+        # get up/down weight from module
+        up_weight = self.lora_up.weight.to(torch.float)
+        down_weight = self.lora_down.weight.to(torch.float)
+
+        # pre-calculated weight
+        if len(down_weight.size()) == 2:
+            # linear
+            weight = self.multiplier * (up_weight @ down_weight) * self.scale
+        elif down_weight.size()[2:4] == (1, 1):
+            # conv2d 1x1
+            weight = (
+                self.multiplier
+                * (up_weight.squeeze(3).squeeze(2) @ down_weight.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
+                * self.scale
+            )
+        else:
+            # conv2d 3x3
+            conved = torch.nn.functional.conv2d(down_weight.permute(1, 0, 2, 3), up_weight).permute(1, 0, 2, 3)
+            weight = self.multiplier * conved * self.scale
+
+        return weight
+
+    def set_region(self, region):
+        self.region = region
+        self.region_mask = None
+
+    def default_forward(self, x):
+        # print("default_forward", self.lora_name, x.size())
+        return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+
+    def forward(self, x):
+        if not self.enabled:
+            return self.org_forward(x)
+
+        if self.network is None or self.network.sub_prompt_index is None:
+            return self.default_forward(x)
+        if not self.regional and not self.use_sub_prompt:
+            return self.default_forward(x)
+
+        if self.regional:
+            return self.regional_forward(x)
+        else:
+            return self.sub_prompt_forward(x)
+
+    def get_mask_for_x(self, x):
+        # calculate size from shape of x
+        if len(x.size()) == 4:
+            h, w = x.size()[2:4]
+            area = h * w
+        else:
+            area = x.size()[1]
+
+        mask = self.network.mask_dic[area]
+        if mask is None:
+            raise ValueError(f"mask is None for resolution {area}")
+        if len(x.size()) != 4:
+            mask = torch.reshape(mask, (1, -1, 1))
+        return mask
+
+    def regional_forward(self, x):
+        if "attn2_to_out" in self.lora_name:
+            return self.to_out_forward(x)
+
+        if self.network.mask_dic is None:  # sub_prompt_index >= 3
+            return self.default_forward(x)
+
+        # apply mask for LoRA result
+        lx = self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+        mask = self.get_mask_for_x(lx)
+        # print("regional", self.lora_name, self.network.sub_prompt_index, lx.size(), mask.size())
+        lx = lx * mask
+
+        x = self.org_forward(x)
+        x = x + lx
+
+        if "attn2_to_q" in self.lora_name and self.network.is_last_network:
+            x = self.postp_to_q(x)
+
+        return x
+
+    def postp_to_q(self, x):
+        # repeat x to num_sub_prompts
+        has_real_uncond = x.size()[0] // self.network.batch_size == 3
+        qc = self.network.batch_size  # uncond
+        qc += self.network.batch_size * self.network.num_sub_prompts  # cond
+        if has_real_uncond:
+            qc += self.network.batch_size  # real_uncond
+
+        query = torch.zeros((qc, x.size()[1], x.size()[2]), device=x.device, dtype=x.dtype)
+        query[: self.network.batch_size] = x[: self.network.batch_size]
+
+        for i in range(self.network.batch_size):
+            qi = self.network.batch_size + i * self.network.num_sub_prompts
+            query[qi : qi + self.network.num_sub_prompts] = x[self.network.batch_size + i]
+
+        if has_real_uncond:
+            query[-self.network.batch_size :] = x[-self.network.batch_size :]
+
+        # print("postp_to_q", self.lora_name, x.size(), query.size(), self.network.num_sub_prompts)
+        return query
+
+    def sub_prompt_forward(self, x):
+        if x.size()[0] == self.network.batch_size:  # if uncond in text_encoder, do not apply LoRA
+            return self.org_forward(x)
+
+        emb_idx = self.network.sub_prompt_index
+        if not self.text_encoder:
+            emb_idx += self.network.batch_size
+
+        # apply sub prompt of X
+        lx = x[emb_idx :: self.network.num_sub_prompts]
+        lx = self.lora_up(self.lora_down(lx)) * self.multiplier * self.scale
+
+        # print("sub_prompt_forward", self.lora_name, x.size(), lx.size(), emb_idx)
+
+        x = self.org_forward(x)
+        x[emb_idx :: self.network.num_sub_prompts] += lx
+
+        return x
+
+    def to_out_forward(self, x):
+        # print("to_out_forward", self.lora_name, x.size(), self.network.is_last_network)
+
+        if self.network.is_last_network:
+            masks = [None] * self.network.num_sub_prompts
+            self.network.shared[self.lora_name] = (None, masks)
+        else:
+            lx, masks = self.network.shared[self.lora_name]
+
+        # call own LoRA
+        x1 = x[self.network.batch_size + self.network.sub_prompt_index :: self.network.num_sub_prompts]
+        lx1 = self.lora_up(self.lora_down(x1)) * self.multiplier * self.scale
+
+        if self.network.is_last_network:
+            lx = torch.zeros(
+                (self.network.num_sub_prompts * self.network.batch_size, *lx1.size()[1:]), device=lx1.device, dtype=lx1.dtype
+            )
+            self.network.shared[self.lora_name] = (lx, masks)
+
+        # print("to_out_forward", lx.size(), lx1.size(), self.network.sub_prompt_index, self.network.num_sub_prompts)
+        lx[self.network.sub_prompt_index :: self.network.num_sub_prompts] += lx1
+        masks[self.network.sub_prompt_index] = self.get_mask_for_x(lx1)
+
+        # if not last network, return x and masks
+        x = self.org_forward(x)
+        if not self.network.is_last_network:
+            return x
+
+        lx, masks = self.network.shared.pop(self.lora_name)
+
+        # if last network, combine separated x with mask weighted sum
+        has_real_uncond = x.size()[0] // self.network.batch_size == self.network.num_sub_prompts + 2
+
+        out = torch.zeros((self.network.batch_size * (3 if has_real_uncond else 2), *x.size()[1:]), device=x.device, dtype=x.dtype)
+        out[: self.network.batch_size] = x[: self.network.batch_size]  # uncond
+        if has_real_uncond:
+            out[-self.network.batch_size :] = x[-self.network.batch_size :]  # real_uncond
+
+        # print("to_out_forward", self.lora_name, self.network.sub_prompt_index, self.network.num_sub_prompts)
+        # for i in range(len(masks)):
+        #     if masks[i] is None:
+        #         masks[i] = torch.zeros_like(masks[-1])
+
+        mask = torch.cat(masks)
+        mask_sum = torch.sum(mask, dim=0) + 1e-4
+        for i in range(self.network.batch_size):
+            # 1枚の画像ごとに処理する
+            lx1 = lx[i * self.network.num_sub_prompts : (i + 1) * self.network.num_sub_prompts]
+            lx1 = lx1 * mask
+            lx1 = torch.sum(lx1, dim=0)
+
+            xi = self.network.batch_size + i * self.network.num_sub_prompts
+            x1 = x[xi : xi + self.network.num_sub_prompts]
+            x1 = x1 * mask
+            x1 = torch.sum(x1, dim=0)
+            x1 = x1 / mask_sum
+
+            x1 = x1 + lx1
+            out[self.network.batch_size + i] = x1
+
+        # print("to_out_forward", x.size(), out.size(), has_real_uncond)
+        return out
+
+
+def parse_block_lr_kwargs(nw_kwargs):
+    down_lr_weight = nw_kwargs.get("down_lr_weight", None)
+    mid_lr_weight = nw_kwargs.get("mid_lr_weight", None)
+    up_lr_weight = nw_kwargs.get("up_lr_weight", None)
+
+    # 以上のいずれにも設定がない場合は無効としてNoneを返す
+    if down_lr_weight is None and mid_lr_weight is None and up_lr_weight is None:
+        return None, None, None
+
+    # extract learning rate weight for each block
+    if down_lr_weight is not None:
+        # if some parameters are not set, use zero
+        if "," in down_lr_weight:
+            down_lr_weight = [(float(s) if s else 0.0) for s in down_lr_weight.split(",")]
+
+    if mid_lr_weight is not None:
+        mid_lr_weight = float(mid_lr_weight)
+
+    if up_lr_weight is not None:
+        if "," in up_lr_weight:
+            up_lr_weight = [(float(s) if s else 0.0) for s in up_lr_weight.split(",")]
+
+    down_lr_weight, mid_lr_weight, up_lr_weight = get_block_lr_weight(
+        down_lr_weight, mid_lr_weight, up_lr_weight, float(nw_kwargs.get("block_lr_zero_threshold", 0.0))
+    )
+
+    return down_lr_weight, mid_lr_weight, up_lr_weight
+
+
+def create_network(
+    multiplier: float,
+    network_dim: Optional[int],
+    network_alpha: Optional[float],
+    unet,
+    neuron_dropout: Optional[float] = None,
+    **kwargs,
+):
+    if network_dim is None:
+        network_dim = 4  # default
+    if network_alpha is None:
+        network_alpha = 1.0
+
+    # extract dim/alpha for conv2d, and block dim
+    conv_dim = kwargs.get("conv_dim", None)
+    conv_alpha = kwargs.get("conv_alpha", None)
+    if conv_dim is not None:
+        conv_dim = int(conv_dim)
+        if conv_alpha is None:
+            conv_alpha = 1.0
+        else:
+            conv_alpha = float(conv_alpha)
+
+    # block dim/alpha/lr
+    block_dims = kwargs.get("block_dims", None)
+    down_lr_weight, mid_lr_weight, up_lr_weight = parse_block_lr_kwargs(kwargs)
+
+    # 以上のいずれかに指定があればblockごとのdim(rank)を有効にする
+    if block_dims is not None or down_lr_weight is not None or mid_lr_weight is not None or up_lr_weight is not None:
+        block_alphas = kwargs.get("block_alphas", None)
+        conv_block_dims = kwargs.get("conv_block_dims", None)
+        conv_block_alphas = kwargs.get("conv_block_alphas", None)
+
+        block_dims, block_alphas, conv_block_dims, conv_block_alphas = get_block_dims_and_alphas(
+            block_dims, block_alphas, network_dim, network_alpha, conv_block_dims, conv_block_alphas, conv_dim, conv_alpha
+        )
+
+        # remove block dim/alpha without learning rate
+        block_dims, block_alphas, conv_block_dims, conv_block_alphas = remove_block_dims_and_alphas(
+            block_dims, block_alphas, conv_block_dims, conv_block_alphas, down_lr_weight, mid_lr_weight, up_lr_weight
+        )
+
+    else:
+        block_alphas = None
+        conv_block_dims = None
+        conv_block_alphas = None
+
+    # rank/module dropout
+    rank_dropout = kwargs.get("rank_dropout", None)
+    if rank_dropout is not None:
+        rank_dropout = float(rank_dropout)
+    module_dropout = kwargs.get("module_dropout", None)
+    if module_dropout is not None:
+        module_dropout = float(module_dropout)
+
+    # すごく引数が多いな ( ^ω^)･･･
+    network = LoRANetwork(
+        unet,
+        multiplier=multiplier,
+        lora_dim=network_dim,
+        alpha=network_alpha,
+        dropout=neuron_dropout,
+        rank_dropout=rank_dropout,
+        module_dropout=module_dropout,
+        conv_lora_dim=conv_dim,
+        conv_alpha=conv_alpha,
+        block_dims=block_dims,
+        block_alphas=block_alphas,
+        conv_block_dims=conv_block_dims,
+        conv_block_alphas=conv_block_alphas,
+        varbose=True,
+    )
+
+    if up_lr_weight is not None or mid_lr_weight is not None or down_lr_weight is not None:
+        network.set_block_lr_weight(up_lr_weight, mid_lr_weight, down_lr_weight)
+
+    return network
+
+
+# このメソッドは外部から呼び出される可能性を考慮しておく
+# network_dim, network_alpha にはデフォルト値が入っている。
+# block_dims, block_alphas は両方ともNoneまたは両方とも値が入っている
+# conv_dim, conv_alpha は両方ともNoneまたは両方とも値が入っている
+def get_block_dims_and_alphas(
+    block_dims, block_alphas, network_dim, network_alpha, conv_block_dims, conv_block_alphas, conv_dim, conv_alpha
+):
+    num_total_blocks = LoRANetwork.NUM_OF_BLOCKS * 2 + 1
+
+    def parse_ints(s):
+        return [int(i) for i in s.split(",")]
+
+    def parse_floats(s):
+        return [float(i) for i in s.split(",")]
+
+    # block_dimsとblock_alphasをパースする。必ず値が入る
+    if block_dims is not None:
+        block_dims = parse_ints(block_dims)
+        assert (
+            len(block_dims) == num_total_blocks
+        ), f"block_dims must have {num_total_blocks} elements / block_dimsは{num_total_blocks}個指定してください"
+    else:
+        print(f"block_dims is not specified. all dims are set to {network_dim} / block_dimsが指定されていません。すべてのdimは{network_dim}になります")
+        block_dims = [network_dim] * num_total_blocks
+
+    if block_alphas is not None:
+        block_alphas = parse_floats(block_alphas)
+        assert (
+            len(block_alphas) == num_total_blocks
+        ), f"block_alphas must have {num_total_blocks} elements / block_alphasは{num_total_blocks}個指定してください"
+    else:
+        print(
+            f"block_alphas is not specified. all alphas are set to {network_alpha} / block_alphasが指定されていません。すべてのalphaは{network_alpha}になります"
+        )
+        block_alphas = [network_alpha] * num_total_blocks
+
+    # conv_block_dimsとconv_block_alphasを、指定がある場合のみパースする。指定がなければconv_dimとconv_alphaを使う
+    if conv_block_dims is not None:
+        conv_block_dims = parse_ints(conv_block_dims)
+        assert (
+            len(conv_block_dims) == num_total_blocks
+        ), f"conv_block_dims must have {num_total_blocks} elements / conv_block_dimsは{num_total_blocks}個指定してください"
+
+        if conv_block_alphas is not None:
+            conv_block_alphas = parse_floats(conv_block_alphas)
+            assert (
+                len(conv_block_alphas) == num_total_blocks
+            ), f"conv_block_alphas must have {num_total_blocks} elements / conv_block_alphasは{num_total_blocks}個指定してください"
+        else:
+            if conv_alpha is None:
+                conv_alpha = 1.0
+            print(
+                f"conv_block_alphas is not specified. all alphas are set to {conv_alpha} / conv_block_alphasが指定されていません。すべてのalphaは{conv_alpha}になります"
+            )
+            conv_block_alphas = [conv_alpha] * num_total_blocks
+    else:
+        if conv_dim is not None:
+            print(
+                f"conv_dim/alpha for all blocks are set to {conv_dim} and {conv_alpha} / すべてのブロックのconv_dimとalphaは{conv_dim}および{conv_alpha}になります"
+            )
+            conv_block_dims = [conv_dim] * num_total_blocks
+            conv_block_alphas = [conv_alpha] * num_total_blocks
+        else:
+            conv_block_dims = None
+            conv_block_alphas = None
+
+    return block_dims, block_alphas, conv_block_dims, conv_block_alphas
+
+
+# 層別学習率用に層ごとの学習率に対する倍率を定義する、外部から呼び出される可能性を考慮しておく
+def get_block_lr_weight(
+    down_lr_weight, mid_lr_weight, up_lr_weight, zero_threshold
+) -> Tuple[List[float], List[float], List[float]]:
+    # パラメータ未指定時は何もせず、今までと同じ動作とする
+    if up_lr_weight is None and mid_lr_weight is None and down_lr_weight is None:
+        return None, None, None
+
+    max_len = LoRANetwork.NUM_OF_BLOCKS  # フルモデル相当でのup,downの層の数
+
+    def get_list(name_with_suffix) -> List[float]:
+        import math
+
+        tokens = name_with_suffix.split("+")
+        name = tokens[0]
+        base_lr = float(tokens[1]) if len(tokens) > 1 else 0.0
+
+        if name == "cosine":
+            return [math.sin(math.pi * (i / (max_len - 1)) / 2) + base_lr for i in reversed(range(max_len))]
+        elif name == "sine":
+            return [math.sin(math.pi * (i / (max_len - 1)) / 2) + base_lr for i in range(max_len)]
+        elif name == "linear":
+            return [i / (max_len - 1) + base_lr for i in range(max_len)]
+        elif name == "reverse_linear":
+            return [i / (max_len - 1) + base_lr for i in reversed(range(max_len))]
+        elif name == "zeros":
+            return [0.0 + base_lr] * max_len
+        else:
+            print(
+                "Unknown lr_weight argument %s is used. Valid arguments:  / 不明なlr_weightの引数 %s が使われました。有効な引数:\n\tcosine, sine, linear, reverse_linear, zeros"
+                % (name)
+            )
+            return None
+
+    if type(down_lr_weight) == str:
+        down_lr_weight = get_list(down_lr_weight)
+    if type(up_lr_weight) == str:
+        up_lr_weight = get_list(up_lr_weight)
+
+    if (up_lr_weight != None and len(up_lr_weight) > max_len) or (down_lr_weight != None and len(down_lr_weight) > max_len):
+        print("down_weight or up_weight is too long. Parameters after %d-th are ignored." % max_len)
+        print("down_weightもしくはup_weightが長すぎます。%d個目以降のパラメータは無視されます。" % max_len)
+        up_lr_weight = up_lr_weight[:max_len]
+        down_lr_weight = down_lr_weight[:max_len]
+
+    if (up_lr_weight != None and len(up_lr_weight) < max_len) or (down_lr_weight != None and len(down_lr_weight) < max_len):
+        print("down_weight or up_weight is too short. Parameters after %d-th are filled with 1." % max_len)
+        print("down_weightもしくはup_weightが短すぎます。%d個目までの不足したパラメータは1で補われます。" % max_len)
+
+        if down_lr_weight != None and len(down_lr_weight) < max_len:
+            down_lr_weight = down_lr_weight + [1.0] * (max_len - len(down_lr_weight))
+        if up_lr_weight != None and len(up_lr_weight) < max_len:
+            up_lr_weight = up_lr_weight + [1.0] * (max_len - len(up_lr_weight))
+
+    if (up_lr_weight != None) or (mid_lr_weight != None) or (down_lr_weight != None):
+        print("apply block learning rate / 階層別学習率を適用します。")
+        if down_lr_weight != None:
+            down_lr_weight = [w if w > zero_threshold else 0 for w in down_lr_weight]
+            print("down_lr_weight (shallower -> deeper, 浅い層->深い層):", down_lr_weight)
+        else:
+            print("down_lr_weight: all 1.0, すべて1.0")
+
+        if mid_lr_weight != None:
+            mid_lr_weight = mid_lr_weight if mid_lr_weight > zero_threshold else 0
+            print("mid_lr_weight:", mid_lr_weight)
+        else:
+            print("mid_lr_weight: 1.0")
+
+        if up_lr_weight != None:
+            up_lr_weight = [w if w > zero_threshold else 0 for w in up_lr_weight]
+            print("up_lr_weight (deeper -> shallower, 深い層->浅い層):", up_lr_weight)
+        else:
+            print("up_lr_weight: all 1.0, すべて1.0")
+
+    return down_lr_weight, mid_lr_weight, up_lr_weight
+
+
+# lr_weightが0のblockをblock_dimsから除外する、外部から呼び出す可能性を考慮しておく
+def remove_block_dims_and_alphas(
+    block_dims, block_alphas, conv_block_dims, conv_block_alphas, down_lr_weight, mid_lr_weight, up_lr_weight
+):
+    # set 0 to block dim without learning rate to remove the block
+    if down_lr_weight != None:
+        for i, lr in enumerate(down_lr_weight):
+            if lr == 0:
+                block_dims[i] = 0
+                if conv_block_dims is not None:
+                    conv_block_dims[i] = 0
+    if mid_lr_weight != None:
+        if mid_lr_weight == 0:
+            block_dims[LoRANetwork.NUM_OF_BLOCKS] = 0
+            if conv_block_dims is not None:
+                conv_block_dims[LoRANetwork.NUM_OF_BLOCKS] = 0
+    if up_lr_weight != None:
+        for i, lr in enumerate(up_lr_weight):
+            if lr == 0:
+                block_dims[LoRANetwork.NUM_OF_BLOCKS + 1 + i] = 0
+                if conv_block_dims is not None:
+                    conv_block_dims[LoRANetwork.NUM_OF_BLOCKS + 1 + i] = 0
+
+    return block_dims, block_alphas, conv_block_dims, conv_block_alphas
+
+
+# 外部から呼び出す可能性を考慮しておく
+def get_block_index(lora_name: str) -> int:
+    block_idx = -1  # invalid lora name
+
+    m = RE_UPDOWN.search(lora_name)
+    if m:
+        g = m.groups()
+        i = int(g[1])
+        j = int(g[3])
+        if g[2] == "resnets":
+            idx = 3 * i + j
+        elif g[2] == "attentions":
+            idx = 3 * i + j
+        elif g[2] == "upsamplers" or g[2] == "downsamplers":
+            idx = 3 * i + 2
+
+        if g[0] == "down":
+            block_idx = 1 + idx  # 0に該当するLoRAは存在しない
+        elif g[0] == "up":
+            block_idx = LoRANetwork.NUM_OF_BLOCKS + 1 + idx
+
+    elif "mid_block_" in lora_name:
+        block_idx = LoRANetwork.NUM_OF_BLOCKS  # idx=12
+
+    return block_idx
+
+
+# Create network from weights for inference, weights are not loaded here (because can be merged)
+def create_network_from_weights(multiplier, file, vae, text_encoder, unet, weights_sd=None, for_inference=False, **kwargs):
+    if weights_sd is None:
+        if os.path.splitext(file)[1] == ".safetensors":
+            from safetensors.torch import load_file, safe_open
+
+            weights_sd = load_file(file)
+        else:
+            weights_sd = torch.load(file, map_location="cpu")
+
+    # get dim/alpha mapping
+    modules_dim = {}
+    modules_alpha = {}
+    for key, value in weights_sd.items():
+        if "." not in key:
+            continue
+
+        lora_name = key.split(".")[0]
+        if "alpha" in key:
+            modules_alpha[lora_name] = value
+        elif "lora_down" in key:
+            dim = value.size()[0]
+            modules_dim[lora_name] = dim
+            # print(lora_name, value.size(), dim)
+
+    # support old LoRA without alpha
+    for key in modules_dim.keys():
+        if key not in modules_alpha:
+            modules_alpha[key] = modules_dim[key]
+
+    module_class = LoRAInfModule if for_inference else LoRAModule
+
+    network = LoRANetwork(
+        text_encoder, unet, multiplier=multiplier, modules_dim=modules_dim, modules_alpha=modules_alpha, module_class=module_class
+    )
+
+    # block lr
+    down_lr_weight, mid_lr_weight, up_lr_weight = parse_block_lr_kwargs(kwargs)
+    if up_lr_weight is not None or mid_lr_weight is not None or down_lr_weight is not None:
+        network.set_block_lr_weight(up_lr_weight, mid_lr_weight, down_lr_weight)
+
+    return network, weights_sd
+
+
+class LoRANetwork(torch.nn.Module):
+    NUM_OF_BLOCKS = 12  # フルモデル相当でのup,downの層の数
+
+    UNET_TARGET_REPLACE_MODULE = ["Transformer2DModel"]
+    UNET_TARGET_REPLACE_MODULE_CONV2D_3X3 = ["ResnetBlock2D", "Downsample2D", "Upsample2D"]
+    LORA_PREFIX_UNET = "lora_unet"
+    
+    def __init__(
+        self,
+        unet,
+        multiplier: float = 1.0,
+        lora_dim: int = 4,
+        alpha: float = 1,
+        dropout: Optional[float] = None,
+        rank_dropout: Optional[float] = None,
+        module_dropout: Optional[float] = None,
+        conv_lora_dim: Optional[int] = None,
+        conv_alpha: Optional[float] = None,
+        block_dims: Optional[List[int]] = None,
+        block_alphas: Optional[List[float]] = None,
+        conv_block_dims: Optional[List[int]] = None,
+        conv_block_alphas: Optional[List[float]] = None,
+        modules_dim: Optional[Dict[str, int]] = None,
+        modules_alpha: Optional[Dict[str, int]] = None,
+        module_class: Type[object] = LoRAModule,
+        varbose: Optional[bool] = False,
+    ) -> None:
+        """
+        LoRA network: すごく引数が多いが、パターンは以下の通り
+        1. lora_dimとalphaを指定
+        2. lora_dim、alpha、conv_lora_dim、conv_alphaを指定
+        3. block_dimsとblock_alphasを指定 :  Conv2d3x3には適用しない
+        4. block_dims、block_alphas、conv_block_dims、conv_block_alphasを指定 : Conv2d3x3にも適用する
+        5. modules_dimとmodules_alphaを指定 (推論用)
+        """
+        super().__init__()
+        self.multiplier = multiplier
+
+        self.lora_dim = lora_dim
+        self.alpha = alpha
+        self.conv_lora_dim = conv_lora_dim
+        self.conv_alpha = conv_alpha
+        self.dropout = dropout
+        self.rank_dropout = rank_dropout
+        self.module_dropout = module_dropout
+
+
+        # create module instances
+        def create_modules(
+            is_unet: bool,
+            root_module: torch.nn.Module,
+            target_replace_modules: List[torch.nn.Module],
+        ) -> List[LoRAModule]:
+            prefix = (
+                self.LORA_PREFIX_UNET
+            )
+            loras = []
+            skipped = []
+            for name, module in root_module.named_modules():
+                if module.__class__.__name__ in target_replace_modules:
+                    for child_name, child_module in module.named_modules():
+                        is_linear = child_module.__class__.__name__ == "Linear"
+                        is_conv2d = child_module.__class__.__name__ == "Conv2d"
+                        is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1)
+
+                        if is_linear or is_conv2d:
+                            lora_name = prefix + "." + name + "." + child_name
+                            lora_name = lora_name.replace(".", "_")
+
+                            dim = None
+                            alpha = None
+
+                            if modules_dim is not None:
+                                # モジュール指定あり
+                                if lora_name in modules_dim:
+                                    dim = modules_dim[lora_name]
+                                    alpha = modules_alpha[lora_name]
+                            elif is_unet and block_dims is not None:
+                                # U-Netでblock_dims指定あり
+                                block_idx = get_block_index(lora_name)
+                                if is_linear or is_conv2d_1x1:
+                                    dim = block_dims[block_idx]
+                                    alpha = block_alphas[block_idx]
+                                elif conv_block_dims is not None:
+                                    dim = conv_block_dims[block_idx]
+                                    alpha = conv_block_alphas[block_idx]
+                            else:
+                                # 通常、すべて対象とする
+                                if is_linear or is_conv2d_1x1:
+                                    dim = self.lora_dim
+                                    alpha = self.alpha
+                                elif self.conv_lora_dim is not None:
+                                    dim = self.conv_lora_dim
+                                    alpha = self.conv_alpha
+
+                            if dim is None or dim == 0:
+                                # skipした情報を出力
+                                if is_linear or is_conv2d_1x1 or (self.conv_lora_dim is not None or conv_block_dims is not None):
+                                    skipped.append(lora_name)
+                                continue
+
+                            lora = module_class(
+                                lora_name,
+                                child_module,
+                                self.multiplier,
+                                dim,
+                                alpha,
+                                dropout=dropout,
+                                rank_dropout=rank_dropout,
+                                module_dropout=module_dropout,
+                            )
+                            loras.append(lora)
+            return loras, skipped
+
+        # extend U-Net target modules if conv2d 3x3 is enabled, or load from weights
+        target_modules = LoRANetwork.UNET_TARGET_REPLACE_MODULE
+        if modules_dim is not None or self.conv_lora_dim is not None or conv_block_dims is not None:
+            target_modules += LoRANetwork.UNET_TARGET_REPLACE_MODULE_CONV2D_3X3
+
+        self.unet_loras, skipped_un = create_modules(True, unet, target_modules)
+
+
+        skipped = skipped_un
+
+        self.up_lr_weight: List[float] = None
+        self.down_lr_weight: List[float] = None
+        self.mid_lr_weight: float = None
+        self.block_lr = False
+
+        # assertion
+        names = set()
+        for lora in self.unet_loras:
+            assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
+            names.add(lora.lora_name)
+
+    def set_multiplier(self, multiplier):
+        self.multiplier = multiplier
+        for lora in self.unet_loras:
+            lora.multiplier = self.multiplier
+
+    def load_weights(self, file):
+        if os.path.splitext(file)[1] == ".safetensors":
+            from safetensors.torch import load_file
+
+            weights_sd = load_file(file)
+        else:
+            weights_sd = torch.load(file, map_location="cpu")
+
+        info = self.load_state_dict(weights_sd, False)
+        return info
+
+    def apply_to(self, unet, apply_unet=True):
+        for lora in self.unet_loras:
+            lora.apply_to()
+            self.add_module(lora.lora_name, lora)
+
+    # マージできるかどうかを返す
+    def is_mergeable(self):
+        return True
+
+    # TODO refactor to common function with apply_to
+    def merge_to(self, text_encoder, unet, weights_sd, dtype, device):
+        apply_text_encoder = apply_unet = False
+        for key in weights_sd.keys():
+            if key.startswith(LoRANetwork.LORA_PREFIX_TEXT_ENCODER):
+                apply_text_encoder = True
+            elif key.startswith(LoRANetwork.LORA_PREFIX_UNET):
+                apply_unet = True
+
+
+        for lora in self.unet_loras:
+            sd_for_lora = {}
+            for key in weights_sd.keys():
+                if key.startswith(lora.lora_name):
+                    sd_for_lora[key[len(lora.lora_name) + 1 :]] = weights_sd[key]
+            lora.merge_to(sd_for_lora, dtype, device)
+
+ 
+    def set_block_lr_weight(
+        self,
+        up_lr_weight: List[float] = None,
+        mid_lr_weight: float = None,
+        down_lr_weight: List[float] = None,
+    ):
+        self.block_lr = True
+        self.down_lr_weight = down_lr_weight
+        self.mid_lr_weight = mid_lr_weight
+        self.up_lr_weight = up_lr_weight
+
+    def get_lr_weight(self, lora: LoRAModule) -> float:
+        lr_weight = 1.0
+        block_idx = get_block_index(lora.lora_name)
+        if block_idx < 0:
+            return lr_weight
+
+        if block_idx < LoRANetwork.NUM_OF_BLOCKS:
+            if self.down_lr_weight != None:
+                lr_weight = self.down_lr_weight[block_idx]
+        elif block_idx == LoRANetwork.NUM_OF_BLOCKS:
+            if self.mid_lr_weight != None:
+                lr_weight = self.mid_lr_weight
+        elif block_idx > LoRANetwork.NUM_OF_BLOCKS:
+            if self.up_lr_weight != None:
+                lr_weight = self.up_lr_weight[block_idx - LoRANetwork.NUM_OF_BLOCKS - 1]
+
+        return lr_weight
+
+    def prepare_optimizer_params(self, unet_lr):
+        self.requires_grad_(True)
+        all_params = []
+
+        def enumerate_params(loras):
+            params = []
+            for lora in loras:
+                params.extend(lora.parameters())
+            return params
+
+
+        if self.unet_loras:
+            if self.block_lr:
+                # 学習率のグラフをblockごとにしたいので、blockごとにloraを分類
+                block_idx_to_lora = {}
+                for lora in self.unet_loras:
+                    idx = get_block_index(lora.lora_name)
+                    if idx not in block_idx_to_lora:
+                        block_idx_to_lora[idx] = []
+                    block_idx_to_lora[idx].append(lora)
+
+                # blockごとにパラメータを設定する
+                for idx, block_loras in block_idx_to_lora.items():
+                    param_data = {"params": enumerate_params(block_loras)}
+
+                    if unet_lr is not None:
+                        param_data["lr"] = unet_lr * self.get_lr_weight(block_loras[0])
+                    if ("lr" in param_data) and (param_data["lr"] == 0):
+                        continue
+                    all_params.append(param_data)
+
+            else:
+                param_data = {"params": enumerate_params(self.unet_loras)}
+                if unet_lr is not None:
+                    param_data["lr"] = unet_lr
+                all_params.append(param_data)
+
+        return all_params
+
+    def enable_gradient_checkpointing(self):
+        # not supported
+        pass
+
+    def prepare_grad_etc(self, unet):
+        self.requires_grad_(True)
+
+    def on_epoch_start(self, unet):
+        self.train()
+
+    def get_trainable_params(self):
+        return self.parameters()
+
+    def save_weights(self, file, dtype, metadata):
+        if metadata is not None and len(metadata) == 0:
+            metadata = None
+
+        state_dict = self.state_dict()
+
+        if dtype is not None:
+            for key in list(state_dict.keys()):
+                v = state_dict[key]
+                v = v.detach().clone().to("cpu").to(dtype)
+                state_dict[key] = v
+
+        if os.path.splitext(file)[1] == ".safetensors":
+            from safetensors.torch import save_file
+
+            if metadata is None:
+                metadata = {}
+            save_file(state_dict, file, metadata)
+        else:
+            torch.save(state_dict, file)
+
+    # mask is a tensor with values from 0 to 1
+    def set_region(self, sub_prompt_index, is_last_network, mask):
+        if mask.max() == 0:
+            mask = torch.ones_like(mask)
+
+        self.mask = mask
+        self.sub_prompt_index = sub_prompt_index
+        self.is_last_network = is_last_network
+
+        for lora in self.unet_loras:
+            lora.set_network(self)
+
+    def set_current_generation(self, batch_size, num_sub_prompts, width, height, shared):
+        self.batch_size = batch_size
+        self.num_sub_prompts = num_sub_prompts
+        self.current_size = (height, width)
+        self.shared = shared
+
+        # create masks
+        mask = self.mask
+        mask_dic = {}
+        mask = mask.unsqueeze(0).unsqueeze(1)  # b(1),c(1),h,w
+        ref_weight = self.text_encoder_loras[0].lora_down.weight if self.text_encoder_loras else self.unet_loras[0].lora_down.weight
+        dtype = ref_weight.dtype
+        device = ref_weight.device
+
+        def resize_add(mh, mw):
+            # print(mh, mw, mh * mw)
+            m = torch.nn.functional.interpolate(mask, (mh, mw), mode="bilinear")  # doesn't work in bf16
+            m = m.to(device, dtype=dtype)
+            mask_dic[mh * mw] = m
+
+        h = height // 8
+        w = width // 8
+        for _ in range(4):
+            resize_add(h, w)
+            if h % 2 == 1 or w % 2 == 1:  # add extra shape if h/w is not divisible by 2
+                resize_add(h + h % 2, w + w % 2)
+            h = (h + 1) // 2
+            w = (w + 1) // 2
+
+        self.mask_dic = mask_dic
+
+    def backup_weights(self):
+        # 重みのバックアップを行う
+        loras: List[LoRAInfModule] = self.text_encoder_loras + self.unet_loras
+        for lora in loras:
+            org_module = lora.org_module_ref[0]
+            if not hasattr(org_module, "_lora_org_weight"):
+                sd = org_module.state_dict()
+                org_module._lora_org_weight = sd["weight"].detach().clone()
+                org_module._lora_restored = True
+
+    def restore_weights(self):
+        # 重みのリストアを行う
+        loras: List[LoRAInfModule] = self.text_encoder_loras + self.unet_loras
+        for lora in loras:
+            org_module = lora.org_module_ref[0]
+            if not org_module._lora_restored:
+                sd = org_module.state_dict()
+                sd["weight"] = org_module._lora_org_weight
+                org_module.load_state_dict(sd)
+                org_module._lora_restored = True
+
+    def pre_calculation(self):
+        # 事前計算を行う
+        loras: List[LoRAInfModule] = self.text_encoder_loras + self.unet_loras
+        for lora in loras:
+            org_module = lora.org_module_ref[0]
+            sd = org_module.state_dict()
+
+            org_weight = sd["weight"]
+            lora_weight = lora.get_weight().to(org_weight.device, dtype=org_weight.dtype)
+            sd["weight"] = org_weight + lora_weight
+            assert sd["weight"].shape == org_weight.shape
+            org_module.load_state_dict(sd)
+
+            org_module._lora_restored = False
+            lora.enabled = False
+
+    def apply_max_norm_regularization(self, max_norm_value, device):
+        downkeys = []
+        upkeys = []
+        alphakeys = []
+        norms = []
+        keys_scaled = 0
+
+        state_dict = self.state_dict()
+        for key in state_dict.keys():
+            if "lora_down" in key and "weight" in key:
+                downkeys.append(key)
+                upkeys.append(key.replace("lora_down", "lora_up"))
+                alphakeys.append(key.replace("lora_down.weight", "alpha"))
+
+        for i in range(len(downkeys)):
+            down = state_dict[downkeys[i]].to(device)
+            up = state_dict[upkeys[i]].to(device)
+            alpha = state_dict[alphakeys[i]].to(device)
+            dim = down.shape[0]
+            scale = alpha / dim
+
+            if up.shape[2:] == (1, 1) and down.shape[2:] == (1, 1):
+                updown = (up.squeeze(2).squeeze(2) @ down.squeeze(2).squeeze(2)).unsqueeze(2).unsqueeze(3)
+            elif up.shape[2:] == (3, 3) or down.shape[2:] == (3, 3):
+                updown = torch.nn.functional.conv2d(down.permute(1, 0, 2, 3), up).permute(1, 0, 2, 3)
+            else:
+                updown = up @ down
+
+            updown *= scale
+
+            norm = updown.norm().clamp(min=max_norm_value / 2)
+            desired = torch.clamp(norm, max=max_norm_value)
+            ratio = desired.cpu() / norm.cpu()
+            sqrt_ratio = ratio**0.5
+            if ratio != 1:
+                keys_scaled += 1
+                state_dict[upkeys[i]] *= sqrt_ratio
+                state_dict[downkeys[i]] *= sqrt_ratio
+            scalednorm = updown.norm() * ratio
+            norms.append(scalednorm.item())
+
+        return keys_scaled, sum(norms) / len(norms), max(norms)

mainrunpodA1111.py

@@ -0,0 +1,501 @@
+import os
+from IPython.display import clear_output
+from subprocess import call, getoutput, Popen, run
+import time
+import ipywidgets as widgets
+import requests
+import sys
+import fileinput
+from torch.hub import download_url_to_file
+from urllib.parse import urlparse, parse_qs, unquote
+import re
+import six
+
+from urllib.request import urlopen, Request
+import tempfile
+from tqdm import tqdm 
+
+
+
+
+def Deps(force_reinstall):
+
+    if not force_reinstall and os.path.exists('/usr/local/lib/python3.10/dist-packages/safetensors'):
+        ntbks()
+        print('Modules and notebooks updated, dependencies already installed')
+        os.environ['TORCH_HOME'] = '/workspace/cache/torch'
+        os.environ['PYTHONWARNINGS'] = 'ignore'
+    else:
+        call('pip install --root-user-action=ignore --disable-pip-version-check --no-deps -qq gdown PyWavelets numpy==1.23.5 accelerate==0.12.0 --force-reinstall', shell=True, stdout=open('/dev/null', 'w'))
+        ntbks()
+        if os.path.exists('deps'):
+            call("rm -r deps", shell=True)
+        if os.path.exists('diffusers'):
+            call("rm -r diffusers", shell=True)
+        call('mkdir deps', shell=True)
+        if not os.path.exists('cache'):
+            call('mkdir cache', shell=True)
+        os.chdir('deps')
+        dwn("https://huggingface.co/TheLastBen/dependencies/resolve/main/rnpddeps-t2.tar.zst", "/workspace/deps/rnpddeps-t2.tar.zst", "Installing dependencies")
+        call('tar -C / --zstd -xf rnpddeps-t2.tar.zst', shell=True, stdout=open('/dev/null', 'w'))
+        call("sed -i 's@~/.cache@/workspace/cache@' /usr/local/lib/python3.10/dist-packages/transformers/utils/hub.py", shell=True)
+        os.chdir('/workspace')
+        call("git clone --depth 1 -q --branch main https://github.com/TheLastBen/diffusers", shell=True, stdout=open('/dev/null', 'w'))
+        #call('pip install --root-user-action=ignore --disable-pip-version-check -qq gradio==3.41.0', shell=True, stdout=open('/dev/null', 'w'))
+        call("rm -r deps", shell=True)
+        os.chdir('/workspace')
+        os.environ['TORCH_HOME'] = '/workspace/cache/torch'
+        os.environ['PYTHONWARNINGS'] = 'ignore'
+        call("sed -i 's@text = _formatwarnmsg(msg)@text =\"\"@g' /usr/lib/python3.10/warnings.py", shell=True)
+        clear_output()
+
+        done()
+
+
+def dwn(url, dst, msg):
+    file_size = None
+    req = Request(url, headers={"User-Agent": "torch.hub"})
+    u = urlopen(req)
+    meta = u.info()
+    if hasattr(meta, 'getheaders'):
+        content_length = meta.getheaders("Content-Length")
+    else:
+        content_length = meta.get_all("Content-Length")
+    if content_length is not None and len(content_length) > 0:
+        file_size = int(content_length[0])
+
+    with tqdm(total=file_size, disable=False, mininterval=0.5,
+              bar_format=msg+' |{bar:20}| {percentage:3.0f}%') as pbar:
+        with open(dst, "wb") as f:
+            while True:
+                buffer = u.read(8192)
+                if len(buffer) == 0:
+                    break
+                f.write(buffer)
+                pbar.update(len(buffer))
+            f.close()
+
+
+def ntbks():
+
+    os.chdir('/workspace')
+    if not os.path.exists('Latest_Notebooks'):
+        call('mkdir Latest_Notebooks', shell=True)
+    else:
+        call('rm -r Latest_Notebooks', shell=True)
+        call('mkdir Latest_Notebooks', shell=True)
+    os.chdir('/workspace/Latest_Notebooks')
+    call('wget -q -i https://huggingface.co/datasets/TheLastBen/RNPD/raw/main/Notebooks.txt', shell=True)
+    call('rm Notebooks.txt', shell=True)
+    os.chdir('/workspace')
+
+
+def repo(Huggingface_token_optional):
+
+    from slugify import slugify
+    from huggingface_hub import HfApi, CommitOperationAdd, create_repo
+    
+    os.chdir('/workspace')
+    if Huggingface_token_optional!="":
+       username = HfApi().whoami(Huggingface_token_optional)["name"]
+       backup=f"https://huggingface.co/datasets/{username}/fast-stable-diffusion/resolve/main/sd_backup_rnpd.tar.zst"
+       headers = {"Authorization": f"Bearer {Huggingface_token_optional}"}
+       response = requests.head(backup, headers=headers)
+       if response.status_code == 302:
+          print('Restoring the SD folder...')
+          open('/workspace/sd_backup_rnpd.tar.zst', 'wb').write(requests.get(backup, headers=headers).content)
+          call('tar --zstd -xf sd_backup_rnpd.tar.zst', shell=True)
+          call('rm sd_backup_rnpd.tar.zst', shell=True)
+       else:
+          print('Backup not found, using a fresh/existing repo...')
+          time.sleep(2)
+          if not os.path.exists('/workspace/sd/stablediffusiond'): #reset later
+             call('wget -q -O sd_mrep.tar.zst https://huggingface.co/TheLastBen/dependencies/resolve/main/sd_mrep.tar.zst', shell=True)
+             call('tar --zstd -xf sd_mrep.tar.zst', shell=True)
+             call('rm sd_mrep.tar.zst', shell=True)        
+          os.chdir('/workspace/sd')
+          if not os.path.exists('stable-diffusion-webui'):
+              call('git clone -q --depth 1 --branch master https://github.com/AUTOMATIC1111/stable-diffusion-webui', shell=True)
+        
+    else:
+        print('Installing/Updating the repo...')
+        os.chdir('/workspace')
+        if not os.path.exists('/workspace/sd/stablediffusiond'): #reset later
+           call('wget -q -O sd_mrep.tar.zst https://huggingface.co/TheLastBen/dependencies/resolve/main/sd_mrep.tar.zst', shell=True)
+           call('tar --zstd -xf sd_mrep.tar.zst', shell=True)
+           call('rm sd_mrep.tar.zst', shell=True)        
+
+        os.chdir('/workspace/sd')
+        if not os.path.exists('stable-diffusion-webui'):
+            call('git clone -q --depth 1 --branch master https://github.com/AUTOMATIC1111/stable-diffusion-webui', shell=True)
+
+
+    os.chdir('/workspace/sd/stable-diffusion-webui/')
+    call('git reset --hard', shell=True)
+    print('')
+    call('git pull', shell=True)
+    os.chdir('/workspace')
+    clear_output()
+    done()
+
+
+
+def mdl(Original_Model_Version, Path_to_MODEL, MODEL_LINK):
+
+    import gdown
+    
+    src=getsrc(MODEL_LINK)  
+
+    if not os.path.exists('/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/SDv1-5.ckpt'):
+        call('ln -s /workspace/auto-models/* /workspace/sd/stable-diffusion-webui/models/Stable-diffusion', shell=True)
+
+    if Path_to_MODEL !='':
+      if os.path.exists(str(Path_to_MODEL)):
+        print('Using the custom model')
+        model=Path_to_MODEL
+      else:
+          print('Wrong path, check that the path to the model is correct')
+
+    elif MODEL_LINK !="":
+
+      if src=='civitai':
+         modelname=get_name(MODEL_LINK, False)
+         model=f'/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/{modelname}'
+         if not os.path.exists(model):
+            dwn(MODEL_LINK, model, 'Downloading the custom model')
+            clear_output()
+         else:
+            print('Model already exists')
+      elif src=='gdrive':
+         modelname=get_name(MODEL_LINK, True)
+         model=f'/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/{modelname}'
+         if not os.path.exists(model):
+            gdown.download(url=MODEL_LINK, output=model, quiet=False, fuzzy=True)
+            clear_output()
+         else:
+            print('Model already exists')
+      else:
+         modelname=os.path.basename(MODEL_LINK)
+         model=f'/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/{modelname}'
+         if not os.path.exists(model):
+            gdown.download(url=MODEL_LINK, output=model, quiet=False, fuzzy=True)
+            clear_output()
+         else:
+            print('Model already exists')
+
+      if os.path.exists(model) and os.path.getsize(model) > 1810671599:
+        print('Model downloaded, using the custom model.')
+      else:
+        call('rm '+model, shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+        print('Wrong link, check that the link is valid')
+
+    else:
+        if Original_Model_Version == "v1.5":
+           model="/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/SDv1-5.ckpt"
+           print('Using the original V1.5 model')
+        elif Original_Model_Version == "v2-512":
+            model='/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/SDv2-512.ckpt'
+            if not os.path.exists('/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/SDv2-512.ckpt'):
+               print('Downloading the V2-512 model...')
+               call('gdown -O '+model+' https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-nonema-pruned.ckpt', shell=True)
+               clear_output()
+            print('Using the original V2-512 model')
+        elif Original_Model_Version == "v2-768":
+           model="/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/SDv2-768.ckpt"
+           print('Using the original V2-768 model')
+        elif Original_Model_Version == "SDXL":
+            model="/workspace/sd/stable-diffusion-webui/models/Stable-diffusion/sd_xl_base_1.0.safetensors"
+            print('Using the original SDXL model')
+
+        else:
+            model="/workspace/sd/stable-diffusion-webui/models/Stable-diffusion"
+            print('Wrong model version, try again')
+    try:
+        model
+    except:
+        model="/workspace/sd/stable-diffusion-webui/models/Stable-diffusion"
+
+    return model
+
+
+
+def loradwn(LoRA_LINK):
+
+    os.makedirs('/workspace/sd/stable-diffusion-webui/models/Lora', exist_ok=True)
+
+    src=getsrc(LoRA_LINK)
+
+    if src=='civitai':
+        modelname=get_name(LoRA_LINK, False)
+        loramodel=f'/workspace/sd/stable-diffusion-webui/models/Lora/{modelname}'
+        if not os.path.exists(loramodel):
+          dwn(LoRA_LINK, loramodel, 'Downloading the LoRA model')
+          clear_output()
+        else:
+          print('Model already exists')
+    elif src=='gdrive':
+        modelname=get_name(LoRA_LINK, True)
+        loramodel=f'/workspace/sd/stable-diffusion-webui/models/Lora/{modelname}'
+        if not os.path.exists(loramodel):
+          gdown.download(url=LoRA_LINK, output=loramodel, quiet=False, fuzzy=True)
+          clear_output()
+        else:
+          print('Model already exists')
+    else:
+        modelname=os.path.basename(LoRA_LINK)
+        loramodel=f'/workspace/sd/stable-diffusion-webui/models/Lora/{modelname}'
+        if not os.path.exists(loramodel):
+          gdown.download(url=LoRA_LINK, output=loramodel, quiet=False, fuzzy=True)
+          clear_output()
+        else:
+          print('Model already exists')
+
+    if os.path.exists(loramodel) :
+      print('LoRA downloaded')
+    else:
+      print('Wrong link, check that the link is valid')
+
+
+
+def CNet(ControlNet_Model, ControlNet_v2_Model):
+    
+    def download(url, model_dir):
+
+        filename = os.path.basename(urlparse(url).path)
+        pth = os.path.abspath(os.path.join(model_dir, filename))
+        if not os.path.exists(pth):
+            print('Downloading: '+os.path.basename(url))
+            download_url_to_file(url, pth, hash_prefix=None, progress=True)
+        else:
+          print(f"The model {filename} already exists")    
+
+    wrngv1=False
+    os.chdir('/workspace/sd/stable-diffusion-webui/extensions')
+    if not os.path.exists("sd-webui-controlnet"):
+      call('git clone https://github.com/Mikubill/sd-webui-controlnet.git', shell=True)
+      os.chdir('/workspace')
+    else:
+      os.chdir('sd-webui-controlnet')
+      call('git reset --hard', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+      call('git pull', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+      os.chdir('/workspace')
+
+    mdldir="/workspace/sd/stable-diffusion-webui/extensions/sd-webui-controlnet/models"
+    for filename in os.listdir(mdldir):
+      if "_sd14v1" in filename:
+        renamed = re.sub("_sd14v1", "-fp16", filename)
+        os.rename(os.path.join(mdldir, filename), os.path.join(mdldir, renamed))
+
+    call('wget -q -O CN_models.txt https://github.com/TheLastBen/fast-stable-diffusion/raw/main/AUTOMATIC1111_files/CN_models.txt', shell=True)
+    call('wget -q -O CN_models_v2.txt https://github.com/TheLastBen/fast-stable-diffusion/raw/main/AUTOMATIC1111_files/CN_models_v2.txt', shell=True)
+      
+    with open("CN_models.txt", 'r') as f:
+        mdllnk = f.read().splitlines()
+    with open("CN_models_v2.txt", 'r') as d:
+        mdllnk_v2 = d.read().splitlines()
+    call('rm CN_models.txt CN_models_v2.txt', shell=True)
+    
+    cfgnames=[os.path.basename(url).split('.')[0]+'.yaml' for url in mdllnk_v2]
+    os.chdir('/workspace/sd/stable-diffusion-webui/extensions/sd-webui-controlnet/models')
+    for name in cfgnames:
+        run(['cp', 'cldm_v21.yaml', name])
+    os.chdir('/workspace')
+
+    if ControlNet_Model == "All" or ControlNet_Model == "all" :     
+      for lnk in mdllnk:
+          download(lnk, mdldir)
+      clear_output()
+
+      
+    elif ControlNet_Model == "15":
+      mdllnk=list(filter(lambda x: 't2i' in x, mdllnk))
+      for lnk in mdllnk:
+          download(lnk, mdldir)
+      clear_output()        
+
+
+    elif ControlNet_Model.isdigit() and int(ControlNet_Model)-1<14 and int(ControlNet_Model)>0:
+      download(mdllnk[int(ControlNet_Model)-1], mdldir)
+      clear_output()
+      
+    elif ControlNet_Model == "none":
+       pass
+       clear_output()
+
+    else:
+      print('Wrong ControlNet V1 choice, try again')
+      wrngv1=True
+
+
+    if ControlNet_v2_Model == "All" or ControlNet_v2_Model == "all" :
+      for lnk_v2 in mdllnk_v2:
+          download(lnk_v2, mdldir)
+      if not wrngv1:
+        clear_output()
+      done()
+
+    elif ControlNet_v2_Model.isdigit() and int(ControlNet_v2_Model)-1<5:
+      download(mdllnk_v2[int(ControlNet_v2_Model)-1], mdldir)
+      if not wrngv1:
+        clear_output()
+      done()
+    
+    elif ControlNet_v2_Model == "none":
+       pass
+       if not wrngv1:
+        clear_output()
+       done()       
+
+    else:
+      print('Wrong ControlNet V2 choice, try again')
+
+
+
+def sd(User, Password, model):
+
+    import gradio
+    
+    gradio.close_all()
+    
+    auth=f"--gradio-auth {User}:{Password}"
+    if User =="" or Password=="":
+      auth=""
+
+    call('wget -q -O /usr/local/lib/python3.10/dist-packages/gradio/blocks.py https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/AUTOMATIC1111_files/blocks.py', shell=True)
+   
+    os.chdir('/workspace/sd/stable-diffusion-webui/modules')
+
+    call("sed -i 's@possible_sd_paths =.*@possible_sd_paths = [\"/workspace/sd/stablediffusion\"]@' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True)
+    call("sed -i 's@\.\.\/@src/@g' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True)
+    call("sed -i 's@src\/generative-models@generative-models@g' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True)
+    
+    call("sed -i 's@\[\"sd_model_checkpoint\"\]@\[\"sd_model_checkpoint\", \"sd_vae\", \"CLIP_stop_at_last_layers\", \"inpainting_mask_weight\", \"initial_noise_multiplier\"\]@g' /workspace/sd/stable-diffusion-webui/modules/shared.py", shell=True)
+    
+    call("sed -i 's@print(\"No module.*@@' /workspace/sd/stablediffusion/ldm/modules/diffusionmodules/model.py", shell=True)
+    os.chdir('/workspace/sd/stable-diffusion-webui')
+    clear_output()
+
+    podid=os.environ.get('RUNPOD_POD_ID')
+    localurl=f"{podid}-3001.proxy.runpod.net"
+
+    for line in fileinput.input('/usr/local/lib/python3.10/dist-packages/gradio/blocks.py', inplace=True):
+      if line.strip().startswith('self.server_name ='):
+          line = f'            self.server_name = "{localurl}"\n'
+      if line.strip().startswith('self.protocol = "https"'):
+          line = '            self.protocol = "https"\n'
+      if line.strip().startswith('if self.local_url.startswith("https") or self.is_colab'):
+          line = ''
+      if line.strip().startswith('else "http"'):
+          line = ''
+      sys.stdout.write(line)
+      
+    if model=="":
+        mdlpth=""
+    else:
+        if os.path.isfile(model):
+            mdlpth="--ckpt "+model
+        else:
+            mdlpth="--ckpt-dir "+model
+
+    configf="--disable-console-progressbars --no-half-vae --disable-safe-unpickle --api --no-download-sd-model --opt-sdp-attention --enable-insecure-extension-access  --skip-version-check --listen --port 3000 "+auth+" "+mdlpth
+
+    return configf
+
+
+
+def save(Huggingface_Write_token):
+    
+    from slugify import slugify
+    from huggingface_hub import HfApi, CommitOperationAdd, create_repo
+    
+    if Huggingface_Write_token=="":
+        print('A huggingface write token is required')
+        
+    else:
+        os.chdir('/workspace')
+        
+        if os.path.exists('sd'):
+        
+            call('tar --exclude="stable-diffusion-webui/models/*/*" --exclude="sd-webui-controlnet/models/*" --zstd -cf sd_backup_rnpd.tar.zst sd', shell=True)
+            api = HfApi()
+            username = api.whoami(token=Huggingface_Write_token)["name"]
+
+            repo_id = f"{username}/{slugify('fast-stable-diffusion')}"
+
+            print("Backing up...")
+
+            operations = [CommitOperationAdd(path_in_repo="sd_backup_rnpd.tar.zst", path_or_fileobj="/workspace/sd_backup_rnpd.tar.zst")]
+
+            create_repo(repo_id,private=True, token=Huggingface_Write_token, exist_ok=True, repo_type="dataset")
+
+            api.create_commit(
+              repo_id=repo_id,
+              repo_type="dataset",
+              operations=operations,
+              commit_message="SD folder Backup",
+              token=Huggingface_Write_token
+            )
+
+            call('rm sd_backup_rnpd.tar.zst', shell=True)
+            clear_output()
+
+            done()
+     
+        else:
+            print('Nothing to backup')
+        
+
+
+
+def getsrc(url):
+
+    parsed_url = urlparse(url)
+    
+    if parsed_url.netloc == 'civitai.com':
+        src='civitai'
+    elif parsed_url.netloc == 'drive.google.com':
+        src='gdrive'
+    elif parsed_url.netloc == 'huggingface.co':
+        src='huggingface'
+    else:
+        src='others'
+    return src
+
+
+
+def get_name(url, gdrive):
+
+    from gdown.download import get_url_from_gdrive_confirmation
+
+    if not gdrive:
+        response = requests.get(url, allow_redirects=False)
+        if "Location" in response.headers:
+            redirected_url = response.headers["Location"]
+            quer = parse_qs(urlparse(redirected_url).query)
+            if "response-content-disposition" in quer:
+                disp_val = quer["response-content-disposition"][0].split(";")
+                for vals in disp_val:
+                    if vals.strip().startswith("filename="):
+                        filenm=unquote(vals.split("=", 1)[1].strip())
+                        return filenm.replace("\"","")
+    else:
+        headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36"}
+        lnk="https://drive.google.com/uc?id={id}&export=download".format(id=url[url.find("/d/")+3:url.find("/view")])
+        res = requests.session().get(lnk, headers=headers, stream=True, verify=True)
+        res = requests.session().get(get_url_from_gdrive_confirmation(res.text), headers=headers, stream=True, verify=True)
+        content_disposition = six.moves.urllib_parse.unquote(res.headers["Content-Disposition"])
+        filenm = re.search(r"filename\*=UTF-8''(.*)", content_disposition).groups()[0].replace(os.path.sep, "_")
+        return filenm   
+
+
+
+
+def done():
+    done = widgets.Button(
+        description='Done!',
+        disabled=True,
+        button_style='success',
+        tooltip='',
+        icon='check'
+    )
+    display(done)

sdxllorarunpod.py

@@ -0,0 +1,1131 @@
+from IPython.display import clear_output
+from subprocess import call, getoutput, Popen
+from IPython.display import display
+import ipywidgets as widgets
+import io
+from PIL import Image, ImageDraw, ImageOps
+import fileinput
+import time
+import os
+from os import listdir
+from os.path import isfile
+import random
+import sys
+from io import BytesIO
+import requests
+from collections import defaultdict
+from math import log, sqrt
+import numpy as np
+import sys
+import fileinput
+from subprocess import check_output
+import six
+import base64
+import re
+
+from urllib.parse import urlparse, parse_qs, unquote
+import urllib.request
+from urllib.request import urlopen, Request
+
+import tempfile
+from tqdm import tqdm
+
+
+
+
+def Deps(force_reinstall):
+
+    if not force_reinstall and os.path.exists('/usr/local/lib/python3.10/dist-packages/safetensors'):
+        ntbks()
+        call('pip install --root-user-action=ignore --disable-pip-version-check -qq diffusers==0.18.1', shell=True, stdout=open('/dev/null', 'w'))
+        print('Modules and notebooks updated, dependencies already installed')        
+        os.environ['TORCH_HOME'] = '/workspace/cache/torch'
+        os.environ['PYTHONWARNINGS'] = 'ignore'
+    else:
+        call('pip install --root-user-action=ignore --disable-pip-version-check --no-deps -qq gdown PyWavelets numpy==1.23.5 accelerate==0.12.0 --force-reinstall', shell=True, stdout=open('/dev/null', 'w'))
+        ntbks()
+        if os.path.exists('deps'):
+            call("rm -r deps", shell=True)
+        if os.path.exists('diffusers'):
+            call("rm -r diffusers", shell=True)
+        call('mkdir deps', shell=True)
+        if not os.path.exists('cache'):
+            call('mkdir cache', shell=True)
+        os.chdir('deps')
+        dwn("https://huggingface.co/TheLastBen/dependencies/resolve/main/rnpddeps-t2.tar.zst", "/workspace/deps/rnpddeps-t2.tar.zst", "Installing dependencies")
+        call('tar -C / --zstd -xf rnpddeps-t2.tar.zst', shell=True, stdout=open('/dev/null', 'w'))
+        call("sed -i 's@~/.cache@/workspace/cache@' /usr/local/lib/python3.10/dist-packages/transformers/utils/hub.py", shell=True)
+        os.chdir('/workspace')
+        call('pip install --root-user-action=ignore --disable-pip-version-check -qq diffusers==0.18.1', shell=True, stdout=open('/dev/null', 'w'))
+        call("git clone --depth 1 -q --branch main https://github.com/TheLastBen/diffusers", shell=True, stdout=open('/dev/null', 'w'))
+        #call('pip install --root-user-action=ignore --disable-pip-version-check -qq gradio==3.41.0', shell=True, stdout=open('/dev/null', 'w'))
+        call("rm -r deps", shell=True)
+        os.chdir('/workspace')
+        os.environ['TORCH_HOME'] = '/workspace/cache/torch'
+        os.environ['PYTHONWARNINGS'] = 'ignore'
+        call("sed -i 's@text = _formatwarnmsg(msg)@text =\"\"@g' /usr/lib/python3.10/warnings.py", shell=True)
+        clear_output()
+
+        done()
+
+
+def dwn(url, dst, msg):
+    file_size = None
+    req = Request(url, headers={"User-Agent": "torch.hub"})
+    u = urlopen(req)
+    meta = u.info()
+    if hasattr(meta, 'getheaders'):
+        content_length = meta.getheaders("Content-Length")
+    else:
+        content_length = meta.get_all("Content-Length")
+    if content_length is not None and len(content_length) > 0:
+        file_size = int(content_length[0])
+
+    with tqdm(total=file_size, disable=False, mininterval=0.5,
+              bar_format=msg+' |{bar:20}| {percentage:3.0f}%') as pbar:
+        with open(dst, "wb") as f:
+            while True:
+                buffer = u.read(8192)
+                if len(buffer) == 0:
+                    break
+                f.write(buffer)
+                pbar.update(len(buffer))
+            f.close()
+
+
+def ntbks():
+
+    os.chdir('/workspace')
+    if not os.path.exists('Latest_Notebooks'):
+        call('mkdir Latest_Notebooks', shell=True)
+    else:
+        call('rm -r Latest_Notebooks', shell=True)
+        call('mkdir Latest_Notebooks', shell=True)
+    os.chdir('/workspace/Latest_Notebooks')
+    call('wget -q -i https://huggingface.co/datasets/TheLastBen/RNPD/raw/main/Notebooks.txt', shell=True)
+    call('rm Notebooks.txt', shell=True)
+    os.chdir('/workspace')
+
+def done():
+    done = widgets.Button(
+        description='Done!',
+        disabled=True,
+        button_style='success',
+        tooltip='',
+        icon='check'
+    )
+    display(done)
+
+
+
+def mdlvxl():
+
+  os.chdir('/workspace')
+
+  if os.path.exists('stable-diffusion-XL') and not os.path.exists('/workspace/stable-diffusion-XL/unet/diffusion_pytorch_model.safetensors'):
+     call('rm -r stable-diffusion-XL', shell=True)
+  if not os.path.exists('stable-diffusion-XL'):
+      print('Downloading SDXL model...')
+      call('mkdir stable-diffusion-XL', shell=True)
+      os.chdir('stable-diffusion-XL')
+      call('git init', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+      call('git lfs install --system --skip-repo', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+      call('git remote add -f origin  https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+      call('git config core.sparsecheckout true', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+      call('echo -e "\nscheduler\ntext_encoder\ntext_encoder_2\ntokenizer\ntokenizer_2\nunet\nvae\nfeature_extractor\nmodel_index.json\n!*.safetensors\n!*.bin\n!*.onnx*\n!*.xml" > .git/info/sparse-checkout', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+      call('git pull origin main', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+      dwn('https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/text_encoder/model.safetensors', 'text_encoder/model.safetensors', '1/4')
+      dwn('https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/text_encoder_2/model.safetensors', 'text_encoder_2/model.safetensors', '2/4')
+      dwn('https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/vae/diffusion_pytorch_model.safetensors', 'vae/diffusion_pytorch_model.safetensors', '3/4')
+      dwn('https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/unet/diffusion_pytorch_model.safetensors', 'unet/diffusion_pytorch_model.safetensors', '4/4')  
+      call('rm -r .git', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+      os.chdir('/workspace')
+      clear_output()
+      while not os.path.exists('/workspace/stable-diffusion-XL/unet/diffusion_pytorch_model.safetensors'):
+            print('Invalid HF token, make sure you have access to the model')
+            time.sleep(8)
+      if os.path.exists('/workspace/stable-diffusion-XL/unet/diffusion_pytorch_model.safetensors'):
+          print('Using SDXL model')
+  else:
+    print('Using SDXL model')
+    
+  call("sed -i 's@\"force_upcast.*@@' /workspace/stable-diffusion-XL/vae/config.json", shell=True)
+
+  
+
+def downloadmodel_hfxl(Path_to_HuggingFace):
+
+  os.chdir('/workspace')
+  if os.path.exists('stable-diffusion-custom'):
+    call("rm -r stable-diffusion-custom", shell=True)
+  clear_output()
+  
+  if os.path.exists('Fast-Dreambooth/token.txt'):
+    with open("Fast-Dreambooth/token.txt") as f:
+       token = f.read()
+    authe=f'https://USER:{token}@'
+  else:
+    authe="https://"
+
+  clear_output()
+  call("mkdir stable-diffusion-custom", shell=True)
+  os.chdir("stable-diffusion-custom")
+  call("git init", shell=True)
+  call("git lfs install --system --skip-repo", shell=True)
+  call('git remote add -f origin '+authe+'huggingface.co/'+Path_to_HuggingFace, shell=True)
+  call("git config core.sparsecheckout true", shell=True)
+  call('echo -e "\nscheduler\ntext_encoder\ntokenizer\nunet\nvae\nfeature_extractor\nmodel_index.json\n!*.fp16.safetensors" > .git/info/sparse-checkout', shell=True)
+  call("git pull origin main", shell=True)
+  if os.path.exists('unet/diffusion_pytorch_model.safetensors'):
+    call("rm -r .git", shell=True)
+    os.chdir('/workspace')
+    clear_output()
+    done()
+  while not os.path.exists('/workspace/stable-diffusion-custom/unet/diffusion_pytorch_model.safetensors'):
+        print('Check the link you provided')
+        os.chdir('/workspace')
+        time.sleep(5)
+  
+        
+
+def downloadmodel_link_xl(MODEL_LINK): 
+
+    import wget
+    import gdown
+    from gdown.download import get_url_from_gdrive_confirmation    
+    
+    def getsrc(url):
+        parsed_url = urlparse(url)
+        if parsed_url.netloc == 'civitai.com':
+            src='civitai'
+        elif parsed_url.netloc == 'drive.google.com':
+            src='gdrive'
+        elif parsed_url.netloc == 'huggingface.co':
+            src='huggingface'
+        else:
+            src='others'
+        return src
+        
+    src=getsrc(MODEL_LINK)
+
+    def get_name(url, gdrive):
+        if not gdrive:
+            response = requests.get(url, allow_redirects=False)
+            if "Location" in response.headers:
+                redirected_url = response.headers["Location"]
+                quer = parse_qs(urlparse(redirected_url).query)
+                if "response-content-disposition" in quer:
+                    disp_val = quer["response-content-disposition"][0].split(";")
+                    for vals in disp_val:
+                        if vals.strip().startswith("filename="):
+                            filenm=unquote(vals.split("=", 1)[1].strip())
+                            return filenm.replace("\"","")
+        else:
+            headers = {"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_10_1) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/39.0.2171.95 Safari/537.36"}
+            lnk="https://drive.google.com/uc?id={id}&export=download".format(id=url[url.find("/d/")+3:url.find("/view")])
+            res = requests.session().get(lnk, headers=headers, stream=True, verify=True)
+            res = requests.session().get(get_url_from_gdrive_confirmation(res.text), headers=headers, stream=True, verify=True)
+            content_disposition = six.moves.urllib_parse.unquote(res.headers["Content-Disposition"])
+            filenm = re.search(r"filename\*=UTF-8''(.*)", content_disposition).groups()[0].replace(os.path.sep, "_")
+            return filenm   
+
+    if src=='civitai':
+       modelname=get_name(MODEL_LINK, False)
+    elif src=='gdrive':
+       modelname=get_name(MODEL_LINK, True)
+    else:
+       modelname=os.path.basename(MODEL_LINK)
+
+
+    os.chdir('/workspace')
+    if src=='huggingface':
+        dwn(MODEL_LINK, modelname,'Downloading the Model')
+    else:
+        call("gdown --fuzzy " +MODEL_LINK+ " -O "+modelname, shell=True)
+    
+    if os.path.exists(modelname):
+      if os.path.getsize(modelname) > 1810671599:
+          
+        print('Converting to diffusers...')    
+        call('python /workspace/diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path '+modelname+' --dump_path stable-diffusion-custom --from_safetensors', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+       
+        if os.path.exists('stable-diffusion-custom/unet/diffusion_pytorch_model.bin'):
+          os.chdir('/workspace')
+          clear_output()
+          done()
+        else:
+            while not os.path.exists('stable-diffusion-custom/unet/diffusion_pytorch_model.bin'):
+              print('Conversion error')
+              os.chdir('/workspace')
+              time.sleep(5)
+    else:
+        while os.path.getsize(modelname) < 1810671599:
+           print('Wrong link, check that the link is valid')
+           os.chdir('/workspace')
+           time.sleep(5)
+
+
+
+def downloadmodel_path_xl(MODEL_PATH):
+
+  import wget
+  os.chdir('/workspace')
+  clear_output() 
+  if os.path.exists(str(MODEL_PATH)):
+  
+    print('Converting to diffusers...')
+    call('python /workspace/diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path '+MODEL_PATH+' --dump_path stable-diffusion-custom --from_safetensors', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+
+    if os.path.exists('stable-diffusion-custom/unet/diffusion_pytorch_model.bin'):
+      clear_output()
+      done()
+    while not os.path.exists('stable-diffusion-custom/unet/diffusion_pytorch_model.bin'):
+      print('Conversion error')
+      os.chdir('/workspace')
+      time.sleep(5)
+  else:
+    while not os.path.exists(str(MODEL_PATH)):
+       print('Wrong path, use the file explorer to copy the path')
+       os.chdir('/workspace')
+       time.sleep(5)
+
+
+
+
+def dls_xlf(Path_to_HuggingFace, MODEL_PATH, MODEL_LINK):
+
+    os.chdir('/workspace')
+    
+    if Path_to_HuggingFace != "":
+      downloadmodel_hfxl(Path_to_HuggingFace)
+      MODEL_NAMExl="/workspace/stable-diffusion-custom"
+
+    elif MODEL_PATH !="":
+        
+      downloadmodel_path_xl(MODEL_PATH)
+      MODEL_NAMExl="/workspace/stable-diffusion-custom"
+
+    elif MODEL_LINK !="":
+
+      downloadmodel_link_xl(MODEL_LINK)
+      MODEL_NAMExl="/workspace/stable-diffusion-custom"
+
+    else:
+        mdlvxl()
+        MODEL_NAMExl="/workspace/stable-diffusion-XL"
+
+    return MODEL_NAMExl    
+    
+
+
+def sess_xl(Session_Name, MODEL_NAMExl):
+    import gdown
+    import wget
+    os.chdir('/workspace')
+    PT=""
+
+    while Session_Name=="":
+      print('Input the Session Name:') 
+      Session_Name=input("")
+    Session_Name=Session_Name.replace(" ","_")
+
+    WORKSPACE='/workspace/Fast-Dreambooth'
+
+    INSTANCE_NAME=Session_Name
+    OUTPUT_DIR="/workspace/models/"+Session_Name
+    SESSION_DIR=WORKSPACE+"/Sessions/"+Session_Name
+    INSTANCE_DIR=SESSION_DIR+"/instance_images"
+    CAPTIONS_DIR=SESSION_DIR+'/captions'
+    MDLPTH=str(SESSION_DIR+"/"+Session_Name+'.safetensors')
+    
+
+    if os.path.exists(str(SESSION_DIR)) and not os.path.exists(MDLPTH):
+        print('Loading session with no previous LoRa model')
+        if MODEL_NAMExl=="":
+            print('No model found, use the "Model Download" cell to download a model.')
+        else:
+            print('Session Loaded, proceed')
+
+    elif not os.path.exists(str(SESSION_DIR)):
+        call('mkdir -p '+INSTANCE_DIR, shell=True)
+        print('Creating session...')
+        if MODEL_NAMExl=="":
+          print('No model found, use the "Model Download" cell to download a model.')
+        else:
+          print('Session created, proceed to uploading instance images')
+          if MODEL_NAMExl=="":
+             print('No model found, use the "Model Download" cell to download a model.')        
+
+    else:
+        print('Session Loaded, proceed')
+        
+        
+    return WORKSPACE, Session_Name, INSTANCE_NAME, OUTPUT_DIR, SESSION_DIR, INSTANCE_DIR, CAPTIONS_DIR, MDLPTH, MODEL_NAMExl
+
+
+
+def uplder(Remove_existing_instance_images, Crop_images, Crop_size, IMAGES_FOLDER_OPTIONAL, INSTANCE_DIR, CAPTIONS_DIR):
+
+    if os.path.exists(INSTANCE_DIR+"/.ipynb_checkpoints"):
+      call('rm -r '+INSTANCE_DIR+'/.ipynb_checkpoints', shell=True)
+
+    uploader = widgets.FileUpload(description="Choose images",accept='image/*, .txt', multiple=True)
+    Upload = widgets.Button(
+        description='Upload',
+        disabled=False,
+        button_style='info', 
+        tooltip='Click to upload the chosen instance images',
+        icon=''
+    )
+
+
+    def up(Upload):
+        with out: 
+            uploader.close()
+            Upload.close()
+            upld(Remove_existing_instance_images, Crop_images, Crop_size, IMAGES_FOLDER_OPTIONAL, INSTANCE_DIR, CAPTIONS_DIR, uploader)
+            done()
+    out=widgets.Output()
+    
+    if IMAGES_FOLDER_OPTIONAL=="":
+      Upload.on_click(up)
+      display(uploader, Upload, out)
+    else:
+       upld(Remove_existing_instance_images, Crop_images, Crop_size, IMAGES_FOLDER_OPTIONAL, INSTANCE_DIR, CAPTIONS_DIR, uploader)
+       done()
+    
+    
+
+def upld(Remove_existing_instance_images, Crop_images, Crop_size, IMAGES_FOLDER_OPTIONAL, INSTANCE_DIR, CAPTIONS_DIR, uploader):
+
+    from tqdm import tqdm
+    if Remove_existing_instance_images:
+        if os.path.exists(str(INSTANCE_DIR)):
+            call("rm -r " +INSTANCE_DIR, shell=True)
+        if os.path.exists(str(CAPTIONS_DIR)):
+            call("rm -r " +CAPTIONS_DIR, shell=True)            
+
+
+    if not os.path.exists(str(INSTANCE_DIR)):
+        call("mkdir -p " +INSTANCE_DIR, shell=True)
+    if not os.path.exists(str(CAPTIONS_DIR)):
+        call("mkdir -p " +CAPTIONS_DIR, shell=True)        
+
+
+    if IMAGES_FOLDER_OPTIONAL !="":
+        if os.path.exists(IMAGES_FOLDER_OPTIONAL+"/.ipynb_checkpoints"):
+          call('rm -r '+IMAGES_FOLDER_OPTIONAL+'/.ipynb_checkpoints', shell=True)    
+    
+        if any(file.endswith('.{}'.format('txt')) for file in os.listdir(IMAGES_FOLDER_OPTIONAL)):
+            call('mv '+IMAGES_FOLDER_OPTIONAL+'/*.txt '+CAPTIONS_DIR, shell=True)
+        if Crop_images:   
+            os.chdir(str(IMAGES_FOLDER_OPTIONAL))
+            call('find . -name "* *" -type f | rename ' "'s/ /-/g'", shell=True)
+            os.chdir('/workspace')    
+            for filename in tqdm(os.listdir(IMAGES_FOLDER_OPTIONAL), bar_format='  |{bar:15}| {n_fmt}/{total_fmt} Uploaded'):
+                extension = filename.split(".")[-1]
+                identifier=filename.split(".")[0]
+                new_path_with_file = os.path.join(INSTANCE_DIR, filename)
+                file = Image.open(IMAGES_FOLDER_OPTIONAL+"/"+filename)
+                file=file.convert("RGB")
+                file=ImageOps.exif_transpose(file)
+                width, height = file.size
+                if file.size !=(Crop_size, Crop_size):
+                    image=crop_image(file, Crop_size)
+                    if extension.upper()=="JPG" or extension.upper()=="jpg":
+                        image[0].save(new_path_with_file, format="JPEG", quality = 100)
+                    else:
+                        image[0].save(new_path_with_file, format=extension.upper())
+                        
+                else:
+                   call("cp \'"+IMAGES_FOLDER_OPTIONAL+"/"+filename+"\' "+INSTANCE_DIR, shell=True)                        
+
+        else:
+            for filename in tqdm(os.listdir(IMAGES_FOLDER_OPTIONAL), bar_format='  |{bar:15}| {n_fmt}/{total_fmt} Uploaded'):
+                call("cp -r " +IMAGES_FOLDER_OPTIONAL+"/. " +INSTANCE_DIR, shell=True)
+
+    elif IMAGES_FOLDER_OPTIONAL =="":
+        up=""  
+        for file in uploader.value:
+          filename = file['name']
+          if filename.split(".")[-1]=="txt":
+            with open(CAPTIONS_DIR+'/'+filename, 'w') as f:
+                f.write(bytes(file['content']).decode())            
+          up=[file for file in uploader.value if not file['name'].endswith('.txt')]
+        if Crop_images:
+            for file in tqdm(up, bar_format='  |{bar:15}| {n_fmt}/{total_fmt} Uploaded'):
+                filename = file['name']
+                img = Image.open(io.BytesIO(file['content']))
+                img=img.convert("RGB")
+                img=ImageOps.exif_transpose(img)
+                extension = filename.split(".")[-1]
+                identifier=filename.split(".")[0]
+
+                if extension.upper()=="JPG" or extension.upper()=="jpg":
+                    img.save(INSTANCE_DIR+"/"+filename, format="JPEG", quality = 100) 
+                else:
+                    img.save(INSTANCE_DIR+"/"+filename, format=extension.upper())                
+                
+                new_path_with_file = os.path.join(INSTANCE_DIR, filename)
+                file = Image.open(new_path_with_file)
+                width, height = file.size    
+                if file.size !=(Crop_size, Crop_size):
+                    image=crop_image(file, Crop_size)
+                    if extension.upper()=="JPG" or extension.upper()=="jpg":
+                        image[0].save(new_path_with_file, format="JPEG", quality = 100) 
+                    else:
+                        image[0].save(new_path_with_file, format=extension.upper())
+
+        else:
+            for file in tqdm(uploader.value, bar_format='  |{bar:15}| {n_fmt}/{total_fmt} Uploaded'):
+                filename = file['name']
+                img = Image.open(io.BytesIO(file['content']))
+                img=img.convert("RGB")
+                extension = filename.split(".")[-1]
+                identifier=filename.split(".")[0]   
+                
+                if extension.upper()=="JPG" or extension.upper()=="jpg":
+                    img.save(INSTANCE_DIR+"/"+filename, format="JPEG", quality = 100) 
+                else:
+                    img.save(INSTANCE_DIR+"/"+filename, format=extension.upper())                  
+   
+
+    os.chdir(INSTANCE_DIR)
+    call('find . -name "* *" -type f | rename ' "'s/ /-/g'", shell=True)
+    os.chdir(CAPTIONS_DIR)
+    call('find . -name "* *" -type f | rename ' "'s/ /-/g'", shell=True)    
+    os.chdir('/workspace')
+
+
+
+
+def caption(CAPTIONS_DIR, INSTANCE_DIR):
+   
+  paths=""
+  out=""
+  widgets_l=""
+  clear_output()
+  def Caption(path):
+      if path!="Select an instance image to caption":
+        
+        name = os.path.splitext(os.path.basename(path))[0]
+        ext=os.path.splitext(os.path.basename(path))[-1][1:]
+        if ext=="jpg" or "JPG":
+          ext="JPEG"        
+
+        if os.path.exists(CAPTIONS_DIR+"/"+name + '.txt'):
+          with open(CAPTIONS_DIR+"/"+name + '.txt', 'r') as f:
+              text = f.read()
+        else:
+          with open(CAPTIONS_DIR+"/"+name + '.txt', 'w') as f:
+              f.write("")
+              with open(CAPTIONS_DIR+"/"+name + '.txt', 'r') as f:
+                  text = f.read()   
+
+        img=Image.open(os.path.join(INSTANCE_DIR,path))
+        img=img.convert("RGB")
+        img=img.resize((420, 420))
+        image_bytes = BytesIO()
+        img.save(image_bytes, format=ext, qualiy=10)
+        image_bytes.seek(0)
+        image_data = image_bytes.read()
+        img= image_data  
+        image = widgets.Image(
+            value=img,
+            width=420,
+            height=420
+        )
+        text_area = widgets.Textarea(value=text, description='', disabled=False, layout={'width': '300px', 'height': '120px'})
+        
+
+        def update_text(text):
+            with open(CAPTIONS_DIR+"/"+name + '.txt', 'w') as f:
+                f.write(text)
+
+        button = widgets.Button(description='Save', button_style='success')
+        button.on_click(lambda b: update_text(text_area.value))
+
+        return widgets.VBox([widgets.HBox([image, text_area, button])])
+
+
+  paths = os.listdir(INSTANCE_DIR)
+  widgets_l = widgets.Select(options=["Select an instance image to caption"]+paths, rows=25)
+
+
+  out = widgets.Output()
+
+  def click(change):
+      with out:
+          out.clear_output()
+          display(Caption(change.new))
+
+  widgets_l.observe(click, names='value')
+  display(widgets.HBox([widgets_l, out]))
+
+
+
+def dbtrainxl(Unet_Training_Epochs, Text_Encoder_Training_Epochs, Unet_Learning_Rate, Text_Encoder_Learning_Rate, dim, Offset_Noise, Resolution, MODEL_NAME, SESSION_DIR, INSTANCE_DIR, CAPTIONS_DIR, External_Captions,  INSTANCE_NAME, Session_Name, OUTPUT_DIR, ofstnselvl, Save_VRAM):
+
+
+    if os.path.exists(INSTANCE_DIR+"/.ipynb_checkpoints"):
+      call('rm -r '+INSTANCE_DIR+'/.ipynb_checkpoints', shell=True)
+    if os.path.exists(CAPTIONS_DIR+"/.ipynb_checkpoints"):
+      call('rm -r '+CAPTIONS_DIR+'/.ipynb_checkpoints', shell=True)
+
+
+    Seed=random.randint(1, 999999)
+    
+    ofstnse=""
+    if Offset_Noise:
+      ofstnse="--offset_noise"
+      
+    GC=''
+    if Save_VRAM:
+        GC='--gradient_checkpointing'
+        
+    extrnlcptn=""
+    if External_Captions:
+      extrnlcptn="--external_captions"      
+
+    precision="fp16"
+
+   
+   
+    def train_only_text(SESSION_DIR, MODEL_NAME, INSTANCE_DIR, OUTPUT_DIR, Seed, Resolution, ofstnse, extrnlcptn, precision, Training_Epochs):
+        print('Training the Text Encoder...')
+        call('accelerate launch /workspace/diffusers/examples/dreambooth/train_dreambooth_rnpd_sdxl_TI.py \
+        '+ofstnse+' \
+        '+extrnlcptn+' \
+        --dim='+str(dim)+' \
+        --ofstnselvl='+str(ofstnselvl)+' \
+        --image_captions_filename \
+        --Session_dir='+SESSION_DIR+' \
+        --pretrained_model_name_or_path='+MODEL_NAME+' \
+        --instance_data_dir='+INSTANCE_DIR+' \
+        --output_dir='+OUTPUT_DIR+' \
+        --captions_dir='+CAPTIONS_DIR+' \
+        --seed='+str(Seed)+' \
+        --resolution='+str(Resolution)+' \
+        --mixed_precision='+str(precision)+' \
+        --train_batch_size=1 \
+        --gradient_accumulation_steps=1 '+GC+ ' \
+        --use_8bit_adam \
+        --learning_rate='+str(Text_Encoder_Learning_Rate)+' \
+        --lr_scheduler="cosine" \
+        --lr_warmup_steps=0 \
+        --num_train_epochs='+str(Training_Epochs), shell=True)
+   
+   
+
+    def train_only_unet(SESSION_DIR, MODEL_NAME, INSTANCE_DIR, OUTPUT_DIR, Seed, Resolution, ofstnse, extrnlcptn, precision, Training_Epochs):
+        print('Training the UNet...')
+        call('accelerate launch /workspace/diffusers/examples/dreambooth/train_dreambooth_rnpd_sdxl_lora.py \
+        '+ofstnse+' \
+        '+extrnlcptn+' \
+        --dim='+str(dim)+' \
+        --ofstnselvl='+str(ofstnselvl)+' \
+        --image_captions_filename \
+        --Session_dir='+SESSION_DIR+' \
+        --pretrained_model_name_or_path='+MODEL_NAME+' \
+        --instance_data_dir='+INSTANCE_DIR+' \
+        --output_dir='+OUTPUT_DIR+' \
+        --captions_dir='+CAPTIONS_DIR+' \
+        --seed='+str(Seed)+' \
+        --resolution='+str(Resolution)+' \
+        --mixed_precision='+str(precision)+' \
+        --train_batch_size=1 \
+        --gradient_accumulation_steps=1 '+GC+ ' \
+        --use_8bit_adam \
+        --learning_rate='+str(Unet_Learning_Rate)+' \
+        --lr_scheduler="cosine" \
+        --lr_warmup_steps=0 \
+        --num_train_epochs='+str(Training_Epochs), shell=True)
+
+      
+
+    if Unet_Training_Epochs!=0:
+        if Text_Encoder_Training_Epochs!=0:
+            train_only_text(SESSION_DIR, MODEL_NAME, INSTANCE_DIR, OUTPUT_DIR, Seed, Resolution, ofstnse, extrnlcptn, precision, Training_Epochs=Text_Encoder_Training_Epochs)
+            clear_output()
+        train_only_unet(SESSION_DIR, MODEL_NAME, INSTANCE_DIR, OUTPUT_DIR, Seed, Resolution, ofstnse, extrnlcptn, precision, Training_Epochs=Unet_Training_Epochs)
+    else  :
+      print('Nothing to do')
+
+
+    if os.path.exists(SESSION_DIR+'/'+Session_Name+'.safetensors'):
+        clear_output()
+        print("DONE, the LoRa model is in the session's folder")
+    else:
+        print("Something went wrong")
+
+
+
+
+def sdcmff(Huggingface_token_optional, MDLPTH, restored):
+
+    from slugify import slugify
+    from huggingface_hub import HfApi, CommitOperationAdd, create_repo
+    
+    os.chdir('/workspace')
+
+    if restored:
+        Huggingface_token_optional=""
+
+    if Huggingface_token_optional!="":
+       username = HfApi().whoami(Huggingface_token_optional)["name"]
+       backup=f"https://huggingface.co/datasets/{username}/fast-stable-diffusion/resolve/main/sdcomfy_backup_rnpd.tar.zst"
+       headers = {"Authorization": f"Bearer {Huggingface_token_optional}"}
+       response = requests.head(backup, headers=headers)
+       if response.status_code == 302:
+          restored=True
+          print('Restoring ComfyUI...')
+          open('/workspace/sdcomfy_backup_rnpd.tar.zst', 'wb').write(requests.get(backup, headers=headers).content)
+          call('tar --zstd -xf sdcomfy_backup_rnpd.tar.zst', shell=True)
+          call('rm sdcomfy_backup_rnpd.tar.zst', shell=True)
+       else:
+          print('Backup not found, using a fresh/existing repo...')
+          time.sleep(2)
+          if not os.path.exists('ComfyUI'):
+              call('git clone -q --depth 1 https://github.com/comfyanonymous/ComfyUI', shell=True)
+    else:
+        print('Installing/Updating the repo...')
+        if not os.path.exists('ComfyUI'):
+            call('git clone -q --depth 1 https://github.com/comfyanonymous/ComfyUI', shell=True)
+
+    os.chdir('ComfyUI')
+    call('git reset --hard', shell=True)
+    print('')
+    call('git pull', shell=True)
+
+    if os.path.exists(MDLPTH):
+        call('ln -s '+MDLPTH+' models/loras', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+        
+    clean_symlinks('models/loras')        
+
+    if not os.path.exists('models/checkpoints/sd_xl_base_1.0.safetensors'):
+        call('ln -s /workspace/auto-models/* models/checkpoints', shell=True)
+
+
+    podid=os.environ.get('RUNPOD_POD_ID')
+    localurl=f"https://{podid}-3001.proxy.runpod.net"
+    call("sed -i 's@print(\"To see the GUI go to: http://{}:{}\".format(address, port))@print(\"\u2714 Connected\")\\n            print(\""+localurl+"\")@' /workspace/ComfyUI/server.py", shell=True)
+    os.chdir('/workspace')
+    
+    return restored 
+
+
+
+
+def test(MDLPTH, User, Password, Huggingface_token_optional, restoreda):
+
+    from slugify import slugify
+    from huggingface_hub import HfApi, CommitOperationAdd, create_repo
+    import gradio
+    
+    gradio.close_all()
+               
+
+    auth=f"--gradio-auth {User}:{Password}"
+    if User =="" or Password=="":
+      auth=""
+
+
+    if restoreda:
+        Huggingface_token_optional=""
+
+    if Huggingface_token_optional!="":
+       username = HfApi().whoami(Huggingface_token_optional)["name"]
+       backup=f"https://huggingface.co/datasets/{username}/fast-stable-diffusion/resolve/main/sd_backup_rnpd.tar.zst"
+       headers = {"Authorization": f"Bearer {Huggingface_token_optional}"}
+       response = requests.head(backup, headers=headers)
+       if response.status_code == 302:
+          restoreda=True
+          print('Restoring the SD folder...')
+          open('/workspace/sd_backup_rnpd.tar.zst', 'wb').write(requests.get(backup, headers=headers).content)
+          call('tar --zstd -xf sd_backup_rnpd.tar.zst', shell=True)
+          call('rm sd_backup_rnpd.tar.zst', shell=True)
+       else:
+          print('Backup not found, using a fresh/existing repo...')
+          time.sleep(2)
+          if not os.path.exists('/workspace/sd/stablediffusiond'): #reset later
+             call('wget -q -O sd_mrep.tar.zst https://huggingface.co/TheLastBen/dependencies/resolve/main/sd_mrep.tar.zst', shell=True)
+             call('tar --zstd -xf sd_mrep.tar.zst', shell=True)
+             call('rm sd_mrep.tar.zst', shell=True)        
+          os.chdir('/workspace/sd')
+          if not os.path.exists('stable-diffusion-webui'):
+              call('git clone -q --depth 1 --branch master https://github.com/AUTOMATIC1111/stable-diffusion-webui', shell=True)            
+        
+    else:
+        print('Installing/Updating the repo...')
+        os.chdir('/workspace')
+        if not os.path.exists('/workspace/sd/stablediffusiond'): #reset later
+           call('wget -q -O sd_mrep.tar.zst https://huggingface.co/TheLastBen/dependencies/resolve/main/sd_mrep.tar.zst', shell=True)
+           call('tar --zstd -xf sd_mrep.tar.zst', shell=True)
+           call('rm sd_mrep.tar.zst', shell=True)        
+
+        os.chdir('/workspace/sd')
+        if not os.path.exists('stable-diffusion-webui'):
+            call('git clone -q --depth 1 --branch master https://github.com/AUTOMATIC1111/stable-diffusion-webui', shell=True)
+
+
+    os.chdir('/workspace/sd/stable-diffusion-webui/')
+    call('git reset --hard', shell=True)
+    print('')
+    call('git pull', shell=True)
+    
+    
+    if os.path.exists(MDLPTH):
+        call('mkdir models/Lora', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+        call('ln -s '+MDLPTH+' models/Lora', shell=True, stdout=open('/dev/null', 'w'), stderr=open('/dev/null', 'w'))
+    
+    if not os.path.exists('models/Stable-diffusion/sd_xl_base_1.0.safetensors'):
+        call('ln -s /workspace/auto-models/* models/Stable-diffusion', shell=True)
+
+    clean_symlinks('models/Lora')
+   
+    os.chdir('/workspace')
+
+
+    call('wget -q -O /usr/local/lib/python3.10/dist-packages/gradio/blocks.py https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/AUTOMATIC1111_files/blocks.py', shell=True)
+   
+    os.chdir('/workspace/sd/stable-diffusion-webui/modules')
+    
+    call("sed -i 's@possible_sd_paths =.*@possible_sd_paths = [\"/workspace/sd/stablediffusion\"]@' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True)
+    call("sed -i 's@\.\.\/@src/@g' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True)
+    call("sed -i 's@src\/generative-models@generative-models@g' /workspace/sd/stable-diffusion-webui/modules/paths.py", shell=True)
+    
+    call("sed -i 's@\[\"sd_model_checkpoint\"\]@\[\"sd_model_checkpoint\", \"sd_vae\", \"CLIP_stop_at_last_layers\", \"inpainting_mask_weight\", \"initial_noise_multiplier\"\]@g' /workspace/sd/stable-diffusion-webui/modules/shared.py", shell=True)
+    call("sed -i 's@print(\"No module.*@@' /workspace/sd/stablediffusion/ldm/modules/diffusionmodules/model.py", shell=True)
+    os.chdir('/workspace/sd/stable-diffusion-webui')
+    clear_output()
+
+    podid=os.environ.get('RUNPOD_POD_ID')
+    localurl=f"{podid}-3001.proxy.runpod.net"
+
+    for line in fileinput.input('/usr/local/lib/python3.10/dist-packages/gradio/blocks.py', inplace=True):
+      if line.strip().startswith('self.server_name ='):
+          line = f'            self.server_name = "{localurl}"\n'
+      if line.strip().startswith('self.protocol = "https"'):
+          line = '            self.protocol = "https"\n'
+      if line.strip().startswith('if self.local_url.startswith("https") or self.is_colab'):
+          line = ''
+      if line.strip().startswith('else "http"'):
+          line = ''
+      sys.stdout.write(line)
+
+
+    configf="--disable-console-progressbars --upcast-sampling --no-half-vae --disable-safe-unpickle --api --opt-sdp-attention --enable-insecure-extension-access --no-download-sd-model  --skip-version-check  --listen --port 3000 --ckpt /workspace/sd/stable-diffusion-webui/models/Stable-diffusion/sd_xl_base_1.0.safetensors "+auth
+    
+    
+    return configf, restoreda
+
+
+
+
+def clean():
+    
+    Sessions=os.listdir("/workspace/Fast-Dreambooth/Sessions")
+
+    s = widgets.Select(
+        options=Sessions,
+        rows=5,
+        description='',
+        disabled=False
+    )
+
+    out=widgets.Output()
+
+    d = widgets.Button(
+        description='Remove',
+        disabled=False,
+        button_style='warning',
+        tooltip='Removet the selected session',
+        icon='warning'
+    )
+
+    def rem(d):
+        with out:
+            if s.value is not None:
+                clear_output()
+                print("THE SESSION "+s.value+" HAS BEEN REMOVED FROM THE STORAGE")
+                call('rm -r /workspace/Fast-Dreambooth/Sessions/'+s.value, shell=True)
+                if os.path.exists('/workspace/models/'+s.value):
+                  call('rm -r /workspace/models/'+s.value, shell=True)
+                s.options=os.listdir("/workspace/Fast-Dreambooth/Sessions")       
+
+
+            else:
+                d.close()
+                s.close()
+                clear_output()
+                print("NOTHING TO REMOVE")
+
+    d.on_click(rem)
+    if s.value is not None:
+        display(s,d,out)
+    else:
+        print("NOTHING TO REMOVE")
+
+
+
+def crop_image(im, size):
+
+  import cv2
+
+  GREEN = "#0F0"
+  BLUE = "#00F"
+  RED = "#F00"    
+
+  def focal_point(im, settings):
+      corner_points = image_corner_points(im, settings) if settings.corner_points_weight > 0 else []
+      entropy_points = image_entropy_points(im, settings) if settings.entropy_points_weight > 0 else []
+      face_points = image_face_points(im, settings) if settings.face_points_weight > 0 else []
+
+      pois = []
+
+      weight_pref_total = 0
+      if len(corner_points) > 0:
+        weight_pref_total += settings.corner_points_weight
+      if len(entropy_points) > 0:
+        weight_pref_total += settings.entropy_points_weight
+      if len(face_points) > 0:
+        weight_pref_total += settings.face_points_weight
+
+      corner_centroid = None
+      if len(corner_points) > 0:
+        corner_centroid = centroid(corner_points)
+        corner_centroid.weight = settings.corner_points_weight / weight_pref_total 
+        pois.append(corner_centroid)
+
+      entropy_centroid = None
+      if len(entropy_points) > 0:
+        entropy_centroid = centroid(entropy_points)
+        entropy_centroid.weight = settings.entropy_points_weight / weight_pref_total
+        pois.append(entropy_centroid)
+
+      face_centroid = None
+      if len(face_points) > 0:
+        face_centroid = centroid(face_points)
+        face_centroid.weight = settings.face_points_weight / weight_pref_total 
+        pois.append(face_centroid)
+
+      average_point = poi_average(pois, settings)
+      
+      return average_point
+
+
+  def image_face_points(im, settings):
+
+      np_im = np.array(im)
+      gray = cv2.cvtColor(np_im, cv2.COLOR_BGR2GRAY)
+
+      tries = [
+        [ f'{cv2.data.haarcascades}haarcascade_eye.xml', 0.01 ],
+        [ f'{cv2.data.haarcascades}haarcascade_frontalface_default.xml', 0.05 ],
+        [ f'{cv2.data.haarcascades}haarcascade_profileface.xml', 0.05 ],
+        [ f'{cv2.data.haarcascades}haarcascade_frontalface_alt.xml', 0.05 ],
+        [ f'{cv2.data.haarcascades}haarcascade_frontalface_alt2.xml', 0.05 ],
+        [ f'{cv2.data.haarcascades}haarcascade_frontalface_alt_tree.xml', 0.05 ],
+        [ f'{cv2.data.haarcascades}haarcascade_eye_tree_eyeglasses.xml', 0.05 ],
+        [ f'{cv2.data.haarcascades}haarcascade_upperbody.xml', 0.05 ]
+      ]
+      for t in tries:
+        classifier = cv2.CascadeClassifier(t[0])
+        minsize = int(min(im.width, im.height) * t[1]) # at least N percent of the smallest side
+        try:
+          faces = classifier.detectMultiScale(gray, scaleFactor=1.1,
+            minNeighbors=7, minSize=(minsize, minsize), flags=cv2.CASCADE_SCALE_IMAGE)
+        except:
+          continue
+
+        if len(faces) > 0:
+          rects = [[f[0], f[1], f[0] + f[2], f[1] + f[3]] for f in faces]
+          return [PointOfInterest((r[0] +r[2]) // 2, (r[1] + r[3]) // 2, size=abs(r[0]-r[2]), weight=1/len(rects)) for r in rects]
+      return []
+
+
+  def image_corner_points(im, settings):
+      grayscale = im.convert("L")
+
+      # naive attempt at preventing focal points from collecting at watermarks near the bottom
+      gd = ImageDraw.Draw(grayscale)
+      gd.rectangle([0, im.height*.9, im.width, im.height], fill="#999")
+
+      np_im = np.array(grayscale)
+
+      points = cv2.goodFeaturesToTrack(
+          np_im,
+          maxCorners=100,
+          qualityLevel=0.04,
+          minDistance=min(grayscale.width, grayscale.height)*0.06,
+          useHarrisDetector=False,
+      )
+
+      if points is None:
+          return []
+
+      focal_points = []
+      for point in points:
+        x, y = point.ravel()
+        focal_points.append(PointOfInterest(x, y, size=4, weight=1/len(points)))
+
+      return focal_points
+
+
+  def image_entropy_points(im, settings):
+      landscape = im.height < im.width
+      portrait = im.height > im.width
+      if landscape:
+        move_idx = [0, 2]
+        move_max = im.size[0]
+      elif portrait:
+        move_idx = [1, 3]
+        move_max = im.size[1]
+      else:
+        return []
+
+      e_max = 0
+      crop_current = [0, 0, settings.crop_width, settings.crop_height]
+      crop_best = crop_current
+      while crop_current[move_idx[1]] < move_max:
+          crop = im.crop(tuple(crop_current))
+          e = image_entropy(crop)
+
+          if (e > e_max):
+            e_max = e
+            crop_best = list(crop_current)
+
+          crop_current[move_idx[0]] += 4
+          crop_current[move_idx[1]] += 4
+
+      x_mid = int(crop_best[0] + settings.crop_width/2)
+      y_mid = int(crop_best[1] + settings.crop_height/2)
+
+      return [PointOfInterest(x_mid, y_mid, size=25, weight=1.0)]
+
+
+  def image_entropy(im):
+      # greyscale image entropy
+      # band = np.asarray(im.convert("L"))
+      band = np.asarray(im.convert("1"), dtype=np.uint8)
+      hist, _ = np.histogram(band, bins=range(0, 256))
+      hist = hist[hist > 0]
+      return -np.log2(hist / hist.sum()).sum()
+
+  def centroid(pois):
+    x = [poi.x for poi in pois]
+    y = [poi.y for poi in pois]
+    return PointOfInterest(sum(x)/len(pois), sum(y)/len(pois))
+
+
+  def poi_average(pois, settings):
+      weight = 0.0
+      x = 0.0
+      y = 0.0
+      for poi in pois:
+          weight += poi.weight
+          x += poi.x * poi.weight
+          y += poi.y * poi.weight
+      avg_x = round(weight and x / weight)
+      avg_y = round(weight and y / weight)
+
+      return PointOfInterest(avg_x, avg_y)
+
+
+  def is_landscape(w, h):
+    return w > h
+
+
+  def is_portrait(w, h):
+    return h > w
+
+
+  def is_square(w, h):
+    return w == h
+
+
+  class PointOfInterest:
+    def __init__(self, x, y, weight=1.0, size=10):
+      self.x = x
+      self.y = y
+      self.weight = weight
+      self.size = size
+
+    def bounding(self, size):
+      return [
+        self.x - size//2,
+        self.y - size//2,
+        self.x + size//2,
+        self.y + size//2
+      ]
+
+  class Settings:
+    def __init__(self, crop_width=512, crop_height=512, corner_points_weight=0.5, entropy_points_weight=0.5, face_points_weight=0.5):
+      self.crop_width = crop_width
+      self.crop_height = crop_height
+      self.corner_points_weight = corner_points_weight
+      self.entropy_points_weight = entropy_points_weight
+      self.face_points_weight = face_points_weight
+
+  settings = Settings(
+      crop_width = size,
+      crop_height = size,
+      face_points_weight = 0.9,
+      entropy_points_weight = 0.15,
+      corner_points_weight = 0.5,
+  )        
+
+  scale_by = 1
+  if is_landscape(im.width, im.height):
+    scale_by = settings.crop_height / im.height
+  elif is_portrait(im.width, im.height):
+    scale_by = settings.crop_width / im.width
+  elif is_square(im.width, im.height):
+    if is_square(settings.crop_width, settings.crop_height):
+      scale_by = settings.crop_width / im.width
+    elif is_landscape(settings.crop_width, settings.crop_height):
+      scale_by = settings.crop_width / im.width
+    elif is_portrait(settings.crop_width, settings.crop_height):
+      scale_by = settings.crop_height / im.height
+
+  im = im.resize((int(im.width * scale_by), int(im.height * scale_by)))
+  im_debug = im.copy()
+
+  focus = focal_point(im_debug, settings)
+
+  # take the focal point and turn it into crop coordinates that try to center over the focal
+  # point but then get adjusted back into the frame
+  y_half = int(settings.crop_height / 2)
+  x_half = int(settings.crop_width / 2)
+
+  x1 = focus.x - x_half
+  if x1 < 0:
+      x1 = 0
+  elif x1 + settings.crop_width > im.width:
+      x1 = im.width - settings.crop_width
+
+  y1 = focus.y - y_half
+  if y1 < 0:
+      y1 = 0
+  elif y1 + settings.crop_height > im.height:
+      y1 = im.height - settings.crop_height
+
+  x2 = x1 + settings.crop_width
+  y2 = y1 + settings.crop_height
+
+  crop = [x1, y1, x2, y2]
+
+  results = []
+
+  results.append(im.crop(tuple(crop)))
+
+  return results
+  
+
+
+def clean_symlinks(path):
+    for item in os.listdir(path):
+        lnk = os.path.join(path, item)
+        if os.path.islink(lnk) and not os.path.exists(os.readlink(lnk)):
+            os.remove(lnk)
+  

train_dreambooth_rnpd_sdxl_lora.py

@@ -0,0 +1,782 @@
+import argparse
+import itertools
+import math
+import os
+from pathlib import Path
+from typing import Optional
+import subprocess
+import sys
+
+import gc
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from torch.utils.data import Dataset
+from transformers import AutoTokenizer, PretrainedConfig
+import bitsandbytes as bnb
+
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import set_seed
+from contextlib import nullcontext
+from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionXLPipeline, UNet2DConditionModel
+from diffusers.optimization import get_scheduler
+from huggingface_hub import HfFolder, Repository, whoami
+from PIL import Image
+from torchvision import transforms
+from tqdm import tqdm
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextConfig, CLIPTextModelWithProjection
+
+from lora_sdxl import *
+
+logger = get_logger(__name__)
+
+
+def import_model_class_from_model_name_or_path(
+    pretrained_model_name_or_path: str, subfolder: str = "text_encoder"
+):
+    text_encoder_config = PretrainedConfig.from_pretrained(
+        pretrained_model_name_or_path,
+        subfolder=subfolder,
+        use_auth_token=True
+    )
+    model_class = text_encoder_config.architectures[0]
+
+    if model_class == "CLIPTextModel":
+        from transformers import CLIPTextModel
+
+        return CLIPTextModel
+    elif model_class == "CLIPTextModelWithProjection":
+        from transformers import CLIPTextModelWithProjection
+
+        return CLIPTextModelWithProjection
+    else:
+        raise ValueError(f"{model_class} is not supported.")
+        
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default=None,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )
+    parser.add_argument(
+        "--tokenizer_name",
+        type=str,
+        default=None,
+        help="Pretrained tokenizer name or path if not the same as model_name",
+    )
+    parser.add_argument(
+        "--instance_data_dir",
+        type=str,
+        default=None,
+        required=True,
+        help="A folder containing the training data of instance images.",
+    )
+    parser.add_argument(
+        "--class_data_dir",
+        type=str,
+        default=None,
+        required=False,
+        help="A folder containing the training data of class images.",
+    )
+    parser.add_argument(
+        "--instance_prompt",
+        type=str,
+        default=None,
+        help="The prompt with identifier specifying the instance",
+    )
+    parser.add_argument(
+        "--class_prompt",
+        type=str,
+        default="",
+        help="The prompt to specify images in the same class as provided instance images.",
+    )
+    parser.add_argument(
+        "--with_prior_preservation",
+        default=False,
+        action="store_true",
+        help="Flag to add prior preservation loss.",
+    )
+    parser.add_argument("--prior_loss_weight", type=float, default=1.0, help="The weight of prior preservation loss.")
+    parser.add_argument(
+        "--num_class_images",
+        type=int,
+        default=100,
+        help=(
+            "Minimal class images for prior preservation loss. If not have enough images, additional images will be"
+            " sampled with class_prompt."
+        ),
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--resolution",
+        type=int,
+        default=512,
+        help=(
+            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
+            " resolution"
+        ),
+    )
+    parser.add_argument(
+        "--center_crop", action="store_true", help="Whether to center crop images before resizing to resolution"
+    )
+    parser.add_argument("--train_text_encoder", action="store_true", help="Whether to train the text encoder")
+    parser.add_argument(
+        "--train_batch_size", type=int, default=4, help="Batch size (per device) for the training dataloader."
+    )
+    parser.add_argument(
+        "--sample_batch_size", type=int, default=4, help="Batch size (per device) for sampling images."
+    )
+    parser.add_argument("--num_train_epochs", type=int, default=1)
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=5e-6,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=(
+            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+            ' "constant", "constant_with_warmup"]'
+        ),
+    )
+    parser.add_argument(
+        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
+    )
+    parser.add_argument(
+        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
+    )
+    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
+    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
+    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
+    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
+    parser.add_argument(
+        "--hub_model_id",
+        type=str,
+        default=None,
+        help="The name of the repository to keep in sync with the local `output_dir`.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=(
+            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
+        ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose"
+            "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+            "and an Nvidia Ampere GPU."
+        ),
+    )
+
+    parser.add_argument(
+        "--save_n_steps",
+        type=int,
+        default=1,
+        help=("Save the model every n global_steps"),
+    )
+    
+    
+    parser.add_argument(
+        "--save_starting_step",
+        type=int,
+        default=1,
+        help=("The step from which it starts saving intermediary checkpoints"),
+    )
+    
+    parser.add_argument(
+        "--stop_text_encoder_training",
+        type=int,
+        default=1000000,
+        help=("The step at which the text_encoder is no longer trained"),
+    )
+
+
+    parser.add_argument(
+        "--image_captions_filename",
+        action="store_true",
+        help="Get captions from filename",
+    )    
+    
+      
+    
+    parser.add_argument(
+        "--Resumetr",
+        type=str,
+        default="False",        
+        help="Resume training info",
+    )    
+    
+ 
+    
+    parser.add_argument(
+        "--Session_dir",
+        type=str,
+        default="",     
+        help="Current session directory",
+    )    
+
+    parser.add_argument(
+        "--external_captions",
+        action="store_true",
+        default=False,        
+        help="Use captions stored in a txt file",
+    )    
+    
+    parser.add_argument(
+        "--captions_dir",
+        type=str,
+        default="",
+        help="The folder where captions files are stored",
+    )        
+
+    parser.add_argument(
+        "--offset_noise",
+        action="store_true",
+        default=False,
+        help="Offset Noise",
+    )
+    
+    parser.add_argument(
+        "--ofstnselvl",
+        type=float,
+        default=0.03,        
+        help="Offset Noise amount",
+    )
+    
+    parser.add_argument(
+        "--resume",
+        action="store_true",
+        default=False,
+        help="resume training",
+    )    
+    
+    parser.add_argument(
+        "--dim",
+        type=int,
+        default=64,        
+        help="LoRa dimension",
+    )    
+
+    args = parser.parse_args()
+    
+    return args
+
+
+
+class DreamBoothDataset(Dataset):
+    """
+    A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
+    It pre-processes the images and the tokenizes prompts.
+    """
+
+    def __init__(
+        self,
+        instance_data_root,
+        args,
+        tokenizers,        
+        text_encoders,
+        size=512,
+        center_crop=False,
+        instance_prompt_hidden_states=None,
+        instance_unet_added_conditions=None,    
+    ):
+        self.size = size
+        self.tokenizers=tokenizers
+        self.text_encoders=text_encoders
+        self.center_crop = center_crop
+        self.instance_prompt_hidden_states = instance_prompt_hidden_states
+        self.instance_unet_added_conditions = instance_unet_added_conditions
+        self.image_captions_filename = None
+
+        self.instance_data_root = Path(instance_data_root)
+        if not self.instance_data_root.exists():
+            raise ValueError("Instance images root doesn't exists.")
+
+        self.instance_images_path = list(Path(instance_data_root).iterdir())
+        self.num_instance_images = len(self.instance_images_path)
+        self._length = self.num_instance_images
+
+        if args.image_captions_filename:
+            self.image_captions_filename = True
+        
+        self.image_transforms = transforms.Compose(
+            [
+                transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
+                transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
+                transforms.ToTensor(),
+                transforms.Normalize([0.5], [0.5]),
+            ]
+        )
+
+    def __len__(self):
+        return self._length
+
+    def __getitem__(self, index, args=parse_args()):
+        example = {}
+        path = self.instance_images_path[index % self.num_instance_images]
+        instance_image = Image.open(path)
+        if not instance_image.mode == "RGB":
+            instance_image = instance_image.convert("RGB")
+
+        if self.image_captions_filename:
+            filename = Path(path).stem
+            
+            pt=''.join([i for i in filename if not i.isdigit()])
+            pt=pt.replace("_"," ")
+            pt=pt.replace("(","")
+            pt=pt.replace(")","")
+            pt=pt.replace("-","")
+            pt=pt.replace("conceptimagedb","")  
+            
+            if args.external_captions:
+              cptpth=os.path.join(args.captions_dir, filename+'.txt')
+              if os.path.exists(cptpth):
+                with open(cptpth, "r") as f:
+                    instance_prompt=f.read()
+              else:
+                instance_prompt=pt
+            else:
+                instance_prompt = pt
+        
+        example["instance_images"] = self.image_transforms(instance_image)
+        with torch.no_grad():
+            example["instance_prompt_ids"], example["instance_added_cond_kwargs"]= compute_embeddings(args, instance_prompt, self.text_encoders, self.tokenizers)
+
+        return example
+
+            
+class PromptDataset(Dataset):
+    "A simple dataset to prepare the prompts to generate class images on multiple GPUs."
+
+    def __init__(self, prompt, num_samples):
+        self.prompt = prompt
+        self.num_samples = num_samples
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, index):
+        example = {}
+        example["prompt"] = self.prompt
+        example["index"] = index
+        return example
+
+
+def encode_prompt(text_encoders, tokenizers, prompt):
+    prompt_embeds_list = []
+
+    for tokenizer, text_encoder in zip(tokenizers, text_encoders):
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+        untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+            removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
+            logger.warning(
+                "The following part of your input was truncated because CLIP can only handle sequences up to"
+                f" {tokenizer.model_max_length} tokens: {removed_text}"
+            )
+            
+        with torch.no_grad():
+            prompt_embeds = text_encoder(
+                text_input_ids.to(text_encoder.device),
+                output_hidden_states=True,
+            )
+
+        # We are only ALWAYS interested in the pooled output of the final text encoder
+        pooled_prompt_embeds = prompt_embeds[0]
+        prompt_embeds = prompt_embeds.hidden_states[-2]
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
+        prompt_embeds_list.append(prompt_embeds)
+
+    prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
+    pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1)
+    return prompt_embeds, pooled_prompt_embeds    
+    
+
+def collate_fn(examples):
+
+    input_ids = [example["instance_prompt_ids"] for example in examples]
+    pixel_values = [example["instance_images"] for example in examples]        
+    add_text_embeds = [example["instance_added_cond_kwargs"]["text_embeds"] for example in examples]
+    add_time_ids = [example["instance_added_cond_kwargs"]["time_ids"] for example in examples]
+
+    pixel_values = torch.stack(pixel_values)
+    pixel_values = pixel_values.to(memory_format=torch.contiguous_format).half()
+
+    input_ids = torch.cat(input_ids, dim=0)
+    add_text_embeds = torch.cat(add_text_embeds, dim=0)
+    add_time_ids = torch.cat(add_time_ids, dim=0)
+
+    batch = {
+        "input_ids": input_ids,
+        "pixel_values": pixel_values,
+        "unet_added_conditions": {"text_embeds": add_text_embeds, "time_ids": add_time_ids},
+    }
+
+    return batch
+
+
+def compute_embeddings(args, prompt, text_encoders, tokenizers):
+    original_size = (args.resolution, args.resolution)
+    target_size = (args.resolution, args.resolution)
+    crops_coords_top_left = (0, 0)
+
+    with torch.no_grad():
+        prompt_embeds, pooled_prompt_embeds = encode_prompt(text_encoders, tokenizers, prompt)
+        add_text_embeds = pooled_prompt_embeds
+
+        # Adapted from pipeline.StableDiffusionXLPipeline._get_add_time_ids
+        add_time_ids = list(original_size + crops_coords_top_left + target_size)
+        add_time_ids = torch.tensor([add_time_ids])
+
+        prompt_embeds = prompt_embeds.to('cuda')
+        add_text_embeds = add_text_embeds.to('cuda')
+        add_time_ids = add_time_ids.to('cuda', dtype=prompt_embeds.dtype)
+        unet_added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+
+    return prompt_embeds, unet_added_cond_kwargs
+    
+    
+class LatentsDataset(Dataset):
+    def __init__(self, latents_cache, text_encoder_cache, cond_cache):
+        self.latents_cache = latents_cache
+        self.text_encoder_cache = text_encoder_cache
+        self.cond_cache = cond_cache
+
+    def __len__(self):
+        return len(self.latents_cache)
+
+    def __getitem__(self, index):
+        return self.latents_cache[index], self.text_encoder_cache[index], self.cond_cache[index]
+    
+    
+
+def main():
+    args = parse_args()
+    logging_dir = Path(args.output_dir, args.logging_dir)
+          
+    accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        mixed_precision=args.mixed_precision,
+        log_with="tensorboard",
+        logging_dir=logging_dir,
+    )
+
+        
+    if args.seed is not None:
+        set_seed(args.seed)
+
+    # Handle the repository creation
+    if accelerator.is_main_process:
+        if args.output_dir is not None:
+            os.makedirs(args.output_dir, exist_ok=True)
+
+    # Load the tokenizers
+    tokenizer_one = AutoTokenizer.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="tokenizer",
+        use_fast=False,
+        use_auth_token=True,
+    )
+    tokenizer_two = AutoTokenizer.from_pretrained(
+        args.pretrained_model_name_or_path,
+        subfolder="tokenizer_2",
+        use_fast=False,
+        use_auth_token=True
+    )
+
+ 
+    
+    # import correct text encoder classes
+    text_encoder_cls_one = import_model_class_from_model_name_or_path(
+        args.pretrained_model_name_or_path, subfolder="text_encoder"
+    )
+    text_encoder_cls_two = import_model_class_from_model_name_or_path(
+        args.pretrained_model_name_or_path, subfolder="text_encoder_2"
+    )
+    
+    # Load scheduler and models
+
+    text_encoder_one = text_encoder_cls_one.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder", use_auth_token=True,
+    )
+    text_encoder_two = text_encoder_cls_two.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="text_encoder_2", use_auth_token=True
+    )
+    vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae", use_auth_token=True)
+    unet = UNet2DConditionModel.from_pretrained(
+        args.pretrained_model_name_or_path, subfolder="unet", use_auth_token=True
+    ) 
+    
+    vae.requires_grad_(False)
+    text_encoder_one.requires_grad_(False)
+    text_encoder_two.requires_grad_(False)
+    unet.requires_grad_(False)
+    text_encoder_one.eval()
+    text_encoder_two.eval()
+    vae.eval()
+    
+    model_path = os.path.join(args.Session_dir, os.path.basename(args.Session_dir) + ".safetensors")
+    network = create_network(1, args.dim, 20000, unet)
+    if args.resume:
+        network.load_weights(model_path)
+
+    def set_diffusers_xformers_flag(model, valid):
+        def fn_recursive_set_mem_eff(module: torch.nn.Module):
+            if hasattr(module, "set_use_memory_efficient_attention_xformers"):
+                module.set_use_memory_efficient_attention_xformers(valid)
+
+            for child in module.children():
+                fn_recursive_set_mem_eff(child)
+
+        fn_recursive_set_mem_eff(model)
+
+    set_diffusers_xformers_flag(unet, True)   
+   
+    network.apply_to(unet, True)
+    trainable_params = network.parameters()
+    
+    tokenizers = [tokenizer_one, tokenizer_two]
+    text_encoders = [text_encoder_one, text_encoder_two]      
+
+
+    if args.gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
+
+    if args.scale_lr:
+        args.learning_rate = (
+            args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
+        )
+
+    optimizer_class = bnb.optim.AdamW8bit
+
+    optimizer = optimizer_class(
+        trainable_params,
+        lr=args.learning_rate,
+        betas=(args.adam_beta1, args.adam_beta2),
+        weight_decay=args.adam_weight_decay,
+        eps=args.adam_epsilon,
+    )
+
+    noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler", use_auth_token=True)
+    
+    train_dataset = DreamBoothDataset(
+        instance_data_root=args.instance_data_dir,
+        tokenizers=tokenizers,
+        text_encoders=text_encoders,
+        size=args.resolution,
+        center_crop=args.center_crop,
+        args=args
+    )
+    
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=args.train_batch_size,
+        shuffle=True,
+        collate_fn=lambda examples: collate_fn(examples),
+    )
+
+
+    # Scheduler and math around the number of training steps.
+    overrode_max_train_steps = False
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if args.max_train_steps is None:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+        overrode_max_train_steps = True
+
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * args.gradient_accumulation_steps,
+        num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+    )
+    
+
+    network, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+        network, optimizer, train_dataloader, lr_scheduler)
+
+    weight_dtype = torch.float32
+    if args.mixed_precision == "fp16":
+        weight_dtype = torch.float16
+    elif args.mixed_precision == "bf16":
+        weight_dtype = torch.bfloat16
+
+    unet.to(accelerator.device, dtype=weight_dtype)
+    vae.to(accelerator.device, dtype=weight_dtype)
+    network.prepare_grad_etc(network)
+
+    
+    latents_cache = []
+    text_encoder_cache = []
+    cond_cache= []
+    for batch in train_dataloader:
+        with torch.no_grad():
+            
+            batch["input_ids"] = batch["input_ids"].to(accelerator.device, non_blocking=True)
+            batch["unet_added_conditions"] = batch["unet_added_conditions"]
+
+            batch["pixel_values"]=(vae.encode(batch["pixel_values"].to(accelerator.device, dtype=weight_dtype)).latent_dist.sample() * vae.config.scaling_factor)
+
+            latents_cache.append(batch["pixel_values"])
+            text_encoder_cache.append(batch["input_ids"])
+            cond_cache.append(batch["unet_added_conditions"])
+
+    train_dataset = LatentsDataset(latents_cache, text_encoder_cache, cond_cache)
+    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1, collate_fn=lambda x: x, shuffle=True)
+
+    del vae, tokenizers, text_encoders
+    gc.collect()
+    torch.cuda.empty_cache()   
+    
+        
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
+    if overrode_max_train_steps:
+        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
+    # Afterwards we recalculate our number of training epochs
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    # We need to initialize the trackers we use, and also store our configuration.
+    # The trackers initializes automatically on the main process.
+    if accelerator.is_main_process:
+        accelerator.init_trackers("dreambooth", config=vars(args))
+
+    def bar(prg):
+       br='|'+'█' * prg + ' ' * (25-prg)+'|'
+       return br
+    
+    # Train!
+    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
+    text_enc_context = nullcontext() if args.train_text_encoder else torch.no_grad()
+    logger.info("***** Running training *****")
+    logger.info(f"  Num examples = {len(train_dataset)}")
+    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
+    logger.info(f"  Num Epochs = {args.num_train_epochs}")
+    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
+    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
+    logger.info(f"  Total optimization steps = {args.max_train_steps}")
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
+    global_step = 0
+
+    for epoch in range(args.num_train_epochs):
+        unet.train()
+        network.train()
+        for step, batch in enumerate(train_dataloader):
+            with accelerator.accumulate(unet):
+
+                with torch.no_grad():
+                    model_input = batch[0][0]
+ 
+                # Sample noise that we'll add to the latents
+                if args.offset_noise:
+                    noise = torch.randn_like(model_input)# + args.ofstnselvl * torch.randn(model_input.shape[0], model_input.shape[1], 1, 1, device=model_input.device)
+                else:
+                    noise = torch.randn_like(model_input)
+
+                bsz = model_input.shape[0]
+
+                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device)
+                timesteps = timesteps.long()
+
+                noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)
+
+                # Predict the noise residual
+                with accelerator.autocast():
+                    model_pred = unet(noisy_model_input, timesteps, batch[0][1], added_cond_kwargs=batch[0][2]).sample
+
+                # Get the target for loss depending on the prediction type
+                target = noise
+
+                loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+                    
+                accelerator.backward(loss)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad(set_to_none=True)
+
+            # Checks if the accelerator has performed an optimization step behind the scenes
+            if accelerator.sync_gradients:
+                progress_bar.update(1)
+                global_step += 1
+                
+            fll=round((global_step*100)/args.max_train_steps)
+            fll=round(fll/4)
+            pr=bar(fll)
+            
+            logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
+            progress_bar.set_postfix(**logs)
+            progress_bar.set_description_str("Progress")
+            accelerator.log(logs, step=global_step)
+
+            if global_step >= args.max_train_steps:
+                break
+
+
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+         network = accelerator.unwrap_model(network)
+    accelerator.end_training()
+    network.save_weights(model_path, torch.float16, None)
+      
+    accelerator.end_training()
+
+if __name__ == "__main__":
+    main()