evan/ai-toolkit
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add support to train MOSS Image Base as an edit model

Browse Source
This commit is contained in:
Jaret Burkett committed 2026-10-16 12:19:06 -06:00
1 parent 0d19b26481
commit b7cd8a13af
7 files changed
+548 -58
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

toolkit/models/registry.py
+1
View File
@@ -58,6 +58,7 @@ ARCH_REGISTRY: Dict[str, dict] = {
"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}},
"moss_image:base_edit": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-Base", "quantize": True, "qtype": "convrot8", "quantize_te": True, "qtype_te": "convrot8", "model_kwargs": {"edit": True, "kv_cache": True, "match_target_res": True, "uncond_skip_layers": [22]}}, "sample": {"width": 1024, "height": 1024, "num_inference_steps": 30, "guidance_scale": 3.0}, "needs_control_image": True},
"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}},