diff --git a/mikazuki/app/api.py b/mikazuki/app/api.py index c54f58a..731ef81 100644 --- a/mikazuki/app/api.py +++ b/mikazuki/app/api.py @@ -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"] diff --git a/mikazuki/global.d.ts b/mikazuki/global.d.ts index 28cad10..7af59dc 100644 --- a/mikazuki/global.d.ts +++ b/mikazuki/global.d.ts @@ -126,4 +126,8 @@ declare namespace Schemastery { type SchemaI = Schemastery.Schemastery; -declare const Schema: Schemastery.Static \ No newline at end of file +declare const Schema: Schemastery.Static + +declare const SHARED_SCHEMAS: Dict + +declare function UpdateSchema(origin: Record, modify?: Record, toDelete?: string[]): Record; diff --git a/mikazuki/schema/lora-flux-master.ts b/mikazuki/schema/lora-flux-master.ts index abac100..d3c0140 100644 --- a/mikazuki/schema/lora-flux-master.ts +++ b/mikazuki/schema/lora-flux-master.ts @@ -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("caption（Tag）选项"), - 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("caption（Tag）选项"), + // 其他选项 + 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 ]); diff --git a/mikazuki/schema/lora-master.ts b/mikazuki/schema/lora-master.ts index 76583b7..10f85e4 100644 --- a/mikazuki/schema/lora-master.ts +++ b/mikazuki/schema/lora-master.ts @@ -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("caption（Tag）选项"), - 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("caption（Tag）选项"), + // 其他选项 + 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 ]); diff --git a/mikazuki/schema/shared.ts b/mikazuki/schema/shared.ts new file mode 100644 index 0000000..764192e --- /dev/null +++ b/mikazuki/schema/shared.ts @@ -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文件填写路径
`--n` 后方为反向提示词
`--w`宽，`--h`高
`--l`: CFG Scale
`--s`: 迭代步数
`--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 +})() \ No newline at end of file diff --git a/mikazuki/tsconfig.json b/mikazuki/tsconfig.json index 9a88d3f..2157858 100644 --- a/mikazuki/tsconfig.json +++ b/mikazuki/tsconfig.json @@ -1,4 +1,8 @@ { + "compilerOptions": { + "target": "ES2020", + "module": "commonjs", + }, "include": [ "**/*.ts" ],