add cli

cli/README.md

@@ -0,0 +1,66 @@
+# Scripts using Stable-Diffusion/Automatic API
+
+*Note*: Start **SD/Automatic** using `python launch.py --api`
+## Generate
+
+Text-to-image with all of the possible parameters  
+Supports upsampling, face restoration and grid creation  
+> python generate.py --help
+
+By default uses parameters from  `generate.json`
+
+Parameters that are not specified will be randomized to some extent:
+- Prompt will be dynamically created from template of random samples: `random.json`
+- Sampler/Scheduler will be randomly picked from available ones
+- CFG Scale set to 5-10
+
+## Train
+
+End-to-end embedding training
+> python train.py --help
+
+Combined pipeline:
+1. Creates embedding  
+2. Extracts images if input is movie  
+3. Preprocesses images  
+4. Runs training  
+
+## Interrogate
+
+Runs CLiP and Booru image interrogation on any provided parameters  
+*(image, list of images, wildcards, folder, etc.)*
+> python interrogate.py
+
+## Promptist
+
+Attempts to beautify the provided prompt  
+> python promptist.py
+
+## Ideas
+
+Generate complex prompt ideas
+> python ideas.py --help
+
+## SDAPI
+
+Utility module that handles async communication to Automatic API endpoints  
+Can be used to manually execute specific commands:
+> python sdapi.py progress  
+> python sdapi.py interrupt
+
+## FFMPEG
+
+Utility module that handles video files  
+Can be used to manually execute specific commands:
+> ffmpeg extract --help  
+> python ffmpeg.py extract --input ~/downloads/vlado.mp4 --output ./vlado --fps 2 --skipstart 3 --skipend 1
+
+## Grid
+
+Utility module to create image grids
+> python grid.py --help
+
+## Bench
+
+Benchmark your Automatic
+> python bench.py

cli/bench.py

@@ -0,0 +1,112 @@
+#!/bin/env python
+"""
+sd api txt2img benchmark
+"""
+import time
+import json
+import asyncio
+import base64
+import io
+import sdapi
+from util import Map, log
+from PIL import Image
+
+options = Map({
+    'restore_faces': False,
+    'prompt': 'photo of two dice on a table',
+    'negative_prompt': 'foggy, blurry',
+    'steps': 20,
+    'batch_size': 1,
+    'n_iter': 1,
+    'seed': -1,
+    'sampler_name': 'Euler a',
+    'cfg_scale': 0,
+    'width': 512,
+    'height': 512
+})
+
+# batch = [1, 1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128]
+batch = [1, 1, 2, 4, 8, 12, 16]
+oom = 0
+
+async def txt2img():
+    t0 = time.perf_counter()
+    data = {}
+    try:
+        data = await sdapi.post('/sdapi/v1/txt2img', options)
+    except:
+        return -1
+    if 'error' in data:
+        return -1
+    info = Map(json.loads(data['info']))
+    log.debug({ 'info': info })
+    for i in range(len(data['images'])):
+        data['images'][i] = Image.open(io.BytesIO(base64.b64decode(data['images'][i].split(',',1)[0])))
+        log.debug({ 'image': data['images'][i].size })
+    t1 = time.perf_counter()
+    return t1 - t0
+
+def memstats():
+    mem = sdapi.getsync('/sdapi/v1/memory')
+    cpu = mem.get('ram', 'unavailable')
+    gpu = mem.get('cuda', 'unavailable')
+    if 'active' in gpu:
+        gpu['session'] = gpu.pop('active')
+    if 'reserved' in gpu:
+        gpu.pop('allocated')
+        gpu.pop('reserved')
+        gpu.pop('inactive')
+    if 'events' in gpu:
+        global oom # pylint: disable=global-statement
+        oom = gpu['events']['oom']
+        gpu.pop('events')
+    return cpu, gpu
+
+def gb(val: float):
+    return round(val / 1024 / 1024 / 1024, 2)
+
+async def main():
+    log.info({ 'benchmark': { 'batch-sizes': batch } })
+    sdapi.quiet = True
+    await sdapi.session()
+    await sdapi.interrupt()
+    opts = await sdapi.get('/sdapi/v1/options')
+    opts = Map(opts)
+    log.info({ 'options': {
+        'resolution': [options.width, options.height],
+        'model': opts.sd_model_checkpoint,
+        'vae': opts.sd_vae,
+        'hypernetwork': opts.sd_hypernetwork,
+        'sampler': options.sampler_name,
+        'clip-stop': opts.CLIP_stop_at_last_layers,
+        'preview': opts.show_progress_every_n_steps
+    } })
+    cpu, gpu = memstats()
+    log.info({ 'system': { 'cpu': cpu, 'gpu': gpu }})
+    for i in range(len(batch)):
+        if oom > 0:
+            continue
+        options['batch_size'] = batch[i]
+        ts = await txt2img()
+        if ts > 0:
+            await asyncio.sleep(0)
+            cpu, gpu = memstats()
+            if i == 0:
+                log.info({ 'warmup': round(ts, 2) })
+            else:
+                peak = gpu['session']['peak'] if 'session' in gpu else 0
+                log.info({ 'batch': batch[i], 'its': round(options.steps / (ts / batch[i]), 2), 'img': round(ts / batch[i], 2), 'wall': round(ts, 2), 'peak': gb(peak), 'oom': oom > 0 })
+        else:
+            await asyncio.sleep(10)
+            cpu, gpu = memstats()
+            log.info({ 'batch': batch[i], 'result': 'error', 'gpu': gpu, 'oom': oom > 0 })
+    if oom > 0:
+        log.info({ 'benchmark': 'ended with oom so you should probably restart your automatic server now' })
+    await sdapi.close()
+
+if __name__ == '__main__':
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        log.warning({ 'interrupted': 'keyboard request' })
+        sdapi.interruptsync()

cli/detect.py

@@ -0,0 +1,42 @@
+#!/bin/env python
+"""
+Detect model type
+
+Works for v1 and v2-base (EPS models), both standard inference and inpainting
+
+But looking at model dumps between EPS and V type models, its only about parametrization, there are no differences in actual model (its just weighted differently without any structural difference)
+So i don't see easy way to auto-detect if model should be run in `EPS` or `V` mode
+Only difference are some calculations in `ldm/models/diffusion/ddpm.py` and by then we already need to know which code-path to trigger
+(maaaybe there could be a way by looking at some cherry-picked base tensors min/max range, but I dont see that as reliable)
+"""
+
+import os
+import sys
+import torch
+
+def signature(model):
+    if model is None:
+        return None
+    try:
+        size = model['state_dict']['model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k.weight'].shape[1]
+        unet = model['state_dict']['model.diffusion_model.input_blocks.0.0.weight'].shape[1]
+    except:
+        return 'unknown'
+    guess = 'v1' if size == 768 else 'v2' # 768 for v1 and 1024 for v2
+    guess += '-inference' if unet == 4 else '-inpainting' # inference models have shorter inputs, 4 for inference 9 for inpainting
+    return guess
+
+def load(file: str):
+    try:
+        model = torch.load(file, map_location='cpu')
+        return model
+    except Exception as err:
+        print(f"Error loading {f}: {err}")
+
+if __name__ == "__main__":
+    sys.argv.pop(0)
+    for f in sys.argv:
+        if os.path.isfile(f):
+            print(f"Model {f} is of type {signature(load(f))}")
+        else:
+            print(f"{f} is not a file")

cli/ffmpeg.py

@@ -0,0 +1,71 @@
+#!/bin/env python
+"""
+use ffmpeg for animation processing
+"""
+import os
+import json
+import subprocess
+import pathlib
+import argparse
+import filetype
+from util import log, Map
+
+
+def probe(src: str):
+    cmd = f"ffprobe -hide_banner -loglevel 0 -print_format json -show_format -show_streams {src}"
+    result = subprocess.run(cmd, shell = True, capture_output = True, text = True, check = True)
+    data = json.loads(result.stdout)
+    i = [x for x in data['streams'] if x["codec_type"] == "video"][0]
+    video = Map({
+        'codec': i['codec_name']+'/'+i['codec_tag_string'],
+        'resolution': [int(i['width']), int(i['height'])],
+        'duration': float(i['duration']),
+        'frames': int(i['nb_frames']),
+        'bitrate': round(float(i['bit_rate']) / 1024),
+    })
+    return video
+
+
+def extract(src: str, dst: str, rate: float = 0.015, fps: float = 0, start = 0, end = 0):
+    images = []
+    if not os.path.isfile(src) or not filetype.is_video(src):
+        log.error({ 'extract': 'input is not movie file' })
+        return
+    dst = dst if dst.endswith('/') else dst + '/'
+
+    video = probe(src)
+    log.info({ 'extract': { 'source': src, **video } })
+
+    ssstart = f' -ss {start}' if start > 0 else ''
+    ssend = f' -to {video.duration - end}' if start > 0 else ''
+    filename = pathlib.Path(src).stem
+    if rate > 0:
+        cmd = f"ffmpeg -hide_banner -y -loglevel info {ssstart} {ssend} -i {src} -filter:v \"select='gt(scene,{rate})',metadata=print\" -vsync vfr -frame_pts 1 {dst}{filename}-%05d.jpg"
+    elif fps > 0:
+        cmd = f"ffmpeg -hide_banner -y -loglevel info {ssstart} {ssend} -i {src} -r {fps} -vsync vfr -frame_pts 1 {dst}{filename}-%05d.jpg"
+    else:
+        log.error({ 'extract': 'requires either rate or fps' })
+        return 0
+    log.debug({ 'extract': cmd })
+    pathlib.Path(dst).mkdir(parents = True, exist_ok = True)
+    result = subprocess.run(cmd, shell = True, capture_output = True, text = True, check = True)
+    for line in result.stderr.split('\n'):
+        if 'pts_time' in line:
+            log.debug({ 'extract': { 'keyframe': line.strip().split(' ')[-1].split(':')[-1] } })
+    images = next(os.walk(dst))[2]
+    log.info({ 'extract': { 'destination': dst, 'keyframes': len(images), 'rate': rate, 'fps': fps } })
+    return len(images)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="ffmpeg pipeline")
+    parser.add_argument("command", choices = ["extract", "animate"])
+    parser.add_argument("--input", type = str, required = True, help="input")
+    parser.add_argument("--output", type = str, required = True, help="output")
+    parser.add_argument("--rate", type = float, default = 0, required = False, help="extraction change rate threshold")
+    parser.add_argument("--fps", type = float, default = 0, required = False, help="extraction frames per second")
+    parser.add_argument("--skipstart", type = float, default = 1, required = False, help="skip time from start of video")
+    parser.add_argument("--skipend", type = float, default = 1, required = False, help="skip time to end of video")
+    params = parser.parse_args()
+    if params.command == "extract":
+        extract(src = params.input, dst = params.output, rate = params.rate, fps = params.fps, start = params.skipstart, end = params.skipend)

