From 4d502e0e153c09b41c1bb7812cff3f580af99e32 Mon Sep 17 00:00:00 2001 From: wochenlong Date: Sat, 10 Oct 2026 11:19:11 +0800 Subject: [PATCH 1/6] feat(musubi): model-family registry and Ideogram 4 LoRA backend (#424) - families.py: per-family spec (scripts, network module, constraints, sampler/cache defaults) so adapter/launcher/preflight/run stay generic - manifest: register ideogram4-lora -> ideogram4; capabilities updated - adapter: family-aware constraints (FP8-only base, no turbo/weighting_scheme, blocks_to_swap<=33, asymmetric-CFG opt-in) and cache-only field filtering - launcher/process/run: family scripts, cache-stage extra flags, per-family tokenizer patching on launch - model_assets: ideogram4-lora assets (FP8 components + tokenizer) and the upstream rotary-embedding patch (pre-v0.3.6 snapshots) - config_import: detect ideogram4-lora configs and route them to their page --- mikazuki/engines/musubi/__init__.py | 4 +- mikazuki/engines/musubi/adapter.py | 154 +++++++-- mikazuki/engines/musubi/families.py | 136 ++++++++ mikazuki/engines/musubi/launcher.py | 20 +- mikazuki/engines/musubi/manifest.py | 5 +- mikazuki/engines/musubi/preflight.py | 36 ++- mikazuki/engines/musubi/routes.py | 12 +- mikazuki/engines/musubi/run.py | 28 +- mikazuki/model_assets.py | 97 +++++- mikazuki/process.py | 29 +- mikazuki/schema/ideogram4-lora.ts | 114 +++++++ mikazuki/utils/config_import.py | 33 ++ tests/test_musubi_ideogram4.py | 465 +++++++++++++++++++++++++++ tests/test_musubi_install.py | 8 + 14 files changed, 1066 insertions(+), 75 deletions(-) create mode 100644 mikazuki/engines/musubi/families.py create mode 100644 mikazuki/schema/ideogram4-lora.ts create mode 100644 tests/test_musubi_ideogram4.py 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..af3b4e6df 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,20 @@ "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", + "ideogram4_timestep_mu", + "ideogram4_timestep_std", # training "max_train_epochs", "max_train_steps", @@ -146,6 +168,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 +232,17 @@ "network_dropout", "scale_weight_norms", "base_weights_multiplier", + "initial_sigma", + "ideogram4_timestep_mu", + "ideogram4_timestep_std", } 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). +CACHE_ONLY_FIELDS = {"vae_dtype", "text_cache_dtype"} + SUBSET_REPEAT_PATTERN = re.compile(r"^(\d+)_(.+)$") @@ -370,7 +403,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 +453,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 +537,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") @@ -491,7 +575,11 @@ def toml_scalar(value: Any) -> str: 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 values.items() + if key not in CACHE_ONLY_FIELDS + ) 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..3c4622a10 --- /dev/null +++ b/mikazuki/engines/musubi/families.py @@ -0,0 +1,136 @@ +"""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 = "建议开启 fp8_base + blocks_to_swap" + 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, + 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..85ba797be 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,11 @@ 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: + warnings.append(f"显存 {dep.vram_total_mb} MB 可能不足以训练 {spec.label}，{spec.vram_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..6db575b23 100644 --- a/mikazuki/engines/musubi/routes.py +++ b/mikazuki/engines/musubi/routes.py @@ -7,7 +7,8 @@ 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, @@ -45,6 +46,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 +67,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 +85,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" diff --git a/mikazuki/engines/musubi/run.py b/mikazuki/engines/musubi/run.py index 69d3a5782..36375a254 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, family_for from mikazuki.engines.runner import RunContext from mikazuki.utils import train_utils @@ -81,21 +82,20 @@ 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" + spec = family_for(getattr(ctx, "variant", "") or DEFAULT_FAMILY) + model_train_type = spec.train_type if not musubi_feature_enabled(): return musubi_disabled_response() ready, failure = musubi_ready_gate() @@ -116,7 +116,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 +141,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 +158,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..63c783bbc 100644 --- a/mikazuki/model_assets.py +++ b/mikazuki/model_assets.py @@ -47,6 +47,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 +115,21 @@ 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. + "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"), + 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"), + 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"), + AssetDef("vae", "VAE（Flux2）", "sd-models/ideogram4/flux2-vae.safetensors", + hf_repo=IDEOGRAM4_REPO, hf_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 +198,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 +262,76 @@ 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 + model.language_model.rotary_emb = type(_rotary)(config.text_config, device=_rotary.inv_freq.device)""" +_IDEOGRAM4_TOKENIZER_LINE = " return AutoTokenizer.from_pretrained(QWEN3_VL_8B_INSTRUCT_REPO_ID)" + + +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, so + offline/China installs never fetch it from the Hub. + 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") + if _TOKENIZER_PATCH_MARK not in patched: + tokenizer_dir = Path(tokenizer_dir) + if not dir_complete(tokenizer_dir): + pass + elif _IDEOGRAM4_TOKENIZER_LINE in patched: + patched = patched.replace( + _IDEOGRAM4_TOKENIZER_LINE, + f' return AutoTokenizer.from_pretrained(r"{tokenizer_dir.as_posix()}") {_TOKENIZER_PATCH_MARK}', + 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 + module.write_text(patched, encoding="utf-8") + 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 +389,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..753f33d04 --- /dev/null +++ b/mikazuki/schema/ideogram4-lora.ts @@ -0,0 +1,114 @@ +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，通常无需修改"), + ideogram4_timestep_mu: Schema.number().step(0.001).description("ideogram4_shift 采样器的 mu（留空使用官方默认值）"), + ideogram4_timestep_std: Schema.number().step(0.001).description("ideogram4_shift 采样器的 std（留空使用官方默认值）"), + 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（官方采样主步 guidance 为 7）'), + sample_seed: Schema.number().default(42).description('种子'), + sample_steps: Schema.number().min(1).max(300).default(20).description('迭代步数（与 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_musubi_ideogram4.py b/tests/test_musubi_ideogram4.py new file mode 100644 index 000000000..669434bdb --- /dev/null +++ b/tests/test_musubi_ideogram4.py @@ -0,0 +1,465 @@ +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 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") + + self.assertEqual(adapted.values["vae_dtype"], "bfloat16") + toml_text = dump_train_toml(adapted.values) + self.assertNotIn("vae_dtype", toml_text) + self.assertNotIn("text_cache_dtype", 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) -> Path: + tokenizer_dir = root / "sd-models" / "ideogram4" / "qwen3-vl-8b-tokenizer" + 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)) + + +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(), + ) + + +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") + + +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"): From eb71c8ce4bdcd23d70dd5e58b99091b478dc6d38 Mon Sep 17 00:00:00 2001 From: wochenlong Date: Sat, 10 Oct 2026 11:19:14 +0800 Subject: [PATCH 2/6] feat(training): Ideogram 4 entry in the training workbench (#424) - new ideogram4-lora schema (FP8 components, ideogram4_shift sampling, caption validation, cache dtype/disk hints, sampler presets) - module table/i18n/router/catalog for model Ideogram 4 on musubi - MusubiGatePage takes the selected schema instead of hardcoding Krea 2; TrainingPage preflights every musubi schema - carry-over guard: vae/turbo/unconditional-DiT/fp8 flags no longer cross Krea 2 and Ideogram 4 pages --- frontend/src/engines/catalog.ts | 2 +- frontend/src/i18n/messages/en-US.ts | 5 +- frontend/src/i18n/messages/zh-CN.ts | 5 +- frontend/src/pages/MusubiGatePage.test.ts | 16 ++++- frontend/src/pages/MusubiGatePage.vue | 13 ++-- frontend/src/pages/TrainingPage.vue | 4 +- frontend/src/pages/TrainingWorkbenchPage.vue | 2 +- frontend/src/router.ts | 2 + frontend/src/training/ideogram4.test.ts | 76 ++++++++++++++++++++ frontend/src/training/modules.test.ts | 7 +- frontend/src/training/modules.ts | 10 ++- frontend/src/training/params.ts | 9 +++ 12 files changed, 134 insertions(+), 17 deletions(-) create mode 100644 frontend/src/training/ideogram4.test.ts 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" > - +