create embedding preview

.gitignore

@@ -18,6 +18,7 @@ __pycache__
 /repositories/**/*
 /extensions/**/*
 /outputs/**/*
+/embeddings/**/*
 /models/**/*
 /train/log/**/*
 /tmp
@@ -27,6 +28,7 @@ __pycache__
 # unexcluded so folders get created
 !/repositories/.placeholder
 !/extensions/.placeholder
+!/embeddings/.placeholder
 !/outputs/.placeholder
 !/models/.placeholder
 !/localizations/.placeholder

cli/modules/preview.py

@@ -1,28 +1,31 @@
 #!/bin/env python
+import os
 import io
 import json
 import base64
-import logging
+from pathlib import Path
 from PIL import Image
 
 from util import Map, log
 from sdapi import getsync, postsync
+from grid import grid
 
+# masks = ['preview-face.jpg', 'preview-body.jpg']
 mask = 'preview-body.jpg'
 template = 'photo of "{name}", {suffix}, high detailed, skin texture, facing camera, 135mm, shot on dslr, 4k, modelshoot style'
-img2img_options = {
+img2img_options = Map({
     'prompt': None,
     'negative_prompt': '',
     'init_images': [],
     'sampler_name': 'DPM2 Karras',
-    'batch_size': 1,
+    'batch_size': 4,
     'n_iter': 1,
     'steps': 30,
     'cfg_scale': 6,
     'width': 512,
     'height': 512,
     'restore_faces': False
-}
+})
 
 def encode(f):
     img = Image.open(f)
@@ -32,28 +35,34 @@ def encode(f):
         encoded = base64.b64encode(values).decode()
         return encoded
 
-def img2img(name: str, suffix: str):
+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.info({ 'preview options': img2img_options })
-    img2img_options['init_images'].append(encode(mask))
+    log.debug({ 'preview options': img2img_options })
+    if len(img2img_options['init_images']) == 0:
+        for i in range(img2img_options.batch_size):
+            img2img_options['init_images'].append(encode(mask))
     data = postsync('/sdapi/v1/img2img', img2img_options)
     if 'error' in data:
         log.error({ 'preview': data['error'], 'reason': data['reason'] })
         return
     info = Map(json.loads(data['info']))
-    log.debug({ 'preview info': info })
     if not 'images' in data:
         log.error({ 'preview': 'no images' })
         return
-    obj = data.copy()
-    del obj['images']
-    log.info({ 'preview': { 'model': options["sd_model_checkpoint"], 'seed': obj['info']['seed'] } })
-    image = Image.open(io.BytesIO(base64.b64decode(data['images'][0].split(",",1)[0])))
-    image.save('test.jpg')
+    fn = os.path.join(cmdflags.embeddings_dir, name + '.preview.png')
+    log.info({ 'preview': { 'name': fn, 'model': options.sd_model_checkpoint, 'seed': info.seed } })
+    images = []
+    for b64 in data['images']:
+        images.append(Image.open(io.BytesIO(base64.b64decode(b64.split(",",1)[0]))))
+    image = grid(images, None, square=True)
+    image.save(fn)
 
 if __name__ == "__main__":
-    # log.setLevel(logging.DEBUG)
     log.info({ 'preview': 'start' })
-    img2img('hanna', 'person, woman, girl, model')
+    cmdflags = getsync('/sdapi/v1/cmd-flags')
+    for f in Path(cmdflags.embeddings_dir).glob('*.pt'):
+        create_preview(f.stem, 'person')

cli/train.py

@@ -31,6 +31,7 @@ from modules.lossrate import gen_loss_rate_str
 from modules.sdapi import close, get, interrupt, post, progress, session
 from modules.process import process_images
 from modules.grid import grid
+from modules.preview import create_preview
 
 
 images = []
@@ -344,6 +345,8 @@ async def pipeline(params):
 
     await plotloss(params)
 
+    create_preview(params.name, params.init)
+
     log.debug({ 'pipeline end' })
     return