update cli

cli/train/latents.py

@@ -44,13 +44,13 @@ options = Map({
 vae = None
 
 
-def get_latents(vae, images, weight_dtype):
+def get_latents(local_vae, images, weight_dtype):
     image_transforms = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.5], [0.5]) ])
     img_tensors = [image_transforms(image) for image in images]
     img_tensors = torch.stack(img_tensors)
     img_tensors = img_tensors.to(device, weight_dtype)
     with torch.no_grad():
-        latents = vae.encode(img_tensors).latent_dist.sample().float().to('cpu').numpy()
+        latents = local_vae.encode(img_tensors).latent_dist.sample().float().to('cpu').numpy()
     return latents
 
 
@@ -58,8 +58,8 @@ def get_npz_filename_wo_ext(data_dir, image_key):
     return os.path.join(data_dir, os.path.splitext(os.path.basename(image_key))[0])
 
 
-def create_vae_latents(params):
-    args = Map({**options, **params})
+def create_vae_latents(local_params):
+    args = Map({**options, **local_params})
     console.log(f'create vae latents args: {args}')
     image_paths = train_util.glob_images(args.input)
     if os.path.exists(args.json):
@@ -73,7 +73,7 @@ def create_vae_latents(params):
         weight_dtype = torch.bfloat16
     else:
         weight_dtype = torch.float32
-    global vae
+    global vae # pylint: disable=global-statement
     if vae is None:
         vae = model_util.load_vae(args.vae, weight_dtype)
         vae.eval()
@@ -142,7 +142,7 @@ def create_vae_latents(params):
 
 
 def unload_vae():
-    global vae
+    global vae # pylint: disable=global-statement
     vae = None
 
 

cli/train/process.py

@@ -1,15 +1,13 @@
+ # pylint: disable=global-statement
 import os
-import sys
 import io
 import math
 import base64
-import pathlib
 import numpy as np
 import mediapipe as mp
 from PIL import Image, ImageOps
 from skimage.metrics import structural_similarity as ssim
 from scipy.stats import beta
-sys.path.append(os.path.join(os.path.dirname(__file__)))
 
 import util
 import sdapi
@@ -23,9 +21,9 @@ all_images_by_type = {}
 
 
 class Result(object):
-    def __init__(self, type: str, input: str, tag: str = None, requested: list = []):
-        self.type = type
-        self.input = input
+    def __init__(self, typ: str, fn: str, tag: str = None, requested: list = []):
+        self.type = typ
+        self.input = fn
         self.output = ''
         self.basename = ''
         self.message = ''
@@ -56,8 +54,8 @@ def detect_dynamicrange(image: Image):
     data = np.asarray(image)
     image = np.float32(data)
     RGB = [0.299, 0.587, 0.114]
-    height, width = image.shape[:2]
-    brightness_image = np.sqrt(image[..., 0] ** 2 * RGB[0] + image[..., 1] ** 2 * RGB[1] + image[..., 2] ** 2 * RGB[2])
+    height, width = image.shape[:2] # pylint: disable=unsubscriptable-object
+    brightness_image = np.sqrt(image[..., 0] ** 2 * RGB[0] + image[..., 1] ** 2 * RGB[1] + image[..., 2] ** 2 * RGB[2]) # pylint: disable=unsubscriptable-object
     hist, _ = np.histogram(brightness_image, bins=256, range=(0, 255))
     img_brightness_pmf = hist / (height * width)
     dist = beta(2, 2)
@@ -264,7 +262,7 @@ def save_image(res: Result, folder: str):
 
 def file(filename: str, folder: str, tag = None, requested = []):
     # initialize result dict
-    res = Result(input = filename, type='unknown', tag=tag, requested = requested)
+    res = Result(fn = filename, typ='unknown', tag=tag, requested = requested)
     # open image
     try:
         res.image = Image.open(filename)

cli/train/sdapi.py

@@ -1,7 +1,5 @@
-import sys
-import json
-import aiohttp
 import asyncio
+import aiohttp
 import requests
 from util import Map
 
@@ -89,9 +87,9 @@ def progress():
 
 
 def options():
-    options = getsync('/sdapi/v1/options')
+    opt = getsync('/sdapi/v1/options')
     flags = getsync('/sdapi/v1/cmd-flags')
-    return { 'options': options, 'flags': flags }
+    return { 'options': opt, 'flags': flags }
 
 
 def shutdown():

cli/train/train.py

