Merge branch 'master' into master

2 files changed, 28 insertions, 22 deletions

diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py
index 12e2f43b..23bb4b6a 100644
--- a/modules/textual_inversion/dataset.py
+++ b/modules/textual_inversion/dataset.py
@@ -24,11 +24,12 @@ class DatasetEntry:
 

 

 class PersonalizedBase(Dataset):

-    def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, device=None, template_file=None, include_cond=False):

-        re_word = re.compile(shared.opts.dataset_filename_word_regex) if len(shared.opts.dataset_filename_word_regex)>0 else None

+    def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, device=None, template_file=None, include_cond=False, batch_size=1):

+        re_word = re.compile(shared.opts.dataset_filename_word_regex) if len(shared.opts.dataset_filename_word_regex) > 0 else None

 

         self.placeholder_token = placeholder_token

 

+        self.batch_size = batch_size

         self.width = width

         self.height = height

         self.flip = transforms.RandomHorizontalFlip(p=flip_p)

@@ -78,14 +79,14 @@ class PersonalizedBase(Dataset):
 

             if include_cond:

                 entry.cond_text = self.create_text(filename_text)

-                entry.cond = cond_model([entry.cond_text]).to(devices.cpu)

+                entry.cond = cond_model([entry.cond_text]).to(devices.cpu).squeeze(0)

 

             self.dataset.append(entry)

-        

+

         assert len(self.dataset) > 1, "No images have been found in the dataset."

-        self.length = len(self.dataset) * repeats

+        self.length = len(self.dataset) * repeats // batch_size

 

-        self.initial_indexes = np.arange(self.length) % len(self.dataset)

+        self.initial_indexes = np.arange(len(self.dataset))

         self.indexes = None

         self.shuffle()

 

@@ -102,13 +103,19 @@ class PersonalizedBase(Dataset):
         return self.length

 

     def __getitem__(self, i):

-        if i % len(self.dataset) == 0:

-            self.shuffle()

+        res = []

+

+        for j in range(self.batch_size):

+            position = i * self.batch_size + j

+            if position % len(self.indexes) == 0:

+                self.shuffle()

+

+            index = self.indexes[position % len(self.indexes)]

+            entry = self.dataset[index]

 

-        index = self.indexes[i % len(self.indexes)]

-        entry = self.dataset[index]

+            if entry.cond is None:

+                entry.cond_text = self.create_text(entry.filename_text)

 

-        if entry.cond is None:

-            entry.cond_text = self.create_text(entry.filename_text)

+            res.append(entry)

 

-        return entry

+        return res

diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index da0d77a0..e754747e 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -199,7 +199,7 @@ def write_loss(log_directory, filename, step, epoch_len, values):
         })

 

 

-def train_embedding(embedding_name, learn_rate, data_root, log_directory, training_width, training_height, steps, create_image_every, save_embedding_every, template_file, save_image_with_stored_embedding, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height):

+def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_directory, training_width, training_height, steps, create_image_every, save_embedding_every, template_file, save_image_with_stored_embedding, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height):

     assert embedding_name, 'embedding not selected'

 

     shared.state.textinfo = "Initializing textual inversion training..."

@@ -231,7 +231,7 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
 

     shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."

     with torch.autocast("cuda"):

-        ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=embedding_name, model=shared.sd_model, device=devices.device, template_file=template_file)

+        ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=embedding_name, model=shared.sd_model, device=devices.device, template_file=template_file, batch_size=batch_size)

 

     hijack = sd_hijack.model_hijack

 

@@ -251,7 +251,7 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
     optimizer = torch.optim.AdamW([embedding.vec], lr=scheduler.learn_rate)

 

     pbar = tqdm.tqdm(enumerate(ds), total=steps-ititial_step)

-    for i, entry in pbar:

+    for i, entries in pbar:

         embedding.step = i + ititial_step

 

         scheduler.apply(optimizer, embedding.step)

@@ -262,10 +262,9 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
             break

 

         with torch.autocast("cuda"):

-            c = cond_model([entry.cond_text])

-

-            x = entry.latent.to(devices.device)

-            loss = shared.sd_model(x.unsqueeze(0), c)[0]

+            c = cond_model([entry.cond_text for entry in entries])

+            x = torch.stack([entry.latent for entry in entries]).to(devices.device)

+            loss = shared.sd_model(x, c)[0]

             del x

 

             losses[embedding.step % losses.shape[0]] = loss.item()

@@ -307,7 +306,7 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
                 p.width = preview_width

                 p.height = preview_height

             else:

-                p.prompt = entry.cond_text

+                p.prompt = entries[0].cond_text

                 p.steps = 20

                 p.width = training_width

                 p.height = training_height

@@ -348,7 +347,7 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
 <p>

 Loss: {losses.mean():.7f}<br/>

 Step: {embedding.step}<br/>

-Last prompt: {html.escape(entry.cond_text)}<br/>

+Last prompt: {html.escape(entries[0].cond_text)}<br/>

 Last saved embedding: {html.escape(last_saved_file)}<br/>

 Last saved image: {html.escape(last_saved_image)}<br/>

 </p>