From dd8a55bcc773e8cdd946800364440e2d5ed6965c Mon Sep 17 00:00:00 2001
From: Anton <antonwiehe@gmail.com>
Date: Tue, 23 Mar 2021 00:03:04 +0100
Subject: [PATCH] Add img, num_cutouts, and center_bias to CLI

---
 big_sleep/big_sleep.py |  4 ++--
 big_sleep/cli.py       | 15 +++++++++++----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/big_sleep/big_sleep.py b/big_sleep/big_sleep.py
index a03399e..3dad9bc 100644
--- 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
index 177d08d..b4dc920 100644
--- 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,7 +27,9 @@ 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__}')
 
@@ -33,7 +37,8 @@ def train(
         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():