From 7d85ac15027f8b261b063ddce0e4b4ca8870286c Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Mon, 23 Jan 2023 13:24:22 -0500
Subject: [PATCH] add preview

---
 .gitignore                                    |  1 +
 cli/modules/interrogate.py                    |  2 ++
 cli/modules/preview.py                        |  5 ++--
 cli/train.py                                  | 27 +++++++++----------
 .../sd-extension-steps-animation              |  2 +-
 5 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/.gitignore b/.gitignore
index 806715328..e633f5daa 100644
--- a/.gitignore
+++ b/.gitignore
@@ -20,6 +20,7 @@ __pycache__
 /outputs/**/*
 /embeddings/**/*
 /models/**/*
+/interrogate/**/*
 /train/log/**/*
 /tmp
 /log
diff --git a/cli/modules/interrogate.py b/cli/modules/interrogate.py
index 0a7c00c15..e9790de42 100755
--- a/cli/modules/interrogate.py
+++ b/cli/modules/interrogate.py
@@ -74,6 +74,8 @@ async def interrogate(f):
     if 'caption' in res:
         for term in res.caption.split(', '):
             term = term.replace('(', '').replace(')', '').split(':')
+            if len(term) < 2:
+                continue
             keywords[term[0]] = term[1]
         keywords = dict(sorted(keywords.items(), key=lambda x:x[1], reverse=True))
         for word in keywords.items():
diff --git a/cli/modules/preview.py b/cli/modules/preview.py
index ba29b097a..4478ac907 100755
--- a/cli/modules/preview.py
+++ b/cli/modules/preview.py
@@ -5,6 +5,7 @@ import json
 import base64
 from pathlib import Path
 from PIL import Image
+from inspect import getsourcefile
 
 from util import Map, log
 from sdapi import getsync, postsync
@@ -38,13 +39,13 @@ def encode(f):
 def create_preview(name: str, suffix: str):
     options = getsync('/sdapi/v1/options')
     cmdflags = getsync('/sdapi/v1/cmd-flags')    
-    print(cmdflags.embeddings_dir)
     img2img_options['prompt'] = template.format(name = name, suffix = suffix)
     log.info({ 'preview prompt': img2img_options['prompt'] })
     log.debug({ 'preview options': img2img_options })
+    mask_path = os.path.join(os.path.dirname(getsourcefile(lambda:0)), mask)
     if len(img2img_options['init_images']) == 0:
         for i in range(img2img_options.batch_size):
-            img2img_options['init_images'].append(encode(mask))
+            img2img_options['init_images'].append(encode(mask_path))
     data = postsync('/sdapi/v1/img2img', img2img_options)
     if 'error' in data:
         log.error({ 'preview': data['error'], 'reason': data['reason'] })
diff --git a/cli/train.py b/cli/train.py
index 7a665344b..a0422b559 100755
--- a/cli/train.py
+++ b/cli/train.py
@@ -69,12 +69,12 @@ async def captions(docs: list):
     log.info({ 'top captions': d })
 
 
-async def cleanup(params):
-    if params.nocleanup:
-        return
+async def preprocess_cleanup(params):
     log.info({ 'preprocess cleanup': params.dst })
     for f in Path(params.dst).glob('*.png'):
         f.unlink()
+    for f in Path(params.dst).glob('*.jpg'):
+        f.unlink()
     for f in Path(params.dst).glob('*.txt'):
         f.unlink()
 
@@ -82,16 +82,6 @@ async def cleanup(params):
 async def preprocess_builtin(params):
     global images # pylint: disable=global-statement
     log.debug({ 'preprocess start' })
-    if os.path.isdir(params.dst):
-        if params.overwrite:
-            log.info({ 'preprocess deleting existing images': params.dst })
-            for f in Path(params.dst).glob('*.png'):
-                f.unlink()
-            for f in Path(params.dst).glob('*.txt'):
-                f.unlink()
-        else:
-            log.error({ 'preprocess output folder already exists': params.dst })
-            return 0
     files = [os.path.join(params.src, f) for f in os.listdir(params.src) if os.path.isfile(os.path.join(params.src, f))]
     candidates = [f for f in files if filetype.is_image(f)]
     not_images = [f for f in files if (not filetype.is_image(f) and not f.endswith('.txt'))]
@@ -146,9 +136,15 @@ async def preprocess(params):
             else:
                 log.error({ 'preprocess video extract': 'no images' })
     elif os.path.isdir(params.src):
+        if params.overwrite:
+            preprocess_cleanup(params)
+        else:
+            log.error({ 'preprocess output folder already exists': params.dst })
+            return 0
+
         if params.preprocess == 'builtin':
             res = await preprocess_builtin(params)
-        if params.preprocess == 'custom':
+        elif params.preprocess == 'custom':
             t0 = time.perf_counter()
             args.preprocess.process_src = params.src
             args.preprocess.process_dst = params.dst
@@ -512,7 +508,8 @@ async def main():
     except Exception as e:
         log.error({ 'exception': e })
     finally:
-        await cleanup(params)
+        if not params.nocleanup:
+            await preprocess_cleanup(params)
         await close()
     return
 
diff --git a/extensions-builtin/sd-extension-steps-animation b/extensions-builtin/sd-extension-steps-animation
index ba6300b7b..b0079fe64 160000
--- a/extensions-builtin/sd-extension-steps-animation
+++ b/extensions-builtin/sd-extension-steps-animation
@@ -1 +1 @@
-Subproject commit ba6300b7b9f4c1640849f391c7b1eed837ae6e3a
+Subproject commit b0079fe64398d3ac96866f145df68bd7f164b14c