Custom Width and Height

1 files changed, 3 insertions, 4 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 = []