pull/608/head
Akegarasu 2024-11-23 11:27:43 +08:00
parent b6d33fec7f
commit 5e9a0e8e9b
No known key found for this signature in database
GPG Key ID: DACA951FEBA569A2
6 changed files with 260 additions and 318 deletions

View File

@ -34,6 +34,7 @@ avaliable_scripts = [
]
avaliable_schemas = []
avaliable_presets = []
trainer_mapping = {
"sd-lora": "./scripts/stable/train_network.py",
@ -67,6 +68,18 @@ async def load_schemas():
})
async def load_presets():
avaliable_presets.clear()
preset_dir = os.path.join(os.getcwd(), "config", "presets")
presets = os.listdir(preset_dir)
for preset_name in presets:
with open(os.path.join(preset_dir, preset_name), encoding="utf-8") as f:
content = f.read()
avaliable_presets.append(toml.loads(content))
@router.post("/run")
async def create_toml_file(request: Request):
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
@ -282,6 +295,17 @@ async def get_all_schemas() -> APIResponse:
})
@router.get("/presets")
async def get_presets() -> APIResponse:
if os.environ.get("MIKAZUKI_SCHEMA_HOT_RELOAD", "0") == "1":
log.info("Hot reloading presets")
await load_presets()
return APIResponseSuccess(data={
"presets": avaliable_presets
})
@router.get("/config/saved_params")
async def get_saved_params() -> APIResponse:
saved_params = app_config["saved_params"]

View File

