diff --git a/README.md b/README.md index 71111b6..3862f50 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t - [inclusionAI/Ming-Image-0.1-Design](https://huggingface.co/inclusionAI/Ming-Image-0.1-Design) (Ming-Image 0.1 Design) - [microsoft/Mage-Flow-Base](https://huggingface.co/microsoft/Mage-Flow-Base) (Mage-Flow) - [OpenMOSS-Team/MOSS-Image-SFT](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT) (MOSS-Image SFT) +- [OpenMOSS-Team/MOSS-Image-Base](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base) (MOSS-Image Base) ### Instruction / Edit - [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) (FLUX.1-Kontext-dev) diff --git a/extensions_built_in/diffusion_models/moss_image/src/checkpoints.py b/extensions_built_in/diffusion_models/moss_image/src/checkpoints.py index 0abd564..552dfe4 100644 --- a/extensions_built_in/diffusion_models/moss_image/src/checkpoints.py +++ b/extensions_built_in/diffusion_models/moss_image/src/checkpoints.py @@ -3,12 +3,14 @@ Weights live in the ComfyUI layout under MODELS_PATH, named like the Comfy-Org/MOSS-Image repack (DiT: the vendor's keys with q/k/v fused). -- ``OpenMOSS-Team/MOSS-Image-SFT`` / ``-Base`` (vendor, the default): a +- ``OpenMOSS-Team/MOSS-Image-SFT`` / ``-Base`` (vendor, the UI defaults): a repack DiT file already on disk wins (best precision for the requested - qtype); otherwise the vendor's diffusers shards are downloaded once and - merged into ``diffusion_models/moss_image_{sft,base}_bf16.safetensors`` in - the Comfy layout. Other hub repos (fine-tunes) merge into a file named - after the repo, never substituted. + qtype; int8 only for convrot). Otherwise Base downloads the repack's bf16 + (byte-identical to the vendor shards); SFT, which the repack has no bf16 + of, downloads the vendor's diffusers shards once and merges them into + ``diffusion_models/moss_image_sft_bf16.safetensors`` in the Comfy layout. + Other hub repos (fine-tunes) merge into a file named after the repo, never + substituted. - ``Comfy-Org/MOSS-Image``: the repack's SFT DiT (int8 convrot only; a local merged bf16 SFT also counts). - ``org/repo/path/file.safetensors``: that file, downloaded to its comfy spot. @@ -64,6 +66,9 @@ COMFY_TRANSFORMER_FILES = { ), } REPO_VARIANT = {SFT_REPO: "sft", BASE_REPO: "base", COMFY_REPO: "sft"} +# vendor repo -> the repack's file holding its weights byte-identical (same +# name and layout the shard merge writes): downloaded instead of merging +VENDOR_MIRROR = {BASE_REPO: "diffusion_models/moss_image_base_bf16.safetensors"} # (repo, file) candidates for the Qwen3.5-9B text encoder COMFY_TEXT_ENCODER_FILES = [ @@ -131,6 +136,13 @@ def comfy_transformer_file( status_fn=status_fn, qtype=qtype, ) + if repo_id in VENDOR_MIRROR: + return resolve_comfy_candidates( + [VENDOR_MIRROR[repo_id]], + repo_id=COMFY_REPO, + hf_token=hf_token, + status_fn=status_fn, + ) rel_path = comfy_transformer_rel(repo_id) found = resolve_comfy_file(rel_path, repo_id, local_only=True) diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index 62141a7..bf6613a 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/extensions_built_in/diffusion_models/ui.tsx @@ -1797,6 +1797,50 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ disableSections: ["network.conv"], additionalSections: ["model.layer_offloading", "model.low_vram"], }, + { + name: "moss_image:base", + label: "MOSS-Image Base", + generateNameOverride: "MOSS-Image Base", + group: "image", + defaults: { + // default updates when [selected, unselected] in the UI + // the vendor repo: Comfy-Org/MOSS-Image base files on disk win, else the + // repack's bf16 (identical weights) is downloaded + "config.process[0].model.name_or_path": [ + "OpenMOSS-Team/MOSS-Image-Base", + defaultNameOrPath, + ], + "config.process[0].model.quantize": [true, false], + "config.process[0].model.qtype": ["convrot8", "qfloat8"], + // Qwen3.5-9B text encoder + "config.process[0].model.quantize_te": [true, false], + "config.process[0].model.qtype_te": ["convrot8", "qfloat8"], + "config.process[0].model.low_vram": [true, false], + "config.process[0].sample.sampler": ["flowmatch", "flowmatch"], + "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"], + "config.process[0].sample.width": [1024, 1024], + "config.process[0].sample.height": [1024, 1024], + // vendor defaults: CFG 3 with block 22 skipped on the negative branch + "config.process[0].sample.guidance_scale": [3, 4], + "config.process[0].sample.sample_steps": [30, 25], + "config.process[0].train.timestep_type": ["shift", "sigmoid"], + // sampling: the CFG negative branch skips these DiT blocks (zero-based) + "config.process[0].model.model_kwargs": [ + { + uncond_skip_layers: [22], + }, + {}, + ], + // keep the LoRA off the block the negative branch skips, so CFG's + // skipped-block contrast stays the base model's + "config.process[0].network.network_kwargs.ignore_if_contains": [ + ["layers.22."], + [], + ], + }, + disableSections: ["network.conv"], + additionalSections: ["model.layer_offloading", "model.low_vram"], + }, { name: "flux2_klein_4b", label: "FLUX.2-klein-base-4B", diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 673387c..8663e52 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -226,6 +226,17 @@ MODEL_TESTS = { # vendor defaults are 50 steps, CFG 3 (negative branch skips block 22) "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0}, }, + "moss_image:base": { + # vendor id; local Comfy-Org base files win, else the repack's bf16 + "model": { + "name_or_path": "OpenMOSS-Team/MOSS-Image-Base", + "quantize": True, + "quantize_te": True, + "qtype": "convrot8", + "qtype_te": "convrot8", + }, + "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0}, + }, "moss_image:comfy": { # the Comfy-Org repack files named outright (a bare repo id prefers any # local bf16): int8 convrot SFT DiT + Qwen3.5-9B TE attach as-is diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index aecef16..67c9623 100644 --- a/toolkit/models/registry.py +++ b/toolkit/models/registry.py @@ -57,6 +57,7 @@ ARCH_REGISTRY: Dict[str, dict] = { "ltx2.5": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}}, "flux3": {"modality": "image", "model": {"name_or_path": "black-forest-labs/flux-3-image-klein-base", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8"}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 28, "guidance_scale": 3.5}}, "moss_image": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-SFT", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8"}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 30, "guidance_scale": 3.0}}, + "moss_image:base": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-Base", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8", "model_kwargs": {"uncond_skip_layers": [22]}}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 30, "guidance_scale": 3.0}}, "ace_step_15": {"modality": "audio", "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0}}, "qwen25_omni": {"modality": "text", "model": {"name_or_path": "ai-toolkit/Qwen2.5-Omni-7B/qwen2_5_omni_7b_convrot8.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 1, "guidance_scale": 1.0}}, "yue2": {"modality": "audio", "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}},