cli/generate.json

@@ -0,0 +1,38 @@
+{
+  "paths":
+  {
+      "root": "/mnt/c/Users/mandi/OneDrive/Generative/Generate",
+      "generate": "image",
+      "upscale": "upscale",
+      "grid": "grid"
+  },
+  "generate":
+  {
+      "restore_faces": false,
+      "prompt": "dynamic",
+      "negative_prompt": "foggy, blurry, blurred, duplicate, ugly, mutilated, mutation, mutated, out of frame, bad anatomy, disfigured, deformed, censored, low res, watermark, text, poorly drawn face, signature",
+      "steps": 20,
+      "batch_size": 1,
+      "n_iter": 1,
+      "seed": -1,
+      "sampler_name": "random",
+      "cfg_scale": 0,
+      "width": 512,
+      "height": 512
+  },
+  "upscale":
+  {
+      "upscaler_1": "SwinIR_4x",
+      "upscaler_2": "None",
+      "upscale_first": false,
+      "upscaling_resize": 0,
+      "gfpgan_visibility": 0,
+      "codeformer_visibility": 0,
+      "codeformer_weight": 0.5
+  },
+  "options":
+  {
+      "sd_model_checkpoint": "mix-berrymix",
+      "sd_vae": "vae-ft-mse-840000-ema-pruned"
+  }
+}

cli/generate.py