@ -126,4 +126,8 @@ declare namespace Schemastery {
type SchemaI<S = any, T = S> = Schemastery.Schemastery<S, T>;
declare const Schema: Schemastery.Static
declare const Schema: Schemastery.Static
declare const SHARED_SCHEMAS: Dict<any>
declare function UpdateSchema(origin: Record<string, any>, modify?: Record<string, any>, toDelete?: string[]): Record<string, any>;

View File

@ -1,14 +1,12 @@
Schema.intersect([
Schema.intersect([
Schema.object({
model_train_type: Schema.string().default("flux-lora").disabled().description("训练种类"),
pretrained_model_name_or_path: Schema.string().role('filepicker', { type: "model-file" }).default("./sd-models/model.safetensors").description("Flux 模型路径"),
ae: Schema.string().role('filepicker', { type: "model-file" }).description("AE 模型文件路径"),
clip_l: Schema.string().role('filepicker', { type: "model-file" }).description("clip_l 模型文件路径"),
t5xxl: Schema.string().role('filepicker', { type: "model-file" }).description("t5xxl 模型文件路径"),
resume: Schema.string().role('filepicker', { type: "folder" }).description("从某个 `save_state` 保存的中断状态继续训练,填写文件路径"),
}).description("训练用模型"),
]),
Schema.object({
model_train_type: Schema.string().default("flux-lora").disabled().description("训练种类"),
pretrained_model_name_or_path: Schema.string().role('filepicker', { type: "model-file" }).default("./sd-models/model.safetensors").description("Flux 模型路径"),
ae: Schema.string().role('filepicker', { type: "model-file" }).description("AE 模型文件路径"),
clip_l: Schema.string().role('filepicker', { type: "model-file" }).description("clip_l 模型文件路径"),
t5xxl: Schema.string().role('filepicker', { type: "model-file" }).description("t5xxl 模型文件路径"),
resume: Schema.string().role('filepicker', { type: "folder" }).description("从某个 `save_state` 保存的中断状态继续训练,填写文件路径"),
}).description("训练用模型"),
Schema.object({
timestep_sampling: Schema.union(["sigma", "uniform", "sigmoid", "shift"]).default("sigmoid").description("时间步采样"),
@ -18,28 +16,21 @@ Schema.intersect([
loss_type: Schema.union(["l1", "l2", "huber", "smooth_l1"]).default("l2").description("损失函数类型"),
guidance_scale: Schema.number().step(0.01).default(1.0).description("CFG 引导缩放"),
t5xxl_max_token_length: Schema.number().step(1).description("T5XXL 最大 token 长度(不填写使用自动)"),
train_t5xxl: Schema.boolean().default(false).description("训练 T5XXL不推荐"),
}).description("Flux 专用参数"),
Schema.object({
train_data_dir: Schema.string().role('filepicker', { type: "folder", internal: "train-dir" }).default("./train/aki").description("训练数据集路径"),
reg_data_dir: Schema.string().role('filepicker', { type: "folder", internal: "train-dir" }).description("正则化数据集路径。默认留空,不使用正则化图像"),
prior_loss_weight: Schema.number().step(0.1).default(1.0).description("正则化 - 先验损失权重"),
resolution: Schema.string().default("768,768").description("训练图片分辨率宽x高。支持非正方形但必须是 64 倍数。"),
enable_bucket: Schema.boolean().default(true).description("启用 arb 桶以允许非固定宽高比的图片"),
min_bucket_reso: Schema.number().default(256).description("arb 桶最小分辨率"),
max_bucket_reso: Schema.number().default(2048).description("arb 桶最大分辨率"),
bucket_reso_steps: Schema.number().default(64).description("arb 桶分辨率划分单位FLUX 必须为 64 的倍数 (FLUX低于64时无法训练)"),
bucket_no_upscale: Schema.boolean().default(true).description("arb 桶不放大图片"),
}).description("数据集设置"),
Schema.object(
UpdateSchema(SHARED_SCHEMAS.RAW.DATASET_SETTINGS, {
resolution: Schema.string().default("768,768").description("训练图片分辨率宽x高。支持非正方形但必须是 64 倍数。"),
enable_bucket: Schema.boolean().default(true).description("启用 arb 桶以允许非固定宽高比的图片"),
min_bucket_reso: Schema.number().default(256).description("arb 桶最小分辨率"),
max_bucket_reso: Schema.number().default(2048).description("arb 桶最大分辨率"),
bucket_reso_steps: Schema.number().default(64).description("arb 桶分辨率划分单位FLUX 需大于 64"),
})
).description("数据集设置"),
Schema.object({
output_name: Schema.string().default("aki").description("模型保存名称"),
output_dir: Schema.string().role('filepicker', { type: "folder" }).default("./output").description("模型保存文件夹"),
save_model_as: Schema.union(["safetensors", "pt", "ckpt"]).default("safetensors").description("模型保存格式"),
save_precision: Schema.union(["fp16", "float", "bf16"]).default("bf16").description("模型保存精度"),
save_every_n_epochs: Schema.number().default(2).description("每 N epoch自动保存一次模型"),
save_state: Schema.boolean().description("保存训练状态 配合 `resume` 参数可以继续从某个状态训练"),
}).description("保存设置"),
// 保存设置
SHARED_SCHEMAS.SAVE_SETTINGS,
Schema.object({
max_train_epochs: Schema.number().min(1).default(20).description("最大训练 epoch轮数"),
@ -50,67 +41,8 @@ Schema.intersect([
network_train_text_encoder_only: Schema.boolean().default(false).description("仅训练文本编码器"),
}).description("训练相关参数"),
Schema.intersect([
Schema.object({
learning_rate: Schema.string().default("1e-4").description("总学习率, 在分开设置 U-Net 与文本编码器学习率后这个值失效。"),
unet_lr: Schema.string().default("5e-4").description("U-Net 学习率"),
text_encoder_lr: Schema.string().default("1e-5").description("文本编码器学习率"),
lr_scheduler: Schema.union([
"linear",
"cosine",
"cosine_with_restarts",
"polynomial",
"constant",
"constant_with_warmup",
]).default("cosine_with_restarts").description("学习率调度器设置"),
lr_warmup_steps: Schema.number().default(0).description('学习率预热步数'),
}).description("学习率与优化器设置"),
Schema.union([
Schema.object({
lr_scheduler: Schema.const('cosine_with_restarts'),
lr_scheduler_num_cycles: Schema.number().default(1).description('重启次数'),
}),
Schema.object({}),
]),
Schema.object({
optimizer_type: Schema.union([
"AdamW",
"AdamW8bit",
"PagedAdamW8bit",
"AdamWScheduleFree",
"Lion",
"Lion8bit",
"PagedLion8bit",
"SGDNesterov",
"SGDNesterov8bit",
"SGDScheduleFree",
"DAdaptation",
"DAdaptAdam",
"DAdaptAdaGrad",
"DAdaptAdanIP",
"DAdaptLion",
"DAdaptSGD",
"AdaFactor",
"Prodigy"
]).default("PagedAdamW8bit").description("优化器设置"),
min_snr_gamma: Schema.number().step(0.1).description("最小信噪比伽马值, 如果启用推荐为 5"),
}),
Schema.union([
Schema.object({
optimizer_type: Schema.const('Prodigy').required(),
prodigy_d0: Schema.string(),
prodigy_d_coef: Schema.string().default("2.0"),
}),
Schema.object({}),
]),
Schema.object({
optimizer_args_custom: Schema.array(String).role('table').description('自定义 optimizer_args一行一个'),
})
]),
// 学习率&优化器设置
SHARED_SCHEMAS.LR_OPTIMIZER,
Schema.intersect([
Schema.object({
@ -134,82 +66,25 @@ Schema.intersect([
]),
]),
Schema.intersect([
Schema.object({
enable_preview: Schema.boolean().default(false).description('启用训练预览图'),
}).description('训练预览图设置'),
// 预览图设置
SHARED_SCHEMAS.PREVIEW_IMAGE,
Schema.union([
Schema.object({
enable_preview: Schema.const(true).required(),
sample_prompts: Schema.string().role('textarea').default(window.__MIKAZUKI__.SAMPLE_PROMPTS_DEFAULT).description(window.__MIKAZUKI__.SAMPLE_PROMPTS_DESCRIPTION),
sample_sampler: Schema.union(["ddim", "pndm", "lms", "euler", "euler_a", "heun", "dpm_2", "dpm_2_a", "dpmsolver", "dpmsolver++", "dpmsingle", "k_lms", "k_euler", "k_euler_a", "k_dpm_2", "k_dpm_2_a"]).default("euler_a").description("生成预览图所用采样器"),
sample_every_n_epochs: Schema.number().default(2).description("每 N 个 epoch 生成一次预览图"),
}),
Schema.object({}),
]),
]),
// 日志设置
SHARED_SCHEMAS.LOG_SETTINGS,
Schema.intersect([
Schema.object({
log_with: Schema.union(["tensorboard", "wandb"]).default("tensorboard").description("日志模块"),
log_prefix: Schema.string().description("日志前缀"),
log_tracker_name: Schema.string().description("日志追踪器名称"),
logging_dir: Schema.string().default("./logs").description("日志保存文件夹"),
}).description('日志设置'),
// caption 选项
// FLUX 去除 max_token_length
Schema.object(UpdateSchema(SHARED_SCHEMAS.RAW.CAPTION_SETTINGS, {}, ["max_token_length"])).description("captionTag选项"),
Schema.union([
Schema.object({
log_with: Schema.const("wandb").required(),
wandb_api_key: Schema.string().required().description("wandb 的 api 密钥"),
}),
Schema.object({}),
]),
]),
// 数据增强
SHARED_SCHEMAS.DATA_ENCHANCEMENT,
Schema.object({
caption_extension: Schema.string().default(".txt").description("Tag 文件扩展名"),
shuffle_caption: Schema.boolean().default(false).description("训练时随机打乱 tokens"),
weighted_captions: Schema.boolean().description("使用带权重的 token不推荐与 shuffle_caption 一同开启"),
keep_tokens: Schema.number().min(0).max(255).step(1).default(0).description("在随机打乱 tokens 时,保留前 N 个不变"),
keep_tokens_separator: Schema.string().description("保留 tokens 时使用的分隔符"),
// max_token_length: Schema.number().default(255).description("最大 token 长度"),
caption_dropout_rate: Schema.number().min(0).step(0.01).description("丢弃全部标签的概率,对一个图片概率不使用 caption 或 class token"),
caption_dropout_every_n_epochs: Schema.number().min(0).max(100).step(1).description("每 N 个 epoch 丢弃全部标签"),
caption_tag_dropout_rate: Schema.number().min(0).step(0.01).description("按逗号分隔的标签来随机丢弃 tag 的概率"),
}).description("captionTag选项"),
// 其他选项
SHARED_SCHEMAS.OTHER,
Schema.object({
color_aug: Schema.boolean().description("颜色改变"),
flip_aug: Schema.boolean().description("图像翻转"),
random_crop: Schema.boolean().description("随机剪裁"),
}).description("数据增强"),
// 速度优化选项
SHARED_SCHEMAS.PRECISION_CACHE_BATCH,
Schema.object({
seed: Schema.number().default(1337).description("随机种子"),
clip_skip: Schema.number().role("slider").min(0).max(12).step(1).default(2).description("CLIP 跳过层数 *玄学*"),
ui_custom_params: Schema.string().role('textarea').description("**危险** 自定义参数,请输入 TOML 格式,将会直接覆盖当前界面内任何参数。实时更新,推荐写完后再粘贴过来"),
}).description("高级设置"),
Schema.object({
mixed_precision: Schema.union(["no", "fp16", "bf16"]).default("bf16").description("训练混合精度, RTX30系列以后也可以指定`bf16`"),
full_fp16: Schema.boolean().description("完全使用 FP16 精度"),
full_bf16: Schema.boolean().description("完全使用 BF16 精度"),
fp8_base: Schema.boolean().default(true).description("对基础模型使用 FP8 精度"),
fp8_base_unet: Schema.boolean().description("仅对 U-Net 使用 FP8 精度CLIP-L不使用"),
no_half_vae: Schema.boolean().description("不使用半精度 VAE"),
sdpa: Schema.boolean().default(true).description("启用 sdpa"),
lowram: Schema.boolean().default(false).description("低内存模式 该模式下会将 U-net、文本编码器、VAE 直接加载到显存中"),
cache_latents: Schema.boolean().default(true).description("缓存图像 latent, 缓存 VAE 输出以减少 VRAM 使用"),
cache_latents_to_disk: Schema.boolean().default(true).description("缓存图像 latent 到磁盘"),
cache_text_encoder_outputs: Schema.boolean().default(true).description("缓存文本编码器的输出,减少显存使用。使用时需要关闭 shuffle_caption"),
cache_text_encoder_outputs_to_disk: Schema.boolean().default(true).description("缓存文本编码器的输出到磁盘"),
persistent_data_loader_workers: Schema.boolean().default(true).description("保留加载训练集的worker减少每个 epoch 之间的停顿。"),
vae_batch_size: Schema.number().min(1).description("vae 编码批量大小"),
}).description("速度优化选项"),
Schema.object({
ddp_timeout: Schema.number().min(0).description("分布式训练超时时间"),
ddp_gradient_as_bucket_view: Schema.boolean(),
}).description("分布式训练"),
// 分布式训练
SHARED_SCHEMAS.DISTRIBUTED_TRAINING
]);

View File

@ -26,25 +26,11 @@ Schema.intersect([
]),
]),
Schema.object({
train_data_dir: Schema.string().role('filepicker', { type: "folder", internal: "train-dir" }).default("./train/aki").description("训练数据集路径"),
reg_data_dir: Schema.string().role('filepicker', { type: "folder", internal: "train-dir" }).description("正则化数据集路径。默认留空,不使用正则化图像"),
prior_loss_weight: Schema.number().step(0.1).default(1.0).description("正则化 - 先验损失权重"),
resolution: Schema.string().default("512,512").description("训练图片分辨率宽x高。支持非正方形但必须是 64 倍数。"),
enable_bucket: Schema.boolean().default(true).description("启用 arb 桶以允许非固定宽高比的图片"),
min_bucket_reso: Schema.number().default(256).description("arb 桶最小分辨率"),
max_bucket_reso: Schema.number().default(1024).description("arb 桶最大分辨率"),
bucket_reso_steps: Schema.number().default(64).description("arb 桶分辨率划分单位SDXL 可以使用 32 (SDXL低于32时失效)"),
}).description("数据集设置"),
// 数据集设置
Schema.object(SHARED_SCHEMAS.RAW.DATASET_SETTINGS).description("数据集设置"),
Schema.object({
output_name: Schema.string().default("aki").description("模型保存名称"),
output_dir: Schema.string().role('filepicker', { type: "folder" }).default("./output").description("模型保存文件夹"),
save_model_as: Schema.union(["safetensors", "pt", "ckpt"]).default("safetensors").description("模型保存格式"),
save_precision: Schema.union(["fp16", "float", "bf16"]).default("fp16").description("模型保存精度"),
save_every_n_epochs: Schema.number().default(2).description("每 N epoch自动保存一次模型"),
save_state: Schema.boolean().description("保存训练状态 配合 `resume` 参数可以继续从某个状态训练"),
}).description("保存设置"),
// 保存设置
SHARED_SCHEMAS.SAVE_SETTINGS,
Schema.object({
max_train_epochs: Schema.number().min(1).default(10).description("最大训练 epoch轮数"),
@ -55,65 +41,8 @@ Schema.intersect([
network_train_text_encoder_only: Schema.boolean().default(false).description("仅训练文本编码器"),
}).description("训练相关参数"),
Schema.intersect([
Schema.object({
learning_rate: Schema.string().default("1e-4").description("总学习率, 在分开设置 U-Net 与文本编码器学习率后这个值失效。"),
unet_lr: Schema.string().default("1e-4").description("U-Net 学习率"),
text_encoder_lr: Schema.string().default("1e-5").description("文本编码器学习率"),
lr_scheduler: Schema.union([
"linear",
"cosine",
"cosine_with_restarts",
"polynomial",
"constant",
"constant_with_warmup",
]).default("cosine_with_restarts").description("学习率调度器设置"),
lr_warmup_steps: Schema.number().default(0).description('学习率预热步数'),
}).description("学习率与优化器设置"),
Schema.union([
Schema.object({
lr_scheduler: Schema.const('cosine_with_restarts'),
lr_scheduler_num_cycles: Schema.number().default(1).description('重启次数'),
}),
Schema.object({}),
]),
Schema.object({
optimizer_type: Schema.union([
"AdamW",
"AdamW8bit",
"PagedAdamW8bit",
"Lion",
"Lion8bit",
"PagedLion8bit",
"SGDNesterov",
"SGDNesterov8bit",
"DAdaptation",
"DAdaptAdam",
"DAdaptAdaGrad",
"DAdaptAdanIP",
"DAdaptLion",
"DAdaptSGD",
"AdaFactor",
"Prodigy"
]).default("AdamW8bit").description("优化器设置"),
min_snr_gamma: Schema.number().step(0.1).description("最小信噪比伽马值, 如果启用推荐为 5"),
}),
Schema.union([
Schema.object({
optimizer_type: Schema.const('Prodigy').required(),
prodigy_d0: Schema.string(),
prodigy_d_coef: Schema.string().default("2.0"),
}),
Schema.object({}),
]),
Schema.object({
optimizer_args_custom: Schema.array(String).role('table').description('自定义 optimizer_args一行一个'),
})
]),
// 学习率&优化器设置
SHARED_SCHEMAS.LR_OPTIMIZER,
Schema.intersect([
Schema.object({
@ -173,88 +102,24 @@ Schema.intersect([
]),
]),
Schema.intersect([
Schema.object({
enable_preview: Schema.boolean().default(false).description('启用训练预览图'),
}).description('训练预览图设置'),
// 预览图设置
SHARED_SCHEMAS.PREVIEW_IMAGE,
Schema.union([
Schema.object({
enable_preview: Schema.const(true).required(),
sample_prompts: Schema.string().role('textarea').default(window.__MIKAZUKI__.SAMPLE_PROMPTS_DEFAULT).description(window.__MIKAZUKI__.SAMPLE_PROMPTS_DESCRIPTION),
sample_sampler: Schema.union(["ddim", "pndm", "lms", "euler", "euler_a", "heun", "dpm_2", "dpm_2_a", "dpmsolver", "dpmsolver++", "dpmsingle", "k_lms", "k_euler", "k_euler_a", "k_dpm_2", "k_dpm_2_a"]).default("euler_a").description("生成预览图所用采样器"),
sample_every_n_epochs: Schema.number().default(2).description("每 N 个 epoch 生成一次预览图"),
}),
Schema.object({}),
]),
]),
// 日志设置
SHARED_SCHEMAS.LOG_SETTINGS,
Schema.intersect([
Schema.object({
log_with: Schema.union(["tensorboard", "wandb"]).default("tensorboard").description("日志模块"),
log_prefix: Schema.string().description("日志前缀"),
log_tracker_name: Schema.string().description("日志追踪器名称"),
logging_dir: Schema.string().default("./logs").description("日志保存文件夹"),
}).description('日志设置'),
// caption 选项
Schema.object(SHARED_SCHEMAS.RAW.CAPTION_SETTINGS).description("captionTag选项"),
Schema.union([
Schema.object({
log_with: Schema.const("wandb").required(),
wandb_api_key: Schema.string().required().description("wandb 的 api 密钥"),
}),
Schema.object({}),
]),
]),
// 数据增强
SHARED_SCHEMAS.DATA_ENCHANCEMENT,
Schema.object({
caption_extension: Schema.string().default(".txt").description("Tag 文件扩展名"),
shuffle_caption: Schema.boolean().default(true).description("训练时随机打乱 tokens"),
weighted_captions: Schema.boolean().description("使用带权重的 token不推荐与 shuffle_caption 一同开启"),
keep_tokens: Schema.number().min(0).max(255).step(1).default(0).description("在随机打乱 tokens 时,保留前 N 个不变"),
keep_tokens_separator: Schema.string().description("保留 tokens 时使用的分隔符"),
max_token_length: Schema.number().default(255).description("最大 token 长度"),
caption_dropout_rate: Schema.number().min(0).step(0.01).description("丢弃全部标签的概率,对一个图片概率不使用 caption 或 class token"),
caption_dropout_every_n_epochs: Schema.number().min(0).max(100).step(1).description("每 N 个 epoch 丢弃全部标签"),
caption_tag_dropout_rate: Schema.number().min(0).step(0.01).description("按逗号分隔的标签来随机丢弃 tag 的概率"),
}).description("captionTag选项"),
// 其他选项
SHARED_SCHEMAS.OTHER,
Schema.object({
noise_offset: Schema.number().step(0.0001).description("在训练中添加噪声偏移来改良生成非常暗或者非常亮的图像,如果启用推荐为 0.1"),
multires_noise_iterations: Schema.number().step(1).description("多分辨率(金字塔)噪声迭代次数 推荐 6-10。无法与 noise_offset 一同启用"),
multires_noise_discount: Schema.number().step(0.01).description("多分辨率(金字塔)衰减率 推荐 0.3-0.8,须同时与上方参数 multires_noise_iterations 一同启用"),
}).description("噪声设置"),
// 速度优化选项
SHARED_SCHEMAS.PRECISION_CACHE_BATCH,
Schema.object({
color_aug: Schema.boolean().description("颜色改变"),
flip_aug: Schema.boolean().description("图像翻转"),
random_crop: Schema.boolean().description("随机剪裁"),
}).description("数据增强"),
Schema.object({
seed: Schema.number().default(1337).description("随机种子"),
clip_skip: Schema.number().role("slider").min(0).max(12).step(1).default(2).description("CLIP 跳过层数 *玄学*"),
ui_custom_params: Schema.string().role('textarea').description("**危险** 自定义参数,请输入 TOML 格式,将会直接覆盖当前界面内任何参数。实时更新,推荐写完后再粘贴过来"),
}).description("高级设置"),
Schema.object({
mixed_precision: Schema.union(["no", "fp16", "bf16"]).default("fp16").description("训练混合精度, RTX30系列以后也可以指定`bf16`"),
full_fp16: Schema.boolean().description("完全使用 FP16 精度"),
full_bf16: Schema.boolean().description("完全使用 BF16 精度"),
fp8_base: Schema.boolean().description("对基础模型使用 FP8 精度"),
no_half_vae: Schema.boolean().description("不使用半精度 VAE"),
xformers: Schema.boolean().default(true).description("启用 xformers"),
sdpa: Schema.boolean().description("启用 sdpa"),
lowram: Schema.boolean().default(false).description("低内存模式 该模式下会将 U-net、文本编码器、VAE 直接加载到显存中"),
cache_latents: Schema.boolean().default(true).description("缓存图像 latent, 缓存 VAE 输出以减少 VRAM 使用"),
cache_latents_to_disk: Schema.boolean().default(true).description("缓存图像 latent 到磁盘"),
cache_text_encoder_outputs: Schema.boolean().description("缓存文本编码器的输出,减少显存使用。使用时需要关闭 shuffle_caption"),
cache_text_encoder_outputs_to_disk: Schema.boolean().description("缓存文本编码器的输出到磁盘"),
persistent_data_loader_workers: Schema.boolean().default(true).description("保留加载训练集的worker减少每个 epoch 之间的停顿。"),
vae_batch_size: Schema.number().min(1).description("vae 编码批量大小"),
}).description("速度优化选项"),
Schema.object({
ddp_timeout: Schema.number().min(0).description("分布式训练超时时间"),
ddp_gradient_as_bucket_view: Schema.boolean(),
}).description("分布式训练"),
// 分布式训练
SHARED_SCHEMAS.DISTRIBUTED_TRAINING
]);

170
mikazuki/schema/shared.ts Normal file
View File

@ -0,0 +1,170 @@
(function () {
const SAMPLE_PROMPTS_DEFAULT = "(masterpiece, best quality:1.2), 1girl, solo, --n lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts,signature, watermark, username, blurry, --w 512 --h 768 --l 7 --s 24 --d 1337"
const SAMPLE_PROMPTS_DESCRIPTION = "预览图生成参数。可填写直接填写参数或单独写入txt文件填写路径<br>`--n` 后方为反向提示词<br>`--w`宽,`--h`高<br>`--l`: CFG Scale<br>`--s`: 迭代步数<br>`--d`: 种子"
let data = {
RAW: {
DATASET_SETTINGS: {
train_data_dir: Schema.string().role('filepicker', { type: "folder", internal: "train-dir" }).default("./train/aki").description("训练数据集路径"),
reg_data_dir: Schema.string().role('filepicker', { type: "folder", internal: "train-dir" }).description("正则化数据集路径。默认留空,不使用正则化图像"),
prior_loss_weight: Schema.number().step(0.1).default(1.0).description("正则化 - 先验损失权重"),
resolution: Schema.string().default("512,512").description("训练图片分辨率宽x高。支持非正方形但必须是 64 倍数。"),
enable_bucket: Schema.boolean().default(true).description("启用 arb 桶以允许非固定宽高比的图片"),
min_bucket_reso: Schema.number().default(256).description("arb 桶最小分辨率"),
max_bucket_reso: Schema.number().default(1024).description("arb 桶最大分辨率"),
bucket_reso_steps: Schema.number().default(64).description("arb 桶分辨率划分单位SDXL 可以使用 32 (SDXL低于32时失效)"),
bucket_no_upscale: Schema.boolean().default(true).description("arb 桶不放大图片"),
},
CAPTION_SETTINGS: {
caption_extension: Schema.string().default(".txt").description("Tag 文件扩展名"),
shuffle_caption: Schema.boolean().default(false).description("训练时随机打乱 tokens"),
weighted_captions: Schema.boolean().description("使用带权重的 token不推荐与 shuffle_caption 一同开启"),
keep_tokens: Schema.number().min(0).max(255).step(1).default(0).description("在随机打乱 tokens 时,保留前 N 个不变"),
keep_tokens_separator: Schema.string().description("保留 tokens 时使用的分隔符"),
max_token_length: Schema.number().default(255).description("最大 token 长度"),
caption_dropout_rate: Schema.number().min(0).step(0.01).description("丢弃全部标签的概率,对一个图片概率不使用 caption 或 class token"),
caption_dropout_every_n_epochs: Schema.number().min(0).max(100).step(1).description("每 N 个 epoch 丢弃全部标签"),
caption_tag_dropout_rate: Schema.number().min(0).step(0.01).description("按逗号分隔的标签来随机丢弃 tag 的概率"),
}
},
SAVE_SETTINGS: Schema.object({
output_name: Schema.string().default("aki").description("模型保存名称"),
output_dir: Schema.string().role('filepicker', { type: "folder" }).default("./output").description("模型保存文件夹"),
save_model_as: Schema.union(["safetensors", "pt", "ckpt"]).default("safetensors").description("模型保存格式"),
save_precision: Schema.union(["fp16", "float", "bf16"]).default("fp16").description("模型保存精度"),
save_every_n_epochs: Schema.number().default(2).description("每 N epoch自动保存一次模型"),
save_state: Schema.boolean().description("保存训练状态 配合 `resume` 参数可以继续从某个状态训练"),
}).description("保存设置"),
LR_OPTIMIZER: Schema.intersect([
Schema.object({
learning_rate: Schema.string().default("1e-4").description("总学习率, 在分开设置 U-Net 与文本编码器学习率后这个值失效。"),
unet_lr: Schema.string().default("1e-4").description("U-Net 学习率"),
text_encoder_lr: Schema.string().default("1e-5").description("文本编码器学习率"),
lr_scheduler: Schema.union([
"linear",
"cosine",
"cosine_with_restarts",
"polynomial",
"constant",
"constant_with_warmup",
]).default("cosine_with_restarts").description("学习率调度器设置"),
lr_warmup_steps: Schema.number().default(0).description('学习率预热步数'),
}).description("学习率与优化器设置"),
Schema.union([
Schema.object({
lr_scheduler: Schema.const('cosine_with_restarts'),
lr_scheduler_num_cycles: Schema.number().default(1).description('重启次数'),
}),
Schema.object({}),
]),
Schema.object({
optimizer_type: Schema.union([
"AdamW",
"AdamW8bit",
"PagedAdamW8bit",
"Lion",
"Lion8bit",
"PagedLion8bit",
"SGDNesterov",
"SGDNesterov8bit",
"DAdaptation",
"DAdaptAdam",
"DAdaptAdaGrad",
"DAdaptAdanIP",
"DAdaptLion",
"DAdaptSGD",
"AdaFactor",
"Prodigy"
]).default("AdamW8bit").description("优化器设置"),
min_snr_gamma: Schema.number().step(0.1).description("最小信噪比伽马值, 如果启用推荐为 5"),
}),
Schema.union([
Schema.object({
optimizer_type: Schema.const('Prodigy').required(),
prodigy_d0: Schema.string(),
prodigy_d_coef: Schema.string().default("2.0"),
}),
Schema.object({}),
]),
Schema.object({
optimizer_args_custom: Schema.array(String).role('table').description('自定义 optimizer_args一行一个'),
})
]),
PREVIEW_IMAGE: Schema.intersect([
Schema.object({
enable_preview: Schema.boolean().default(false).description('启用训练预览图'),
}).description('训练预览图设置'),
Schema.union([
Schema.object({
enable_preview: Schema.const(true).required(),
sample_prompts: Schema.string().role('textarea').default(SAMPLE_PROMPTS_DEFAULT).description(SAMPLE_PROMPTS_DESCRIPTION),
sample_sampler: Schema.union(["ddim", "pndm", "lms", "euler", "euler_a", "heun", "dpm_2", "dpm_2_a", "dpmsolver", "dpmsolver++", "dpmsingle", "k_lms", "k_euler", "k_euler_a", "k_dpm_2", "k_dpm_2_a"]).default("euler_a").description("生成预览图所用采样器"),
sample_every_n_epochs: Schema.number().default(2).description("每 N 个 epoch 生成一次预览图"),
}),
Schema.object({}),
]),
]),
LOG_SETTINGS: Schema.intersect([
Schema.object({
log_with: Schema.union(["tensorboard", "wandb"]).default("tensorboard").description("日志模块"),
log_prefix: Schema.string().description("日志前缀"),
log_tracker_name: Schema.string().description("日志追踪器名称"),
logging_dir: Schema.string().default("./logs").description("日志保存文件夹"),
}).description('日志设置'),
Schema.union([
Schema.object({
log_with: Schema.const("wandb").required(),
wandb_api_key: Schema.string().required().description("wandb 的 api 密钥"),
}),
Schema.object({}),
]),
]),
DATA_ENCHANCEMENT: Schema.object({
color_aug: Schema.boolean().description("颜色改变"),
flip_aug: Schema.boolean().description("图像翻转"),
random_crop: Schema.boolean().description("随机剪裁"),
}).description("数据增强"),
OTHER: Schema.object({
seed: Schema.number().default(1337).description("随机种子"),
clip_skip: Schema.number().role("slider").min(0).max(12).step(1).default(2).description("CLIP 跳过层数 *玄学*"),
ui_custom_params: Schema.string().role('textarea').description("**危险** 自定义参数,请输入 TOML 格式,将会直接覆盖当前界面内任何参数。实时更新,推荐写完后再粘贴过来"),
}).description("其他设置"),
PRECISION_CACHE_BATCH: Schema.object({
mixed_precision: Schema.union(["no", "fp16", "bf16"]).default("bf16").description("训练混合精度, RTX30系列以后也可以指定`bf16`"),
full_fp16: Schema.boolean().description("完全使用 FP16 精度"),
full_bf16: Schema.boolean().description("完全使用 BF16 精度"),
fp8_base: Schema.boolean().default(true).description("对基础模型使用 FP8 精度"),
fp8_base_unet: Schema.boolean().description("仅对 U-Net 使用 FP8 精度CLIP-L不使用"),
no_half_vae: Schema.boolean().description("不使用半精度 VAE"),
sdpa: Schema.boolean().default(true).description("启用 sdpa"),
lowram: Schema.boolean().default(false).description("低内存模式 该模式下会将 U-net、文本编码器、VAE 直接加载到显存中"),
cache_latents: Schema.boolean().default(true).description("缓存图像 latent, 缓存 VAE 输出以减少 VRAM 使用"),
cache_latents_to_disk: Schema.boolean().default(true).description("缓存图像 latent 到磁盘"),
cache_text_encoder_outputs: Schema.boolean().default(true).description("缓存文本编码器的输出,减少显存使用。使用时需要关闭 shuffle_caption"),
cache_text_encoder_outputs_to_disk: Schema.boolean().default(true).description("缓存文本编码器的输出到磁盘"),
persistent_data_loader_workers: Schema.boolean().default(true).description("保留加载训练集的worker减少每个 epoch 之间的停顿。"),
vae_batch_size: Schema.number().min(1).description("vae 编码批量大小"),
}).description("速度优化选项"),
DISTRIBUTED_TRAINING: Schema.object({
ddp_timeout: Schema.number().min(0).description("分布式训练超时时间"),
ddp_gradient_as_bucket_view: Schema.boolean(),
}).description("分布式训练"),
}
return data
})()

View File

@ -1,4 +1,8 @@
{
"compilerOptions": {
"target": "ES2020",
"module": "commonjs",
},
"include": [
"**/*.ts"
],