add preview

.gitignore

@@ -20,6 +20,7 @@ __pycache__
 /outputs/**/*
 /embeddings/**/*
 /models/**/*
+/interrogate/**/*
 /train/log/**/*
 /tmp
 /log

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():

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'] })

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