mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-03-12 20:42:45 -07:00
45582d129d
with XSeg model you can train your own mask segmentator of dst(and src) faces that will be used in merger for whole_face. Instead of using a pretrained model (which does not exist), you control which part of faces should be masked. Workflow is not easy, but at the moment it is the best solution for obtaining the best quality of whole_face's deepfakes using minimum effort without rotoscoping in AfterEffects. new scripts: XSeg) data_dst edit.bat XSeg) data_dst merge.bat XSeg) data_dst split.bat XSeg) data_src edit.bat XSeg) data_src merge.bat XSeg) data_src split.bat XSeg) train.bat Usage: unpack dst faceset if packed run XSeg) data_dst split.bat this scripts extracts (previously saved) .json data from jpg faces to use in label tool. run XSeg) data_dst edit.bat new tool 'labelme' is used use polygon (CTRL-N) to mask the face name polygon "1" (one symbol) as include polygon name polygon "0" (one symbol) as exclude polygon 'exclude polygons' will be applied after all 'include polygons' Hot keys: ctrl-N create polygon ctrl-J edit polygon A/D navigate between frames ctrl + mousewheel image zoom mousewheel vertical scroll alt+mousewheel horizontal scroll repeat for 10/50/100 faces, you don't need to mask every frame of dst, only frames where the face is different significantly, for example: closed eyes changed head direction changed light the more various faces you mask, the more quality you will get Start masking from the upper left area and follow the clockwise direction. Keep the same logic of masking for all frames, for example: the same approximated jaw line of the side faces, where the jaw is not visible the same hair line Mask the obstructions using polygon with name "0". run XSeg) data_dst merge.bat this script merges .json data of polygons into jpg faces, therefore faceset can be sorted or packed as usual. run XSeg) train.bat train the model Check the faces of 'XSeg dst faces' preview. if some faces have wrong or glitchy mask, then repeat steps: split run edit find these glitchy faces and mask them merge train further or restart training from scratch Restart training of XSeg model is only possible by deleting all 'model\XSeg_*' files. If you want to get the mask of the predicted face in merger, you should repeat the same steps for src faceset. New mask modes available in merger for whole_face: XSeg-prd - XSeg mask of predicted face -> faces from src faceset should be labeled XSeg-dst - XSeg mask of dst face -> faces from dst faceset should be labeled XSeg-prd*XSeg-dst - the smallest area of both if workspace\model folder contains trained XSeg model, then merger will use it, otherwise you will get transparent mask by using XSeg-* modes. Some screenshots: label tool: https://i.imgur.com/aY6QGw1.jpg trainer : https://i.imgur.com/NM1Kn3s.jpg merger : https://i.imgur.com/glUzFQ8.jpg example of the fake using 13 segmented dst faces : https://i.imgur.com/wmvyizU.gifv
599 lines
25 KiB
Python
599 lines
25 KiB
Python
import colorsys
|
|
import inspect
|
|
import json
|
|
import operator
|
|
import os
|
|
import pickle
|
|
import shutil
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from core import imagelib
|
|
from core.interact import interact as io
|
|
from core.leras import nn
|
|
from samplelib import SampleGeneratorBase
|
|
from core import pathex
|
|
from core.cv2ex import *
|
|
|
|
|
|
class ModelBase(object):
|
|
def __init__(self, is_training=False,
|
|
saved_models_path=None,
|
|
training_data_src_path=None,
|
|
training_data_dst_path=None,
|
|
pretraining_data_path=None,
|
|
pretrained_model_path=None,
|
|
no_preview=False,
|
|
force_model_name=None,
|
|
force_gpu_idxs=None,
|
|
cpu_only=False,
|
|
debug=False,
|
|
force_model_class_name=None,
|
|
**kwargs):
|
|
self.is_training = is_training
|
|
self.saved_models_path = saved_models_path
|
|
self.training_data_src_path = training_data_src_path
|
|
self.training_data_dst_path = training_data_dst_path
|
|
self.pretraining_data_path = pretraining_data_path
|
|
self.pretrained_model_path = pretrained_model_path
|
|
self.no_preview = no_preview
|
|
self.debug = debug
|
|
|
|
self.model_class_name = model_class_name = Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]
|
|
|
|
if force_model_class_name is None:
|
|
if force_model_name is not None:
|
|
self.model_name = force_model_name
|
|
else:
|
|
while True:
|
|
# gather all model dat files
|
|
saved_models_names = []
|
|
for filepath in pathex.get_file_paths(saved_models_path):
|
|
filepath_name = filepath.name
|
|
if filepath_name.endswith(f'{model_class_name}_data.dat'):
|
|
saved_models_names += [ (filepath_name.split('_')[0], os.path.getmtime(filepath)) ]
|
|
|
|
# sort by modified datetime
|
|
saved_models_names = sorted(saved_models_names, key=operator.itemgetter(1), reverse=True )
|
|
saved_models_names = [ x[0] for x in saved_models_names ]
|
|
|
|
if len(saved_models_names) != 0:
|
|
io.log_info ("Choose one of saved models, or enter a name to create a new model.")
|
|
io.log_info ("[r] : rename")
|
|
io.log_info ("[d] : delete")
|
|
io.log_info ("")
|
|
for i, model_name in enumerate(saved_models_names):
|
|
s = f"[{i}] : {model_name} "
|
|
if i == 0:
|
|
s += "- latest"
|
|
io.log_info (s)
|
|
|
|
inp = io.input_str(f"", "0", show_default_value=False )
|
|
model_idx = -1
|
|
try:
|
|
model_idx = np.clip ( int(inp), 0, len(saved_models_names)-1 )
|
|
except:
|
|
pass
|
|
|
|
if model_idx == -1:
|
|
if len(inp) == 1:
|
|
is_rename = inp[0] == 'r'
|
|
is_delete = inp[0] == 'd'
|
|
|
|
if is_rename or is_delete:
|
|
if len(saved_models_names) != 0:
|
|
|
|
if is_rename:
|
|
name = io.input_str(f"Enter the name of the model you want to rename")
|
|
elif is_delete:
|
|
name = io.input_str(f"Enter the name of the model you want to delete")
|
|
|
|
if name in saved_models_names:
|
|
|
|
if is_rename:
|
|
new_model_name = io.input_str(f"Enter new name of the model")
|
|
|
|
for filepath in pathex.get_paths(saved_models_path):
|
|
filepath_name = filepath.name
|
|
|
|
model_filename, remain_filename = filepath_name.split('_', 1)
|
|
if model_filename == name:
|
|
|
|
if is_rename:
|
|
new_filepath = filepath.parent / ( new_model_name + '_' + remain_filename )
|
|
filepath.rename (new_filepath)
|
|
elif is_delete:
|
|
filepath.unlink()
|
|
continue
|
|
|
|
self.model_name = inp
|
|
else:
|
|
self.model_name = saved_models_names[model_idx]
|
|
|
|
else:
|
|
self.model_name = io.input_str(f"No saved models found. Enter a name of a new model", "new")
|
|
self.model_name = self.model_name.replace('_', ' ')
|
|
break
|
|
|
|
|
|
self.model_name = self.model_name + '_' + self.model_class_name
|
|
else:
|
|
self.model_name = force_model_class_name
|
|
|
|
self.iter = 0
|
|
self.options = {}
|
|
self.loss_history = []
|
|
self.sample_for_preview = None
|
|
self.choosed_gpu_indexes = None
|
|
|
|
model_data = {}
|
|
self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
|
|
if self.model_data_path.exists():
|
|
io.log_info (f"Loading {self.model_name} model...")
|
|
model_data = pickle.loads ( self.model_data_path.read_bytes() )
|
|
self.iter = model_data.get('iter',0)
|
|
if self.iter != 0:
|
|
self.options = model_data['options']
|
|
self.loss_history = model_data.get('loss_history', [])
|
|
self.sample_for_preview = model_data.get('sample_for_preview', None)
|
|
self.choosed_gpu_indexes = model_data.get('choosed_gpu_indexes', None)
|
|
|
|
if self.is_first_run():
|
|
io.log_info ("\nModel first run.")
|
|
|
|
self.device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(suggest_best_multi_gpu=True)) \
|
|
if not cpu_only else nn.DeviceConfig.CPU()
|
|
|
|
nn.initialize(self.device_config)
|
|
|
|
####
|
|
self.default_options_path = saved_models_path / f'{self.model_class_name}_default_options.dat'
|
|
self.default_options = {}
|
|
if self.default_options_path.exists():
|
|
try:
|
|
self.default_options = pickle.loads ( self.default_options_path.read_bytes() )
|
|
except:
|
|
pass
|
|
|
|
self.choose_preview_history = False
|
|
self.batch_size = self.load_or_def_option('batch_size', 1)
|
|
#####
|
|
|
|
io.input_skip_pending()
|
|
self.on_initialize_options()
|
|
|
|
if self.is_first_run():
|
|
# save as default options only for first run model initialize
|
|
self.default_options_path.write_bytes( pickle.dumps (self.options) )
|
|
|
|
self.autobackup_hour = self.options.get('autobackup_hour', 0)
|
|
self.write_preview_history = self.options.get('write_preview_history', False)
|
|
self.target_iter = self.options.get('target_iter',0)
|
|
self.random_flip = self.options.get('random_flip',True)
|
|
|
|
self.on_initialize()
|
|
self.options['batch_size'] = self.batch_size
|
|
|
|
if self.is_training:
|
|
self.preview_history_path = self.saved_models_path / ( f'{self.get_model_name()}_history' )
|
|
self.autobackups_path = self.saved_models_path / ( f'{self.get_model_name()}_autobackups' )
|
|
|
|
if self.write_preview_history or io.is_colab():
|
|
if not self.preview_history_path.exists():
|
|
self.preview_history_path.mkdir(exist_ok=True)
|
|
else:
|
|
if self.iter == 0:
|
|
for filename in pathex.get_image_paths(self.preview_history_path):
|
|
Path(filename).unlink()
|
|
|
|
if self.generator_list is None:
|
|
raise ValueError( 'You didnt set_training_data_generators()')
|
|
else:
|
|
for i, generator in enumerate(self.generator_list):
|
|
if not isinstance(generator, SampleGeneratorBase):
|
|
raise ValueError('training data generator is not subclass of SampleGeneratorBase')
|
|
|
|
self.update_sample_for_preview(choose_preview_history=self.choose_preview_history)
|
|
|
|
if self.autobackup_hour != 0:
|
|
self.autobackup_start_time = time.time()
|
|
|
|
if not self.autobackups_path.exists():
|
|
self.autobackups_path.mkdir(exist_ok=True)
|
|
|
|
io.log_info( self.get_summary_text() )
|
|
|
|
def update_sample_for_preview(self, choose_preview_history=False, force_new=False):
|
|
if self.sample_for_preview is None or choose_preview_history or force_new:
|
|
if choose_preview_history and io.is_support_windows():
|
|
io.log_info ("Choose image for the preview history. [p] - next. [enter] - confirm.")
|
|
wnd_name = "[p] - next. [enter] - confirm."
|
|
io.named_window(wnd_name)
|
|
io.capture_keys(wnd_name)
|
|
choosed = False
|
|
while not choosed:
|
|
self.sample_for_preview = self.generate_next_samples()
|
|
preview = self.get_static_preview()
|
|
io.show_image( wnd_name, (preview*255).astype(np.uint8) )
|
|
|
|
while True:
|
|
key_events = io.get_key_events(wnd_name)
|
|
key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (0,0,False,False,False)
|
|
if key == ord('\n') or key == ord('\r'):
|
|
choosed = True
|
|
break
|
|
elif key == ord('p'):
|
|
break
|
|
|
|
try:
|
|
io.process_messages(0.1)
|
|
except KeyboardInterrupt:
|
|
choosed = True
|
|
|
|
io.destroy_window(wnd_name)
|
|
else:
|
|
self.sample_for_preview = self.generate_next_samples()
|
|
|
|
try:
|
|
self.get_static_preview()
|
|
except:
|
|
self.sample_for_preview = self.generate_next_samples()
|
|
|
|
self.last_sample = self.sample_for_preview
|
|
|
|
def load_or_def_option(self, name, def_value):
|
|
options_val = self.options.get(name, None)
|
|
if options_val is not None:
|
|
return options_val
|
|
|
|
def_opt_val = self.default_options.get(name, None)
|
|
if def_opt_val is not None:
|
|
return def_opt_val
|
|
|
|
return def_value
|
|
|
|
def ask_override(self):
|
|
return self.is_training and self.iter != 0 and io.input_in_time ("Press enter in 2 seconds to override model settings.", 5 if io.is_colab() else 2 )
|
|
|
|
def ask_autobackup_hour(self, default_value=0):
|
|
default_autobackup_hour = self.options['autobackup_hour'] = self.load_or_def_option('autobackup_hour', default_value)
|
|
self.options['autobackup_hour'] = io.input_int(f"Autobackup every N hour", default_autobackup_hour, add_info="0..24", help_message="Autobackup model files with preview every N hour. Latest backup located in model/<>_autobackups/01")
|
|
|
|
def ask_write_preview_history(self, default_value=False):
|
|
default_write_preview_history = self.load_or_def_option('write_preview_history', default_value)
|
|
self.options['write_preview_history'] = io.input_bool(f"Write preview history", default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
|
|
|
|
if self.options['write_preview_history']:
|
|
if io.is_support_windows():
|
|
self.choose_preview_history = io.input_bool("Choose image for the preview history", False)
|
|
elif io.is_colab():
|
|
self.choose_preview_history = io.input_bool("Randomly choose new image for preview history", False, help_message="Preview image history will stay stuck with old faces if you reuse the same model on different celebs. Choose no unless you are changing src/dst to a new person")
|
|
|
|
def ask_target_iter(self, default_value=0):
|
|
default_target_iter = self.load_or_def_option('target_iter', default_value)
|
|
self.options['target_iter'] = max(0, io.input_int("Target iteration", default_target_iter))
|
|
|
|
def ask_random_flip(self):
|
|
default_random_flip = self.load_or_def_option('random_flip', True)
|
|
self.options['random_flip'] = io.input_bool("Flip faces randomly", default_random_flip, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.")
|
|
|
|
def ask_batch_size(self, suggest_batch_size=None):
|
|
default_batch_size = self.load_or_def_option('batch_size', suggest_batch_size or self.batch_size)
|
|
self.options['batch_size'] = self.batch_size = max(0, io.input_int("Batch_size", default_batch_size, help_message="Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."))
|
|
|
|
|
|
#overridable
|
|
def on_initialize_options(self):
|
|
pass
|
|
|
|
#overridable
|
|
def on_initialize(self):
|
|
'''
|
|
initialize your models
|
|
|
|
store and retrieve your model options in self.options['']
|
|
|
|
check example
|
|
'''
|
|
pass
|
|
|
|
#overridable
|
|
def onSave(self):
|
|
#save your models here
|
|
pass
|
|
|
|
#overridable
|
|
def onTrainOneIter(self, sample, generator_list):
|
|
#train your models here
|
|
|
|
#return array of losses
|
|
return ( ('loss_src', 0), ('loss_dst', 0) )
|
|
|
|
#overridable
|
|
def onGetPreview(self, sample):
|
|
#you can return multiple previews
|
|
#return [ ('preview_name',preview_rgb), ... ]
|
|
return []
|
|
|
|
#overridable if you want model name differs from folder name
|
|
def get_model_name(self):
|
|
return self.model_name
|
|
|
|
#overridable , return [ [model, filename],... ] list
|
|
def get_model_filename_list(self):
|
|
return []
|
|
|
|
#overridable
|
|
def get_MergerConfig(self):
|
|
#return predictor_func, predictor_input_shape, MergerConfig() for the model
|
|
raise NotImplementedError
|
|
|
|
def get_pretraining_data_path(self):
|
|
return self.pretraining_data_path
|
|
|
|
def get_target_iter(self):
|
|
return self.target_iter
|
|
|
|
def is_reached_iter_goal(self):
|
|
return self.target_iter != 0 and self.iter >= self.target_iter
|
|
|
|
def get_previews(self):
|
|
return self.onGetPreview ( self.last_sample )
|
|
|
|
def get_static_preview(self):
|
|
return self.onGetPreview (self.sample_for_preview)[0][1] #first preview, and bgr
|
|
|
|
def save(self):
|
|
Path( self.get_summary_path() ).write_text( self.get_summary_text() )
|
|
|
|
self.onSave()
|
|
|
|
model_data = {
|
|
'iter': self.iter,
|
|
'options': self.options,
|
|
'loss_history': self.loss_history,
|
|
'sample_for_preview' : self.sample_for_preview,
|
|
'choosed_gpu_indexes' : self.choosed_gpu_indexes,
|
|
}
|
|
pathex.write_bytes_safe (self.model_data_path, pickle.dumps(model_data) )
|
|
|
|
if self.autobackup_hour != 0:
|
|
diff_hour = int ( (time.time() - self.autobackup_start_time) // 3600 )
|
|
|
|
if diff_hour > 0 and diff_hour % self.autobackup_hour == 0:
|
|
self.autobackup_start_time += self.autobackup_hour*3600
|
|
self.create_backup()
|
|
|
|
def create_backup(self):
|
|
io.log_info ("Creating backup...", end='\r')
|
|
|
|
if not self.autobackups_path.exists():
|
|
self.autobackups_path.mkdir(exist_ok=True)
|
|
|
|
bckp_filename_list = [ self.get_strpath_storage_for_file(filename) for _, filename in self.get_model_filename_list() ]
|
|
bckp_filename_list += [ str(self.get_summary_path()), str(self.model_data_path) ]
|
|
|
|
for i in range(24,0,-1):
|
|
idx_str = '%.2d' % i
|
|
next_idx_str = '%.2d' % (i+1)
|
|
|
|
idx_backup_path = self.autobackups_path / idx_str
|
|
next_idx_packup_path = self.autobackups_path / next_idx_str
|
|
|
|
if idx_backup_path.exists():
|
|
if i == 24:
|
|
pathex.delete_all_files(idx_backup_path)
|
|
else:
|
|
next_idx_packup_path.mkdir(exist_ok=True)
|
|
pathex.move_all_files (idx_backup_path, next_idx_packup_path)
|
|
|
|
if i == 1:
|
|
idx_backup_path.mkdir(exist_ok=True)
|
|
for filename in bckp_filename_list:
|
|
shutil.copy ( str(filename), str(idx_backup_path / Path(filename).name) )
|
|
|
|
previews = self.get_previews()
|
|
plist = []
|
|
for i in range(len(previews)):
|
|
name, bgr = previews[i]
|
|
plist += [ (bgr, idx_backup_path / ( ('preview_%s.jpg') % (name)) ) ]
|
|
|
|
for preview, filepath in plist:
|
|
preview_lh = ModelBase.get_loss_history_preview(self.loss_history, self.iter, preview.shape[1], preview.shape[2])
|
|
img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8)
|
|
cv2_imwrite (filepath, img )
|
|
|
|
def debug_one_iter(self):
|
|
images = []
|
|
for generator in self.generator_list:
|
|
for i,batch in enumerate(next(generator)):
|
|
if len(batch.shape) == 4:
|
|
images.append( batch[0] )
|
|
|
|
return imagelib.equalize_and_stack_square (images)
|
|
|
|
def generate_next_samples(self):
|
|
sample = []
|
|
for generator in self.generator_list:
|
|
if generator.is_initialized():
|
|
sample.append ( generator.generate_next() )
|
|
else:
|
|
sample.append ( [] )
|
|
self.last_sample = sample
|
|
return sample
|
|
|
|
def train_one_iter(self):
|
|
|
|
iter_time = time.time()
|
|
losses = self.onTrainOneIter()
|
|
iter_time = time.time() - iter_time
|
|
|
|
self.loss_history.append ( [float(loss[1]) for loss in losses] )
|
|
|
|
if self.iter % 10 == 0:
|
|
plist = []
|
|
|
|
if io.is_colab():
|
|
previews = self.get_previews()
|
|
for i in range(len(previews)):
|
|
name, bgr = previews[i]
|
|
plist += [ (bgr, self.get_strpath_storage_for_file('preview_%s.jpg' % (name) ) ) ]
|
|
|
|
if self.write_preview_history:
|
|
plist += [ (self.get_static_preview(), str (self.preview_history_path / ('%.6d.jpg' % (self.iter))) ) ]
|
|
|
|
for preview, filepath in plist:
|
|
preview_lh = ModelBase.get_loss_history_preview(self.loss_history, self.iter, preview.shape[1], preview.shape[2])
|
|
img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8)
|
|
cv2_imwrite (filepath, img )
|
|
|
|
self.iter += 1
|
|
|
|
return self.iter, iter_time
|
|
|
|
def pass_one_iter(self):
|
|
self.generate_next_samples()
|
|
|
|
def finalize(self):
|
|
nn.close_session()
|
|
|
|
def is_first_run(self):
|
|
return self.iter == 0
|
|
|
|
def is_debug(self):
|
|
return self.debug
|
|
|
|
def set_batch_size(self, batch_size):
|
|
self.batch_size = batch_size
|
|
|
|
def get_batch_size(self):
|
|
return self.batch_size
|
|
|
|
def get_iter(self):
|
|
return self.iter
|
|
|
|
def set_iter(self, iter):
|
|
self.iter = iter
|
|
self.loss_history = self.loss_history[:iter]
|
|
|
|
def get_loss_history(self):
|
|
return self.loss_history
|
|
|
|
def set_training_data_generators (self, generator_list):
|
|
self.generator_list = generator_list
|
|
|
|
def get_training_data_generators (self):
|
|
return self.generator_list
|
|
|
|
def get_model_root_path(self):
|
|
return self.saved_models_path
|
|
|
|
def get_strpath_storage_for_file(self, filename):
|
|
return str( self.saved_models_path / ( self.get_model_name() + '_' + filename) )
|
|
|
|
def get_summary_path(self):
|
|
return self.get_strpath_storage_for_file('summary.txt')
|
|
|
|
def get_summary_text(self):
|
|
###Generate text summary of model hyperparameters
|
|
#Find the longest key name and value string. Used as column widths.
|
|
width_name = max([len(k) for k in self.options.keys()] + [17]) + 1 # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration"
|
|
width_value = max([len(str(x)) for x in self.options.values()] + [len(str(self.get_iter())), len(self.get_model_name())]) + 1 # Single space buffer to right edge
|
|
if len(self.device_config.devices) != 0: #Check length of GPU names
|
|
width_value = max([len(device.name)+1 for device in self.device_config.devices] + [width_value])
|
|
width_total = width_name + width_value + 2 #Plus 2 for ": "
|
|
|
|
summary_text = []
|
|
summary_text += [f'=={" Model Summary ":=^{width_total}}=='] # Model/status summary
|
|
summary_text += [f'=={" "*width_total}==']
|
|
summary_text += [f'=={"Model name": >{width_name}}: {self.get_model_name(): <{width_value}}=='] # Name
|
|
summary_text += [f'=={" "*width_total}==']
|
|
summary_text += [f'=={"Current iteration": >{width_name}}: {str(self.get_iter()): <{width_value}}=='] # Iter
|
|
summary_text += [f'=={" "*width_total}==']
|
|
|
|
summary_text += [f'=={" Model Options ":-^{width_total}}=='] # Model options
|
|
summary_text += [f'=={" "*width_total}==']
|
|
for key in self.options.keys():
|
|
summary_text += [f'=={key: >{width_name}}: {str(self.options[key]): <{width_value}}=='] # self.options key/value pairs
|
|
summary_text += [f'=={" "*width_total}==']
|
|
|
|
summary_text += [f'=={" Running On ":-^{width_total}}=='] # Training hardware info
|
|
summary_text += [f'=={" "*width_total}==']
|
|
if len(self.device_config.devices) == 0:
|
|
summary_text += [f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='] # cpu_only
|
|
else:
|
|
for device in self.device_config.devices:
|
|
summary_text += [f'=={"Device index": >{width_name}}: {device.index: <{width_value}}=='] # GPU hardware device index
|
|
summary_text += [f'=={"Name": >{width_name}}: {device.name: <{width_value}}=='] # GPU name
|
|
vram_str = f'{device.total_mem_gb:.2f}GB' # GPU VRAM - Formated as #.## (or ##.##)
|
|
summary_text += [f'=={"VRAM": >{width_name}}: {vram_str: <{width_value}}==']
|
|
summary_text += [f'=={" "*width_total}==']
|
|
summary_text += [f'=={"="*width_total}==']
|
|
summary_text = "\n".join (summary_text)
|
|
return summary_text
|
|
|
|
@staticmethod
|
|
def get_loss_history_preview(loss_history, iter, w, c):
|
|
loss_history = np.array (loss_history.copy())
|
|
|
|
lh_height = 100
|
|
lh_img = np.ones ( (lh_height,w,c) ) * 0.1
|
|
|
|
if len(loss_history) != 0:
|
|
loss_count = len(loss_history[0])
|
|
lh_len = len(loss_history)
|
|
|
|
l_per_col = lh_len / w
|
|
plist_max = [ [ max (0.0, loss_history[int(col*l_per_col)][p],
|
|
*[ loss_history[i_ab][p]
|
|
for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
|
|
]
|
|
)
|
|
for p in range(loss_count)
|
|
]
|
|
for col in range(w)
|
|
]
|
|
|
|
plist_min = [ [ min (plist_max[col][p], loss_history[int(col*l_per_col)][p],
|
|
*[ loss_history[i_ab][p]
|
|
for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
|
|
]
|
|
)
|
|
for p in range(loss_count)
|
|
]
|
|
for col in range(w)
|
|
]
|
|
|
|
plist_abs_max = np.mean(loss_history[ len(loss_history) // 5 : ]) * 2
|
|
|
|
for col in range(0, w):
|
|
for p in range(0,loss_count):
|
|
point_color = [1.0]*c
|
|
point_color[0:3] = colorsys.hsv_to_rgb ( p * (1.0/loss_count), 1.0, 1.0 )
|
|
|
|
ph_max = int ( (plist_max[col][p] / plist_abs_max) * (lh_height-1) )
|
|
ph_max = np.clip( ph_max, 0, lh_height-1 )
|
|
|
|
ph_min = int ( (plist_min[col][p] / plist_abs_max) * (lh_height-1) )
|
|
ph_min = np.clip( ph_min, 0, lh_height-1 )
|
|
|
|
for ph in range(ph_min, ph_max+1):
|
|
lh_img[ (lh_height-ph-1), col ] = point_color
|
|
|
|
lh_lines = 5
|
|
lh_line_height = (lh_height-1)/lh_lines
|
|
for i in range(0,lh_lines+1):
|
|
lh_img[ int(i*lh_line_height), : ] = (0.8,)*c
|
|
|
|
last_line_t = int((lh_lines-1)*lh_line_height)
|
|
last_line_b = int(lh_lines*lh_line_height)
|
|
|
|
lh_text = 'Iter: %d' % (iter) if iter != 0 else ''
|
|
|
|
lh_img[last_line_t:last_line_b, 0:w] += imagelib.get_text_image ( (last_line_b-last_line_t,w,c), lh_text, color=[0.8]*c )
|
|
return lh_img
|