File size: 6,929 Bytes
34097e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 |
import os
import os.path
from modules import shared
import modules.scripts as scripts
from scripts import model_util, util
from scripts.model_util import MAX_MODEL_COUNT
LORA_TRAIN_METADATA_NAMES = {
"ss_session_id": "Session ID",
"ss_training_started_at": "Training started at",
"ss_output_name": "Output name",
"ss_learning_rate": "Learning rate",
"ss_text_encoder_lr": "Text encoder LR",
"ss_unet_lr": "UNet LR",
"ss_num_train_images": "# of training images",
"ss_num_reg_images": "# of reg images",
"ss_num_batches_per_epoch": "Batches per epoch",
"ss_num_epochs": "Total epochs",
"ss_epoch": "Epoch",
"ss_batch_size_per_device": "Batch size/device",
"ss_total_batch_size": "Total batch size",
"ss_gradient_checkpointing": "Gradient checkpointing",
"ss_gradient_accumulation_steps": "Gradient accum. steps",
"ss_max_train_steps": "Max train steps",
"ss_lr_warmup_steps": "LR warmup steps",
"ss_lr_scheduler": "LR scheduler",
"ss_network_module": "Network module",
"ss_network_dim": "Network dim",
"ss_network_alpha": "Network alpha",
"ss_mixed_precision": "Mixed precision",
"ss_full_fp16": "Full FP16",
"ss_v2": "V2",
"ss_resolution": "Resolution",
"ss_clip_skip": "Clip skip",
"ss_max_token_length": "Max token length",
"ss_color_aug": "Color aug",
"ss_flip_aug": "Flip aug",
"ss_random_crop": "Random crop",
"ss_shuffle_caption": "Shuffle caption",
"ss_cache_latents": "Cache latents",
"ss_enable_bucket": "Enable bucket",
"ss_min_bucket_reso": "Min bucket reso.",
"ss_max_bucket_reso": "Max bucket reso.",
"ss_seed": "Seed",
"ss_keep_tokens": "Keep tokens",
"ss_dataset_dirs": "Dataset dirs.",
"ss_reg_dataset_dirs": "Reg dataset dirs.",
"ss_sd_model_name": "SD model name",
"ss_vae_name": "VAE name",
"ss_training_comment": "Comment",
}
xy_grid = None # XY Grid module
script_class = None # additional_networks scripts.Script class
axis_params = [{}] * MAX_MODEL_COUNT
def update_axis_params(i, module, model):
axis_params[i] = {"module": module, "model": model}
def get_axis_model_choices(i):
module = axis_params[i].get("module", "None")
model = axis_params[i].get("model", "None")
if module == "LoRA":
if model != "None":
sort_by = shared.opts.data.get("additional_networks_sort_models_by", "name")
return ["None"] + model_util.get_model_list(module, model, "", sort_by)
return [f"select `Model {i+1}` in `Additional Networks`. models in same folder for selected one will be shown here."]
def update_script_args(p, value, arg_idx):
global script_class
for s in scripts.scripts_txt2img.alwayson_scripts:
if isinstance(s, script_class):
args = list(p.script_args)
# print(f"Changed arg {arg_idx} from {args[s.args_from + arg_idx - 1]} to {value}")
args[s.args_from + arg_idx] = value
p.script_args = tuple(args)
break
def confirm_models(p, xs):
for x in xs:
if x in ["", "None"]:
continue
if not model_util.find_closest_lora_model_name(x):
raise RuntimeError(f"Unknown LoRA model: {x}")
def apply_module(p, x, xs, i):
update_script_args(p, True, 0) # set Enabled to True
update_script_args(p, x, 2 + 4 * i) # enabled, separate_weights, ({module}, model, weight_unet, weight_tenc), ...
def apply_model(p, x, xs, i):
name = model_util.find_closest_lora_model_name(x)
update_script_args(p, True, 0)
update_script_args(p, name, 3 + 4 * i) # enabled, separate_weights, (module, {model}, weight_unet, weight_tenc), ...
def apply_weight(p, x, xs, i):
update_script_args(p, True, 0)
update_script_args(p, x, 4 + 4 * i) # enabled, separate_weights, (module, model, {weight_unet, weight_tenc}), ...
update_script_args(p, x, 5 + 4 * i)
def apply_weight_unet(p, x, xs, i):
update_script_args(p, True, 0)
update_script_args(p, x, 4 + 4 * i) # enabled, separate_weights, (module, model, {weight_unet}, weight_tenc), ...
def apply_weight_tenc(p, x, xs, i):
update_script_args(p, True, 0)
update_script_args(p, x, 5 + 4 * i) # enabled, separate_weights, (module, model, weight_unet, {weight_tenc}), ...
def format_lora_model(p, opt, x):
global xy_grid
model = model_util.find_closest_lora_model_name(x)
if model is None or model.lower() in ["", "none"]:
return "None"
value = xy_grid.format_value(p, opt, model)
model_path = model_util.lora_models.get(model)
metadata = model_util.read_model_metadata(model_path, "LoRA")
if not metadata:
return value
metadata_names = util.split_path_list(shared.opts.data.get("additional_networks_xy_grid_model_metadata", ""))
if not metadata_names:
return value
for name in metadata_names:
name = name.strip()
if name in metadata:
formatted_name = LORA_TRAIN_METADATA_NAMES.get(name, name)
value += f"\n{formatted_name}: {metadata[name]}, "
return value.strip(" ").strip(",")
def initialize(script):
global xy_grid, script_class
xy_grid = None
script_class = script
for scriptDataTuple in scripts.scripts_data:
if os.path.basename(scriptDataTuple.path) == "xy_grid.py" or os.path.basename(scriptDataTuple.path) == "xyz_grid.py":
xy_grid = scriptDataTuple.module
for i in range(MAX_MODEL_COUNT):
model = xy_grid.AxisOption(
f"AddNet Model {i+1}",
str,
lambda p, x, xs, i=i: apply_model(p, x, xs, i),
format_lora_model,
confirm_models,
cost=0.5,
choices=lambda i=i: get_axis_model_choices(i),
)
weight = xy_grid.AxisOption(
f"AddNet Weight {i+1}",
float,
lambda p, x, xs, i=i: apply_weight(p, x, xs, i),
xy_grid.format_value_add_label,
None,
cost=0.5,
)
weight_unet = xy_grid.AxisOption(
f"AddNet UNet Weight {i+1}",
float,
lambda p, x, xs, i=i: apply_weight_unet(p, x, xs, i),
xy_grid.format_value_add_label,
None,
cost=0.5,
)
weight_tenc = xy_grid.AxisOption(
f"AddNet TEnc Weight {i+1}",
float,
lambda p, x, xs, i=i: apply_weight_tenc(p, x, xs, i),
xy_grid.format_value_add_label,
None,
cost=0.5,
)
xy_grid.axis_options.extend([model, weight, weight_unet, weight_tenc])
|