@@ -46,8 +46,6 @@ lycoris_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir
 sys.path.append(lycoris_path)
 import train_network
 
-print('HERE6')
-
 # globals
 args = None
 valid_steps = ['original', 'face', 'body', 'blur', 'range', 'upscale', 'restore', 'interrogate', 'resize', 'square', 'segment']
@@ -69,26 +67,29 @@ def mem_stats():
 
 def parse_args():
     global args # pylint: disable=global-statement
-    parser = argparse.ArgumentParser(description = 'train')
-    # basic section
-    parser.add_argument('--type', type=str, choices=['embedding', 'lora', 'lycoris', 'dreambooth'], default=None, required=True, help='training type')
-    parser.add_argument('--name', type=str, default=None, required=True, help='output filename')
-    parser.add_argument('--overwrite', default = False, action='store_true', help = "overwrite existing training, default: %(default)s")
-    parser.add_argument('--tag', type=str, default='person', required=False, help='primary tags, default: %(default)s')
-    parser.add_argument('--input', type=str, default=None, required=True, help='input folder with training images')
-    parser.add_argument('--output', type=str, default='', required=False, help='where to store processed images, default is system temp/train')
-    parser.add_argument('--process', type=str, default='original,interrogate,resize,square', required=False, help=f'list of possible processing steps: {valid_steps}, default: %(default)s')
+    parser = argparse.ArgumentParser(description = 'Train')
 
-    # global params
-    parser.add_argument('--gradient', type=int, default=1, required=False, help='gradient accumulation steps, default: %(default)s')
-    parser.add_argument('--steps', type=int, default=2500, required=False, help='training steps, default: %(default)s')
-    parser.add_argument('--batch', type=int, default=1, required=False, help='batch size, default: %(default)s')
-    parser.add_argument('--lr', type=float, default=1e-04, required=False, help='model learning rate, default: %(default)s')
-    parser.add_argument('--dim', type=int, default=40, required=False, help='network dimension or number of vectors, default: %(default)s')
+    group_main = parser.add_argument_group('Main')
+    group_main.add_argument('--type', type=str, choices=['embedding', 'lora', 'lycoris', 'dreambooth'], default=None, required=True, help='training type')
+    group_main.add_argument('--name', type=str, default=None, required=True, help='output filename')
+    group_main.add_argument('--overwrite', default = False, action='store_true', help = "overwrite existing training, default: %(default)s")
+    group_main.add_argument('--tag', type=str, default='person', required=False, help='primary tags, default: %(default)s')
+
+    group_data = parser.add_argument_group('Dataset')
+    group_data.add_argument('--input', type=str, default=None, required=True, help='input folder with training images')
+    group_data.add_argument('--output', type=str, default='', required=False, help='where to store processed images, default is system temp/train')
+    group_data.add_argument('--process', type=str, default='original,interrogate,resize,square', required=False, help=f'list of possible processing steps: {valid_steps}, default: %(default)s')
+
+    group_train = parser.add_argument_group('Train')
+    group_train.add_argument('--gradient', type=int, default=1, required=False, help='gradient accumulation steps, default: %(default)s')
+    group_train.add_argument('--steps', type=int, default=2500, required=False, help='training steps, default: %(default)s')
+    group_train.add_argument('--batch', type=int, default=1, required=False, help='batch size, default: %(default)s')
+    group_train.add_argument('--lr', type=float, default=1e-04, required=False, help='model learning rate, default: %(default)s')
+    group_train.add_argument('--dim', type=int, default=40, required=False, help='network dimension or number of vectors, default: %(default)s')
 
     # lora params
-    parser.add_argument('--repeats', type=int, default=10, required=False, help='number of repeats per image, default: %(default)s')
-    parser.add_argument('--alpha', type=float, default=0, required=False, help='alpha for weights scaling, default: dim/2')
+    group_train.add_argument('--repeats', type=int, default=10, required=False, help='number of repeats per image, default: %(default)s')
+    group_train.add_argument('--alpha', type=float, default=0, required=False, help='alpha for weights scaling, default: dim/2')
 
     args = parser.parse_args()
 

cli/train/util.py

@@ -42,8 +42,8 @@ def get_memory():
     return Map(mem)
 
 
-class Map(dict):
-    __slots__ = ('__dict__')
+class Map(dict): # pylint: disable=C0205
+    __slots__ = ('__dict__') # pylint: disable=C0325
     def __init__(self, *args, **kwargs):
         super(Map, self).__init__(*args, **kwargs)
         for arg in args: