Add img, num_cutouts, and center_bias to CLI

lucidrains · Mar 22, 2021 · dd8a55b · dd8a55b
1 parent 5146f10
commit dd8a55b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/big_sleep/big_sleep.py b/big_sleep/big_sleep.py
@@ -316,8 +316,8 @@ def __init__(
         save_best = False,
         experimental_resample = False,
         ema_decay = 0.99,
-        num_cutouts=128,
-        center_bias=False,
+        num_cutouts = 128,
+        center_bias = False,
     ):
         super().__init__()
 

diff --git a/big_sleep/cli.py b/big_sleep/cli.py
@@ -4,8 +4,10 @@
 from pathlib import Path
 from .version import __version__;
 
+
 def train(
-    text,
+    text=None,
+    img=None,
     text_min="",
     lr = .07,
     image_size = 512,
@@ -25,15 +27,18 @@ def train(
     class_temperature = 2.,
     save_best = False,
     experimental_resample = False,
-    ema_decay = 0.5
+    ema_decay = 0.5,
+    num_cutouts = 128,
+    center_bias = False,
 ):
     print(f'Starting up... v{__version__}')
 
     if random:
         seed = rnd.randint(0, 1e6)
 
     imagine = Imagine(
-        text,
+        text=text,
+        img=img,
         text_min=text_min,
         lr = lr,
         image_size = image_size,
@@ -51,7 +56,9 @@ def train(
         save_date_time = save_date_time,
         save_best = save_best,
         experimental_resample = experimental_resample,
-        ema_decay = ema_decay
+        ema_decay = ema_decay,
+        num_cutouts = num_cutouts,
+        center_bias = center_bias,
     )
 
     if not overwrite and imagine.filename.exists():