diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5c6041bef..260f06f38 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,13 @@
- **修复 Git 更新收走模型和训练集（#356）**：停止自动 stash，自动补齐旧包缺失的程序文件；新包保留完整浅克隆工作树，并在打包前执行数据保留测试。已经受影响的用户请保留整个旧目录（含 `.git`），参见 [恢复说明](https://github.com/wochenlong/lora-scripts-next/issues/356)。
+### Musubi-tuner 新增 Ideogram 4（#424）
+
+- Musubi 引擎支持 **Ideogram 4 LoRA**：训练页新增模型入口（`/lora/ideogram4.html`），三阶段任务（缓存 latents → 缓存文本编码器输出 → 训练）与任务队列、日志、预览沿用现有能力
+- 组件只发**量化权重（FP8）**、许可为 **非商用**，整合包不预置；「训练用模型」区提供 conditional / unconditional DiT、Qwen3-VL-8B 文本编码器、Flux2 VAE 与 tokenizer 的清单与下载入口
+- 上游快照补丁：修正文本编码器 rotary embedding 在 `to_empty()` 后未初始化的问题（上游 v0.3.6 修复），并把 tokenizer 指向本地目录，离线可用
+- 引擎能力、预检与配置导入识别 `ideogram4-lora`（跨页导入自动跳转到对应训练页），详见 [Ideogram 4 训练指南](docs/ideogram4-musubi.md)
+
### 数据集标签编辑（#423）
- **编辑页重做**：默认只展示图片墙，选中图片后才展开编辑面板
diff --git a/README-zh.md b/README-zh.md
index 0440b6d3d..682b9e9da 100644
--- a/README-zh.md
+++ b/README-zh.md
@@ -36,7 +36,7 @@ Next Trainer 把数据集管理、打标、标签编辑、训练配置与任务
| --- | --- | --- |
| Kohya | SD 1.5、SDXL、Flux、Anima | 支持范围随模型而异；已接入独立环境管理 |
| Anima Fast | Anima 2B / 2.9B LoRA | 独立运行环境；两种规格有明确的底模路径入口 |
-| Musubi | Krea 2 LoRA | 可选安装；Linux 多卡见专项指南 |
+| Musubi | Krea 2 / Ideogram 4 LoRA | 可选安装；Ideogram 4 为非商用许可的量化权重，需自行下载；Linux 多卡见专项指南 |
| DiffSynth | Qwen-Image-2.1 文生图 / 图像编辑 LoRA | 当前为单卡 BF16；Edit 使用 Batch 1 |
| AI Toolkit | SDXL、Flux.1 Dev、Klein 4B / 9B、Krea 2 RAW、Anima、Qwen-Image-2.1 LoRA | 可选安装、独立环境；Klein 支持 base / distilled，Klein 与 Qwen 提供图像编辑入口 |
diff --git a/README.md b/README.md
index 0938f6369..35380772d 100644
--- a/README.md
+++ b/README.md
@@ -36,7 +36,7 @@ Settings are organised as **UI, Training Engines, API, Plugins, Advanced** and *
| --- | --- | --- |
| Kohya | SD 1.5, SDXL, Flux, Anima | Capabilities vary by model; standalone environment management is integrated |
| Anima Fast | Anima 2B / 2.9B LoRA | Isolated runtime with explicit base-model path selection |
-| Musubi | Krea 2 LoRA | Optional installation; see the Linux multi-GPU guide |
+| Musubi | Krea 2 / Ideogram 4 LoRA | Optional installation; Ideogram 4 ships non-commercial quantized weights you fetch yourself; see the Linux multi-GPU guide |
| DiffSynth | Qwen-Image-2.1 text-to-image / image-editing LoRA | Single-GPU BF16; Edit currently uses batch size 1 |
| AI Toolkit | SDXL, Flux.1 Dev, Klein 4B / 9B, Krea 2 RAW, Anima, Qwen-Image-2.1 LoRA | Optional installation with an isolated runtime; Klein supports base / distilled variants, and Klein and Qwen offer image-editing entry points |
diff --git a/docs/dev-progress.md b/docs/dev-progress.md
index 1cbd5ce8f..faa2b1cec 100644
--- a/docs/dev-progress.md
+++ b/docs/dev-progress.md
@@ -8,6 +8,7 @@
| 方向 | 已集成内容 | 边界 |
| --- | --- | --- |
+| Musubi / Ideogram 4（#424） | Musubi 引擎新增 Ideogram 4 LoRA：训练页模型入口、三阶段缓存与训练、条件化/非条件化 DiT 与采样参数、配置导入识别；本仓补丁修正上游文本编码器 rotary 未初始化问题并支持本地 tokenizer | 官方仅发 FP8 量化权重且为**非商用**许可，整合包不预置，需自行下载；LoRA 不能合并回 FP8 基座；上游为 experimental，真机验收见 #424 |
| 标签编辑与中文释义（#423） | 编辑页重做（图片墙优先、标签/自由两种编辑模式、批量操作与历史记录）；内置 Danbooru 中文词典（约 33 万条），翻译来源为词库 / 本地模型 / API；翻译设置与 API 配置进入全局设置 | 词库可完全离线；本地模型与 API 来源依赖用户自行安装或配置；全局 API 配置的存储统一仍在推进，打标页当前只选择来源 |
| AI Toolkit（#399 及验收修复） | 独立环境、多模型 LoRA 入口、本地模型组件、文生图与图像编辑、统一任务和预览 | Qwen-Image-2.1 两种模式各通过 3 步真机冒烟及 20 步预览采样；不等于全模型、浏览器全流程或长期稳定性验收 |
| 引擎列表（#398） | 搜索、安装状态筛选、置顶与排序 | 顺序按浏览器及访问地址保存，不同步到其他设备，也不改变默认训练引擎 |
@@ -55,6 +56,13 @@ AI Toolkit 本轮修复包括 GPU 显示名称转设备编号、Windows 中文
Anima 底模选择适用于 Kohya LoRA、Anima Fast LoRA 和 Anima 全量微调。Anima Fast 另有明确的训练时长模式及验证集拆分配置；引擎参数与限制见 [Anima Fast 指南](anima-fast.md)。
+### Musubi / Ideogram 4
+
+- Musubi 引擎的训练类型从 Krea 2 扩展为 Krea 2 / Ideogram 4 两个模型族：训练页模型入口、三阶段缓存与训练、模型族专属参数与约束（量化-only 基座、非对称 CFG 采样开关、blocks_to_swap 上限 33）。
+- 组件只发 FP8 量化权重（非商用许可），整合包不预置；「训练用模型」区提供清单与下载入口，tokenizer 建议本地放置以支持离线。
+- 本仓补丁修正上游快照中文本编码器 rotary embedding 未初始化的问题（上游 v0.3.6 修复）；文档见 [Ideogram 4 训练指南](ideogram4-musubi.md)。
+- 真机验收（3 步冒烟、缓存复用、显存实测）在 [#424](https://github.com/wochenlong/lora-scripts-next/issues/424) 跟踪，不代表已完成。
+
### DiffSynth / Qwen-Image-2.1
- 支持完整模型目录或受支持的 BF16 DiT、文本编码器与 VAE 组件，Processor 资源自动管理。
diff --git a/docs/ideogram4-musubi.md b/docs/ideogram4-musubi.md
new file mode 100644
index 000000000..284057809
--- /dev/null
+++ b/docs/ideogram4-musubi.md
@@ -0,0 +1,79 @@
+# Ideogram 4（Musubi-Tuner）训练指南
+
+Ideogram 4 通过 **Musubi-Tuner** 引擎接入，训练入口：训练页选择 **模型 = Ideogram 4，引擎 = Musubi-Tuner，目标 = LoRA**，或直接访问 `/lora/ideogram4.html`。
+
+当前支持 **LoRA 训练**；上游尚未支持全量微调，也 **不支持把 LoRA 合并回 FP8 底模**。
+
+## 许可与权重（务必先读）
+
+- 模型许可为 **Ideogram Non-Commercial**：仅限非商业用途，微调同样受此限制；禁止使用其输出训练与 Ideogram 竞争的模型。
+- 官方仅发布 **量化权重（FP8 / NVFP4）**，没有 BF16/FP16 底模。DiT 始终以预量化 FP8 加载作为冻结基座，LoRA 以计算精度运行——这是唯一的工作模式，不是可选项。
+- 本项目的整合包 **不预置这些权重**。请在「训练用模型」区下载或手动放置，组件取自 Comfy-Org 单文件布局：
+
+| 组件 | 文件 |
+| --- | --- |
+| conditional DiT | `ideogram4_fp8_scaled.safetensors` |
+| unconditional DiT（非对称 CFG，可选） | `ideogram4_unconditional_fp8_scaled.safetensors` |
+| 文本编码器（Qwen3-VL-8B，FP8） | `qwen3vl_8b_fp8_scaled.safetensors` |
+| VAE（Flux2） | `flux2-vae.safetensors` |
+| tokenizer（Qwen3-VL-8B） | 目录，建议本地放置 |
+
+默认路径：`sd-models/ideogram4/`，tokenizer 目录为 `sd-models/ideogram4/qwen3-vl-8b-tokenizer`。
+
+### 需要多少磁盘空间
+
+官方只发量化权重，**有 FP8 版本，而且 FP8 就是我们默认使用、也是 musubi 训练脚本唯一支持的加载格式**。Comfy-Org 镜像上主要的文件体积（2026-10 实测）：
+
+| 文件 | 体积 | 是否必需 |
+| --- | --- | --- |
+| `ideogram4_fp8_scaled.safetensors`（conditional DiT） | 8.64 GB | 必需 |
+| `qwen3vl_8b_fp8_scaled.safetensors`（Qwen3-VL-8B 文本编码器） | 9.86 GB | 必需 |
+| `flux2-vae.safetensors` | 0.31 GB | 必需 |
+| tokenizer（Qwen3-VL-8B 目录） | 约 0.02 GB | 必需 |
+| `ideogram4_unconditional_fp8_scaled.safetensors` | 8.64 GB | 可选（非对称 CFG 推理 / 显式开启的训练采样） |
+
+镜像里还有 **NVFP4**（DiT 5.11 GB + 文本编码器 5.87 GB）与 **int8_convrot**（8.93 GB）变体，体积更小；我们的默认资源清单与上游文档都按 FP8 单文件路径配置，其它量化格式 **未做过验证**。
+
+配套空间按需预留：
+
+- **权重最低要求**：约 **18.8 GB**（必需四项）；需要无条件 DiT 再 +8.6 GB（合计约 27.5 GB）。
+- **文本缓存**：每 token 存 53248 通道——BF16 下 128 token ≈ 13.6 MB/图，1000 张约 **13.6 GB**；`text_cache_dtype=fp8_e4m3fn` 可减半。
+- **latent 缓存**：每图 MB 级，可忽略。
+- **引擎环境**：若本机还没装 musubi-tuner（含 torch 等），安装另需约 **8–12 GB**。
+
+**建议预留 60–80 GB**（权重 + 一份缓存 + 引擎与产物余量）；已有 musubi 环境且数据集只有几百张图时，40 GB 也够。下载期间会同时在目标目录写入一份文件，不要贴着剩余空间下限操作。
+
+
+## 运行流程与显存
+
+训练按三阶段任务组执行：**缓存 latents → 缓存文本编码器输出 → 训练**。缓存阶段固定使用第一张 GPU。
+
+- **文本缓存很大**：每个 token 存 53248 通道，BF16 下 128 token ≈ 13.6 MB/图（1k 图 ≈ 13.6 GB），256 token 翻倍，512 token ≈ 54.5 GB。可把 `text_cache_dtype` 设为 `fp8_e4m3fn` 减半，但属于精度权衡。
+- 基座为 9.3B FP8 DiT + Qwen3-VL-8B 文本编码器。显存不足时按顺序尝试：`blocks_to_swap`（上限 **33**）、`gradient_checkpointing`、关闭训练中采样预览。
+- 预检会在显存低于 16 GB 时提示，并检查已安装的 musubi-tuner 快照是否包含 Ideogram 4 脚本；缺失时请在「设置 → 训练引擎」修复/重装。
+
+## 参数要点
+
+- **时间步采样**：默认 `ideogram4_shift`（官方对齐的分辨率感知采样器），通常无需修改。上游另有 `ideogram4_timestep_mu/std` 两个兼容参数，但明确忽略其取值，界面不提供。
+- **损失**：纯 MSE flow matching，`weighting_scheme` 只能是 `none`（界面不提供该选项，导入其它值时会被忽略并提示）。
+- **caption**：官方用结构化 JSON（`high_level_description` / `style_description` / `compositional_deconstruction`）。纯文本也能训练；开启 `validate_caption_structure` 会校验结构，配合 `warn_on_caption_issues` 可只警告不中断。
+- **LoRA 目标**：仅训练 conditional transformer（`attention.qkv`、`attention.o`、`feed_forward.w1/w2/w3`）。
+- **注意力**：`sdpa` / `flash_attn` / `xformers` 可用；**SageAttention 不可用**（head 维度 256 超出其 int8/fp8 核支持范围）。
+- 精度：`mixed_precision` 仅 bf16；`dit_dtype` 默认 bfloat16（底模始终 FP8 加载）。
+
+## 训练中采样
+
+- 需要同时给出 `text_encoder` 与 `vae` 路径。
+- 采样调度由 `sampler_preset` 决定：`V4_DEFAULT_20`（18+2 步，默认）、`V4_QUALITY_48`（45+3 步）、`V4_TURBO_12`（11+1 步）；`initial_sigma` 默认 1.004。
+- 默认只用 conditional DiT 采样（LoRA 只挂在 conditional DiT 上）。若要使用官方非对称 CFG，需要提供 unconditional DiT 并开启对应开关。
+- 官方结构化 JSON prompt 也可用于采样，界面默认给出一个三键示例。
+
+## 本仓的上游补丁
+
+我们固定的上游快照（2026-08-13）中，`ideogram4_utils.py` 的文本编码器加载在 `to_empty()` 之后 **没有重建 rotary embedding**，非持久化的 `inv_freq` 缓冲会保持未初始化，导致文本缓存是垃圾值。上游在 v0.3.6 修复。本仓不升级 pin，而是在启动训练前对该文件打等价补丁（幂等、写入前做语法校验），同时把 tokenizer 指向本地目录以支持离线运行。
+
+## 已知限制
+
+- 上游标注该模型支持为 **experimental**；上游 issue 报告过文本编码器预缓存阶段 OOM 与 resume 变慢，长训练前请先小步验证。
+- 只能 LoRA；LoRA 不能合并回 FP8 权重，生成时通过 forward hook 加载，结果可能与把 LoRA 合并进权重的工具（如 ComfyUI）不同。
+- 真机验收（3 步冒烟 + 缓存复用）在 [#424](https://github.com/wochenlong/lora-scripts-next/issues/424) 跟踪。
diff --git a/frontend/src/engines/catalog.ts b/frontend/src/engines/catalog.ts
index 1ee3a4e40..2fba17632 100644
--- a/frontend/src/engines/catalog.ts
+++ b/frontend/src/engines/catalog.ts
@@ -75,7 +75,7 @@ export const ENGINE_CATALOG: readonly EngineDefinition[] = [
requiresGpu: true,
managesRuntime: true,
mark: "M",
- tags: ["lora", "krea2", "nvidia"],
+ tags: ["lora", "krea2", "ideogram4", "nvidia"],
version: "plugin",
updatedAt: "2026-08",
},
diff --git a/frontend/src/i18n/messages/en-US.ts b/frontend/src/i18n/messages/en-US.ts
index 17bd7e188..4a15428a2 100644
--- a/frontend/src/i18n/messages/en-US.ts
+++ b/frontend/src/i18n/messages/en-US.ts
@@ -49,6 +49,7 @@ export default {
flux: "Flux",
lumina: "Lumina 2",
krea2: "Krea 2",
+ ideogram4: "Ideogram 4",
klein: "Klein (FLUX.2)",
"qwen-image-21": "Qwen-Image-2.1",
},
@@ -208,6 +209,7 @@ export default {
"flux-lora": { title: "Flux LoRA", area: "Flux · Kohya-ss · LoRA" },
"lumina2-lora": { title: "Lumina 2 LoRA", area: "Lumina 2 · Kohya-ss · LoRA" },
"krea2-lora": { title: "Krea 2 LoRA", area: "Krea 2 · Musubi-Tuner · LoRA" },
+ "ideogram4-lora": { title: "Ideogram 4 LoRA", area: "Ideogram 4 · Musubi-Tuner · LoRA" },
"qwen-image-21-lora": { title: "Qwen-Image-2.1 LoRA", area: "Qwen-Image-2.1 · DiffSynth-Studio · LoRA" },
"ai-toolkit-sdxl-lora": { title: "SDXL LoRA", area: "SDXL · AI Toolkit · LoRA" },
"ai-toolkit-flux-lora": { title: "Flux LoRA", area: "Flux · AI Toolkit · LoRA" },
@@ -711,6 +713,7 @@ export default {
flux: "Flux",
anima: "Anima",
krea2: "Krea 2",
+ ideogram4: "Ideogram 4",
klein: "Klein",
"qwen-image-21": "Qwen-Image-2.1",
nvidia: "NVIDIA",
@@ -806,7 +809,7 @@ export default {
},
musubi: {
name: "Musubi-Tuner",
- summary: "Optional engine for Krea 2 and other image-model LoRA.",
+ summary: "Optional engine for Krea 2, Ideogram 4 and other image-model LoRA.",
sizeHint: "Large download (usually several GB); extra disk space and download time required",
},
"ai-toolkit": {
diff --git a/frontend/src/i18n/messages/zh-CN.ts b/frontend/src/i18n/messages/zh-CN.ts
index 873a24fad..6e1dec2d2 100644
--- a/frontend/src/i18n/messages/zh-CN.ts
+++ b/frontend/src/i18n/messages/zh-CN.ts
@@ -49,6 +49,7 @@ export default {
flux: "Flux",
lumina: "Lumina 2",
krea2: "Krea 2",
+ ideogram4: "Ideogram 4",
klein: "Klein (FLUX.2)",
"qwen-image-21": "Qwen-Image-2.1",
},
@@ -208,6 +209,7 @@ export default {
"flux-lora": { title: "Flux LoRA", area: "Flux · Kohya-ss · LoRA" },
"lumina2-lora": { title: "Lumina 2 LoRA", area: "Lumina 2 · Kohya-ss · LoRA" },
"krea2-lora": { title: "Krea 2 LoRA", area: "Krea 2 · Musubi-Tuner · LoRA" },
+ "ideogram4-lora": { title: "Ideogram 4 LoRA", area: "Ideogram 4 · Musubi-Tuner · LoRA" },
"qwen-image-21-lora": { title: "Qwen-Image-2.1 LoRA", area: "Qwen-Image-2.1 · DiffSynth-Studio · LoRA" },
"ai-toolkit-sdxl-lora": { title: "SDXL LoRA", area: "SDXL · AI Toolkit · LoRA" },
"ai-toolkit-flux-lora": { title: "Flux LoRA", area: "Flux · AI Toolkit · LoRA" },
@@ -711,6 +713,7 @@ export default {
flux: "Flux",
anima: "Anima",
krea2: "Krea 2",
+ ideogram4: "Ideogram 4",
klein: "Klein",
"qwen-image-21": "Qwen-Image-2.1",
nvidia: "NVIDIA",
@@ -806,7 +809,7 @@ export default {
},
musubi: {
name: "Musubi-Tuner",
- summary: "可选训练引擎，支持 Krea 2 等图像模型 LoRA。",
+ summary: "可选训练引擎，支持 Krea 2、Ideogram 4 等图像模型 LoRA。",
sizeHint: "体积较大（通常数 GB），需额外磁盘与下载时间",
},
"ai-toolkit": {
diff --git a/frontend/src/pages/MusubiGatePage.test.ts b/frontend/src/pages/MusubiGatePage.test.ts
index d6b25e269..a76452892 100644
--- a/frontend/src/pages/MusubiGatePage.test.ts
+++ b/frontend/src/pages/MusubiGatePage.test.ts
@@ -21,15 +21,17 @@ const installing: EngineStatus = { id: "musubi", state: "installing", featureEna
const notInstalled: EngineStatus = { id: "musubi", state: "not_installed", featureEnabled: true }
const trainingPageStub = {
- template: '
',
+ props: ["schemaName", "title", "area"],
+ template: '
',
}
const routerLinkStub = {
props: ["to"],
template: '',
}
-function mountPage() {
+function mountPage(props: Record = {}) {
return mount(MusubiGatePage, {
+ props,
global: { plugins: [i18n], stubs: { TrainingPage: trainingPageStub, RouterLink: routerLinkStub } },
})
}
@@ -49,10 +51,20 @@ describe("MusubiGatePage", () => {
expect(status).toHaveBeenCalledWith("musubi")
expect(setInterval).not.toHaveBeenCalled()
expect(wrapper.find(".training-page-stub").exists()).toBe(true)
+ expect(wrapper.get(".training-page-stub").attributes("data-schema")).toBe("krea2-lora")
expect(wrapper.get('[data-testid="engine-ready-chip"]').text()).toContain("训练环境准备就绪")
wrapper.unmount()
})
+ it("forwards the selected musubi schema to the training form", async () => {
+ status.mockResolvedValue(ready)
+ const wrapper = mountPage({ schemaName: "ideogram4-lora" })
+ await flushPromises()
+
+ expect(wrapper.get(".training-page-stub").attributes("data-schema")).toBe("ideogram4-lora")
+ wrapper.unmount()
+ })
+
it("shows the install view instead of the form when musubi is not installed", async () => {
status.mockResolvedValue(notInstalled)
const wrapper = mountPage()
diff --git a/frontend/src/pages/MusubiGatePage.vue b/frontend/src/pages/MusubiGatePage.vue
index 6af3c7d4b..9e0db1fe8 100644
--- a/frontend/src/pages/MusubiGatePage.vue
+++ b/frontend/src/pages/MusubiGatePage.vue
@@ -8,7 +8,10 @@ import EngineStatusBar from "../components/EngineStatusBar.vue"
import TrainingPage from "./TrainingPage.vue"
import { SCHEMA_META } from "../training/modules"
-withDefaults(defineProps<{ bare?: boolean }>(), { bare: false })
+const props = withDefaults(defineProps<{ bare?: boolean; schemaName?: string }>(), {
+ bare: false,
+ schemaName: "krea2-lora",
+})
const status = ref({ id: "musubi", state: "unknown", featureEnabled: true })
const loading = ref(false)
@@ -20,9 +23,9 @@ let logSource: EventSource | undefined
let progressSource: EventSource | undefined
const { t } = useI18n()
-const meta = SCHEMA_META["krea2-lora"]
-const title = computed(() => t(meta.titleKey))
-const area = computed(() => t(meta.areaKey))
+const meta = computed(() => SCHEMA_META[props.schemaName] ?? SCHEMA_META["krea2-lora"])
+const title = computed(() => t(meta.value.titleKey))
+const area = computed(() => t(meta.value.areaKey))
const ready = computed(() => status.value.state === "ready")
const working = computed(() => ["installing", "auditing"].includes(status.value.state))
const label = computed(() => {
@@ -133,7 +136,7 @@ onBeforeUnmount(() => {
v-if="ready"
:title="title"
:area="area"
- schema-name="krea2-lora"
+ :schema-name="schemaName"
:bare="bare"
>
diff --git a/frontend/src/pages/TrainingPage.vue b/frontend/src/pages/TrainingPage.vue
index 987b98e52..90669cd1a 100644
--- a/frontend/src/pages/TrainingPage.vue
+++ b/frontend/src/pages/TrainingPage.vue
@@ -13,7 +13,7 @@ import { applyReadonlyDefaults, cloneFormModel, cloneFormValue, createDefaultMod
import { loadTrainingSchema } from "../schema/loader"
import { buildTrainingConfig, checkTrainingConfig, hydrateImportedConfig, pickCarryOverFields, sanitizePersistedDraft } from "../training/params"
import { QWEN_VALIDATION_FIELDS, validateQwenConfig } from "../training/qwenValidation"
-import { isAiToolkitSchema, moduleForSchema, moduleForTrainType } from "../training/modules"
+import { isAiToolkitSchema, isMusubiSchema, moduleForSchema, moduleForTrainType } from "../training/modules"
import { collectDatasetCheckTargets } from "../training/datasetCheck"
import { datasetsApi, type DatasetValidation, type ValidateFinding } from "../api/datasets"
import { copyText } from "../utils/clipboard"
@@ -331,7 +331,7 @@ async function submit() {
if (!preflight.ok) throw new Error(preflight.errors?.join("\n") || t("training.submitConfirm.preflightFail"))
preflight.warnings?.forEach((warning) => ElMessage.warning(warning))
}
- if (props.schemaName === "krea2-lora") {
+ if (isMusubiSchema(props.schemaName)) {
const preflight = await trainingApi.musubiPreflight(output.value)
if (!preflight.ok) throw new Error(preflight.errors?.join("\n") || t("training.submitConfirm.preflightFail"))
preflight.warnings?.forEach((warning) => ElMessage.warning(warning))
diff --git a/frontend/src/pages/TrainingWorkbenchPage.vue b/frontend/src/pages/TrainingWorkbenchPage.vue
index 38829bb33..dcfd157cc 100644
--- a/frontend/src/pages/TrainingWorkbenchPage.vue
+++ b/frontend/src/pages/TrainingWorkbenchPage.vue
@@ -118,7 +118,7 @@ watch([model, engine, target], () => {
-
+
diff --git a/frontend/src/router.ts b/frontend/src/router.ts
index 846145f30..aacaeed5f 100644
--- a/frontend/src/router.ts
+++ b/frontend/src/router.ts
@@ -55,6 +55,8 @@ const routes: RouteRecordRaw[] = [
{ path: "/lora/anima-finetune.html", redirect: () => trainingQuery("anima", "kohya", "finetune") },
{ path: "/lora/krea2.html", redirect: () => trainingQuery("krea2", "musubi", "lora") },
{ path: "/lora/krea2", redirect: () => trainingQuery("krea2", "musubi", "lora") },
+ { path: "/lora/ideogram4.html", redirect: () => trainingQuery("ideogram4", "musubi", "lora") },
+ { path: "/lora/ideogram4", redirect: () => trainingQuery("ideogram4", "musubi", "lora") },
{ path: "/lora/klein", redirect: () => trainingQuery("klein", "ai-toolkit", "lora") },
{ path: "/dreambooth/index.html", redirect: () => trainingQuery("sdxl", "kohya", "finetune") },
{ path: "/tagger.html", redirect: "/dataset/tagger" },
diff --git a/frontend/src/training/ideogram4.test.ts b/frontend/src/training/ideogram4.test.ts
new file mode 100644
index 000000000..ae5a99ba6
--- /dev/null
+++ b/frontend/src/training/ideogram4.test.ts
@@ -0,0 +1,87 @@
+import { readFileSync } from "node:fs"
+import { resolve } from "node:path"
+import { describe, expect, it } from "vitest"
+import { createDefaultModel, executeSchemaSources } from "../schema/adapter"
+import { buildTrainingConfig, pickCarryOverFields } from "./params"
+
+const source = readFileSync(resolve(process.cwd(), "../mikazuki/schema/ideogram4-lora.ts"), "utf8")
+const schema = executeSchemaSources([{ name: "ideogram4-lora", hash: "test", schema: source }], "ideogram4-lora")
+
+describe("Ideogram 4 schema defaults", () => {
+ it("points at the FP8 quantized components", () => {
+ expect(createDefaultModel(schema)).toMatchObject({
+ model_train_type: "ideogram4-lora",
+ dit: "./sd-models/ideogram4/ideogram4_fp8_scaled.safetensors",
+ text_encoder: "./sd-models/ideogram4/qwen3vl_8b_fp8_scaled.safetensors",
+ vae: "./sd-models/ideogram4/flux2-vae.safetensors",
+ output_name: "next-ideogram4-lora",
+ resolution: "1024,1024",
+ })
+ })
+
+ it("uses the official sampler and precision defaults", () => {
+ expect(createDefaultModel(schema)).toMatchObject({
+ timestep_sampling: "ideogram4_shift",
+ sampler_preset: "V4_DEFAULT_20",
+ initial_sigma: 1.004,
+ sample_cfg: 7,
+ mixed_precision: "bf16",
+ dit_dtype: "bfloat16",
+ text_cache_dtype: "bf16",
+ blocks_to_swap: 0,
+ network_dim: 32,
+ network_alpha: 32,
+ use_unconditional_dit_for_lora_sampling: false,
+ })
+ })
+
+ it("keeps Krea 2 only switches out of the form", () => {
+ const model = createDefaultModel(schema)
+ for (const key of ["turbo_dit", "turbo_dit_cache", "fp8_base", "fp8_scaled", "weighting_scheme"]) {
+ expect(model).not.toHaveProperty(key)
+ }
+ })
+})
+
+describe("Ideogram 4 config serialization", () => {
+ it("locks the train type and carries the sampler defaults", () => {
+ const model = { ...createDefaultModel(schema), model_train_type: "krea2-lora" }
+ const config = buildTrainingConfig(model, "ideogram4-lora")
+
+ expect(config.model_train_type).toBe("ideogram4-lora")
+ expect(config.sampler_preset).toBe("V4_DEFAULT_20")
+ expect(config.sample_cfg).toBe(7)
+ expect(config).not.toHaveProperty("fp8_base")
+ })
+
+ it("does not inherit Krea 2 model paths when switching pages", () => {
+ const defaults = createDefaultModel(schema)
+ const carried = pickCarryOverFields(
+ {
+ dit: "./sd-models/krea2/krea2.safetensors",
+ vae: "./sd-models/krea2/qwen_image_vae.safetensors",
+ text_encoder: "./sd-models/krea2/qwen3_vl_4b.safetensors",
+ fp8_base: true,
+ fp8_scaled: true,
+ learning_rate: "2e-4",
+ },
+ defaults,
+ )
+
+ expect(carried.learning_rate).toBe("2e-4")
+ for (const key of ["dit", "vae", "text_encoder", "fp8_base", "fp8_scaled"]) {
+ expect(carried).not.toHaveProperty(key)
+ }
+ })
+
+ it("still carries an external VAE between kohya pages", () => {
+ const kohyaDefaults = { model_train_type: "sdxl-lora", vae: "", fp8_base: false }
+ const carried = pickCarryOverFields(
+ { vae: "./models/vae.safetensors", fp8_base: true, learning_rate: "2e-4" },
+ kohyaDefaults,
+ )
+
+ expect(carried.vae).toBe("./models/vae.safetensors")
+ expect(carried.fp8_base).toBe(true)
+ })
+})
diff --git a/frontend/src/training/modules.test.ts b/frontend/src/training/modules.test.ts
index ad75abb83..54f652fcf 100644
--- a/frontend/src/training/modules.test.ts
+++ b/frontend/src/training/modules.test.ts
@@ -78,10 +78,13 @@ describe("training module mapping", () => {
expect(resolveModule("anima", "musubi", "lora")).toBeUndefined()
})
- it("limits musubi engine to the krea2 model", () => {
+ it("limits musubi engine to the krea2 and ideogram4 models", () => {
expect(isEngineSupported("krea2", "musubi")).toBe(true)
+ expect(isEngineSupported("ideogram4", "musubi")).toBe(true)
+ expect(resolveModule("krea2", "musubi", "lora")?.schemaName).toBe("krea2-lora")
+ expect(resolveModule("ideogram4", "musubi", "lora")?.schemaName).toBe("ideogram4-lora")
for (const model of TRAINING_MODELS) {
- if (model === "krea2") continue
+ if (model === "krea2" || model === "ideogram4") continue
expect(isEngineSupported(model, "musubi")).toBe(false)
}
})
diff --git a/frontend/src/training/modules.ts b/frontend/src/training/modules.ts
index 544a72769..073ecf86f 100644
--- a/frontend/src/training/modules.ts
+++ b/frontend/src/training/modules.ts
@@ -1,6 +1,6 @@
import type { FormModel } from "../schema/adapter"
-export type TrainingModel = "anima" | "sd15" | "sdxl" | "flux" | "lumina" | "krea2" | "klein" | "qwen-image-21"
+export type TrainingModel = "anima" | "sd15" | "sdxl" | "flux" | "lumina" | "krea2" | "ideogram4" | "klein" | "qwen-image-21"
export type TrainingEngine = "kohya" | "anima-fast" | "musubi" | "ai-toolkit" | "diffsynth"
export type TrainingTarget = "lora" | "finetune"
@@ -17,7 +17,7 @@ export interface TrainingModule {
legacyStorageKey?: string
}
-export const TRAINING_MODELS: readonly TrainingModel[] = ["anima", "sd15", "sdxl", "flux", "lumina", "krea2", "klein", "qwen-image-21"]
+export const TRAINING_MODELS: readonly TrainingModel[] = ["anima", "sd15", "sdxl", "flux", "lumina", "krea2", "ideogram4", "klein", "qwen-image-21"]
export const TRAINING_ENGINES: readonly TrainingEngine[] = ["kohya", "anima-fast", "musubi", "ai-toolkit", "diffsynth"]
export const TRAINING_TARGETS: readonly TrainingTarget[] = ["lora", "finetune"]
@@ -42,6 +42,7 @@ export const TRAINING_MODULES: readonly TrainingModule[] = [
{ model: "flux", engine: "kohya", target: "lora", schemaName: "flux-lora" },
{ model: "lumina", engine: "kohya", target: "lora", schemaName: "lumina2-lora" },
{ model: "krea2", engine: "musubi", target: "lora", schemaName: "krea2-lora" },
+ { model: "ideogram4", engine: "musubi", target: "lora", schemaName: "ideogram4-lora" },
{ model: "klein", engine: "ai-toolkit", target: "lora", schemaName: "klein-lora" },
{ model: "sdxl", engine: "ai-toolkit", target: "lora", schemaName: "ai-toolkit-sdxl-lora" },
{ model: "flux", engine: "ai-toolkit", target: "lora", schemaName: "ai-toolkit-flux-lora" },
@@ -67,6 +68,7 @@ export const SCHEMA_META: Record
"flux-lora": { titleKey: "training.schemas.flux-lora.title", areaKey: "training.schemas.flux-lora.area" },
"lumina2-lora": { titleKey: "training.schemas.lumina2-lora.title", areaKey: "training.schemas.lumina2-lora.area" },
"krea2-lora": { titleKey: "training.schemas.krea2-lora.title", areaKey: "training.schemas.krea2-lora.area" },
+ "ideogram4-lora": { titleKey: "training.schemas.ideogram4-lora.title", areaKey: "training.schemas.ideogram4-lora.area" },
"klein-lora": { titleKey: "training.schemas.klein-lora.title", areaKey: "training.schemas.klein-lora.area" },
}
@@ -117,3 +119,7 @@ export function firstSupportedTarget(model: TrainingModel, engine: TrainingEngin
export function isAiToolkitSchema(name: string): boolean {
return TRAINING_MODULES.some(module => module.engine === "ai-toolkit" && module.schemaName === name)
}
+
+export function isMusubiSchema(name: string): boolean {
+ return TRAINING_MODULES.some(module => module.engine === "musubi" && module.schemaName === name)
+}
diff --git a/frontend/src/training/params.ts b/frontend/src/training/params.ts
index 3ad3cf606..fca94931c 100644
--- a/frontend/src/training/params.ts
+++ b/frontend/src/training/params.ts
@@ -43,6 +43,26 @@ export const CROSS_SCHEMA_DENY_KEYS = [
"model_input_mode", "model_variant", "model_path", "model_config_dir", "dit_path", "text_encoder_path", "vae_path", "training_task", "control_data_dirs",
] as const
+/**
+ * Asset paths that are only unsafe across *model-family* pages (musubi/klein):
+ * Krea 2's Turbo DiT, Ideogram 4's asymmetric-CFG DiT and the fp8 pairing flags
+ * must not leak between those pages. Kohya pages legitimately share an external
+ * VAE and their own fp8 switches, so these stay out of the global deny list.
+ */
+const MODEL_ASSET_DENY_TRAIN_TYPES = new Set([
+ "krea2-lora", "ideogram4-lora", "klein-lora", "klein-4b-lora", "klein-9b-lora",
+])
+const MODEL_ASSET_DENY_KEYS = ["vae", "turbo_dit", "turbo_dit_cache", "unconditional_dit", "fp8_base", "fp8_scaled"] as const
+
+/** Deny set for a carry-over / autosave-sanitize boundary, keyed on the target page. */
+export function denyKeysForDefaults(defaults: FormModel): Set {
+ const deny = new Set(CROSS_SCHEMA_DENY_KEYS)
+ if (MODEL_ASSET_DENY_TRAIN_TYPES.has(String(defaults.model_train_type ?? ""))) {
+ for (const key of MODEL_ASSET_DENY_KEYS) deny.add(key)
+ }
+ return deny
+}
+
/** schemaName → locked model_train_type for pages that own a single train type. */
export const SCHEMA_TRAIN_TYPES: Record = {
"qwen-image-21-lora": "qwen-image-21-lora",
@@ -52,6 +72,7 @@ export const SCHEMA_TRAIN_TYPES: Record = {
"flux-lora": "flux-lora",
"lumina2-lora": "lumina2-lora",
"krea2-lora": "krea2-lora",
+ "ideogram4-lora": "ideogram4-lora",
}
export interface ParamDiagnostics {
@@ -65,7 +86,7 @@ export function pickCarryOverFields(
defaults: FormModel,
fieldDefaults?: FormModel,
): FormModel {
- const deny = new Set(CROSS_SCHEMA_DENY_KEYS)
+ const deny = denyKeysForDefaults(defaults)
const carried: FormModel = {}
for (const [key, value] of Object.entries(carry)) {
if (deny.has(key)) continue
@@ -88,7 +109,7 @@ export function sanitizePersistedDraft(saved: FormModel, defaults: FormModel): F
if (expected === undefined || draft.model_train_type === undefined || draft.model_train_type === expected) {
return draft
}
- for (const key of CROSS_SCHEMA_DENY_KEYS) delete draft[key]
+ for (const key of denyKeysForDefaults(defaults)) delete draft[key]
return draft
}
diff --git a/mikazuki/engines/musubi/__init__.py b/mikazuki/engines/musubi/__init__.py
index e7972815e..769a94a0a 100644
--- a/mikazuki/engines/musubi/__init__.py
+++ b/mikazuki/engines/musubi/__init__.py
@@ -1,6 +1,8 @@
"""Independent musubi-tuner backend integration for lora-scripts-next.
-First supported training type: Krea 2 LoRA (krea2-lora).
+Supported training types: Krea 2 LoRA (krea2-lora) and Ideogram 4 LoRA
+(ideogram4-lora). TRAIN_TYPE stays the pack's primary/legacy type for callers
+that only need one; multi-type callers should read manifest.TRAIN_TYPES.
"""
TRAIN_TYPE = "krea2-lora"
diff --git a/mikazuki/engines/musubi/adapter.py b/mikazuki/engines/musubi/adapter.py
index 012304f7b..3e498ade0 100644
--- a/mikazuki/engines/musubi/adapter.py
+++ b/mikazuki/engines/musubi/adapter.py
@@ -6,10 +6,17 @@
import re
import sys
+from .families import (
+ DEFAULT_FAMILY,
+ FAMILIES,
+ family_for,
+)
from .settings import RuntimeConfig
-NETWORK_MODULE = "musubi_tuner.networks.lora_krea2"
+# Legacy alias for the pack's original Krea 2 module; family-aware code should
+# read FamilySpec.network_module instead.
+NETWORK_MODULE = FAMILIES[DEFAULT_FAMILY].network_module
UI_ONLY_FIELDS = {
"model_train_type",
@@ -37,6 +44,7 @@
"vae",
"text_encoder",
"turbo_dit",
+ "unconditional_dit",
"network_weights",
"base_weights",
"sample_prompts",
@@ -53,6 +61,18 @@
"turbo_dit_cache",
"fp8_base",
"fp8_scaled",
+ # ideogram4 (Ideogram 4): FP8-only base, optional asymmetric-CFG DiT and
+ # caption / sampling controls.
+ "unconditional_dit",
+ "use_unconditional_dit_for_lora_sampling",
+ "dit_dtype",
+ "vae_dtype",
+ "text_cache_dtype",
+ "validate_caption_structure",
+ "warn_on_caption_issues",
+ "log_loss_stats",
+ "sampler_preset",
+ "initial_sigma",
# training
"max_train_epochs",
"max_train_steps",
@@ -146,6 +166,10 @@
"turbo_dit_cache",
"fp8_base",
"fp8_scaled",
+ "use_unconditional_dit_for_lora_sampling",
+ "validate_caption_structure",
+ "warn_on_caption_issues",
+ "log_loss_stats",
"gradient_checkpointing",
"persistent_data_loader_workers",
"use_pinned_memory_for_block_swap",
@@ -206,10 +230,17 @@
"network_dropout",
"scale_weight_norms",
"base_weights_multiplier",
+ "initial_sigma",
}
DATASET_GENERAL_KEYS = {"resolution", "caption_extension", "batch_size", "enable_bucket", "bucket_no_upscale"}
+# Cache-stage flags: forwarded to the caching scripts as CLI flags, but never
+# written into the train TOML (the train script only shares a subset of args).
+# NB: ``vae_dtype`` is also a legitimate train arg (the sampling VAE load), so it
+# stays in the TOML on purpose.
+CACHE_ONLY_FIELDS = {"text_cache_dtype"}
+
SUBSET_REPEAT_PATTERN = re.compile(r"^(\d+)_(.+)$")
@@ -370,7 +401,13 @@ def build_dataset_config(source: dict[str, Any], runtime: RuntimeConfig) -> dict
return {"general": general, "datasets": datasets}
-def adapt_config(source: dict[str, Any], runtime: RuntimeConfig, run_id: str) -> AdaptedConfig:
+def adapt_config(
+ source: dict[str, Any],
+ runtime: RuntimeConfig,
+ run_id: str,
+ family: str = DEFAULT_FAMILY,
+) -> AdaptedConfig:
+ spec = family_for(family)
warnings: list[str] = []
dataset = build_dataset_config(source, runtime)
@@ -414,42 +451,82 @@ def adapt_config(source: dict[str, Any], runtime: RuntimeConfig, run_id: str) ->
if custom_optimizer_args:
values["optimizer_args"] = normalize_kv_args([*values.get("optimizer_args", []), *custom_optimizer_args])
- for field_name, label in (("dit", "DiT 模型路径"), ("vae", "VAE 模型路径"), ("text_encoder", "Qwen3-VL 文本编码器路径")):
+ for field_name, label in (
+ ("dit", f"{spec.label} DiT 模型路径"),
+ ("vae", "VAE 模型路径"),
+ ("text_encoder", spec.text_encoder_label),
+ ):
if is_empty(values.get(field_name)):
raise AdapterError(f"缺少 {label} ({field_name})")
- # Krea 2 DiT is trained in bf16; fp16 is not supported by the trainer.
+ # Both trainers run bf16; fp16 is not supported by either.
mixed = str(values.get("mixed_precision", "") or "").strip().lower()
if mixed in {"", "none", "null"}:
values["mixed_precision"] = "bf16"
elif mixed == "fp16":
values["mixed_precision"] = "bf16"
- warnings.append("Krea 2 训练仅支持 bf16，已将 mixed_precision 从 fp16 改为 bf16")
+ warnings.append(f"{spec.label} 训练仅支持 bf16，已将 mixed_precision 从 fp16 改为 bf16")
elif mixed != "bf16":
- raise AdapterError(f"mixed_precision={mixed} 不被 Krea 2 支持，请使用 bf16")
-
- # Krea 2: fp8_base and fp8_scaled must be paired. Trainer asserts
- # ``fp8_scaled requires fp8_base``; plain fp8_base without scaled is also rejected.
- # Prefer both ON when either is requested (product default). Both OFF is allowed.
- fp8_base = bool(values.get("fp8_base"))
- fp8_scaled = bool(values.get("fp8_scaled"))
- if fp8_base or fp8_scaled:
- if not (fp8_base and fp8_scaled):
- values["fp8_base"] = True
- values["fp8_scaled"] = True
- warnings.append("Krea 2 的 fp8 须同时开启 fp8_base 与 fp8_scaled，已自动成对开启")
-
- # RAW-train / Turbo-sample constraints (see musubi krea2_train_network.py).
- if not is_empty(values.get("turbo_dit")):
- if int_value(values.get("blocks_to_swap")) > 0:
+ raise AdapterError(f"mixed_precision={mixed} 不被 {spec.label} 支持，请使用 bf16")
+
+ if spec.supports_fp8_pair:
+ # Krea 2: fp8_base and fp8_scaled must be paired. Trainer asserts
+ # ``fp8_scaled requires fp8_base``; plain fp8_base without scaled is also rejected.
+ # Prefer both ON when either is requested (product default). Both OFF is allowed.
+ fp8_base = bool(values.get("fp8_base"))
+ fp8_scaled = bool(values.get("fp8_scaled"))
+ if fp8_base or fp8_scaled:
+ if not (fp8_base and fp8_scaled):
+ values["fp8_base"] = True
+ values["fp8_scaled"] = True
+ warnings.append("Krea 2 的 fp8 须同时开启 fp8_base 与 fp8_scaled，已自动成对开启")
+ else:
+ # Ideogram 4's base is always FP8; upstream takes no --fp8_base/--fp8_scaled.
+ dropped = [key for key in ("fp8_base", "fp8_scaled") if values.pop(key, None) is not None]
+ if dropped:
+ warnings.append(f"{spec.label} 基底固定为 FP8，无需 {'/'.join(dropped)}，已移除")
+
+ if spec.supports_turbo_dit:
+ # RAW-train / Turbo-sample constraints (see musubi krea2_train_network.py).
+ if not is_empty(values.get("turbo_dit")):
+ if int_value(values.get("blocks_to_swap")) > 0:
+ raise AdapterError(
+ "turbo_dit 与 blocks_to_swap 互斥：block swap 的 offloader 会绕过外部权重切换，"
+ "导致 RAW/Turbo 权重混合。请关闭 blocks_to_swap，或去掉 turbo_dit 改用 RAW 采样"
+ )
+ if is_empty(values.get("sample_prompts")):
+ warnings.append("已设置 turbo_dit 但未配置 sample_prompts；Turbo 仅用于训练中采样预览")
+ elif values.get("turbo_dit_cache"):
+ raise AdapterError("turbo_dit_cache 需要同时设置 turbo_dit（Turbo DiT 模型路径）")
+ else:
+ for key in ("turbo_dit", "turbo_dit_cache"):
+ if values.pop(key, None) is not None:
+ warnings.append(f"{spec.label} 不支持 Krea 2 的 {key}，已忽略")
+
+ if spec.blocks_to_swap_max and int_value(values.get("blocks_to_swap")) > spec.blocks_to_swap_max:
+ raise AdapterError(
+ f"blocks_to_swap={int_value(values.get('blocks_to_swap'))} 超出 {spec.label} 上限 "
+ f"{spec.blocks_to_swap_max}"
+ )
+
+ if not spec.supports_fp8_pair:
+ # Ideogram 4 uses plain MSE flow matching; upstream rejects other schemes.
+ scheme = str(values.get("weighting_scheme", "") or "").strip()
+ if scheme and scheme != "none":
+ values.pop("weighting_scheme", None)
+ warnings.append(f"{spec.label} 仅支持 weighting_scheme=none，已忽略 {scheme}")
+ # Asymmetric CFG is inference-only unless explicitly opted in for LoRA sampling.
+ unconditional = values.get("unconditional_dit")
+ opt_in = bool(values.get("use_unconditional_dit_for_lora_sampling"))
+ if opt_in and is_empty(unconditional):
raise AdapterError(
- "turbo_dit 与 blocks_to_swap 互斥：block swap 的 offloader 会绕过外部权重切换，"
- "导致 RAW/Turbo 权重混合。请关闭 blocks_to_swap，或去掉 turbo_dit 改用 RAW 采样"
+ "use_unconditional_dit_for_lora_sampling 需要同时设置 unconditional_dit（unconditional DiT 路径）"
+ )
+ if not opt_in and not is_empty(unconditional):
+ warnings.append(
+ f"{spec.label} 训练中采样默认只用 conditional DiT；unconditional_dit 需与 "
+ "use_unconditional_dit_for_lora_sampling 同时开启才会用于训练采样"
)
- if is_empty(values.get("sample_prompts")):
- warnings.append("已设置 turbo_dit 但未配置 sample_prompts；Turbo 仅用于训练中采样预览")
- elif values.get("turbo_dit_cache"):
- raise AdapterError("turbo_dit_cache 需要同时设置 turbo_dit（Turbo DiT 模型路径）")
# torch.compile needs Triton, which is unavailable on Windows.
if sys.platform == "win32" and values.get("compile"):
@@ -458,14 +535,19 @@ def adapt_config(source: dict[str, Any], runtime: RuntimeConfig, run_id: str) ->
values.pop(key, None)
warnings.append("compile 在 Windows 上不可用（依赖仅限 Linux 的 Triton），已自动关闭")
- values["network_module"] = NETWORK_MODULE
- if not is_empty(source.get("network_module")) and source["network_module"] != NETWORK_MODULE:
- warnings.append(f"network_module={source['network_module']} 已替换为 {NETWORK_MODULE}")
+ values["network_module"] = spec.network_module
+ if not is_empty(source.get("network_module")) and source["network_module"] != spec.network_module:
+ warnings.append(f"network_module={source['network_module']} 已替换为 {spec.network_module}")
- values.setdefault("network_dim", 32)
- values.setdefault("network_alpha", 32)
+ values.setdefault("network_dim", spec.default_network_dim)
+ values.setdefault("network_alpha", spec.default_network_alpha)
values.setdefault("log_with", "tensorboard")
+ if not is_empty(values.get("sample_prompts")):
+ for key, default in spec.sampler_defaults.items():
+ if is_empty(values.get(key)):
+ values[key] = default
+
output_dir = source.get("output_dir")
values["output_dir"] = resolve_path(output_dir, runtime.lora_next_root) if not is_empty(output_dir) else runtime.output_dir.as_posix()
logging_dir = source.get("logging_dir")
@@ -490,8 +572,16 @@ def toml_scalar(value: Any) -> str:
return f'"{escaped}"'
+def train_toml_values(values: dict[str, Any]) -> dict[str, Any]:
+ """Values that belong in the train TOML (drops cache-stage-only flags)."""
+ return {key: value for key, value in values.items() if key not in CACHE_ONLY_FIELDS}
+
+
def dump_train_toml(values: dict[str, Any]) -> str:
- return "".join(f"{key} = {toml_scalar(value)}\n" for key, value in values.items())
+ return "".join(
+ f"{key} = {toml_scalar(value)}\n"
+ for key, value in train_toml_values(values).items()
+ )
def dump_dataset_toml(dataset: dict[str, Any]) -> str:
diff --git a/mikazuki/engines/musubi/families.py b/mikazuki/engines/musubi/families.py
new file mode 100644
index 000000000..610117b87
--- /dev/null
+++ b/mikazuki/engines/musubi/families.py
@@ -0,0 +1,137 @@
+"""Model-family registry for the musubi-tuner engine pack.
+
+A family pairs one upstream musubi-tuner architecture with its entry scripts,
+LoRA network module and parameter constraints. The pack currently ships:
+
+- ``krea2`` — Krea 2 LoRA (RAW/Turbo DiT, fp8 pair)
+- ``ideogram4`` — Ideogram 4 LoRA (FP8-only base, Qwen3-VL-8B text encoder,
+ asymmetric CFG for inference)
+
+Everything family-specific lives here so adapter/launcher/preflight/run stay
+generic. ``manifest.TRAIN_TYPES`` maps train type -> family name.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass(frozen=True)
+class FamilySpec:
+ family: str
+ train_type: str
+ label: str
+ network_module: str
+ cache_latents_script: str
+ cache_text_encoder_script: str
+ train_script: str
+ text_encoder_label: str = "Qwen3-VL 文本编码器路径"
+ # Krea 2 only: optional Turbo DiT used for sampling; mutually exclusive with
+ # blocks_to_swap.
+ supports_turbo_dit: bool = False
+ # Krea 2 only: fp8_base/fp8_scaled must be enabled together. Ideogram 4's
+ # base is always FP8, so the flags must not leak into its TOML.
+ supports_fp8_pair: bool = False
+ # Upper bound for blocks_to_swap as documented upstream (0 = no limit).
+ blocks_to_swap_max: int = 0
+ # Krea 2 refuses to start without the local Qwen3-VL tokenizer directory
+ # (its encoder is patched to it). Ideogram 4 falls back to the Hub id when
+ # the directory is absent, so it only warns.
+ require_local_tokenizer: bool = True
+ vram_hint_mb: int = 12000
+ vram_hint: str = ""
+ default_network_dim: int = 32
+ default_network_alpha: int = 32
+ # Config keys forwarded to the cache stages as CLI flags.
+ cache_latents_args: tuple[str, ...] = ()
+ cache_text_encoder_args: tuple[str, ...] = ()
+ # Values appended to the train TOML as defaults when sampling is enabled.
+ sampler_defaults: dict[str, Any] = field(default_factory=dict)
+ # Defaults applied before prompt generation (UI-only sample keys).
+ sample_defaults: dict[str, Any] = field(default_factory=dict)
+ # Overrides applied on top of sample_defaults when a Turbo/aux DiT is set.
+ turbo_sample_defaults: dict[str, Any] = field(default_factory=dict)
+
+
+KREA2 = FamilySpec(
+ family="krea2",
+ train_type="krea2-lora",
+ label="Krea 2",
+ network_module="musubi_tuner.networks.lora_krea2",
+ cache_latents_script="krea2_cache_latents.py",
+ cache_text_encoder_script="krea2_cache_text_encoder_outputs.py",
+ train_script="krea2_train_network.py",
+ text_encoder_label="Qwen3-VL-4B 文本编码器路径",
+ supports_turbo_dit=True,
+ supports_fp8_pair=True,
+ vram_hint="建议开启 fp8_base + blocks_to_swap",
+ sample_defaults={"sample_cfg": 4.5, "sample_steps": 28},
+ turbo_sample_defaults={"sample_cfg": 1, "sample_steps": 8},
+)
+
+IDEOGRAM4 = FamilySpec(
+ family="ideogram4",
+ train_type="ideogram4-lora",
+ label="Ideogram 4",
+ network_module="musubi_tuner.networks.lora_ideogram4",
+ cache_latents_script="ideogram4_cache_latents.py",
+ cache_text_encoder_script="ideogram4_cache_text_encoder_outputs.py",
+ train_script="ideogram4_train_network.py",
+ text_encoder_label="Qwen3-VL-8B 文本编码器路径（qwen3vl_8b_fp8_scaled）",
+ blocks_to_swap_max=33,
+ require_local_tokenizer=False,
+ vram_hint_mb=16000,
+ vram_hint="建议提高 blocks_to_swap（上限 33）并开启 gradient_checkpointing，必要时关闭训练中采样预览",
+ cache_latents_args=("vae_dtype",),
+ cache_text_encoder_args=("text_cache_dtype", "validate_caption_structure", "warn_on_caption_issues"),
+ sampler_defaults={
+ "sampler_preset": "V4_DEFAULT_20",
+ "initial_sigma": 1.004,
+ },
+ sample_defaults={
+ "sample_cfg": 7,
+ "sample_steps": 20,
+ },
+)
+
+FAMILIES: dict[str, FamilySpec] = {spec.family: spec for spec in (KREA2, IDEOGRAM4)}
+DEFAULT_FAMILY = KREA2.family
+TRAIN_TYPE_TO_FAMILY: dict[str, str] = {spec.train_type: spec.family for spec in FAMILIES.values()}
+
+
+def family_for(family: str | None) -> FamilySpec:
+ """Resolve a family name; unknown/empty values fall back to Krea 2."""
+ return FAMILIES.get(str(family or "").strip(), KREA2)
+
+
+def family_for_variant(variant: str | None) -> FamilySpec:
+ """Resolve the registry variant (``manifest.TRAIN_TYPES`` value)."""
+ return family_for(variant)
+
+
+def family_for_train_type(train_type: str | None) -> FamilySpec:
+ """Resolve a page-level train type (``model_train_type``) to its family."""
+ return family_for(TRAIN_TYPE_TO_FAMILY.get(str(train_type or "").strip()))
+
+
+def cache_extra_args(
+ spec: FamilySpec,
+ stage: str,
+ values: dict[str, Any],
+) -> list[str]:
+ """CLI flags for a cache stage, taken from the adapted config values."""
+ keys = spec.cache_latents_args if stage == "cache_latents" else spec.cache_text_encoder_args
+ args: list[str] = []
+ for key in keys:
+ raw = values.get(key)
+ if raw is None or raw is False:
+ continue
+ if raw is True:
+ args.append(f"--{key}")
+ continue
+ text = str(raw).strip()
+ if not text:
+ continue
+ args.extend((f"--{key}", text))
+ return args
diff --git a/mikazuki/engines/musubi/launcher.py b/mikazuki/engines/musubi/launcher.py
index a3938937d..73ea339ae 100644
--- a/mikazuki/engines/musubi/launcher.py
+++ b/mikazuki/engines/musubi/launcher.py
@@ -10,6 +10,7 @@
except ModuleNotFoundError: # Python 3.10 runtime
import toml as tomllib # type: ignore
+from .families import DEFAULT_FAMILY, family_for
from .settings import RuntimeConfig
_VALID_ACCELERATE_MIXED_PRECISION = {"no", "fp16", "bf16", "fp8"}
@@ -83,14 +84,18 @@ def build_cache_latents_spec(
task_id: str,
gpu_ids: list[str] | None = None,
skip_existing: bool = True,
+ *,
+ family: str = DEFAULT_FAMILY,
+ extra_args: list[str] | None = None,
) -> LaunchSpec:
command = [
str(runtime.python),
- str(runtime.musubi_root / "krea2_cache_latents.py"),
+ str(runtime.musubi_root / family_for(family).cache_latents_script),
"--dataset_config",
str(dataset_toml),
"--vae",
vae,
+ *(extra_args or []),
]
if skip_existing:
command.append("--skip_existing")
@@ -108,14 +113,18 @@ def build_cache_text_encoder_spec(
task_id: str,
gpu_ids: list[str] | None = None,
skip_existing: bool = True,
+ *,
+ family: str = DEFAULT_FAMILY,
+ extra_args: list[str] | None = None,
) -> LaunchSpec:
command = [
str(runtime.python),
- str(runtime.musubi_root / "krea2_cache_text_encoder_outputs.py"),
+ str(runtime.musubi_root / family_for(family).cache_text_encoder_script),
"--dataset_config",
str(dataset_toml),
"--text_encoder",
text_encoder,
+ *(extra_args or []),
]
if skip_existing:
command.append("--skip_existing")
@@ -133,14 +142,15 @@ def build_train_spec(
gpu_ids: list[str] | None = None,
*,
cpu_threads: int = 1,
+ family: str = DEFAULT_FAMILY,
) -> LaunchSpec:
- """Build musubi Krea2 train argv via ``accelerate launch`` (multi-GPU capable).
+ """Build musubi train argv via ``accelerate launch`` (multi-GPU capable).
- Upstream musubi-tuner documents ``accelerate launch … krea2_train_network.py``.
+ Upstream musubi-tuner documents ``accelerate launch … _train_network.py``.
We use the musubi venv's ``python -m accelerate.commands.launch`` so the
correct Accelerate/Torch stack is used (not the Kohya/main env).
"""
- script = runtime.musubi_root / "krea2_train_network.py"
+ script = runtime.musubi_root / family_for(family).train_script
launch_opts = [
"--num_cpu_threads_per_process",
str(cpu_threads),
diff --git a/mikazuki/engines/musubi/manifest.py b/mikazuki/engines/musubi/manifest.py
index a6251409c..4c78fb928 100644
--- a/mikazuki/engines/musubi/manifest.py
+++ b/mikazuki/engines/musubi/manifest.py
@@ -5,6 +5,7 @@
TRAIN_TYPES = {
"krea2-lora": "krea2",
+ "ideogram4-lora": "ideogram4",
}
UPSTREAM = {
@@ -20,9 +21,9 @@
FEATURE_FLAG_ENV = "LORA_ENABLE_MUSUBI"
CAPABILITIES = {
- "model_families": ["krea2"],
+ "model_families": ["krea2", "ideogram4"],
"tasks": ["lora"],
- "variants": ["krea2"],
+ "variants": ["krea2", "ideogram4"],
}
PATCHES = []
diff --git a/mikazuki/engines/musubi/preflight.py b/mikazuki/engines/musubi/preflight.py
index 3c17c8e98..e5a607eb4 100644
--- a/mikazuki/engines/musubi/preflight.py
+++ b/mikazuki/engines/musubi/preflight.py
@@ -7,6 +7,7 @@
import os
import subprocess
+from .families import DEFAULT_FAMILY, family_for
from .settings import RuntimeConfig
@@ -116,18 +117,32 @@ def run_preflight(
runtime: RuntimeConfig,
dataset: dict[str, Any] | None = None,
probe: DependencyProbe = default_dependency_probe,
+ family: str = DEFAULT_FAMILY,
) -> PreflightResult:
"""Validate adapted musubi config values (post-adapter, absolute paths)."""
+ spec = family_for(family)
errors: list[str] = []
warnings: list[str] = []
facts: dict[str, Any] = {
"musubi_root": str(runtime.musubi_root),
"python": str(runtime.python),
+ "model_family": spec.family,
}
if not runtime.python.is_file():
errors.append(f"musubi-tuner venv python 不存在: {runtime.python}")
+ missing_scripts = [
+ name
+ for name in (spec.cache_latents_script, spec.cache_text_encoder_script, spec.train_script)
+ if not (runtime.musubi_root / name).is_file()
+ ]
+ if missing_scripts:
+ errors.append(
+ f"已安装的 musubi-tuner 缺少 {spec.label} 脚本（{', '.join(missing_scripts)}），"
+ "请在「设置 → 训练引擎」对 musubi-tuner 执行修复/重装"
+ )
+
for field_name in ("dit", "vae", "text_encoder"):
raw = values.get(field_name)
if raw is None or str(raw).strip() == "":
@@ -135,16 +150,25 @@ def run_preflight(
elif not Path(str(raw)).is_file():
errors.append(f"模型文件不存在: {field_name}={raw}")
if values.get("text_encoder"):
- from mikazuki.model_assets import dir_complete, krea2_tokenizer_dir
+ from mikazuki.model_assets import dir_complete, tokenizer_dir_for
- tokenizer_dir = krea2_tokenizer_dir(runtime.lora_next_root)
+ tokenizer_dir = tokenizer_dir_for(spec.train_type, runtime.lora_next_root)
if not dir_complete(tokenizer_dir):
- errors.append(
+ message = (
f"Qwen3-VL tokenizer 文件不在 {tokenizer_dir}，请先在「训练用模型」区探测并下载 tokenizer"
)
+ if spec.require_local_tokenizer:
+ errors.append(message)
+ else:
+ warnings.append(
+ f"{message}；否则训练时将尝试从 Hugging Face 下载"
+ )
turbo = values.get("turbo_dit")
if turbo is not None and str(turbo).strip() and not Path(str(turbo)).is_file():
errors.append(f"Turbo DiT 文件不存在: turbo_dit={turbo}")
+ unconditional = values.get("unconditional_dit")
+ if unconditional is not None and str(unconditional).strip() and not Path(str(unconditional)).is_file():
+ errors.append(f"unconditional DiT 文件不存在: unconditional_dit={unconditional}")
images: list[Path] = []
for entry in (dataset or {}).get("datasets", []):
@@ -175,11 +199,12 @@ def run_preflight(
major, minor = (int(part) for part in dep.transformers_version.split(".")[:2])
if (major, minor) < (4, 57):
errors.append(
- f"Krea 2 的 Qwen3-VL 文本编码器需要 transformers>=4.57，当前为 {dep.transformers_version}"
+ f"{spec.label} 的 Qwen3-VL 文本编码器需要 transformers>=4.57，当前为 {dep.transformers_version}"
)
except (ValueError, TypeError):
warnings.append(f"无法解析 transformers 版本: {dep.transformers_version}")
- if dep.vram_total_mb and dep.vram_total_mb < 12000:
- warnings.append(f"显存 {dep.vram_total_mb} MB 可能不足以训练 Krea 2，建议开启 fp8_base + blocks_to_swap")
+ if dep.vram_total_mb and dep.vram_total_mb < spec.vram_hint_mb:
+ hint = f"，{spec.vram_hint}" if spec.vram_hint else ""
+ warnings.append(f"显存 {dep.vram_total_mb} MB 可能不足以训练 {spec.label}{hint}")
return PreflightResult(ok=not errors, errors=errors, warnings=warnings, facts=facts)
diff --git a/mikazuki/engines/musubi/routes.py b/mikazuki/engines/musubi/routes.py
index 9d7496537..05e9767ee 100644
--- a/mikazuki/engines/musubi/routes.py
+++ b/mikazuki/engines/musubi/routes.py
@@ -7,12 +7,14 @@
from mikazuki.app.models import APIResponseFail, APIResponseSuccess
from mikazuki.engines.musubi import TRAIN_TYPE as MUSUBI_TRAIN_TYPE
-from mikazuki.engines.musubi.manifest import UPSTREAM
+from mikazuki.engines.musubi.families import family_for_train_type
+from mikazuki.engines.musubi.manifest import TRAIN_TYPES, UPSTREAM
from mikazuki.engines.musubi.adapter import (
AdapterError as MusubiAdapterError,
adapt_config as adapt_musubi_config,
dump_dataset_toml as dump_musubi_dataset_toml,
dump_train_toml as dump_musubi_train_toml,
+ train_toml_values as musubi_train_toml_values,
)
from mikazuki.engines.musubi.environment import start_install_task as start_musubi_install_task
from mikazuki.engines.musubi.extension_state import (
@@ -45,6 +47,7 @@ async def status():
runtime = musubi_runtime()
data["feature_enabled"] = musubi_feature_enabled()
data["train_type"] = MUSUBI_TRAIN_TYPE
+ data["train_types"] = sorted(TRAIN_TYPES)
data["runtime"] = {
"musubi_root": str(runtime.musubi_root),
"python": str(runtime.python),
@@ -65,11 +68,12 @@ async def preflight(config: dict):
repair_musubi_layout_venv(musubi_default_layout(Path.cwd()))
runtime = musubi_runtime()
run_id = f"{datetime.now().strftime('%Y%m%d-%H%M%S')}-musubi"
+ family = family_for_train_type(config.get("model_train_type"))
try:
- adapted = adapt_musubi_config(config, runtime, run_id)
+ adapted = adapt_musubi_config(config, runtime, run_id, family=family.family)
except MusubiAdapterError as exc:
return APIResponseFail(message=str(exc))
- result = run_musubi_preflight(adapted.values, runtime, adapted.dataset)
+ result = run_musubi_preflight(adapted.values, runtime, adapted.dataset, family=family.family)
result.warnings = [*adapted.warnings, *result.warnings]
if result.ok:
return APIResponseSuccess(data=result.as_dict())
@@ -82,12 +86,13 @@ async def dry_run(config: dict):
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
autosave_dir = os.path.join(os.getcwd(), "config", "autosave")
os.makedirs(autosave_dir, exist_ok=True)
+ family = family_for_train_type(config.get("model_train_type"))
config.pop("gpu_ids", None)
config.pop("model_train_type", None)
runtime = musubi_runtime()
run_id = f"{timestamp}-musubi"
try:
- adapted = adapt_musubi_config(config, runtime, run_id)
+ adapted = adapt_musubi_config(config, runtime, run_id, family=family.family)
except MusubiAdapterError as exc:
return APIResponseFail(message=str(exc))
toml_file_path = Path(autosave_dir) / f"{run_id}.toml"
@@ -98,7 +103,7 @@ async def dry_run(config: dict):
return APIResponseSuccess(data={
"toml_path": str(toml_file_path),
"dataset_toml_path": str(dataset_file_path),
- "config": adapted.values,
+ "config": musubi_train_toml_values(adapted.values),
"dataset": adapted.dataset,
"warnings": adapted.warnings,
})
diff --git a/mikazuki/engines/musubi/run.py b/mikazuki/engines/musubi/run.py
index 69d3a5782..82568fce7 100644
--- a/mikazuki/engines/musubi/run.py
+++ b/mikazuki/engines/musubi/run.py
@@ -31,6 +31,7 @@
discover_runtime as discover_musubi_runtime,
feature_enabled as musubi_feature_enabled,
)
+from mikazuki.engines.musubi.families import DEFAULT_FAMILY, FAMILIES, family_for
from mikazuki.engines.runner import RunContext
from mikazuki.utils import train_utils
@@ -81,21 +82,27 @@ def musubi_fail_from_preflight(result):
return APIResponseFail(message=message, data=result.as_dict())
-def musubi_apply_sample_defaults(config: dict) -> None:
- """Krea 2 sampling defaults: 1024px; Turbo schedule wants CFG off + few steps."""
+def musubi_apply_sample_defaults(config: dict, spec) -> None:
+ """Family sampling defaults: 1024px; Krea 2 Turbo wants CFG off + few steps."""
config.setdefault("sample_width", 1024)
config.setdefault("sample_height", 1024)
- turbo = str(config.get("turbo_dit", "") or "").strip()
- if turbo:
- config.setdefault("sample_cfg", 1)
- config.setdefault("sample_steps", 8)
- else:
- config.setdefault("sample_cfg", 4.5)
- config.setdefault("sample_steps", 28)
+ defaults = dict(spec.sample_defaults)
+ if spec.turbo_sample_defaults and str(config.get("turbo_dit", "") or "").strip():
+ defaults.update(spec.turbo_sample_defaults)
+ for key, value in defaults.items():
+ config.setdefault(key, value)
def handle_run(config: dict, ctx: RunContext):
- model_train_type = "krea2-lora"
+ variant = str(getattr(ctx, "variant", "") or "").strip() or DEFAULT_FAMILY
+ if variant not in FAMILIES:
+ # Never silently fall back to Krea 2: a wrong family would launch the
+ # wrong scripts against the user's weights.
+ return APIResponseFail(
+ message=f"未注册的 musubi 模型族: {variant}（可用: {', '.join(sorted(FAMILIES))}）"
+ )
+ spec = family_for(variant)
+ model_train_type = spec.train_type
if not musubi_feature_enabled():
return musubi_disabled_response()
ready, failure = musubi_ready_gate()
@@ -116,7 +123,7 @@ def handle_run(config: dict, ctx: RunContext):
config["sample_prompts"] = prompt_file
train_utils.normalize_sample_prompt_file(prompt_file)
elif should_generate_sample_prompts(config):
- musubi_apply_sample_defaults(config)
+ musubi_apply_sample_defaults(config, spec)
try:
positive_prompt, sample_prompts_arg = get_sample_prompts(config=config, model_train_type=model_train_type)
if positive_prompt is not None and train_utils.is_promopt_like(sample_prompts_arg):
@@ -141,8 +148,8 @@ def handle_run(config: dict, ctx: RunContext):
sanitize_config(config)
run_id = f"{ctx.timestamp}-musubi"
- adapted = adapt_musubi_config(config, runtime, run_id)
- preflight = run_musubi_preflight(adapted.values, runtime, adapted.dataset)
+ adapted = adapt_musubi_config(config, runtime, run_id, family=spec.family)
+ preflight = run_musubi_preflight(adapted.values, runtime, adapted.dataset, family=spec.family)
if not preflight.ok:
return musubi_fail_from_preflight(preflight)
toml_file_path = Path(ctx.autosave_dir) / f"{run_id}.toml"
@@ -158,7 +165,7 @@ def handle_run(config: dict, ctx: RunContext):
"warnings": [*adapted.warnings, *preflight.warnings],
}
return process.run_musubi_train(
- str(toml_file_path), runtime, adapted.values, ctx.gpu_ids, metadata=metadata
+ str(toml_file_path), runtime, adapted.values, ctx.gpu_ids, metadata=metadata, family=spec.family
)
except MusubiAdapterError as exc:
return APIResponseFail(message=str(exc))
diff --git a/mikazuki/model_assets.py b/mikazuki/model_assets.py
index c0eaa3b80..85608ee83 100644
--- a/mikazuki/model_assets.py
+++ b/mikazuki/model_assets.py
@@ -18,6 +18,7 @@
from pathlib import Path
from typing import Any, Callable
import os
+import re
try:
import tomllib
@@ -47,6 +48,8 @@ class AssetDef:
KREA2_REPO = "Comfy-Org/Krea-2"
QWEN3_VL_REPO = "Qwen/Qwen3-VL-4B-Instruct"
+IDEOGRAM4_REPO = "Comfy-Org/Ideogram-4"
+IDEOGRAM4_TOKENIZER_REPO = "Qwen/Qwen3-VL-8B-Instruct"
TOKENIZER_FILES = ["tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt"]
TOKENIZER_REQUIRED = ("tokenizer.json", "tokenizer_config.json")
@@ -113,6 +116,27 @@ def _klein_assets(variant: str, repo: str, dit_file: str, te_repo: str, te_dir:
AssetDef("tokenizer", "Qwen3-VL tokenizer（目录）", "sd-models/krea2/qwen3-vl-tokenizer", kind="dir",
hf_repo=QWEN3_VL_REPO, ms_repo=QWEN3_VL_REPO),
),
+ # Ideogram 4 ships quantized-only weights (FP8) under the Ideogram
+ # Non-Commercial license; the pack must not bundle them, users fetch them.
+ # ModelScope carries the same Comfy-Org mirror layout (verified 2026-10-10),
+ # so both download sources work out of the box.
+ "ideogram4-lora": (
+ AssetDef("dit", "Ideogram 4 conditional DiT（FP8 底模，非商用许可）", "sd-models/ideogram4/ideogram4_fp8_scaled.safetensors",
+ hf_repo=IDEOGRAM4_REPO, hf_file="diffusion_models/ideogram4_fp8_scaled.safetensors",
+ ms_repo=IDEOGRAM4_REPO, ms_file="diffusion_models/ideogram4_fp8_scaled.safetensors"),
+ AssetDef("unconditional_dit", "Ideogram 4 unconditional DiT（非对称 CFG，可选）",
+ "sd-models/ideogram4/ideogram4_unconditional_fp8_scaled.safetensors", optional=True,
+ hf_repo=IDEOGRAM4_REPO, hf_file="diffusion_models/ideogram4_unconditional_fp8_scaled.safetensors",
+ ms_repo=IDEOGRAM4_REPO, ms_file="diffusion_models/ideogram4_unconditional_fp8_scaled.safetensors"),
+ AssetDef("text_encoder", "文本编码器（Qwen3-VL-8B，FP8）", "sd-models/ideogram4/qwen3vl_8b_fp8_scaled.safetensors",
+ hf_repo=IDEOGRAM4_REPO, hf_file="text_encoders/qwen3vl_8b_fp8_scaled.safetensors",
+ ms_repo=IDEOGRAM4_REPO, ms_file="text_encoders/qwen3vl_8b_fp8_scaled.safetensors"),
+ AssetDef("vae", "VAE（Flux2）", "sd-models/ideogram4/flux2-vae.safetensors",
+ hf_repo=IDEOGRAM4_REPO, hf_file="vae/flux2-vae.safetensors",
+ ms_repo=IDEOGRAM4_REPO, ms_file="vae/flux2-vae.safetensors"),
+ AssetDef("tokenizer", "Qwen3-VL-8B tokenizer（目录）", "sd-models/ideogram4/qwen3-vl-8b-tokenizer", kind="dir",
+ hf_repo=IDEOGRAM4_TOKENIZER_REPO, ms_repo=IDEOGRAM4_TOKENIZER_REPO),
+ ),
}
@@ -181,13 +205,17 @@ def asset_dir_complete(asset: AssetDef, path: Path) -> bool:
return True
-def krea2_tokenizer_dir(project_root: Path) -> Path:
- for asset in manifest_for("krea2-lora"):
+def tokenizer_dir_for(train_type: str, project_root: Path) -> Path:
+ for asset in manifest_for(train_type):
if asset.key == "tokenizer":
return _target_path(asset.default_path, project_root)
return _target_path("sd-models/krea2/qwen3-vl-tokenizer", project_root)
+def krea2_tokenizer_dir(project_root: Path) -> Path:
+ return tokenizer_dir_for("krea2-lora", project_root)
+
+
KREA2_ENCODER_REL = Path("src/musubi_tuner/krea2/krea2_encoder.py")
_TOKENIZER_PATCH_MARK = "# mikazuki: patched tokenizer path"
_TOKENIZER_LINES = (
@@ -241,6 +269,87 @@ def patch_krea2_tokenizer_everywhere(project_root: Path, log: Callable[[str], No
return patched
+IDEOGRAM4_UTILS_REL = Path("src/musubi_tuner/ideogram4/ideogram4_utils.py")
+_ROTARY_PATCH_MARK = "# mikazuki: rebuild rotary embedding after to_empty"
+_ROTARY_ANCHOR = " _materialize_meta_tensors(model)"
+_ROTARY_PATCH = """ _materialize_meta_tensors(model)
+ # mikazuki: rebuild rotary embedding after to_empty (upstream #1139) — the
+ # non-persistent inv_freq buffer is not in the checkpoint, so to_empty()
+ # leaves it uninitialized and the cached text encoder outputs turn to junk.
+ _rotary = model.language_model.rotary_emb
+ try:
+ model.language_model.rotary_emb = type(_rotary)(config.text_config, device=_rotary.inv_freq.device)
+ except TypeError: # older/newer transformers builds may not accept the device kwarg
+ model.language_model.rotary_emb = type(_rotary)(config.text_config)"""
+_IDEOGRAM4_TOKENIZER_LINE = " return AutoTokenizer.from_pretrained(QWEN3_VL_8B_INSTRUCT_REPO_ID)"
+# Original Hub call or a previous local-directory patch: either shape can be
+# rewritten, so the patch can follow a moved/re-downloaded tokenizer directory
+# (and fall back to the Hub id when the local directory disappears).
+_IDEOGRAM4_TOKENIZER_CALL_RE = re.compile(
+ r'(?m)^\s*return AutoTokenizer\.from_pretrained\((?:r"[^"]*"|QWEN3_VL_8B_INSTRUCT_REPO_ID)\)(?:[^\n]*)$'
+)
+
+
+def _ideogram4_tokenizer_call(tokenizer_dir: Path) -> str:
+ if dir_complete(tokenizer_dir):
+ return f' return AutoTokenizer.from_pretrained(r"{tokenizer_dir.as_posix()}") {_TOKENIZER_PATCH_MARK}'
+ return _IDEOGRAM4_TOKENIZER_LINE
+
+
+def patch_ideogram4_source(source_root: Path, tokenizer_dir: Path, log: Callable[[str], None] = print) -> bool:
+ """Fix the Ideogram 4 text encoder in the musubi source tree.
+
+ Two idempotent source patches:
+ 1. rebuild the Qwen3-VL rotary embedding after ``to_empty()`` (upstream #1139);
+ 2. point the tokenizer at the local directory when it is complete, and
+ re-point (or revert to the Hub id) when that directory changes.
+ Refuses to write a result that does not compile.
+ """
+ module = Path(source_root) / IDEOGRAM4_UTILS_REL
+ if not module.is_file():
+ return False
+ text = module.read_text(encoding="utf-8")
+ patched = text
+ if _ROTARY_PATCH_MARK not in patched:
+ if _ROTARY_ANCHOR in patched:
+ patched = patched.replace(_ROTARY_ANCHOR, _ROTARY_PATCH, 1)
+ else:
+ log(f"[patch] rotary anchor not found in {module}; skipped")
+ desired_call = _ideogram4_tokenizer_call(Path(tokenizer_dir))
+ if _IDEOGRAM4_TOKENIZER_CALL_RE.search(patched):
+ patched = _IDEOGRAM4_TOKENIZER_CALL_RE.sub(desired_call, patched, count=1)
+ else:
+ log(f"[patch] tokenizer load line not found in {module}; skipped")
+ if patched == text:
+ return False
+ try:
+ compile(patched, str(module), "exec")
+ except SyntaxError as exc:
+ log(f"[patch] result would not compile ({exc}); {module} left unchanged")
+ return False
+ # Write atomically: a partial file would break the installed snapshot.
+ tmp = module.with_name(module.name + ".mikazuki-tmp")
+ tmp.write_text(patched, encoding="utf-8")
+ os.replace(tmp, module)
+ log(f"[patch] ideogram4 rotary rebuild + tokenizer -> {tokenizer_dir}")
+ return True
+
+
+def patch_ideogram4_everywhere(project_root: Path, log: Callable[[str], None] = print) -> bool:
+ """Patch every known musubi source root (installed extension, vendor, upstream cache)."""
+ from mikazuki.engines.musubi.extension_state import default_layout as musubi_default_layout
+
+ tokenizer_dir = tokenizer_dir_for("ideogram4-lora", project_root)
+ patched = False
+ for root in (
+ musubi_default_layout(project_root).source,
+ Path(project_root) / "vendor" / "musubi-tuner",
+ Path(project_root) / ".cache" / "musubi" / "upstream",
+ ):
+ patched = patch_ideogram4_source(root, tokenizer_dir, log) or patched
+ return patched
+
+
def check_assets(train_type: str, values: dict[str, Any], project_root: Path) -> list[dict]:
items = []
for asset in manifest_for(train_type):
@@ -298,6 +407,8 @@ def download_assets(
log(f"[done] {asset.label}")
if train_type == "krea2-lora" and asset.key == "tokenizer":
patch_krea2_tokenizer_everywhere(project_root, log)
+ elif train_type == "ideogram4-lora" and asset.key == "tokenizer":
+ patch_ideogram4_everywhere(project_root, log)
continue
target = _target_path(str(item.get("path") or asset.default_path), project_root)
target.parent.mkdir(parents=True, exist_ok=True)
diff --git a/mikazuki/process.py b/mikazuki/process.py
index 03094b6b5..38f432879 100644
--- a/mikazuki/process.py
+++ b/mikazuki/process.py
@@ -293,10 +293,17 @@ def run_musubi_train(toml_path: str,
runtime,
values: dict,
gpu_ids: Optional[list] = None,
- metadata: Optional[dict] = None):
- """Launch a musubi-tuner Krea 2 run: cache latents -> cache TE outputs -> train."""
- from mikazuki.model_assets import krea2_tokenizer_dir, patch_krea2_tokenizer_path
+ metadata: Optional[dict] = None,
+ family: Optional[str] = None):
+ """Launch a musubi-tuner run: cache latents -> cache TE outputs -> train."""
+ from mikazuki.engines.musubi.families import cache_extra_args, family_for
+ from mikazuki.model_assets import (
+ patch_ideogram4_source,
+ patch_krea2_tokenizer_path,
+ tokenizer_dir_for,
+ )
+ spec = family_for(family)
log.info(f"musubi-tuner training started with config file / musubi 训练开始，使用配置文件: {toml_path}")
if gpu_ids:
log.info(f"Using GPU(s) / 使用 GPU: {gpu_ids}")
@@ -307,20 +314,26 @@ def run_musubi_train(toml_path: str,
gpu_ids[0],
gpu_ids[0],
)
- patch_krea2_tokenizer_path(runtime.musubi_root, krea2_tokenizer_dir(runtime.lora_next_root), log=log.info)
+ tokenizer_dir = tokenizer_dir_for(spec.train_type, runtime.lora_next_root)
+ if spec.family == "krea2":
+ patch_krea2_tokenizer_path(runtime.musubi_root, tokenizer_dir, log=log.info)
+ else:
+ patch_ideogram4_source(runtime.musubi_root, tokenizer_dir, log=log.info)
train_task_id = str(uuid.uuid4())
dataset_toml = Path(str(values["dataset_config"]))
cache_latents_spec = build_cache_latents_spec(
- runtime, dataset_toml, str(values["vae"]), f"{train_task_id}-cache_latents", gpu_ids
+ runtime, dataset_toml, str(values["vae"]), f"{train_task_id}-cache_latents", gpu_ids,
+ family=spec.family, extra_args=cache_extra_args(spec, "cache_latents", values),
)
cache_te_spec = build_cache_text_encoder_spec(
- runtime, dataset_toml, str(values["text_encoder"]), f"{train_task_id}-cache_text_encoder", gpu_ids
+ runtime, dataset_toml, str(values["text_encoder"]), f"{train_task_id}-cache_text_encoder", gpu_ids,
+ family=spec.family, extra_args=cache_extra_args(spec, "cache_text_encoder", values),
)
- train_spec = build_train_spec(runtime, Path(toml_path), train_task_id, gpu_ids)
+ train_spec = build_train_spec(runtime, Path(toml_path), train_task_id, gpu_ids, family=spec.family)
base_metadata = {
"backend": "musubi",
- "train_type": "krea2-lora",
+ "train_type": spec.train_type,
"config_path": str(Path(toml_path).resolve()),
"dataset_config": str(dataset_toml.resolve()),
"musubi_root": str(runtime.musubi_root),
diff --git a/mikazuki/schema/ideogram4-lora.ts b/mikazuki/schema/ideogram4-lora.ts
new file mode 100644
index 000000000..eec2df130
--- /dev/null
+++ b/mikazuki/schema/ideogram4-lora.ts
@@ -0,0 +1,112 @@
+Schema.intersect([
+ Schema.object({
+ model_train_type: Schema.string().default("ideogram4-lora").disabled().description("训练种类"),
+ dit: Schema.string().role('filepicker', { type: "model-file" }).default("./sd-models/ideogram4/ideogram4_fp8_scaled.safetensors").description("Ideogram 4 conditional DiT（FP8 量化底模，官方仅发量化权重；非商用许可，需自行下载）"),
+ text_encoder: Schema.string().role('filepicker', { type: "model-file" }).default("./sd-models/ideogram4/qwen3vl_8b_fp8_scaled.safetensors").description("Qwen3-VL-8B 文本编码器路径（qwen3vl_8b_fp8_scaled）"),
+ vae: Schema.string().role('filepicker', { type: "model-file" }).default("./sd-models/ideogram4/flux2-vae.safetensors").description("VAE 模型路径（Flux2 VAE）"),
+ unconditional_dit: Schema.string().role('filepicker', { type: "model-file" }).description("unconditional DiT 路径（可选）。非对称 CFG 推理需要；训练中采样默认只用 conditional DiT，除非在下方开启对应开关"),
+ use_unconditional_dit_for_lora_sampling: Schema.boolean().default(false).description("训练中采样使用官方非对称 CFG（需要上面的 unconditional DiT 路径）"),
+ }).description("训练用模型"),
+
+ Schema.object({
+ timestep_sampling: Schema.union(["ideogram4_shift", "sigma", "uniform", "sigmoid", "shift", "flux_shift"]).default("ideogram4_shift").description("时间步采样。Ideogram 4 官方对齐的取样器为 ideogram4_shift，通常无需修改"),
+ min_timestep: Schema.number().min(0).max(1000).step(1).description("最小时间步（0-1000，留空不限制）"),
+ max_timestep: Schema.number().min(0).max(1000).step(1).description("最大时间步（0-1000，留空不限制）"),
+ validate_caption_structure: Schema.boolean().default(false).description("校验官方结构化 JSON caption（high_level_description / style_description / compositional_deconstruction）。纯文本 caption 默认也接受"),
+ warn_on_caption_issues: Schema.boolean().default(false).description("caption 校验失败时仅警告，不中断"),
+ log_loss_stats: Schema.boolean().default(false).description("输出 prediction / target 诊断统计（排查发散时使用）"),
+ }).description("Ideogram 4 专用参数"),
+
+ Schema.object({
+ train_data_dir: Schema.string().role('filepicker', { type: "folder", internal: "train-dir" }).default("./train/aki").description("训练数据集路径（子目录按 Kohya 约定命名为 重复次数_概念名）"),
+ resolution: Schema.string().default("1024,1024").description("训练图片分辨率，宽x高。支持非正方形，但必须是 16 倍数。"),
+ enable_bucket: Schema.boolean().default(true).description("启用 arb 桶以允许非固定宽高比的图片"),
+ bucket_no_upscale: Schema.boolean().default(false).description("arb 桶不放大图片"),
+ caption_extension: Schema.string().default(".txt").description("Tag 文件扩展名（官方为结构化 JSON caption；纯文本同样可用）"),
+ }).description("数据集设置"),
+
+ Schema.object({
+ output_name: Schema.string().default("next-ideogram4-lora").description("模型保存名称（Next Trainer · Ideogram 4 默认；建议按角色/风格自行改名）"),
+ output_dir: Schema.string().role('filepicker', { type: "folder" }).default("./output").description("模型保存文件夹"),
+ save_precision: Schema.union(["fp16", "float", "bf16"]).default("bf16").description("模型保存精度"),
+ save_every_n_epochs: Schema.number().default(2).description("每 N epoch（轮）自动保存一次模型"),
+ save_every_n_steps: Schema.number().min(1).description("每 N 步自动保存一次模型（与 save_every_n_epochs 二选一即可）"),
+ save_state: Schema.boolean().default(false).description("保存训练状态"),
+ }).description("保存设置"),
+
+ Schema.object({
+ max_train_epochs: Schema.number().min(1).default(16).description("最大训练 epoch（轮数）"),
+ max_train_steps: Schema.number().min(1).description("最大训练步数（设置了 epoch 时由 epoch 推导，可不填）"),
+ train_batch_size: Schema.number().min(1).default(1).description("批量大小, 越高显存占用越高"),
+ gradient_checkpointing: Schema.boolean().default(true).description("梯度检查点（9.3B 底模建议开启）"),
+ gradient_accumulation_steps: Schema.number().min(1).default(1).description("梯度累加步数"),
+ seed: Schema.number().default(42).description("随机种子"),
+ }).description("训练相关参数"),
+
+ Schema.object({
+ learning_rate: Schema.string().default("1e-4").description("学习率"),
+ lr_scheduler: Schema.union([
+ "linear",
+ "cosine",
+ "cosine_with_restarts",
+ "polynomial",
+ "constant",
+ "constant_with_warmup",
+ ]).default("constant").description("学习率调度器设置"),
+ lr_warmup_steps: Schema.number().default(0).description("学习率预热步数"),
+ optimizer_type: Schema.union(["AdamW", "AdamW8bit", "Adafactor"]).default("AdamW8bit").description("优化器设置"),
+ optimizer_args_custom: Schema.array(String).role('table').description('自定义 optimizer_args，一行一个，例如 weight_decay=0.01'),
+ max_grad_norm: Schema.number().step(0.01).default(1.0).description("梯度裁剪阈值，0 为不裁剪"),
+ }).description("学习率与优化器设置"),
+
+ Schema.object({
+ network_dim: Schema.number().min(1).default(32).description("网络维度（rank）"),
+ network_alpha: Schema.number().min(1).default(32).description("常用值：等于 network_dim 或 network_dim*1/2"),
+ network_dropout: Schema.number().step(0.01).default(0).description('dropout 概率'),
+ network_weights: Schema.string().role('filepicker', { type: "model-file" }).description("从已有的 LoRA 模型上继续训练，填写路径"),
+ scale_weight_norms: Schema.number().step(0.01).min(0).description("最大范数正则化。如果使用，推荐为 1"),
+ network_args_custom: Schema.array(String).role('table').description("自定义 network_args，一行一个。例如 exclude_patterns=['.*\\.mlp\\..*'] 只训练注意力层"),
+ }).description("网络设置"),
+
+ Schema.object({
+ mixed_precision: Schema.union(["bf16"]).default("bf16").description("训练混合精度（Ideogram 4 仅支持 bf16）"),
+ dit_dtype: Schema.union(["bfloat16", "float16", "float32"]).default("bfloat16").description("DiT 计算精度（底模权重始终以 FP8 加载，这里是计算 dtype）"),
+ vae_dtype: Schema.union(["bfloat16", "float16", "float32"]).default("bfloat16").description("VAE 计算精度（缓存 latents 阶段）"),
+ text_cache_dtype: Schema.union(["bf16", "fp8_e4m3fn"]).default("bf16").description("文本编码器缓存精度。每 token 需存 53248 通道：bf16 约 13.6 MB/图（1k 图 ≈ 13.6 GB），fp8_e4m3fn 可减半，但会损失精度"),
+ blocks_to_swap: Schema.number().min(0).max(33).step(1).default(0).description("交换到内存的 DiT block 数量以节省显存（Ideogram 4 上限 33）"),
+ sdpa: Schema.boolean().default(true).description("启用 sdpa"),
+ sage_attn: Schema.boolean().default(false).description("启用 SageAttention。Ideogram 4 的 head 维度为 256，不在 SageAttention 支持范围内，会失败"),
+ flash_attn: Schema.boolean().default(false).description("启用 Flash Attention（需要额外安装）"),
+ split_attn: Schema.boolean().default(false).description("注意力分拆计算，省显存但更慢"),
+ persistent_data_loader_workers: Schema.boolean().default(true).description("保留加载训练集的 worker，减少每个 epoch 之间的停顿。"),
+ max_data_loader_n_workers: Schema.number().min(0).default(8).description("数据加载器进程数"),
+ }).description("显存与精度设置"),
+
+ Schema.object({
+ enable_preview: Schema.boolean().default(false).description("启用训练中采样预览图（需要 text_encoder 与 vae 路径）"),
+ sample_every_n_epochs: Schema.number().default(2).description("每 N 个 epoch 生成一次预览图"),
+ sample_at_first: Schema.boolean().default(false).description("训练开始前先采样一次（用于验证模型与提示词配置）"),
+ positive_prompts: Schema.string().role('textarea').default('{"high_level_description": "a photo of a subject", "style_description": "natural lighting", "compositional_deconstruction": "centered subject, plain background"}').description("Prompt。官方推荐结构化 JSON caption（high_level_description / style_description / compositional_deconstruction）"),
+ negative_prompts: Schema.string().role('textarea').description("Negative Prompt（官方非对称 CFG 路径忽略该参数）"),
+ sampler_preset: Schema.union(["V4_DEFAULT_20", "V4_QUALITY_48", "V4_TURBO_12"]).default("V4_DEFAULT_20").description("采样调度：V4_DEFAULT_20（18+2 步）/ V4_QUALITY_48（45+3 步）/ V4_TURBO_12（11+1 步）"),
+ initial_sigma: Schema.number().step(0.001).default(1.004).description("首个去噪 sigma（官方默认 1.004）"),
+ sample_width: Schema.number().default(1024).description('预览图宽'),
+ sample_height: Schema.number().default(1024).description('预览图高'),
+ sample_cfg: Schema.number().min(1).max(30).default(7).description('CFG Scale。Ideogram 4 预览的步数与引导由 sampler_preset 决定，此项不生效'),
+ sample_seed: Schema.number().default(42).description('种子'),
+ sample_steps: Schema.number().min(1).max(300).default(20).description('迭代步数。Ideogram 4 预览的步数由 sampler_preset 决定，此项不生效'),
+ prompt_file: Schema.string().role('textarea').description('预览图 Prompt 文件路径。填写后将采用文件内的 prompt，而下方的选项将失效。'),
+ }).description("采样预览设置"),
+
+ 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("日志保存文件夹"),
+ wandb_api_key: Schema.string().description("wandb 的 api 密钥（log_with 选 wandb 时必填）"),
+ }).description("日志设置"),
+
+ Schema.object({
+ ui_custom_params: Schema.string().role('textarea').description("**危险** 自定义参数，请输入 TOML 格式，将会直接覆盖当前界面内任何参数。实时更新建议写完后再粘贴过来"),
+ }).description("其他设置"),
+])
diff --git a/mikazuki/utils/config_import.py b/mikazuki/utils/config_import.py
index 427100c71..693b0b9fc 100644
--- a/mikazuki/utils/config_import.py
+++ b/mikazuki/utils/config_import.py
@@ -7,11 +7,13 @@
from dataclasses import dataclass, field
from typing import Any
+from mikazuki.engines.musubi.families import IDEOGRAM4 as _IDEOGRAM4_SPEC, KREA2 as _KREA2_SPEC
from mikazuki.utils.train_utils import ensure_enable_preview_flag
ANIMA_TRAIN_TYPES = frozenset({"anima-lora", "sd3-lora"})
ANIMA_FAST_TRAIN_TYPES = frozenset({"anima-lora-fast"})
MUSUBI_TRAIN_TYPES = frozenset({"krea2-lora"})
+IDEOGRAM4_TRAIN_TYPES = frozenset({"ideogram4-lora"})
FLUX_TRAIN_TYPES = frozenset({"flux-lora", "flux-finetune"})
LUMINA_TRAIN_TYPES = frozenset({"lumina-lora"})
SDXL_TRAIN_TYPES = frozenset({"sdxl-lora", "sdxl-finetune"})
@@ -63,6 +65,16 @@
"fp8_scaled",
})
+IDEOGRAM4_CONFIG_MARKERS = frozenset({
+ "unconditional_dit",
+ "use_unconditional_dit_for_lora_sampling",
+ "sampler_preset",
+ "text_cache_dtype",
+ "validate_caption_structure",
+ "warn_on_caption_issues",
+ "log_loss_stats",
+})
+
LUMINA_CONFIG_MARKERS = frozenset({
"gemma2",
})
@@ -118,6 +130,11 @@
(re.compile(r"qwen3[-_]?vl", re.I), "文本编码器路径含 qwen3-vl"),
)
+IDEOGRAM4_PATH_RULES: tuple[tuple[re.Pattern[str], str], ...] = (
+ (re.compile(r"/ideogram4?/", re.I), "模型路径含 /ideogram/"),
+ (re.compile(r"ideogram", re.I), "模型路径含 ideogram"),
+)
+
LUMINA_PATH_RULES: tuple[tuple[re.Pattern[str], str], ...] = (
(re.compile(r"/lumina/", re.I), "模型路径含 /lumina/"),
(re.compile(r"gemma2", re.I), "路径含 gemma2"),
@@ -168,6 +185,12 @@
"accepted": MUSUBI_TRAIN_TYPES,
"default_train_type": "krea2-lora",
},
+ "ideogram4-lora": {
+ "label": "Ideogram 4 LoRA 训练",
+ "path": "/lora/ideogram4.html",
+ "accepted": IDEOGRAM4_TRAIN_TYPES,
+ "default_train_type": "ideogram4-lora",
+ },
"lumina-lora": {
"label": "Lumina LoRA 训练",
"path": "/lora/lumina.html",
@@ -206,6 +229,7 @@
"sd3-lora": {"path": "/lora/sd3.html", "label": "Anima LoRA 训练"},
"anima-lora-fast": {"path": "/lora/anima-fast.html", "label": "Anima Fast 训练"},
"krea2-lora": {"path": "/lora/krea2.html", "label": "Krea 2 LoRA 训练"},
+ "ideogram4-lora": {"path": "/lora/ideogram4.html", "label": "Ideogram 4 LoRA 训练"},
"flux-lora": {"path": "/lora/flux.html", "label": "Flux LoRA 训练"},
"flux-finetune": {"path": "/lora/flux.html", "label": "Flux 训练"},
"lumina-lora": {"path": "/lora/lumina.html", "label": "Lumina LoRA 训练"},
@@ -339,9 +363,18 @@ def analyze_train_type(config: dict) -> TrainTypeAnalysis:
config,
marker_keys=MUSUBI_CONFIG_MARKERS,
path_rules=MUSUBI_PATH_RULES,
+ network_modules=frozenset({_KREA2_SPEC.network_module}),
)
families.append(("krea2-lora", musubi_score, musubi_reasons))
+ ideogram4_score, ideogram4_reasons = _score_family(
+ config,
+ marker_keys=IDEOGRAM4_CONFIG_MARKERS,
+ path_rules=IDEOGRAM4_PATH_RULES,
+ network_modules=frozenset({_IDEOGRAM4_SPEC.network_module}),
+ )
+ families.append(("ideogram4-lora", ideogram4_score, ideogram4_reasons))
+
flux_score, flux_reasons = _score_family(
config,
marker_keys=FLUX_CONFIG_MARKERS,
diff --git a/tests/test_engine_registry.py b/tests/test_engine_registry.py
index db6e6eede..1f032449a 100644
--- a/tests/test_engine_registry.py
+++ b/tests/test_engine_registry.py
@@ -23,6 +23,7 @@ def test_discovers_migrated_packs():
def test_train_type_mapping():
mapping = registry.train_type_map()
assert mapping["krea2-lora"] == ("musubi", "krea2")
+ assert mapping["ideogram4-lora"] == ("musubi", "ideogram4")
assert mapping["anima-lora-fast"] == ("anima-fast", "anima")
assert mapping["sd-lora"] == ("kohya", "sd15")
assert mapping["flux-finetune"] == ("kohya", "flux")
@@ -35,6 +36,9 @@ def test_resolve_train_type_roundtrip():
pack, variant = registry.resolve_train_type("krea2-lora")
assert pack.engine_id == "musubi"
assert variant == "krea2"
+ pack, variant = registry.resolve_train_type("ideogram4-lora")
+ assert pack.engine_id == "musubi"
+ assert variant == "ideogram4"
assert registry.resolve_train_type("no-such-type") is None
diff --git a/tests/test_musubi_ideogram4.py b/tests/test_musubi_ideogram4.py
new file mode 100644
index 000000000..4029ff975
--- /dev/null
+++ b/tests/test_musubi_ideogram4.py
@@ -0,0 +1,530 @@
+from __future__ import annotations
+
+import tempfile
+import unittest
+from pathlib import Path
+
+from mikazuki.engines.musubi.adapter import AdapterError, adapt_config, dump_train_toml
+from mikazuki.engines.musubi.families import (
+ IDEOGRAM4,
+ KREA2,
+ cache_extra_args,
+ family_for_train_type,
+ family_for_variant,
+)
+from mikazuki.engines.musubi.launcher import (
+ build_cache_latents_spec,
+ build_cache_text_encoder_spec,
+ build_train_spec,
+)
+from mikazuki.engines.musubi.manifest import CAPABILITIES, TRAIN_TYPES
+from mikazuki.engines.musubi.preflight import ProbeFacts, run_preflight
+from mikazuki.engines.musubi.settings import RuntimeConfig
+from mikazuki.model_assets import check_assets, manifest_for, patch_ideogram4_source, tokenizer_dir_for
+from mikazuki.utils.config_import import PAGE_SPECS, TRAIN_TYPE_TARGETS, analyze_train_type
+
+
+def make_runtime(root: Path) -> RuntimeConfig:
+ return RuntimeConfig(
+ musubi_root=(root / "vendor" / "musubi-tuner").resolve(),
+ python=(root / "vendor" / "musubi-tuner" / ".venv" / "Scripts" / "python.exe").resolve(),
+ lora_next_root=root.resolve(),
+ output_dir=(root / "output" / "musubi").resolve(),
+ logging_dir=(root / "logs" / "musubi").resolve(),
+ cache_dir=(root / ".cache" / "musubi").resolve(),
+ )
+
+
+def ideogram4_config(root: Path) -> dict:
+ return {
+ "model_train_type": "ideogram4-lora",
+ "dit": "./sd-models/ideogram4/ideogram4_fp8_scaled.safetensors",
+ "vae": "./sd-models/ideogram4/flux2-vae.safetensors",
+ "text_encoder": "./sd-models/ideogram4/qwen3vl_8b_fp8_scaled.safetensors",
+ "train_data_dir": str(root / "train"),
+ "resolution": "1024,1024",
+ "train_batch_size": 1,
+ "learning_rate": "1e-4",
+ "mixed_precision": "bf16",
+ "max_train_epochs": 16,
+ }
+
+
+def write_runtime_files(runtime: RuntimeConfig) -> None:
+ runtime.musubi_root.mkdir(parents=True, exist_ok=True)
+ runtime.python.parent.mkdir(parents=True, exist_ok=True)
+ runtime.python.write_text("", encoding="utf-8")
+ for name in (
+ IDEOGRAM4.cache_latents_script,
+ IDEOGRAM4.cache_text_encoder_script,
+ IDEOGRAM4.train_script,
+ KREA2.cache_latents_script,
+ KREA2.cache_text_encoder_script,
+ KREA2.train_script,
+ ):
+ (runtime.musubi_root / name).write_text("", encoding="utf-8")
+
+
+def probe_with(vram_mb: int = 24576) -> object:
+ def probe(_runtime: RuntimeConfig) -> ProbeFacts:
+ return ProbeFacts(
+ python_version="3.12.0",
+ torch_version="2.6.0",
+ cuda_available=True,
+ cuda_version="12.8",
+ gpu_name="test-gpu",
+ vram_total_mb=vram_mb,
+ transformers_version="4.57.6",
+ )
+
+ return probe
+
+
+class FamilyRegistryTests(unittest.TestCase):
+ def test_manifest_registers_both_train_types(self):
+ self.assertEqual(TRAIN_TYPES["krea2-lora"], "krea2")
+ self.assertEqual(TRAIN_TYPES["ideogram4-lora"], "ideogram4")
+ self.assertIn("ideogram4", CAPABILITIES["model_families"])
+ self.assertIn("ideogram4", CAPABILITIES["variants"])
+
+ def test_variant_and_train_type_resolution(self):
+ self.assertIs(family_for_variant("ideogram4"), IDEOGRAM4)
+ self.assertIs(family_for_train_type("ideogram4-lora"), IDEOGRAM4)
+ self.assertIs(family_for_train_type("krea2-lora"), KREA2)
+ # Unknown/empty values keep the pack's original Krea 2 default.
+ self.assertIs(family_for_train_type("nope"), KREA2)
+ self.assertIs(family_for_variant(None), KREA2)
+
+ def test_family_scripts_and_modules(self):
+ self.assertEqual(IDEOGRAM4.network_module, "musubi_tuner.networks.lora_ideogram4")
+ self.assertEqual(IDEOGRAM4.cache_latents_script, "ideogram4_cache_latents.py")
+ self.assertEqual(IDEOGRAM4.train_script, "ideogram4_train_network.py")
+ self.assertFalse(IDEOGRAM4.supports_fp8_pair)
+ self.assertEqual(IDEOGRAM4.blocks_to_swap_max, 33)
+
+
+class AdapterIdeogram4Tests(unittest.TestCase):
+ def adapt(self, root: Path, **overrides):
+ config = ideogram4_config(root)
+ config.update(overrides)
+ return adapt_config(config, make_runtime(root), "run-id", family="ideogram4")
+
+ def test_network_module_and_sampler_defaults(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ (root / "train").mkdir()
+ adapted = self.adapt(root, sample_prompts=str(root / "prompts.txt"))
+
+ self.assertEqual(adapted.values["network_module"], "musubi_tuner.networks.lora_ideogram4")
+ self.assertEqual(adapted.values["sampler_preset"], "V4_DEFAULT_20")
+ self.assertAlmostEqual(adapted.values["initial_sigma"], 1.004)
+ self.assertEqual(adapted.values["network_dim"], 32)
+
+ def test_no_sampler_defaults_without_prompts(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ (root / "train").mkdir()
+ adapted = self.adapt(root)
+
+ self.assertNotIn("sampler_preset", adapted.values)
+ self.assertNotIn("initial_sigma", adapted.values)
+
+ def test_fp8_and_turbo_flags_are_dropped(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ (root / "train").mkdir()
+ adapted = self.adapt(
+ root,
+ fp8_base=True,
+ fp8_scaled=True,
+ turbo_dit="./sd-models/krea2/krea2-turbo.safetensors",
+ turbo_dit_cache=True,
+ )
+
+ for key in ("fp8_base", "fp8_scaled", "turbo_dit", "turbo_dit_cache"):
+ self.assertNotIn(key, adapted.values)
+ warnings = "\n".join(adapted.warnings)
+ self.assertIn("Ideogram 4", warnings)
+ self.assertIn("FP8", warnings)
+ self.assertIn("turbo_dit", warnings)
+
+ def test_unsupported_fields_raise(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ (root / "train").mkdir()
+ with self.assertRaises(AdapterError):
+ self.adapt(root, blocks_to_swap=40)
+
+ def test_weighting_scheme_is_ignored(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ (root / "train").mkdir()
+ adapted = self.adapt(root, weighting_scheme="sigmoid")
+
+ self.assertNotIn("weighting_scheme", adapted.values)
+ self.assertIn("weighting_scheme", "\n".join(adapted.warnings))
+
+ def test_unconditional_dit_requires_optin_for_sampling(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ (root / "train").mkdir()
+ unconditional = root / "uncond.safetensors"
+ unconditional.write_text("", encoding="utf-8")
+ adapted = self.adapt(root, unconditional_dit=str(unconditional), use_unconditional_dit_for_lora_sampling=True)
+ informative = self.adapt(root, unconditional_dit=str(unconditional))
+ with self.assertRaises(AdapterError):
+ self.adapt(root, use_unconditional_dit_for_lora_sampling=True)
+
+ self.assertTrue(adapted.values["use_unconditional_dit_for_lora_sampling"])
+ self.assertIn("unconditional_dit", informative.values)
+ self.assertIn("use_unconditional_dit_for_lora_sampling", "\n".join(informative.warnings))
+
+ def test_fp16_is_coerced_to_bf16(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ (root / "train").mkdir()
+ adapted = self.adapt(root, mixed_precision="fp16")
+
+ self.assertEqual(adapted.values["mixed_precision"], "bf16")
+ self.assertIn("Ideogram 4", "\n".join(adapted.warnings))
+
+ def test_krea2_fp8_pair_behaviour_is_unchanged(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ (root / "train").mkdir()
+ config = {
+ "model_train_type": "krea2-lora",
+ "dit": "./sd-models/krea2/krea2.safetensors",
+ "vae": "./sd-models/krea2/qwen_image_vae.safetensors",
+ "text_encoder": "./sd-models/krea2/qwen3_vl_4b.safetensors",
+ "train_data_dir": str(root / "train"),
+ "fp8_base": True,
+ }
+ adapted = adapt_config(config, make_runtime(root), "run-id", family="krea2")
+
+ self.assertTrue(adapted.values["fp8_base"])
+ self.assertTrue(adapted.values["fp8_scaled"])
+
+ def test_missing_model_field_uses_family_label(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ (root / "train").mkdir()
+ with self.assertRaises(AdapterError) as ctx:
+ self.adapt(root, text_encoder="")
+
+ self.assertIn("Qwen3-VL-8B", str(ctx.exception))
+
+ def test_cache_extra_args_per_stage(self):
+ values = {
+ "vae_dtype": "bfloat16",
+ "text_cache_dtype": "fp8_e4m3fn",
+ "validate_caption_structure": True,
+ "warn_on_caption_issues": False,
+ }
+ self.assertEqual(cache_extra_args(IDEOGRAM4, "cache_latents", values), ["--vae_dtype", "bfloat16"])
+ self.assertEqual(
+ cache_extra_args(IDEOGRAM4, "cache_text_encoder", values),
+ ["--text_cache_dtype", "fp8_e4m3fn", "--validate_caption_structure"],
+ )
+ self.assertEqual(cache_extra_args(KREA2, "cache_text_encoder", values), [])
+
+ def test_cache_only_fields_stay_out_of_the_train_toml(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ (root / "train").mkdir()
+ adapted = self.adapt(
+ root,
+ vae_dtype="bfloat16",
+ text_cache_dtype="fp8_e4m3fn",
+ validate_caption_structure=True,
+ dit_dtype="bfloat16",
+ )
+ toml_text = dump_train_toml(adapted.values)
+
+ # text_cache_dtype is cache-stage only; the rest must reach the train TOML
+ # (vae_dtype also drives the sampling VAE load upstream).
+ self.assertNotIn("text_cache_dtype", toml_text)
+ for key in ("vae_dtype", "validate_caption_structure", "dit_dtype"):
+ self.assertIn(key, toml_text)
+
+
+class LauncherIdeogram4Tests(unittest.TestCase):
+ def test_script_names_follow_family(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ runtime = make_runtime(root)
+ dataset_toml = root / "dataset.toml"
+ train_toml = root / "train.toml"
+ train_toml.write_text('mixed_precision = "bf16"\n', encoding="utf-8")
+
+ latents = build_cache_latents_spec(runtime, dataset_toml, "/models/vae.safetensors", "t", family="ideogram4")
+ te = build_cache_text_encoder_spec(runtime, dataset_toml, "/models/te.safetensors", "t", family="ideogram4")
+ train = build_train_spec(runtime, train_toml, "t", family="ideogram4")
+ krea2_train = build_train_spec(runtime, train_toml, "t", family="krea2")
+
+ self.assertTrue(latents.command[1].endswith("ideogram4_cache_latents.py"))
+ self.assertTrue(te.command[1].endswith("ideogram4_cache_text_encoder_outputs.py"))
+ self.assertTrue(any(part.endswith("ideogram4_train_network.py") for part in train.command))
+ self.assertTrue(any(part.endswith("krea2_train_network.py") for part in krea2_train.command))
+
+ def test_cache_stage_extra_args_are_appended(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ runtime = make_runtime(root)
+ latents = build_cache_latents_spec(
+ runtime,
+ root / "dataset.toml",
+ "/models/vae.safetensors",
+ "t",
+ family="ideogram4",
+ extra_args=["--vae_dtype", "bfloat16"],
+ )
+
+ self.assertIn("--vae_dtype", latents.command)
+ self.assertEqual(latents.command[latents.command.index("--vae_dtype") + 1], "bfloat16")
+
+
+class PreflightIdeogram4Tests(unittest.TestCase):
+ def test_missing_scripts_error_mentions_family_and_repair(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ runtime = make_runtime(root)
+ runtime.musubi_root.mkdir(parents=True)
+ runtime.python.parent.mkdir(parents=True)
+ runtime.python.write_text("", encoding="utf-8")
+ result = run_preflight({}, runtime, None, probe_with(), family="ideogram4")
+
+ self.assertFalse(result.ok)
+ joined = "\n".join(result.errors)
+ self.assertIn("Ideogram 4", joined)
+ self.assertIn("修复", joined)
+
+ def test_tokenizer_missing_warns_for_ideogram4(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ runtime = make_runtime(root)
+ write_runtime_files(runtime)
+ models = root / "models"
+ models.mkdir()
+ values = {}
+ for key in ("dit", "vae", "text_encoder"):
+ path = models / f"{key}.safetensors"
+ path.write_text("", encoding="utf-8")
+ values[key] = str(path)
+ result = run_preflight(values, runtime, None, probe_with(), family="ideogram4")
+
+ self.assertTrue(result.ok, result.errors)
+ joined = "\n".join(result.warnings)
+ self.assertIn("tokenizer", joined.lower())
+ self.assertIn("Hugging Face", joined)
+
+ def test_low_vram_hint_is_family_specific(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ runtime = make_runtime(root)
+ write_runtime_files(runtime)
+ models = root / "models"
+ models.mkdir()
+ values = {}
+ for key in ("dit", "vae", "text_encoder"):
+ path = models / f"{key}.safetensors"
+ path.write_text("", encoding="utf-8")
+ values[key] = str(path)
+ result = run_preflight(values, runtime, None, probe_with(vram_mb=8192), family="ideogram4")
+
+ joined = "\n".join(result.warnings)
+ self.assertIn(IDEOGRAM4.vram_hint.split("，")[0], joined)
+
+ def test_unconditional_dit_must_exist(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ runtime = make_runtime(root)
+ write_runtime_files(runtime)
+ models = root / "models"
+ models.mkdir()
+ values = {}
+ for key in ("dit", "vae", "text_encoder"):
+ path = models / f"{key}.safetensors"
+ path.write_text("", encoding="utf-8")
+ values[key] = str(path)
+ values["unconditional_dit"] = str(models / "missing.safetensors")
+ result = run_preflight(values, runtime, None, probe_with(), family="ideogram4")
+
+ self.assertFalse(result.ok)
+ self.assertIn("unconditional", "\n".join(result.errors))
+
+
+class Ideogram4PatchTests(unittest.TestCase):
+ ROTARY_ANCHOR = " _materialize_meta_tensors(model)"
+ TOKENIZER_LINE = " return AutoTokenizer.from_pretrained(QWEN3_VL_8B_INSTRUCT_REPO_ID)"
+
+ def _write_module(self, root: Path) -> Path:
+ module = root / "src" / "musubi_tuner" / "ideogram4" / "ideogram4_utils.py"
+ module.parent.mkdir(parents=True, exist_ok=True)
+ module.write_text(
+ "def load_ideogram4_text_encoder(config, state_dict):\n"
+ " model = build(config)\n"
+ f"{self.ROTARY_ANCHOR}\n"
+ " return model\n"
+ "\n"
+ "\n"
+ "def load_ideogram4_tokenizer():\n"
+ f"{self.TOKENIZER_LINE}\n",
+ encoding="utf-8",
+ )
+ return module
+
+ def _write_tokenizer_dir(self, root: Path, name: str = "qwen3-vl-8b-tokenizer") -> Path:
+ tokenizer_dir = root / "sd-models" / "ideogram4" / name
+ tokenizer_dir.mkdir(parents=True, exist_ok=True)
+ (tokenizer_dir / "tokenizer.json").write_text("{}", encoding="utf-8")
+ (tokenizer_dir / "tokenizer_config.json").write_text("{}", encoding="utf-8")
+ return tokenizer_dir
+
+ def test_patch_sets_rotary_and_tokenizer(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ module = self._write_module(root)
+ tokenizer_dir = self._write_tokenizer_dir(root)
+
+ first = patch_ideogram4_source(root, tokenizer_dir, log=lambda _line: None)
+ second = patch_ideogram4_source(root, tokenizer_dir, log=lambda _line: None)
+ text = module.read_text(encoding="utf-8")
+ compile(text, str(module), "exec")
+
+ self.assertTrue(first)
+ self.assertFalse(second, "patch must be idempotent")
+ self.assertIn("rotary_emb = type(_rotary)(config.text_config", text)
+ self.assertIn(tokenizer_dir.as_posix(), text)
+ self.assertNotIn(self.TOKENIZER_LINE, text)
+
+ def test_patch_without_tokenizer_dir_only_fixes_rotary(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ module = self._write_module(root)
+
+ patched = patch_ideogram4_source(root, root / "missing-tokenizer", log=lambda _line: None)
+ text = module.read_text(encoding="utf-8")
+
+ self.assertTrue(patched)
+ self.assertIn("_rotary = model.language_model.rotary_emb", text)
+ self.assertIn(self.TOKENIZER_LINE, text)
+
+ def test_patch_skips_missing_source(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ self.assertFalse(patch_ideogram4_source(root, root / "nope", log=lambda _line: None))
+
+ def test_module_without_anchors_is_left_untouched(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ module = root / "src" / "musubi_tuner" / "ideogram4" / "ideogram4_utils.py"
+ module.parent.mkdir(parents=True, exist_ok=True)
+ module.write_text("def other():\n return 1\n", encoding="utf-8")
+ original = module.read_text(encoding="utf-8")
+
+ patched = patch_ideogram4_source(root, root / "nope", log=lambda _line: None)
+ after = module.read_text(encoding="utf-8")
+
+ self.assertFalse(patched)
+ self.assertEqual(after, original)
+
+ def test_tokenizer_patch_follows_and_reverts_with_the_directory(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ module = self._write_module(root)
+ first_dir = self._write_tokenizer_dir(root, "first")
+ second_dir = self._write_tokenizer_dir(root, "second")
+
+ patch_ideogram4_source(root, first_dir, log=lambda _line: None)
+ text_first = module.read_text(encoding="utf-8")
+ patch_ideogram4_source(root, second_dir, log=lambda _line: None)
+ text_second = module.read_text(encoding="utf-8")
+
+ (second_dir / "tokenizer.json").unlink()
+ patch_ideogram4_source(root, second_dir, log=lambda _line: None)
+ text_reverted = module.read_text(encoding="utf-8")
+
+ self.assertIn(first_dir.as_posix(), text_first)
+ self.assertIn(second_dir.as_posix(), text_second)
+ self.assertNotIn(first_dir.as_posix(), text_second)
+ self.assertIn(self.TOKENIZER_LINE, text_reverted)
+ self.assertIn("rotary_emb = type(_rotary)(config.text_config", text_reverted)
+
+
+class Ideogram4AssetsTests(unittest.TestCase):
+ def test_asset_manifest(self):
+ assets = {asset.key: asset for asset in manifest_for("ideogram4-lora")}
+
+ self.assertEqual(
+ assets["dit"].hf_file, "diffusion_models/ideogram4_fp8_scaled.safetensors"
+ )
+ self.assertEqual(
+ assets["unconditional_dit"].hf_file,
+ "diffusion_models/ideogram4_unconditional_fp8_scaled.safetensors",
+ )
+ self.assertTrue(assets["unconditional_dit"].optional)
+ self.assertEqual(assets["vae"].hf_file, "vae/flux2-vae.safetensors")
+ self.assertEqual(assets["tokenizer"].kind, "dir")
+ self.assertIn("非商用", assets["dit"].label)
+
+ def test_tokenizer_dir_resolution(self):
+ with tempfile.TemporaryDirectory() as td:
+ root = Path(td)
+ self.assertEqual(
+ tokenizer_dir_for("ideogram4-lora", root),
+ (root / "sd-models" / "ideogram4" / "qwen3-vl-8b-tokenizer").resolve(),
+ )
+
+ def test_both_download_sources_are_configured(self):
+ with tempfile.TemporaryDirectory() as td:
+ items = {item["key"]: item for item in check_assets("ideogram4-lora", {}, Path(td))}
+
+ # The assets dialog defaults to ModelScope; a missing mirror makes the
+ # whole batch fail, so every required asset needs both sources.
+ for key in ("dit", "unconditional_dit", "text_encoder", "vae", "tokenizer"):
+ self.assertTrue(items[key]["sources"]["huggingface"], key)
+ self.assertTrue(items[key]["sources"]["modelscope"], key)
+
+
+class Ideogram4ConfigImportTests(unittest.TestCase):
+ def test_detects_ideogram4_config(self):
+ analysis = analyze_train_type(
+ {
+ "dit": "/models/sd-models/ideogram4/ideogram4_fp8_scaled.safetensors",
+ "text_encoder": "/models/sd-models/ideogram4/qwen3vl_8b_fp8_scaled.safetensors",
+ "network_module": "musubi_tuner.networks.lora_ideogram4",
+ "sampler_preset": "V4_DEFAULT_20",
+ }
+ )
+ self.assertEqual(analysis.train_type, "ideogram4-lora")
+
+ def test_still_detects_krea2_config(self):
+ analysis = analyze_train_type(
+ {
+ "dit": "/models/sd-models/krea2/krea2.safetensors",
+ "text_encoder": "/models/sd-models/krea2/qwen3_vl_4b.safetensors",
+ "network_module": "musubi_tuner.networks.lora_krea2",
+ "fp8_scaled": True,
+ }
+ )
+ self.assertEqual(analysis.train_type, "krea2-lora")
+
+ def test_page_spec_and_target_registered(self):
+ self.assertEqual(PAGE_SPECS["ideogram4-lora"]["default_train_type"], "ideogram4-lora")
+ self.assertEqual(TRAIN_TYPE_TARGETS["ideogram4-lora"]["path"], "/lora/ideogram4.html")
+
+
+class DispatchGuardTests(unittest.TestCase):
+ def test_unknown_family_variant_fails_loudly(self):
+ from mikazuki.engines.musubi.run import handle_run
+ from mikazuki.engines.runner import RunContext
+
+ result = handle_run({}, RunContext(timestamp="t", autosave_dir=".", variant="ideogram4x"))
+
+ self.assertEqual(result.status, "fail")
+ self.assertIn("未注册的 musubi 模型族", result.message)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/test_musubi_install.py b/tests/test_musubi_install.py
index b6adef4b2..b3b843876 100644
--- a/tests/test_musubi_install.py
+++ b/tests/test_musubi_install.py
@@ -484,6 +484,14 @@ def test_relative_musubi_python_stays_under_venv(self):
class PreflightTests(unittest.TestCase):
def _values_and_dataset(self, root: Path) -> tuple[dict, dict]:
+ musubi_root = root / "vendor" / "musubi-tuner"
+ musubi_root.mkdir(parents=True, exist_ok=True)
+ for name in (
+ "krea2_cache_latents.py",
+ "krea2_cache_text_encoder_outputs.py",
+ "krea2_train_network.py",
+ ):
+ (musubi_root / name).write_text("", encoding="utf-8")
models = root / "models"
models.mkdir(exist_ok=True)
for name in ("dit.safetensors", "vae.safetensors", "te.safetensors"):