diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index 209d66b..d8be53f 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/extensions_built_in/diffusion_models/ui.tsx @@ -1841,6 +1841,59 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ disableSections: ["network.conv"], additionalSections: ["model.layer_offloading", "model.low_vram"], }, + { + name: "moss_image:edit", + label: "MOSS-Image SFT [Edit Training]", + generateNameOverride: "MOSS-Image SFT Edit", + group: "experimental", + defaults: { + // default updates when [selected, unselected] in the UI + "config.process[0].model.name_or_path": [ + "OpenMOSS-Team/MOSS-Image-SFT", + defaultNameOrPath, + ], + "config.process[0].model.quantize": [true, false], + "config.process[0].model.qtype": ["convrot8", "qfloat8"], + // Qwen3.5-9B text encoder + vision tower (references ride in the prompt) + "config.process[0].model.quantize_te": [true, false], + "config.process[0].model.qtype_te": ["convrot8", "qfloat8"], + "config.process[0].model.low_vram": [true, false], + // prompt embeds depend on the control images, so the TE stays loaded + "config.process[0].train.unload_text_encoder": [false, 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], + "config.process[0].sample.guidance_scale": [3, 4], + "config.process[0].sample.sample_steps": [30, 25], + "config.process[0].train.timestep_type": ["shift", "sigmoid"], + // edit: references as clean t=0 latents + in the prompt; kv_cache: + // references attend only to themselves, so sampling computes their K/V + // once (must be trained that way) + "config.process[0].model.model_kwargs": [ + { + edit: true, + kv_cache: true, + match_target_res: true, + uncond_skip_layers: [22], + }, + {}, + ], + "config.process[0].network.network_kwargs.ignore_if_contains": [ + ["layers.22."], + [], + ], + }, + disableSections: ["network.conv", "train.unload_text_encoder"], + additionalSections: [ + "datasets.multi_control_paths", + "sample.multi_ctrl_imgs", + "model.low_vram", + "model.layer_offloading", + "model.qie.match_target_res", + "model.model_kwargs.kv_cache", + ], + }, { name: "moss_image:base_edit", label: "MOSS-Image Base [Edit Training]", diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index deb815c..d46ea60 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -237,6 +237,20 @@ MODEL_TESTS = { }, "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0}, }, + "moss_image:edit": { + # edit LoRA path: vision-tower prompt + t=0 reference latents, isolated + # references with their K/V cached across sampling steps + "model": { + "name_or_path": "OpenMOSS-Team/MOSS-Image-SFT", + "quantize": True, + "quantize_te": True, + "qtype": "convrot8", + "qtype_te": "convrot8", + "model_kwargs": {"edit": True, "kv_cache": True, "match_target_res": True}, + }, + "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0}, + "needs_control_image": True, + }, "moss_image:base_edit": { # edit LoRA path: vision-tower prompt + t=0 reference latents, isolated # references with their K/V cached across sampling steps diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index b858878..a243626 100644 --- a/toolkit/models/registry.py +++ b/toolkit/models/registry.py @@ -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:edit": {"modality": "image", "model": {"name_or_path": "OpenMOSS-Team/MOSS-Image-SFT", "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}, "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}},