@@ -0,0 +1,355 @@
+#!/bin/env python
+# pylint: disable=no-member
+"""
+generate batches of images from prompts and upscale them
+
+params: run with `--help`
+
+default workflow runs infinite loop and prints stats when interrupted:
+1. choose random scheduler lookup all available and pick one
+2. generate dynamic prompt based on styles, embeddings, places, artists, suffixes
+3. beautify prompt
+4. generate 3x3 images
+5. create image grid
+6. upscale images with face restoration
+"""
+
+import argparse
+import asyncio
+import base64
+import io
+import json
+import logging
+import math
+import os
+import pathlib
+import secrets
+import time
+import sys
+from random import randrange
+
+from PIL import Image
+from PIL.ExifTags import TAGS
+from PIL.TiffImagePlugin import ImageFileDirectory_v2
+from sdapi import close, get, interrupt, post, session
+from util import Map, log
+
+sd = {}
+random = {}
+stats = Map({ "images": 0, "wall": 0, "generate": 0, "upscale": 0 })
+avg = {}
+
+
+def grid(data):
+    if len(data.image) > 1:
+        w, h = data.image[0].size
+        rows = round(math.sqrt(len(data.image)))
+        cols = math.ceil(len(data.image) / rows)
+        image = Image.new('RGB', size = (cols * w, rows * h), color = 'black')
+        for i, img in enumerate(data.image):
+            image.paste(img, box=(i % cols * w, i // cols * h))
+        short = data.info.prompt[:min(len(data.info.prompt), 96)] # limit prompt part of filename to 96 chars
+        name = "{seed:0>9}-{short}.jpg".format(short = short, seed = data.info.all_seeds[0]) # pylint: disable=consider-using-f-string
+        f = os.path.join(sd.paths.root, sd.paths.grid, name)
+        log.info({ "grid": { "name": f, "size": image.size, "images": len(data.image) } })
+        image.save(f, "JPEG", exif = exif(data.info, None, "grid"), optimize = True, quality = 70)
+        return image
+
+
+def exif(info, i = None, op = "generate"):
+    seed = [info.all_seeds[i]] if len(info.all_seeds) > 0 and i is not None else info.all_seeds # always returns list
+    seed = ', '.join([str(x) for x in seed]) # int list to str list to single str
+    template = "{prompt} | negative {negative_prompt} | seed {s} | steps {steps} | cfgscale {cfg_scale} | sampler {sampler_name} | batch {batch_size} | timestamp {job_timestamp} | model {model} | vae {vae}".format(s = seed, model = sd.options["sd_model_checkpoint"], vae = sd.options["sd_vae"], **info) # pylint: disable=consider-using-f-string
+    if op == "upscale":
+        template += ' | faces gfpgan' if sd.upscale.gfpgan_visibility > 0 else ''
+        template += ' | faces codeformer' if sd.upscale.codeformer_visibility > 0 else ''
+        template += ' | upscale {resize}x {upscaler}'.format(resize = sd.upscale.upscaling_resize, upscaler = sd.upscale.upscaler_1) if sd.upscale.upscaler_1 != "None" else '' # pylint: disable=consider-using-f-string
+        template += ' | upscale {resize}x {upscaler}'.format(resize = sd.upscale.upscaling_resize, upscaler = sd.upscale.upscaler_2) if sd.upscale.upscaler_2 != "None" else '' # pylint: disable=consider-using-f-string
+    if op == "grid":
+        template += ' | grid {num}'.format(num = sd.generate.batch_size * sd.generate.n_iter) # pylint: disable=consider-using-f-string
+    ifd = ImageFileDirectory_v2()
+    exif_stream = io.BytesIO()
+    _TAGS = dict(((v, k) for k, v in TAGS.items())) # enumerate possible exif tags
+    ifd[_TAGS["ImageDescription"]] = template
+    ifd.save(exif_stream)
+    val = b"Exif\x00\x00" + exif_stream.getvalue()
+    return val
+
+
+def prompt(params): # generate dynamic prompt or use one if provided
+    sd.generate.prompt = params.prompt if params.prompt != "dynamic" else secrets.choice(random.prompts)
+    embedding = params.embedding if params.embedding != 'random' else secrets.choice(random.embeddings)
+    sd.generate.prompt = sd.generate.prompt.replace('<embedding>', embedding)
+    artist = params.artist if params.artist != 'random' else secrets.choice(random.artists)
+    sd.generate.prompt = sd.generate.prompt.replace('<artist>', artist)
+    style = params.style if params.style != 'random' else secrets.choice(random.styles)
+    sd.generate.prompt = sd.generate.prompt.replace('<style>', style)
+    suffix = params.suffix if params.suffix != 'random' else secrets.choice(random.suffixes)
+    sd.generate.prompt = sd.generate.prompt.replace('<suffix>', suffix)
+    place = params.suffix if params.suffix != 'random' else secrets.choice(random.places)
+    sd.generate.prompt = sd.generate.prompt.replace('<place>', place)
+    if params.prompts or params.debug:
+        log.info({ 'random initializers': random })
+    if params.prompt == "dynamic":
+        log.info({ 'dynamic prompt': sd.generate.prompt })
+    return sd.generate.prompt
+
+
+def sampler(params, options): # find sampler
+    if params.sampler == 'random':
+        sd.generate.sampler_name = secrets.choice(options.samplers)
+        log.info({ 'random sampler': sd.generate.sampler_name })
+    else:
+        found = [i for i in options.samplers if i.startswith(params.sampler)]
+        if len(found) == 0:
+            log.error({ 'sampler error': sd.generate.sampler_name, 'available': options.samplers})
+            exit()
+        sd.generate.sampler_name = found[0]
+    return sd.generate.sampler_name
+
+
+async def generate(prompt = None): # pylint: disable=redefined-outer-name
+    if prompt is not None:
+        sd.generate.prompt = prompt
+    log.info({ 'generate': sd.generate })
+    names = []
+    b64s = []
+    images = []
+    info = Map({})
+    data = await post('/sdapi/v1/txt2img', sd.generate)
+    if 'error' in data:
+        log.error({ 'generate': data['error'], 'reason': data['reason'] })
+        return Map({})
+    info = Map(json.loads(data['info']))
+    log.debug({ 'info': info })
+    images = data['images']
+    short = info.prompt[:min(len(info.prompt), 96)] # limit prompt part of filename to 64 chars
+    for i in range(len(images)):
+        b64s.append(images[i])
+        images[i] = Image.open(io.BytesIO(base64.b64decode(images[i].split(",",1)[0])))
+        name = "{seed:0>9}-{short}.jpg".format(short = short, seed = info.all_seeds[i]) # pylint: disable=consider-using-f-string
+        f = os.path.join(sd.paths.root, sd.paths.generate, name)
+        names.append(f)
+        log.info({ "image": { "name": f, "size": images[i].size } })
+        images[i].save(f, "JPEG", exif = exif(info, i), optimize = True, quality = 70)
+    return Map({ "name": names, "image": images, "b64": b64s, "info": info })
+
+
+async def upscale(data):
+    data.upscaled = []
+    if sd.upscale.upscaling_resize <=1:
+        return data
+    sd.upscale.image = ""
+    log.info({ 'upscale': sd.upscale })
+    for i in range(len(data.image)):
+        f = data.name[i].replace(sd.paths.generate, sd.paths.upscale)
+        sd.upscale.image = data.b64[i]
+        res = await post('/sdapi/v1/extra-single-image', sd.upscale)
+        image = Image.open(io.BytesIO(base64.b64decode(res['image'].split(",",1)[0])))
+        data.upscaled.append(image)
+        log.info({ "image": { "name": f, "size": image.size } })
+        image.save(f, "JPEG", exif = exif(data.info, i, "upscale"), optimize = True, quality = 70)
+    return data
+
+
+async def init():
+    """
+    import torch
+    log.info({ "torch": torch.__version__, "available": torch.cuda.is_available() })
+    current_device = torch.cuda.current_device()
+    mem_free, mem_total = torch.cuda.mem_get_info()
+    log.info({ "cuda": torch.version.cuda, "available": torch.cuda.is_available(), "arch": torch.cuda.get_arch_list(), "device": torch.cuda.get_device_name(current_device), "memory": { "free": round(mem_free / 1024 / 1024), "total": (mem_total / 1024 / 1024) } })
+    """
+    options = Map({})
+    options.flags = await get('/sdapi/v1/cmd-flags')
+    log.debug({ 'flags': options.flags })
+    data = await get('/sdapi/v1/sd-models')
+    options.models = [obj["title"] for obj in data]
+    log.debug({ 'registered models': options.models })
+    found = [i for i in options.models if i.startswith(sd.options.sd_model_checkpoint)]
+    if len(found) == 0:
+        log.error({ 'model error': sd.generate.sd_model_checkpoint, 'available': options.models})
+        exit()
+    sd.options.sd_model_checkpoint = found[0]
+    data = await get('/sdapi/v1/samplers')
+    options.samplers = [obj["name"] for obj in data]
+    log.debug({ 'registered samplers': options.samplers })
+    data = await get('/sdapi/v1/upscalers')
+    options.upscalers = [obj["name"] for obj in data]
+    log.debug({ 'registered upscalers': options.upscalers })
+    data = await get('/sdapi/v1/face-restorers')
+    options.restorers = [obj["name"] for obj in data]
+    log.debug({ 'registered face restorers': options.restorers })
+    await interrupt()
+    await post('/sdapi/v1/options', sd.options)
+    options.options = await get('/sdapi/v1/options')
+    log.info({ 'target models': { 'diffuser': options.options["sd_model_checkpoint"], 'vae': options.options["sd_vae"] } })
+    log.info({ 'paths': sd.paths })
+    options.queue = await get('/queue/status')
+    log.info({ 'queue': options.queue })
+    pathlib.Path(sd.paths.root).mkdir(parents = True, exist_ok = True)
+    pathlib.Path(os.path.join(sd.paths.root, sd.paths.generate)).mkdir(parents = True, exist_ok = True)
+    pathlib.Path(os.path.join(sd.paths.root, sd.paths.upscale)).mkdir(parents = True, exist_ok = True)
+    pathlib.Path(os.path.join(sd.paths.root, sd.paths.grid)).mkdir(parents = True, exist_ok = True)
+    return options
+
+
+def args(): # parse cmd arguments
+    global sd # pylint: disable=global-statement
+    global random # pylint: disable=global-statement
+    parser = argparse.ArgumentParser(description = "sd pipeline")
+    parser.add_argument("--config", type = str, default = 'generate.json', required = False, help = "configuration file")
+    parser.add_argument("--random", type = str, default = 'random.json', required = False, help = "prompt file with randomized sections")
+    parser.add_argument("--max", type = int, default = 1, required = False, help = "maximum number of generated images")
+    parser.add_argument("--prompt", type = str, default = 'dynamic', required = False, help = "prompt")
+    parser.add_argument("--negative", type = str, default = '', required = False, help = "negative prompt")
+    parser.add_argument("--artist", type = str, default = 'random', required = False, help = "artist style, used to guide dynamic prompt when prompt is not provided")
+    parser.add_argument("--embedding", type = str, default = 'random', required = False, help = "use embedding, used to guide dynamic prompt when prompt is not provided")
+    parser.add_argument("--style", type = str, default = 'random', required = False, help = "image style, used to guide dynamic prompt when prompt is not provided")
+    parser.add_argument("--suffix", type = str, default = 'random', required = False, help = "style suffix, used to guide dynamic prompt when prompt is not provided")
+    parser.add_argument("--place", type = str, default = 'random', required = False, help = "place locator, used to guide dynamic prompt when prompt is not provided")
+    parser.add_argument('--faces', default = False, action='store_true', help = "restore faces during upscaling")
+    parser.add_argument("--steps", type = int, default = 0, required = False, help = "number of steps")
+    parser.add_argument("--batch", type = int, default = 0, required = False, help = "batch size, limited by gpu vram")
+    parser.add_argument("--n", type = int, default = 0, required = False, help = "number of iterations")
+    parser.add_argument("--cfg", type = int, default = 0, required = False, help = "classifier free guidance scale")
+    parser.add_argument("--sampler", type = str, default = 'random', required = False, help = "sampler")
+    parser.add_argument("--seed", type = int, default = 0, required = False, help = "seed, default is random")
+    parser.add_argument("--upscale", type = int, default = 0, required = False, help = "upscale factor, disabled if 0")
+    parser.add_argument("--model", type = str, default = '', required = False, help = "diffusion model")
+    parser.add_argument("--vae", type = str, default = '', required = False, help = "vae model")
+    parser.add_argument("--path", type = str, default = '', required = False, help = "output path")
+    parser.add_argument("--width", type = int, default = 0, required = False, help = "width")
+    parser.add_argument("--height", type = int, default = 0, required = False, help = "height")
+    parser.add_argument("--beautify", default = False, action='store_true', help = "beautify prompt")
+    parser.add_argument('--prompts', default = False, action='store_true', help = "print dynamic prompt templates")
+    parser.add_argument('--debug', default = False, action='store_true', help = "print extra debug information")
+    params = parser.parse_args()
+    if params.debug:
+        log.setLevel(logging.DEBUG)
+        log.debug({ 'debug': True })
+    log.debug({ 'args': params.__dict__ })
+    home = pathlib.Path(sys.argv[0]).parent
+    if os.path.isfile(params.config):
+        try:
+            with open(params.config, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                sd = Map(data)
+                log.debug({ 'config': sd })
+        except Exception as e:
+            log.error({ 'config error': params.config, 'exception': e })
+            exit()
+    elif os.path.isfile(os.path.join(home, params.config)):
+        try:
+            with open(os.path.join(home, params.config), 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                sd = Map(data)
+                log.debug({ 'config': sd })
+        except Exception as e:
+            log.error({ 'config error': params.config, 'exception': e })
+            exit()
+    else:
+        log.error({ 'config file not found': params.config})
+        exit()
+    if params.prompt == 'dynamic':
+        log.info({ 'prompt template': params.random })
+        if os.path.isfile(params.random):
+            try:
+                with open(params.random, 'r', encoding='utf-8') as f:
+                    data = json.load(f)
+                    random = Map(data)
+                    log.debug({ 'random template': sd })
+            except:
+                log.error({ 'random template error': params.random})
+                exit()
+        elif os.path.isfile(os.path.join(home, params.random)):
+            try:
+                with open(os.path.join(home, params.random), 'r', encoding='utf-8') as f:
+                    data = json.load(f)
+                    random = Map(data)
+                    log.debug({ 'random template': sd })
+            except:
+                log.error({ 'random template error': params.random})
+                exit()
+        else:
+            log.error({ 'random template file not found': params.random})
+            exit()
+        _dynamic = prompt(params)
+
+    sd.paths.root = params.path if params.path != '' else sd.paths.root
+    sd.generate.restore_faces = params.faces if params.faces is not None else sd.generate.restore_faces
+    sd.generate.seed = params.seed if params.seed > 0 else sd.generate.seed
+    sd.generate.sampler_name = params.sampler if params.sampler != 'random' else sd.generate.sampler_name
+    sd.generate.batch_size = params.batch if params.batch > 0 else sd.generate.batch_size
+    sd.generate.negative_prompt = params.negative if params.negative != '' else sd.generate.negative_prompt
+    sd.generate.cfg_scale = params.cfg if params.cfg > 0 else sd.generate.cfg_scale
+    sd.generate.n_iter = params.n if params.n > 0 else sd.generate.n_iter
+    sd.generate.width = params.width if params.width > 0 else sd.generate.width
+    sd.generate.height = params.height if params.height > 0 else sd.generate.height
+    sd.generate.steps = params.steps if params.steps > 0 else sd.generate.steps
+    sd.upscale.upscaling_resize = params.upscale if params.upscale > 0 else sd.upscale.upscaling_resize
+    sd.upscale.codeformer_visibility = 1 if params.faces else sd.upscale.codeformer_visibility
+    sd.options.sd_vae = params.vae if params.vae != '' else sd.options.sd_vae
+    sd.options.sd_model_checkpoint = params.model if params.model != '' else sd.options.sd_model_checkpoint
+    sd.upscale.upscaler_1 = "SwinIR_4x" if params.upscale > 1 else sd.upscale.upscaler_1
+    if sd.generate.cfg_scale == 0:
+        sd.generate.cfg_scale = randrange(5, 10)
+    return params
+
+
+async def main():
+    params = args()
+    sess = await session()
+    if sess is None:
+        await close()
+        exit()
+    options = await init()
+    iteration = 0
+    while True:
+        iteration += 1
+        log.info('')
+        log.info({ 'iteration': iteration, 'batch': sd.generate.batch_size, 'n': sd.generate.n_iter, 'total': sd.generate.n_iter * sd.generate.batch_size })
+        dynamic = prompt(params)
+        if params.beautify:
+            try:
+                from promptist import beautify # pylint: disable=import-outside-toplevel
+                sd.generate.prompt = beautify(dynamic)
+            except Exception as e:
+                log.error({ 'beautify': e })
+        scheduler = sampler(params, options)
+        t0 = time.perf_counter()
+        data = await generate() # generate returns list of images
+        if not 'image' in data:
+            break
+        stats.images += len(data.image)
+        t1 = time.perf_counter()
+        avg[scheduler] = (t1 - t0) / len(data.image)
+        stats.generate += t1 - t0
+        _image = grid(data)
+        data = await upscale(data)
+        t2 = time.perf_counter()
+        stats.upscale += t2 - t1
+        stats.wall += t2 - t0
+        log.info({ "time" : { "generate": t1 - t0, "average": (t1 - t0) / len(data.image), "upscale": t2 - t1 } })
+        if params.max != 0 and stats.images >= params.max:
+            break
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        asyncio.run(interrupt())
+        asyncio.run(close())
+        log.info({ "interrupt": True })
+    finally:
+        log.info({ "sampler performance": avg })
+        log.info({ "stats" : stats })
+        asyncio.run(close())
+"""
+    except Exception as e:
+        log.info({ "sampler performance": avg })
+        log.info({ "stats": stats })
+        log.critical({ "exception": e })
+        exit()
+"""

cli/grid.py

@@ -0,0 +1,115 @@
+#!/bin/env python
+"""
+Create image grid
+"""
+
+import os
+import argparse
+import math
+import logging
+from pathlib import Path
+
+import filetype
+from PIL import Image, ImageDraw, ImageFont
+from util import log
+
+params = None
+
+
+def wrap(text: str, font: ImageFont.ImageFont, length: int):
+    lines = ['']
+    for word in text.split():
+        line = f'{lines[-1]} {word}'.strip()
+        if font.getlength(line) <= length:
+            lines[-1] = line
+        else:
+            lines.append(word)
+    return '\n'.join(lines)
+
+
+def grid(images, labels = None, width = 0, height = 0): # pylint: disable=redefined-outer-name
+    if params.horizontal:
+        rows = 1
+    elif params.vertical:
+        rows = len(images)
+    elif params.square:
+        rows = round(math.sqrt(len(images)))
+    else:
+        rows = math.floor(math.sqrt(len(images)))
+    cols = math.ceil(len(images) / rows)
+    size = [0, 0]
+    if width == 0:
+        w = max([i.size[0] for i in images])
+        size[0] = cols * w + cols * params.border
+    else:
+        size[0] = width
+        w = round(width / cols)
+    if height == 0:
+        h = max([i.size[1] for i in images])
+        size[1] = rows * h + rows * params.border
+    else:
+        size[1] = height
+        h = round(height / rows)
+    size = tuple(size)
+    image = Image.new('RGB', size = size, color = 'black') # pylint: disable=redefined-outer-name
+    font = ImageFont.truetype('DejaVuSansMono', round(w / 20))
+    for i, img in enumerate(images): # pylint: disable=redefined-outer-name
+        x = (i % cols * w) + (i % cols * params.border)
+        y = (i // cols * h) + (i // cols * params.border)
+        img.thumbnail((w, h), Image.HAMMING)
+        image.paste(img, box=(x, y))
+        if labels is not None and len(images) == len(labels):
+            ctx = ImageDraw.Draw(image)
+            label = wrap(labels[i], font, w)
+            ctx.text((x + 1 + round(w / 200), y + 1 + round(w / 200)), label, font = font, fill = (0, 0, 0))
+            ctx.text((x, y), label, font = font, fill = (255, 255, 255))
+    log.info({ 'grid': { 'images': len(images), 'rows': rows, 'cols': cols, 'cell': [w, h] } })
+    return image
+
+
+if __name__ == '__main__':
+    log.info({ 'create grid' })
+    parser = argparse.ArgumentParser(description='image grid utility')
+    parser.add_argument("--square", default = False, action='store_true', help = "create square grid")
+    parser.add_argument("--horizontal", default = False, action='store_true', help = "create horizontal grid")
+    parser.add_argument("--vertical", default = False, action='store_true', help = "create vertical grid")
+    parser.add_argument("--width", type = int, default = 0, required = False, help = "fixed grid width")
+    parser.add_argument("--height", type = int, default = 0, required = False, help = "fixed grid height")
+    parser.add_argument("--border", type = int, default = 0, required = False, help = "image border")
+    parser.add_argument('--debug', default = False, action='store_true', help = "print extra debug information")
+    parser.add_argument('output', type = str)
+    parser.add_argument('input', type = str, nargs = '*')
+    params = parser.parse_args()
+    output = params.output if params.output.lower().endswith('.jpg') else params.output + '.jpg'
+    if params.debug:
+        log.setLevel(logging.DEBUG)
+        log.debug({ 'debug': True })
+    log.debug({ 'args': params.__dict__ })
+    images = []
+    labels = []
+    for f in params.input:
+        path = Path(f)
+        if path.is_dir():
+            files = [os.path.join(f, file) for file in os.listdir(f) if os.path.isfile(os.path.join(f, file))]
+        elif path.is_file():
+            files = [f]
+        else:
+            log.warning({ 'grid not a valid file/folder', f})
+            continue
+        files.sort()
+        for file in files:
+            if not filetype.is_image(file):
+                continue
+            log.debug(file)
+            img = Image.open(file)
+            # img.verify()
+            images.append(img)
+            fp = Path(file)
+            labels.append(fp.stem)
+    # log.info({ 'folder': path.parent, 'labels': labels })
+    if len(images) > 0:
+        image = grid(images, labels, params.width, params.height)
+        image.save(output, 'JPEG', optimize = True, quality = 60)
+        log.info({ 'grid': { 'file': output, 'size': list(image.size) } })
+    else:
+        log.info({ 'grid': 'nothing to do' })

cli/ideas.py

@@ -0,0 +1,60 @@
+#!/bin/env python
+"""
+generate prompt ideas
+model from: <https://huggingface.co/FredZhang7/distilgpt2-stable-diffusion-v2>
+"""
+
+import logging
+import argparse
+from util import log
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+
+tokenizer = None
+model = None
+
+def prompt(text: str, temp: float = 0.9, top: int = 8, penalty: float = 1.2, alpha: float = 0.6, num: int = 5, length: int = 80):
+    global tokenizer, model # pylint: disable=global-statement
+    if tokenizer is None:
+        tokenizer = GPT2Tokenizer.from_pretrained('distilgpt2')
+        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+    if model is None:
+        model = GPT2LMHeadModel.from_pretrained('FredZhang7/distilgpt2-stable-diffusion-v2')
+    input_ids = tokenizer(text, return_tensors='pt').input_ids
+    output = model.generate(input_ids,
+        do_sample = True,
+        temperature = temp,
+        top_k = top,
+        max_length = length,
+        num_return_sequences = num,
+        repetition_penalty = penalty,
+        penalty_alpha = alpha,
+        no_repeat_ngram_size = 1,
+        early_stopping = True
+    )
+    outputs = []
+    for i in range(len(output)):
+        outputs.append(tokenizer.decode(output[i], skip_special_tokens=True))
+    return outputs
+
+
+if __name__ == "__main__": # create & train test embedding when used from cli
+    log.info({ 'idea': 'generate prompts' })
+    parser = argparse.ArgumentParser(description='idea: generate prompts')
+    parser.add_argument("--temp", type = float, default = 0.9, required = False, help = "higher temperature produces more diverse results with a higher risk of less coherent text, default: %(default)s")
+    parser.add_argument("--top", type = int, default = 8, required = False, help = "number of tokens to sample from at each step, default: %(default)s")
+    parser.add_argument("--penalty", type = float, default = 1.2, required = False, help = "penalty value for each repetition of a token, default: %(default)s")
+    parser.add_argument("--alpha", type = float, default = 0.6, required = False, help = "penalty alpha value, default: %(default)s")
+    parser.add_argument("--num", type = int, default = 10, required = False, help = "number of results to generate, default: %(default)s")
+    parser.add_argument("--length", type = int, default = 85, required = False, help = "maximum number of output tokens, default: %(default)s")
+    parser.add_argument('--debug', default = False, action='store_true', help = "print extra debug information, default: %(default)s")
+    parser.add_argument('text', type = str, nargs = '*')
+    params = parser.parse_args()
+    if params.debug:
+        log.setLevel(logging.DEBUG)
+        log.debug({ 'debug': True })
+    log.debug({ 'args': params.__dict__ })
+    sentence = ' '.join(params.text)
+    res = prompt(text = sentence, temp = params.temp, top = params.top, penalty = params.penalty, alpha = params.alpha, num = params.num, length = params.length)
+    log.info({ 'ideas for':  sentence })
+    for line in res:
+        log.info(line)

cli/interrogate.py

@@ -0,0 +1,90 @@
+#!/bin/env python
+"""
+use clip to interrogate image(s)
+"""
+
+import io
+import base64
+import sys
+import os
+import asyncio
+import filetype
+from PIL import Image
+from util import log, Map
+import sdapi
+
+
+def decode(encoding):
+    if encoding.startswith("data:image/"):
+        encoding = encoding.split(";")[1].split(",")[1]
+    return Image.open(io.BytesIO(base64.b64decode(encoding)))
+
+
+def encode(f):
+    image = Image.open(f)
+    exif = image.getexif()
+    with io.BytesIO() as stream:
+        image.save(stream, 'JPEG', exif = exif)
+        values = stream.getvalue()
+        encoded = base64.b64encode(values).decode()
+        return encoded
+
+
+async def interrogate(f):
+    if not filetype.is_image(f):
+        log.info({ 'interrogate skip': f })
+        return
+    json = Map({ 'image': encode(f) })
+    log.info({ 'interrogate': f })
+    # run clip
+    json.model = 'clip'
+    res = await sdapi.post('/sdapi/v1/interrogate', json)
+    # res = sdapi.postsync('/sdapi/v1/interrogate', json)
+    caption = ""
+    style = ""
+    if 'caption' in res:
+        caption = res.caption
+        log.info({ 'interrogate caption': caption })
+        if ', by' in caption:
+            style = caption.split(', by')[1].strip()
+            log.info({ 'interrogate style': style })
+    else:
+        log.error({ 'interrogate clip error': res })
+    # run booru
+    json.model = 'deepdanbooru'
+    res = await sdapi.post('/sdapi/v1/interrogate', json)
+    # res = sdapi.postsync('/sdapi/v1/interrogate', json)
+    keywords = {}
+    if 'caption' in res:
+        for term in res.caption.split(', '):
+            term = term.replace('(', '').replace(')', '').split(':')
+            keywords[term[0]] = term[1]
+        keywords = dict(sorted(keywords.items(), key=lambda x:x[1], reverse=True))
+        log.info({ 'interrogate keywords': keywords })
+    else:
+        log.error({ 'interrogate booru error': res })
+    return caption, keywords, style
+
+
+async def main():
+    sys.argv.pop(0)
+    await sdapi.session()
+    if len(sys.argv) == 0:
+        log.error({ 'interrogate': 'no files specified' })
+    for arg in sys.argv:
+        if os.path.exists(arg):
+            if os.path.isfile(arg):
+                await interrogate(arg)
+            elif os.path.isdir(arg):
+                for root, _dirs, files in os.walk(arg):
+                    for f in files:
+                        await interrogate(os.path.join(root, f))
+            else:
+                log.error({ 'interrogate unknown file type': arg })
+        else:
+            log.error({ 'interrogate file missing': arg })
+    await sdapi.close()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

cli/losschart.py

@@ -0,0 +1,132 @@
+#!/bin/env python
+
+import io
+import os
+import sys
+import json
+import pathlib
+import logging
+import numpy as np
+import scipy as sp
+from PIL import Image, ImageFont, ImageDraw
+from matplotlib import pyplot as plt
+from util import log, Map
+
+def settings(logdir: str, name: str):
+    filename = os.path.join(logdir, name, 'settings.json')
+    # shutil.copyfile(filename, os.path.join(logdir, f"{name}.json"))
+    with open(filename, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+        data = Map(data)
+        log.debug({ 'settings': data })
+    return data
+
+
+def plot(logdir: str, name: str):
+    f = os.path.join(logdir, name, 'train.csv')
+    if not os.path.isfile(f):
+        log.debug({ 'train log missing': f })
+        return
+    name = pathlib.Path(f).parent.name
+    # shutil.copyfile(f, os.path.join(logdir, f"{name}.csv"))
+    img = os.path.join(logdir, f"{name}.png")
+
+    step, loss, rate = plt.np.loadtxt(f, delimiter = ',', skiprows = 1, usecols = [0, 3, 4], unpack = True)
+    d = settings(logdir, name)
+    try:
+        log.debug({ 'loss plot': name, 'output': img, 'data': f, 'records': len(step) })
+    except:
+        return # no data
+    if len(step) < 5:
+        return
+
+    plt.rcParams.update({'font.variant':'small-caps'})
+    plt.rc('axes', edgecolor='gray')
+    plt.rc('font', size=10)
+    plt.rc('font', variant='small-caps')
+    plt.grid(color='gray', linewidth=1, axis='both', alpha=0.5)
+    plt.rcParams['figure.figsize'] = [14, 6]
+    plt.subplots_adjust(right=1000)
+
+    fig, ax1 = plt.subplots()
+    fig.set_facecolor('black')
+
+    ax1.set_facecolor(color = (0.1, 0.1, 0.1, 0.5))
+    ax1.tick_params(axis='x', labelcolor='white')
+    ax1.set_xlabel('step'.upper(), color='white')
+    ax1.set_xlim(0, d.steps)
+    ax1.set_ylim(0, 0.5)
+    ax1.set_axisbelow(True)
+    ax1.xaxis.grid(color='gray', linestyle='dashed')
+    ax1.yaxis.grid(color='gray', linestyle='dashed')
+
+    # loss with additional interpolated values to smooth out the curve
+    ax1.set_ylabel('loss'.upper(), color='gray')
+    ax1.plot(step, loss, 'go')
+    spline = sp.interpolate.make_interp_spline(step, loss)
+    x_ = np.linspace(min(step), max(step), num = len(step * 3), endpoint = True, retstep = False, dtype = int, axis = 0)
+    y_ = spline(x_)
+    ax1.plot(x_, y_, color='gray')
+    ax1.tick_params(axis='y', labelcolor='gray')
+
+    # moving average
+    window = 10
+    if len(loss) > window:
+        ma = []
+        for ind in range(window - 1):
+            ma.insert(0, np.nan)
+        for ind in range(len(loss) - window + 1):
+            ma.append(np.mean(loss[ind:ind+window]))
+        ax1.plot(step, ma, color="maroon", linewidth=5)
+
+    # learning rate
+    ax2 = ax1.twinx()
+    ax2.set_ylabel('learning rate'.upper(), color='cyan')
+    ax2.plot(step, rate, color='cyan', linewidth=3, linestyle='dashed')
+    ax2.tick_params(axis='y', labelcolor='cyan')
+
+    # create chart and convert to pil
+    fig.tight_layout()
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    pltimg = Image.open(buf)
+    size = (pltimg.size[0], pltimg.size[1] + 240)
+    image = Image.new('RGB', size = size, color = (206, 100, 0))
+    font = ImageFont.truetype('DejaVuSansMono', 18)
+    image.paste(pltimg, box=(0, 240))
+    buf.close()
+
+    # text
+    textl = f"""NAME: {d.embedding_name.upper()}
+IMAGES: {d.num_of_dataset_images}
+VECTORS: {d.num_vectors_per_token}
+STEPS: {d.steps}
+BATCH-SIZE: {d.batch_size}
+GRADIENT-STEP: {d.gradient_step}
+LEARN-RATE: {d.learn_rate}
+SAMPLING-METHOD: {d.latent_sampling_method}
+MODEL: {d.model_name.upper()}
+"""
+
+    minval = f"{round(np.min(loss), 4)} @ {round(step[np.argmin(loss)])}"
+    maxval = f"{round(np.max(loss), 4)} @ {round(step[np.argmax(loss)])}"
+    textr = f"""{d.datetime}
+LOSS: {round(loss[-1], 4)}
+MIN: {minval}
+MAX: {maxval}
+"""
+    ctx = ImageDraw.Draw(image)
+    ctx.text((8, 8), textl, font = font, fill = (255, 255, 255), spacing = 8)
+    ctx.text((image.size[0] - 220, 8), textr, font = font, fill = (255, 255, 255))
+
+    image.save(img)
+
+
+if __name__ == "__main__":
+    log.setLevel(logging.DEBUG)
+    if len(sys.argv) == 2:
+        arg = sys.argv[1]
+        log.debug({ 'args': arg })
+        plot(os.path.dirname(arg), os.path.basename(arg))
+    else:
+        log.debug({ 'error': 'specify embedding name'})

cli/promptist.py

@@ -0,0 +1,38 @@
+#!/bin/env python
+"""
+use microsoft promptist to beautify prompt
+- <https://huggingface.co/spaces/microsoft/Promptist>
+"""
+
+import sys
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from util import log
+
+def load_prompter():
+    model = AutoModelForCausalLM.from_pretrained("microsoft/Promptist") # pylint: disable=redefined-outer-name
+    tokenizer = AutoTokenizer.from_pretrained("gpt2") # pylint: disable=redefined-outer-name
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "left"
+    return model, tokenizer
+
+
+model, tokenizer = load_prompter()
+
+
+def beautify(plain_text):
+    input_ids = tokenizer(plain_text.strip() + " Rephrase:", return_tensors = "pt").input_ids
+    eos_id = tokenizer.eos_token_id
+    outputs = model.generate(input_ids, do_sample = False, max_new_tokens = 75, num_beams = 8, num_return_sequences = 8, eos_token_id = eos_id, pad_token_id = eos_id, length_penalty = -1.0)
+    output_texts = tokenizer.batch_decode(outputs, skip_special_tokens = True)
+    texts = []
+    for output_text in output_texts:
+        texts.append(output_text.replace(plain_text+" Rephrase:", "").strip())
+    longest = max(texts, key = len)
+    log.info({ 'beautified': longest })
+    return longest
+
+if __name__ == "__main__": # create & train test embedding when used from cli
+    sys.argv.pop(0)
+    text = ' '.join(sys.argv)
+    log.info({ 'prompt': text })
+    output = beautify(text)

cli/random.json

@@ -0,0 +1,31 @@
+{
+  "prompts": [
+    "<style> of <embedding> <place>, high detailed, by <artist>, <suffix>"
+  ],
+  "negative": [
+    "watermark, fog, clouds, blurry, duplicate, deformed, mutation"
+  ],
+  "places": [
+    "standing in the city", "on a spaceship", "in fantasy landscape", "on a shore", "in a forest", "in winter wonderland"
+  ],
+  "embeddings": [
+    "man", "man next to a beautiful girl", "man next to a car", "beautiful girl", "sexy naked girl", "cute girl holding a flower", "beautiful robot",
+    "young korean girl with medium-length white hair", "monster", "pin up girl",
+    "man vlado", "beutiful girl ana", "man lee", "beautiful girl abby"
+  ],
+  "artists": [
+    "John Salminen", "Greg Rutkowski", "Akihiko Yoshida", "Alejandro Burdisio", "Artgerm", "Patrick Brown", "Walt Disney", "Neal Adams", "Jeremy Chong",
+    "Chris Rallis", "Roy Lichtenstein", "Claude Monet", "Jon Whitcomb", "Pablo Picasso", "Raymond Leech", "Tom Lovell", "Noriyoshi Ohrai", "Shingei",
+    "Helmut Newton", "Maciej Kuciara", "Daniel F. Gerhartz", "Stephan Martinière", "Magali Villeneuve", "Carne Griffiths", "Alberto Seveso",
+    "Vincent Van Gogh", "WLOP", "Frank Xavier Leyendecker", "Peter Lindbergh", "Nick Gentry", "Howard Chandler Christy", "Raphael", "Henri Matisse"
+  ],
+  "styles": [
+    "illustration", "painting", "portrait", "photograph", "drawing", "sketch", "pencil sketch", "3d render", "cartoon", "anime", "scribbles", "pop art",
+    "ink painting", "steampunk illustration", "dc comics illustration", "marvel comics", "vray render", "photoillustration", "pixar", "marble sculpture",
+    "bronze sculpture", "christmas theme"
+  ],
+  "suffixes": [
+    "cinematic lighting", "artstation", "fineart", "cinematic", "photorealistic", "soft light", "sharp focus", "bokeh", "dreamlike", "semirealism",
+    "colorful", "black and white", "intricate", "elegant"
+  ]
+}

cli/sdapi.py

@@ -0,0 +1,148 @@
+#!/bin/env python
+"""
+helper methods that creates HTTP session with managed connection pool
+provides async HTTP get/post methods and several helper methods
+"""
+
+import asyncio
+import logging
+import sys
+
+import aiohttp
+import requests
+from util import Map, log
+
+sd_url = "http://127.0.0.1:7860" # automatic1111 api url root
+use_session = True
+timeout = aiohttp.ClientTimeout(total = None, sock_connect = 10, sock_read = None) # default value is 5 minutes, we need longer for training
+
+sess = None
+quiet = False
+
+async def result(req):
+    if req.status != 200:
+        if not quiet:
+            log.error({ 'request error': req.status, 'reason': req.reason, 'url': req.url })
+        if not use_session and sess is not None:
+            await sess.close()
+        return Map({ 'error': req.status, 'reason': req.reason, 'url': req.url })
+    else:
+        json = await req.json()
+        if type(json) == list:
+            res = json
+        elif json is None:
+            res = {}
+        else:
+            res = Map(json)
+        log.debug({ 'request': req.status, 'url': req.url, 'reason': req.reason, 'result': res })
+        return res
+
+
+def resultsync(req: requests.Response):
+    if req.status_code != 200:
+        if not quiet:
+            log.error({ 'request error': req.status_code, 'reason': req.reason, 'url': req.url })
+        return Map({ 'error': req.status_code, 'reason': req.reason, 'url': req.url })
+    else:
+        json = req.json()
+        if type(json) == list:
+            res = json
+        elif json is None:
+            res = {}
+        else:
+            res = Map(json)
+        log.debug({ 'request': req.status_code, 'url': req.url, 'reason': req.reason, 'result': res })
+        return res
+
+
+async def get(endpoint: str, json: dict = None):
+    global sess # pylint: disable=global-statement
+    sess = sess if sess is not None else await session()
+    async with sess.get(url = endpoint, json = json) as req:
+        res = await result(req)
+        return res
+
+
+def getsync(endpoint: str, json: dict = None):
+    req = requests.get(f'{sd_url}{endpoint}', json = json) # pylint: disable=missing-timeout
+    res = resultsync(req)
+    return res
+
+
+async def post(endpoint: str, json: dict = None):
+    global sess # pylint: disable=global-statement
+    sess = sess if sess is not None else await session()
+    async with sess.post(url = endpoint, json = json) as req:
+        res = await result(req)
+        return res
+
+
+def postsync(endpoint: str, json: dict = None):
+    req = requests.post(f'{sd_url}{endpoint}', json = json) # pylint: disable=missing-timeout
+    res = resultsync(req)
+    return res
+
+
+async def interrupt():
+    res = await get('/sdapi/v1/progress?skip_current_image=true')
+    if 'state' in res and res.state.job_count > 0:
+        log.debug({ 'interrupt': res.state })
+        res = await post('/sdapi/v1/interrupt')
+        await asyncio.sleep(1)
+        return res
+    else:
+        log.debug({ 'interrupt': 'idle' })
+        return { 'interrupt': 'idle' }
+
+
+def interruptsync():
+    res = getsync('/sdapi/v1/progress?skip_current_image=true')
+    if 'state' in res and res.state.job_count > 0:
+        log.debug({ 'interrupt': res.state })
+        res = postsync('/sdapi/v1/interrupt')
+        return res
+    else:
+        log.debug({ 'interrupt': 'idle' })
+        return { 'interrupt': 'idle' }
+
+
+async def progress():
+    res = await get('/sdapi/v1/progress?skip_current_image=true')
+    log.debug({ 'progress': res })
+    return res
+
+
+async def session():
+    global sess # pylint: disable=global-statement
+    time = aiohttp.ClientTimeout(total = None, sock_connect = 10, sock_read = None) # default value is 5 minutes, we need longer for training
+    sess = aiohttp.ClientSession(timeout = time, base_url = sd_url)
+    log.debug({ 'sdapi': 'session created', 'endpoint': sd_url })
+    """
+    sess = await aiohttp.ClientSession(timeout = timeout).__aenter__()
+    try:
+        async with sess.get(url = f'{sd_url}/') as req:
+            log.debug({ 'sdapi': 'session created', 'endpoint': sd_url })
+    except Exception as e:
+        log.error({ 'sdapi': e })
+        await asyncio.sleep(0)
+        await sess.__aexit__(None, None, None)
+        sess = None
+    return sess
+    """
+    return sess
+
+
+async def close():
+    if sess is not None:
+        await asyncio.sleep(0)
+        await sess.__aexit__(None, None, None)
+        log.debug({ 'sdapi': 'session closed', 'endpoint': sd_url })
+
+
+if __name__ == "__main__":
+    log.setLevel(logging.DEBUG)
+    if 'interrupt' in sys.argv:
+        asyncio.run(interrupt())
+    if 'progress' in sys.argv:
+        asyncio.run(progress())
+    asyncio.run(close())

cli/train.json

@@ -0,0 +1,67 @@
+{
+  "training_model": "sd-v15-runwayml.ckpt",
+  "extract_video": {
+    "rate": 0,
+    "fps": 2,
+    "skipstart": 0,
+    "skipend": 0
+  },
+  "create_embedding": {
+      "name": "test",
+      "num_vectors_per_token": 1,
+      "overwrite_old": false,
+      "init_text": "*"
+  },
+  "preprocess": {
+      "process_src": "",
+      "process_dst": "",
+      "process_width": 512,
+      "process_height": 512,
+      "process_flip": false,
+      "process_split": false,
+      "process_caption": true,
+      "process_caption_deepbooru": false,
+      "preprocess_txt_action": "ignore",
+      "process_focal_crop": true,
+      "process_focal_crop_face_weight": 0.9,
+      "process_focal_crop_entropy_weight": 0.3,
+      "process_focal_crop_edges_weight": 0.5,
+      "process_focal_crop_debug": false,
+      "split_threshold": 0.5,
+      "overlap_ratio": 0.2
+  },
+  "train_embedding": {
+      "embedding_name": "",
+      "learn_rate": "0.010:10, 0.008:20, 0.006:40, 0.004:80, 0.002:120, 0.001:160, 0.0005:200",
+      "batch_size": 2,
+      "steps": 200,
+      "data_root": "",
+      "log_directory": "train/log",
+      "template_filename": "subject_filewords.txt",
+      "gradient_step": 10,
+      "training_width": 512,
+      "training_height": 512,
+      "shuffle_tags": false,
+      "tag_drop_out": 0,
+      "clip_grad_mode": "disabled",
+      "clip_grad_value": "0.1",
+      "latent_sampling_method": "deterministic",
+      "create_image_every": 10,
+      "save_embedding_every": 10,
+      "save_image_with_stored_embedding": false,
+      "preview_from_txt2img": false,
+      "preview_prompt": "",
+      "preview_negative_prompt": "blurry, duplicate, ugly, deformed, low res, watermark, text",
+      "preview_steps": 20,
+      "preview_sampler_index": 0,
+      "preview_cfg_scale": 6,
+      "preview_seed": -1,
+      "preview_width": 512,
+      "preview_height": 512,
+      "varsize": false
+  },
+      "create_hypernetwork": {
+  },
+      "train_hypernetwork": {
+  }
+}

cli/train.md

@@ -0,0 +1,23 @@
+# learning notes
+
+## Using accumulation: 1
+
+- steps: 4500
+- batch': 2
+- accumulation: 1
+- rate: "0.005000:100, 0.002500:300, 0.001000:600, 0.000500:1000, 0.000250:1500, 0.000100:2100, 0.000050:2800, 0.000025:3600, 0.000010:4500",
+
+## Using accumulation: 10
+
+- steps: 200
+- batch': 2
+- accumulation: 10
+- learning-rate: '0.010:10, 0.008:20, 0.006:40, 0.004:80, 0.002:120, 0.001:160, 0.0005:200'
+
+## Train
+
+> train.py --name laurentaylor-v4 --src ~/generative/Input/laurentaylor/ --init person,woman,girl,model --overwrite
+
+## Prompt
+
+    a medium shot photo of "dreamkelly", extremely detailed 8k wallpaper, intricate, high detail, dramatic, modelshoot style

cli/train.py

@@ -0,0 +1,429 @@
+#!/bin/env python
+# pylint: disable=no-member
+"""
+simple implementation of training api: `/sdapi/v1/train`
+- supports: create embedding, image preprocess, train embedding (with all known parameters)
+- does not (yet) support: create hyper-network, train hyper-network
+- compatible with progress api: `/sdapi/v1/progress`
+- if interrupted, auto-continues from last known step
+- create and preprocess executed as sync jobs
+- train is executed as async job with progress monitoring
+"""
+
+import argparse
+import asyncio
+import logging
+import math
+import os
+import sys
+import pathlib
+import time
+import json
+from pathlib import Path
+
+import filetype
+from ffmpeg import extract
+from sdapi import close, get, interrupt, post, progress, session
+from util import Map, log
+from losschart import plot
+from PIL import Image
+
+images = []
+use_pbar = False
+
+args = {}
+options = None
+cmdflags = None
+
+
+async def plotloss(params):
+    logdir = os.path.abspath(os.path.join(cmdflags['embeddings_dir'], '../train/log'))
+    try:
+        plot(logdir, params.name)
+    except Exception as err:
+        log.warning({ 'loss chart error': err })
+
+
+async def captions(docs: list):
+    exclude = ['a', 'in', 'on', 'out', 'at', 'the', 'and', 'with', 'next', 'to', 'it', 'for', 'of', 'into', 'that']
+    d = dict()
+    for f in docs:
+        text = open(f, 'r', encoding='utf-8')
+        for line in text:
+            line = line.strip()
+            line = line.lower()
+            words = line.split(" ")
+            for word in words:
+                if word in exclude:
+                    continue
+                d[word] = d[word] + 1 if word in d else 1
+    pairs = ((value, key) for (key,value) in d.items())
+    sort = sorted(pairs, reverse = True)
+    if len(sort) > 10:
+        del sort[10:]
+    d = {k: v for v, k in sort}
+    log.info({ 'top captions': d })
+
+
+async def cleanup(params):
+    if params.nocleanup:
+        return
+    log.info({ 'cleanup deleting preprocessed images': params.dst })
+    for f in Path(params.dst).glob('*.png'):
+        f.unlink()
+    for f in Path(params.dst).glob('*.txt'):
+        f.unlink()
+
+
+async def preprocess(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
+    if os.path.isdir(params.src):
+        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'))]
+        images = []
+        low_res = []
+        for f in candidates:
+            img = Image.open(f)
+            mp = (img.size[0] * img.size[1]) / 1024 / 1024
+            if mp < 1 or img.size[0] < 512 or img.size[1] < 512:
+                low_res.append(f)
+                os.rename(f, f + '.skip')
+            else:
+                images.append(f)
+        log.debug({ 'preprocess skipping': not_images })
+        log.debug({ 'preprocess low res': low_res })
+        args.preprocess.process_src = params.src
+        args.preprocess.process_dst = params.dst
+        log.debug({ 'preprocess args': args.preprocess })
+        _res = await post('/sdapi/v1/preprocess', json = args.preprocess)
+        processed = [os.path.join(params.dst, f) for f in os.listdir(params.dst) if os.path.isfile(os.path.join(params.dst, f))]
+        processed_imgs = [f for f in processed if f.endswith('.png')]
+        processed_docs = [f for f in processed if f.endswith('.txt')]
+        log.info({ 'preprocess': {
+            'source': params.src,
+            'destination': params.dst,
+            'files': len(files),
+            'images': len(images),
+            'processed': len(processed_imgs),
+            'captions': len(processed_docs),
+            'skipped': len(not_images),
+            'low-res': len(low_res) }
+        })
+        if len(processed_docs) > 0:
+            await captions(processed_docs)
+        return len(processed_imgs)
+    elif os.path.isfile(params.src):
+        if not filetype.is_video(params.src):
+            kind = filetype.guess(params.src)
+            log.error({ 'preprocess error': { 'not a valid movie file': params.src, 'guess': kind } })
+        else:
+            extract_dst = os.path.join(params.dst, 'extract')
+            log.debug({ 'preprocess args': args.extract_video })
+            images = extract(params.src, extract_dst, rate = args.extract_video.rate, fps = args.extract_video.fps, start = args.extract_video.skipstart, end = args.extract_video.skipend) # extract keyframes from movie
+            if images > 0:
+                params.src = extract_dst
+                processed_count = await preprocess(params) # call again but now with keyframes
+                return processed_count
+            else:
+                log.error({ 'preprocess video extract': 'no images' })
+                return 0
+    else:
+        log.error({ 'preprocess error': { 'not a valid input': params.src } })
+        return 0
+
+
+async def check(params):
+    log.debug({ 'setting options' })
+    global options # pylint: disable=global-statement
+    options = await get('/sdapi/v1/options')
+
+    options['training_xattention_optimizations'] = False
+    options['training_image_repeats_per_epoch'] = 1
+
+    log.debug({ 'check model': args.training_model })
+    if len(args.training_model) > 0 and not options['sd_model_checkpoint'].startswith(args.training_model):
+        models = await get('/sdapi/v1/sd-models')
+        models = [obj["title"] for obj in models]
+        found = [i for i in models if i.startswith(args.training_model)]
+        if len(found) == 0:
+            log.error({ 'model not found': args.training_model, 'available': models })
+            exit()
+        else:
+            log.warning({ 'switching model': found[0] })
+        options['sd_model_checkpoint'] = found[0]
+
+    log.debug({ 'check embedding': params.name })
+    global cmdflags # pylint: disable=global-statement
+    cmdflags = await get('/sdapi/v1/cmd-flags')
+
+    lst = os.path.join(cmdflags['embeddings_dir'])
+    log.debug({ 'embeddings folder': lst })
+    path = Path(cmdflags['embeddings_dir']).glob(f'{params.name}.pt*')
+    matches = [f for f in path]
+    for match in matches:
+        if params.overwrite:
+            log.info({ 'delete embedding': match.name })
+            os.remove(os.path.join(cmdflags['embeddings_dir'], match.name))
+        else:
+            log.error({ 'embedding exists': match.name })
+            await close()
+            exit()
+    logdir = os.path.abspath(os.path.join(cmdflags['embeddings_dir'], '../train/log', params.name))
+    f = os.path.join(logdir, 'train.csv')
+    if os.path.isfile(f):
+        if params.overwrite:
+            log.info({ 'delete training log': f })
+            os.remove(os.path.join(logdir, 'train.csv'))
+        else:
+            log.warning({ 'training log exists': f })
+    f = os.path.join(logdir, '..', params.name, '.png')
+    if os.path.isfile(f):
+        if params.overwrite:
+            log.info({ 'delete training graph': f })
+            os.remove(f)
+
+    log.debug({ 'options': 'update' })
+    await post('/sdapi/v1/options', options)
+    return
+
+
+async def create(params):
+    log.debug({ 'create start' })
+    if not os.path.isdir(args.preprocess.process_dst):
+        log.error({ 'train source not found': args.preprocess.process_dst })
+        exit()
+    if params.vectors == -1: # dynamically determine number of vectors depending on number of input images
+        if len(images) <= 20:
+            vectors = 2
+        elif len(images) <= 100:
+            vectors = 4
+        else:
+            vectors = 6
+    if os.path.exists(params.name) and os.path.isfile(params.name):
+        log.info({ 'deleting existing embedding': { 'name': params.name } })
+        os.remove(params.name)
+    args.create_embedding.name = params.name
+    args.create_embedding.init_text = params.init
+    args.create_embedding.num_vectors_per_token = vectors
+    log.debug({ 'create args': args.create_embedding })
+    res = await post('/sdapi/v1/create/embedding', args.create_embedding)
+    log.info({ 'create embedding': { 'name': params.name, 'init': params.init, 'vectors': vectors, 'message': res.info } })
+    log.debug({ 'create end' })
+    return params.name
+
+
+async def train(params):
+    log.debug({ 'train start' })
+    args.train_embedding.embedding_name = params.name
+    args.train_embedding.data_root = args.preprocess.process_dst
+    if params.accumulation > -1:
+        args.train_embedding.gradient_step = params.accumulation
+    log.info({ 'train': {
+        'name': params.name,
+        'source': args.preprocess.process_dst,
+        'steps': args.train_embedding.steps,
+        'batch': args.train_embedding.batch_size,
+        'accumulation': args.train_embedding.gradient_step,
+        'learning-rate': args.train_embedding.learn_rate }
+    })
+    log.debug({ 'train args': args.train_embedding })
+    res = await post('/sdapi/v1/train/embedding', args.train_embedding)
+    log.debug({ 'train end': res.info })
+    return
+
+
+async def pipeline(params):
+    log.debug({ 'pipeline start' })
+
+    # interrupt
+    await interrupt()
+
+    # preproceess
+    if not params.skippreprocess:
+        num = await preprocess(params)
+        if num == 0:
+            log.warning({ 'preprocess': 'no resulting images'})
+            return
+    else:
+        args.preprocess.process_dst = params.src
+
+    # create embedding
+    name = await create(params)
+    if not params.name in name:
+        log.error({ 'create embedding failed': name })
+        return
+
+    # train embedding
+    await train(params)
+
+    await plotloss(params)
+
+    log.debug({ 'pipeline end' })
+    return
+
+
+async def monitor(params):
+    step = 0
+    t0 = time.perf_counter()
+    t1 = time.perf_counter()
+    finished = 0
+    while True:
+        await asyncio.sleep(10)
+        res = await progress()
+        if (res.state.job_count == params.steps and res.state.job_no >= res.state.job_count) or (res.eta_relative < 0) or (res.interrupted) or (res.state.job_count == 0): # need exit case if interrupted or failed
+            if res.interrupted:
+                log.info({ 'monitor interrupted': { 'embedding': params.name } })
+                break # exit for monitor job
+            else:
+                finished += 1
+            if finished >= 2: # do it more than once since preprocessing job can finish just in time for monitor to finish
+                log.info({ 'monitor finished': { 'embedding': params.name } })
+                break
+        else:
+            if res.state.job_no == 0:
+                step = 0
+                t0 = time.perf_counter()
+                t1 = time.perf_counter()
+            try:
+                if 'Loss:' in res.textinfo:
+                    text = res.textinfo.split('<br/>')[0].split()
+                    loss = float(text[-1])
+                else:
+                    loss = -1
+            except:
+                loss = -1
+            if math.isnan(loss):
+                log.error({ 'monitor': { 'progress': round(100 * res.progress), 'embedding': params.name, 'eta': round(res.eta_relative), 'step': res.state.job_no, 'steps': res.state.job_count, 'loss': 'nan' } })
+                await interrupt()
+            else:
+                elapsed = t1 - t0
+                log.info({ 'monitor': {
+                    'job': res.state.job,
+                    'progress': round(100 * res.progress),
+                    'embedding': params.name,
+                    'epoch': (1 + res.state.job_no // len(images)) if len(images) > 0 else 'n/a',
+                    'step': res.state.job_no,
+                    'steps': res.state.job_count,
+                    'loss': loss if loss > -1 else 'n/a',
+                    'total': round(1.0 * elapsed * res.state.job_count / res.state.job_no) if res.state.job_no > 0 and t1 != t0 else 'n/a',
+                    'elapsed': round(elapsed),
+                    'remaining': round(res.eta_relative),
+                    'it/s': round((res.state.job_no - step) / (time.perf_counter() - t1), 2) }
+                })
+                if step % 10 == 0:
+                    await plotloss(params)
+                step = res.state.job_no
+            t1 = time.perf_counter()
+    return
+
+
+async def main():
+    parser = argparse.ArgumentParser(description="sd train pipeline")
+    parser.add_argument("--config", type = str, default = 'train.json', required = False, help = "configuration file, default: %(default)s")
+    parser.add_argument("--name", type = str, required = True, help = "embedding name, set to auto to use src folder name")
+    parser.add_argument("--src", type = str, required = True, help = "source image folder or movie file")
+    parser.add_argument("--init", type = str, default = "person", required = False, help = "initialization class, default: %(default)s")
+    parser.add_argument("--dst", type = str, default = "/tmp", required = False, help = "destination image folder for processed images, default: %(default)s")
+    parser.add_argument("--steps", type = int, default = -1, required = False, help = "training steps, default: %(default)s")
+    parser.add_argument("--vectors", type = int, default = -1, required = False, help = "number of vectors per token, default: dynamic")
+    parser.add_argument("--batch", type = int, default = -1, required = False, help = "batch size, default: %(default)s")
+    parser.add_argument("--rate", type = str, default = "", required = False, help = "learning rate, default: dynamic")
+    parser.add_argument("--accumulation", type = int, default = 10, required = False, help = "accumulate gradient over n steps, default: dynamic")
+    parser.add_argument("--type", choices = ['subject', 'style'], default = 'subject', required = False, help = "subject or style, default: %(default)s")
+    parser.add_argument('--overwrite', default = False, action='store_true', help = "overwrite existing embedding, default: %(default)s")
+    parser.add_argument('--skipcaption', default = False, action='store_true', help = "do not auto-generate captions, default: %(default)s")
+    parser.add_argument('--skippreprocess', default = False, action='store_true', help = "skip preprocessing, default: %(default)s")
+    parser.add_argument("--skipstart", type = float, default = -1, required = False, help = "if processing video skip first n seconds, default: %(default)s")
+    parser.add_argument("--skipend", type = float, default = -1, required = False, help = "if processing video skip last n seconds, default: %(default)s")
+    parser.add_argument('--nocleanup', default = False, action='store_true', help = "skip cleanup after completion, default: %(default)s")
+    parser.add_argument('--debug', default = False, action='store_true', help = "print extra debug information, default: %(default)s")
+    params = parser.parse_args()
+    if params.debug:
+        log.setLevel(logging.DEBUG)
+        log.debug({ 'debug': True })
+    log.debug({ 'args': params.__dict__ })
+    home = pathlib.Path(sys.argv[0]).parent
+    global args # pylint: disable=global-statement
+    if os.path.isfile(params.config):
+        try:
+            with open(params.config, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                args = Map(data) # pylint: disable=redefined-outer-name
+                log.debug({ 'config': args })
+        except Exception as e:
+            log.error({ 'config error': params.config, 'exception': e })
+            exit()
+    elif os.path.isfile(os.path.join(home, params.config)):
+        try:
+            with open(os.path.join(home, params.config), 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                args = Map(data) # pylint: disable=redefined-outer-name
+                log.debug({ 'config': args })
+        except Exception as e:
+            log.error({ 'config error': params.config, 'exception': e })
+            exit()
+    else:
+        log.error({ 'config file not found': params.config})
+        exit()
+    if params.skipstart > -1:
+        args.extract_video.skipstart = params.skipstart
+    if params.skipend > -1:
+        args.extract_video.skipend = params.skipend
+    if params.steps > -1:
+        args.train_embedding.steps = params.steps
+    if params.batch > -1:
+        args.train_embedding.batch_size = params.batch
+    if params.rate != '':
+        args.train_embedding.learn_rate = params.rate
+    if params.type == 'subject':
+        if params.skipcaption:
+            args.train_embedding.template_filename = 'subject.txt'
+            args.preprocess.process_caption = False
+        else:
+            args.train_embedding.template_filename = 'subject_filewords.txt'
+    if params.type == 'style':
+        if params.skipcaption:
+            args.train_embedding.template_filename = 'style.txt'
+            args.preprocess.process_caption = False
+        else:
+            args.train_embedding.template_filename = 'style_filewords.txt'
+    if params.name == 'auto':
+        params.name = pathlib.PurePath(params.src).name
+        log.info({ 'training name': params.name })
+    if params.dst == "/tmp":
+        params.dst = os.path.join("/tmp/train", params.name)
+    log.info({ 'args': params.__dict__ })
+    params.src = os.path.abspath(params.src)
+    params.dst = os.path.abspath(params.dst)
+    try:
+        await session()
+        await check(params)
+        a = asyncio.create_task(pipeline(params))
+        b = asyncio.create_task(monitor(params))
+        await asyncio.gather(a, b) # wait for both pipeline and monitor to finish
+    except Exception as e:
+        log.error({ 'exception': e })
+    finally:
+        await cleanup(params)
+        await close()
+    return
+
+if __name__ == "__main__": # create & train test embedding when used from cli
+    log.info({ 'train embedding' })
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        log.warning({ 'interrupted': 'keyboard request' })
+        # asyncio.run(interrupt())

cli/util.py

@@ -0,0 +1,49 @@
+#!/bin/env python
+"""
+generic helper methods
+"""
+
+import logging
+
+logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s')
+log = logging.getLogger("sd")
+
+class Map(dict):
+    def __init__(self, *args, **kwargs):
+        super(Map, self).__init__(*args, **kwargs)
+        for arg in args:
+            if isinstance(arg, dict):
+                for k, v in arg.items():
+                    if isinstance(v, dict):
+                        v = Map(v)
+                    if isinstance(v, list):
+                        self.__convert(v)
+                    self[k] = v
+        if kwargs:
+            for k, v in kwargs.items():
+                if isinstance(v, dict):
+                    v = Map(v)
+                elif isinstance(v, list):
+                    self.__convert(v)
+                self[k] = v
+    def __convert(self, v):
+        for elem in range(0, len(v)): # pylint: disable=consider-using-enumerate
+            if isinstance(v[elem], dict):
+                v[elem] = Map(v[elem])
+            elif isinstance(v[elem], list):
+                self.__convert(v[elem])
+    def __getattr__(self, attr):
+        return self.get(attr)
+    def __setattr__(self, key, value):
+        self.__setitem__(key, value)
+    def __setitem__(self, key, value):
+        super(Map, self).__setitem__(key, value)
+        self.__dict__.update({key: value})
+    def __delattr__(self, item):
+        self.__delitem__(item)
+    def __delitem__(self, key):
+        super(Map, self).__delitem__(key)
+        del self.__dict__[key]
+
+if __name__ == "__main__":
+    pass

cli/versions.py

@@ -0,0 +1,33 @@
+#!/bin/env python
+"""
+print module versions
+"""
+
+import importlib
+import pkg_resources
+
+modules = [
+    'diffusers', 'xformers', 'tokenizers', 'accelerate', 'safetensors'
+]
+
+def get_torch():
+    try:
+        torch = importlib.import_module('torch')
+        print('torch:', { 'version': torch.__version__ })
+        print('cuda:', { 'available': torch.cuda.is_available(), 'version': torch.version.cuda, 'arch': torch.cuda.get_arch_list() })
+        print('device:', { 'name': torch.cuda.get_device_name(torch.cuda.current_device()) })
+    except Exception as err:
+        print('torch:', { 'error': err })
+
+
+def version(name: str):
+    try:
+        ver = pkg_resources.get_distribution(name).version
+        print(f"{name}: {ver}")
+    except Exception as err:
+        print(f"{name} error: {err}")
+
+if __name__ == "__main__": # create & train test embedding when used from cli
+    get_torch()
+    for module in modules:
+        version(module)