Custom Width and Height

4 files changed, 26 insertions, 23 deletions

diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py
index acc4ce59..bcf772d2 100644
--- a/modules/textual_inversion/dataset.py
+++ b/modules/textual_inversion/dataset.py
@@ -15,13 +15,12 @@ re_tag = re.compile(r"[a-zA-Z][_\w\d()]+")
 

 

 class PersonalizedBase(Dataset):

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

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

 

         self.placeholder_token = placeholder_token

 

-        self.size = size

-        self.width = size

-        self.height = size

+        self.width = width

+        self.height = height

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

 

         self.dataset = []

diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index b3de6fd7..d7efdef2 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -7,8 +7,9 @@ import tqdm
 from modules import shared, images

 

 

-def preprocess(process_src, process_dst, process_size, process_flip, process_split, process_caption):

-    size = process_size

+def preprocess(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption):

+    width = process_width

+    height = process_height

     src = os.path.abspath(process_src)

     dst = os.path.abspath(process_dst)

 

@@ -55,23 +56,23 @@ def preprocess(process_src, process_dst, process_size, process_flip, process_spl
         is_wide = ratio < 1 / 1.35

 

         if process_split and is_tall:

-            img = img.resize((size, size * img.height // img.width))

+            img = img.resize((width, height * img.height // img.width))

 

-            top = img.crop((0, 0, size, size))

+            top = img.crop((0, 0, width, height))

             save_pic(top, index)

 

-            bot = img.crop((0, img.height - size, size, img.height))

+            bot = img.crop((0, img.height - height, width, img.height))

             save_pic(bot, index)

         elif process_split and is_wide:

-            img = img.resize((size * img.width // img.height, size))

+            img = img.resize((width * img.width // img.height, height))

 

-            left = img.crop((0, 0, size, size))

+            left = img.crop((0, 0, width, height))

             save_pic(left, index)

 

-            right = img.crop((img.width - size, 0, img.width, size))

+            right = img.crop((img.width - width, 0, img.width, height))

             save_pic(right, index)

         else:

-            img = images.resize_image(1, img, size, size)

+            img = images.resize_image(1, img, width, height)

             save_pic(img, index)

 

         shared.state.nextjob()

diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index 769682ea..5965c5a0 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -6,7 +6,6 @@ import torch
 import tqdm

 import html

 import datetime

-import math

 

 

 from modules import shared, devices, sd_hijack, processing, sd_models

@@ -157,7 +156,7 @@ def create_embedding(name, num_vectors_per_token, init_text='*'):
     return fn

 

 

-def train_embedding(embedding_name, learn_rate, data_root, log_directory, training_size, steps, num_repeats, create_image_every, save_embedding_every, template_file):

+def train_embedding(embedding_name, learn_rate, data_root, log_directory, training_width, training_height, steps, num_repeats, create_image_every, save_embedding_every, template_file):

     assert embedding_name, 'embedding not selected'

 

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

@@ -183,7 +182,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, size=training_size, repeats=num_repeats, 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=num_repeats, placeholder_token=embedding_name, model=shared.sd_model, device=devices.device, template_file=template_file)

 

     hijack = sd_hijack.model_hijack

 

@@ -227,7 +226,7 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
             loss.backward()

             optimizer.step()

 

-        epoch_num = math.floor(embedding.step / epoch_len)

+        epoch_num = embedding.step // epoch_len

         epoch_step = embedding.step - (epoch_num * epoch_len) + 1

 

         pbar.set_description(f"[Epoch {epoch_num}: {epoch_step}/{epoch_len}]loss: {losses.mean():.7f}")

@@ -243,8 +242,8 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
                 sd_model=shared.sd_model,

                 prompt=text,

                 steps=20,

-                height=training_size,

-                width=training_size,

+				height=training_height,

+				width=training_width,

                 do_not_save_grid=True,

                 do_not_save_samples=True,

             )

diff --git a/modules/ui.py b/modules/ui.py
index f821fd8d..8c06ad7c 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -1029,7 +1029,8 @@ def create_ui(wrap_gradio_gpu_call):
 

                     process_src = gr.Textbox(label='Source directory')

                     process_dst = gr.Textbox(label='Destination directory')

-                    process_size = gr.Slider(minimum=64, maximum=2048, step=64, label="Size (width and height)", value=512)

+                    process_width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)

+                    process_height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)

 

                     with gr.Row():

                         process_flip = gr.Checkbox(label='Create flipped copies')

@@ -1050,7 +1051,8 @@ def create_ui(wrap_gradio_gpu_call):
                     dataset_directory = gr.Textbox(label='Dataset directory', placeholder="Path to directory with input images")

                     log_directory = gr.Textbox(label='Log directory', placeholder="Path to directory where to write outputs", value="textual_inversion")

                     template_file = gr.Textbox(label='Prompt template file', value=os.path.join(script_path, "textual_inversion_templates", "style_filewords.txt"))

-                    training_size = gr.Slider(minimum=64, maximum=2048, step=64, label="Size (width and height)", value=512)

+                    training_width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512)

+                    training_height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512)

                     steps = gr.Number(label='Max steps', value=100000, precision=0)

                     num_repeats = gr.Number(label='Number of repeats for a single input image per epoch', value=100, precision=0)

                     create_image_every = gr.Number(label='Save an image to log directory every N steps, 0 to disable', value=500, precision=0)

@@ -1095,7 +1097,8 @@ def create_ui(wrap_gradio_gpu_call):
             inputs=[

                 process_src,

                 process_dst,

-                process_size,

+                process_width,

+                process_height,

                 process_flip,

                 process_split,

                 process_caption,

@@ -1114,7 +1117,8 @@ def create_ui(wrap_gradio_gpu_call):
                 learn_rate,

                 dataset_directory,

                 log_directory,

-                training_size,

+                training_width,

+                training_height,

                 steps,

                 num_repeats,

                 create_image_every,