sd3 support

pull/628/head
Akegarasu 2025-02-06 21:47:27 +08:00
parent c8975ad17e
commit 3ec04aeb9d
No known key found for this signature in database
GPG Key ID: DACA951FEBA569A2
3 changed files with 96 additions and 5 deletions

View File

@ -94,11 +94,11 @@ def get_sample_prompts(config: dict) -> Tuple[Optional[str], str]:
positive_prompts = config.pop('positive_prompts', None)
negative_prompts = config.pop('negative_prompts', '')
sample_width = config.pop('sample_width', 512) # 默认宽度 512
sample_height = config.pop('sample_height', 512) # 默认高度 512
sample_cfg = config.pop('sample_cfg', 7) # 默认 CFG 值 7.5
sample_seed = config.pop('sample_seed', 2333) # 默认随机种子 42
sample_steps = config.pop('sample_steps', 24) # 默认步数 50
sample_width = config.pop('sample_width', 512)
sample_height = config.pop('sample_height', 512)
sample_cfg = config.pop('sample_cfg', 7)
sample_seed = config.pop('sample_seed', 2333)
sample_steps = config.pop('sample_steps', 24)
randomly_choice_prompt = config.pop('randomly_choice_prompt', False)
if randomly_choice_prompt:

View File

@ -0,0 +1,90 @@
Schema.intersect([
Schema.object({
model_train_type: Schema.string().default("sd3-lora").disabled().description("训练种类"),
pretrained_model_name_or_path: Schema.string().role('filepicker', { type: "model-file" }).default("./sd-models/model.safetensors").description("SD3 模型路径"),
clip_l: Schema.string().role('filepicker', { type: "model-file" }).description("clip_l 模型文件路径"),
clip_g: Schema.string().role('filepicker', { type: "model-file" }).description("clip_g 模型文件路径"),
t5xxl: Schema.string().role('filepicker', { type: "model-file" }).description("t5xxl 模型文件路径"),
resume: Schema.string().role('filepicker', { type: "folder" }).description("从某个 `save_state` 保存的中断状态继续训练,填写文件路径"),
}).description("训练用模型"),
Schema.object({
t5xxl_max_token_length: Schema.number().step(1).description("T5XXL 最大 token 长度(不填写使用自动)"),
train_t5xxl: Schema.boolean().default(false).description("训练 T5XXL不推荐"),
}).description("SD3 专用参数"),
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 桶分辨率划分单位"),
})
).description("数据集设置"),
// 保存设置
SHARED_SCHEMAS.SAVE_SETTINGS,
Schema.object({
max_train_epochs: Schema.number().min(1).default(20).description("最大训练 epoch轮数"),
train_batch_size: Schema.number().min(1).default(1).description("批量大小, 越高显存占用越高"),
gradient_checkpointing: Schema.boolean().default(true).description("梯度检查点"),
gradient_accumulation_steps: Schema.number().min(1).default(1).description("梯度累加步数"),
network_train_unet_only: Schema.boolean().default(true).description("仅训练 U-Net"),
network_train_text_encoder_only: Schema.boolean().default(false).description("仅训练文本编码器"),
}).description("训练相关参数"),
// 学习率&优化器设置
SHARED_SCHEMAS.LR_OPTIMIZER,
Schema.intersect([
Schema.object({
network_module: Schema.union(["networks.lora_sd3", "lycoris.kohya"]).default("networks.lora_sd3").description("训练网络模块"),
network_weights: Schema.string().role('filepicker').description("从已有的 LoRA 模型上继续训练,填写路径"),
network_dim: Schema.number().min(1).default(4).description("网络维度,常用 4~128不是越大越好, 低dim可以降低显存占用"),
network_alpha: Schema.number().min(1).default(1).description("常用值:等于 network_dim 或 network_dim*1/2 或 1。使用较小的 alpha 需要提升学习率"),
network_args_custom: Schema.array(String).role('table').description('自定义 network_args一行一个'),
enable_base_weight: Schema.boolean().default(false).description('启用基础权重(差异炼丹)'),
}).description("网络设置"),
// lycoris 参数
SHARED_SCHEMAS.LYCORIS_MAIN,
SHARED_SCHEMAS.LYCORIS_LOKR,
SHARED_SCHEMAS.NETWORK_OPTION_BASEWEIGHT,
]),
// 预览图设置
SHARED_SCHEMAS.PREVIEW_IMAGE,
// 日志设置
SHARED_SCHEMAS.LOG_SETTINGS,
// caption 选项
// FLUX 去除 max_token_length
Schema.object(UpdateSchema(SHARED_SCHEMAS.RAW.CAPTION_SETTINGS, {}, ["max_token_length"])).description("captionTag选项"),
// 噪声设置
SHARED_SCHEMAS.NOISE_SETTINGS,
// 数据增强
SHARED_SCHEMAS.DATA_ENCHANCEMENT,
// 其他选项
SHARED_SCHEMAS.OTHER,
// 速度优化选项
Schema.object(
UpdateSchema(SHARED_SCHEMAS.RAW.PRECISION_CACHE_BATCH, {
fp8_base: Schema.boolean().default(true).description("对基础模型使用 FP8 精度"),
fp8_base_unet: Schema.boolean().description("仅对 U-Net 使用 FP8 精度CLIP-L不使用"),
sdpa: Schema.boolean().default(true).description("启用 sdpa"),
cache_text_encoder_outputs: Schema.boolean().default(true).description("缓存文本编码器的输出,减少显存使用。使用时需要关闭 shuffle_caption"),
cache_text_encoder_outputs_to_disk: Schema.boolean().default(true).description("缓存文本编码器的输出到磁盘"),
}, ["xformers"])
).description("速度优化选项"),
// 分布式训练
SHARED_SCHEMAS.DISTRIBUTED_TRAINING
]);

View File

@ -123,6 +123,7 @@
"constant_with_warmup",
]).default("cosine_with_restarts").description("学习率调度器设置"),
lr_warmup_steps: Schema.number().default(0).description('学习率预热步数'),
loss_type: Schema.union(["l1", "l2", "huber", "smooth_l1"]).description("损失函数类型"),
}).description("学习率与优化器设置"),
Schema.union([