From e05a3a4f2cc9bbbdd93407ad295f19edb6f89876 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Fri, 16 Oct 2026 08:09:45 -0600 Subject: [PATCH] Add support for MOSS Image --- README.md | 2 + .../diffusion_models/__init__.py | 2 + .../diffusion_models/moss_image/__init__.py | 3 + .../diffusion_models/moss_image/moss_image.py | 400 ++++++++++++++++++ .../moss_image/src/__init__.py | 0 .../moss_image/src/checkpoints.py | 274 ++++++++++++ .../moss_image/src/pipeline.py | 143 +++++++ .../moss_image/src/transformer.py | 340 +++++++++++++++ extensions_built_in/diffusion_models/ui.tsx | 43 ++ .../inference_engine/latent_preview.py | 1 + testing/test_model_loading.py | 24 ++ toolkit/models/registry.py | 1 + toolkit/models/v2/text_encoders/qwen35.py | 56 +++ version.py | 2 +- 14 files changed, 1290 insertions(+), 1 deletion(-) create mode 100644 extensions_built_in/diffusion_models/moss_image/__init__.py create mode 100644 extensions_built_in/diffusion_models/moss_image/moss_image.py create mode 100644 extensions_built_in/diffusion_models/moss_image/src/__init__.py create mode 100644 extensions_built_in/diffusion_models/moss_image/src/checkpoints.py create mode 100644 extensions_built_in/diffusion_models/moss_image/src/pipeline.py create mode 100644 extensions_built_in/diffusion_models/moss_image/src/transformer.py create mode 100644 toolkit/models/v2/text_encoders/qwen35.py diff --git a/README.md b/README.md index 14823bc..71111b6 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t - [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev) (FLUX.2) - [black-forest-labs/FLUX.2-klein-base-4B](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B) (FLUX.2-klein-base-4B) - [black-forest-labs/FLUX.2-klein-base-9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B) (FLUX.2-klein-base-9B) +- [black-forest-labs/flux-3-image-klein-base](https://huggingface.co/black-forest-labs/flux-3-image-klein-base) (FLUX.3 Image Klein) - [ostris/Flex.1-alpha](https://huggingface.co/ostris/Flex.1-alpha) (Flex.1) - [ostris/Flex.2-preview](https://huggingface.co/ostris/Flex.2-preview) (Flex.2) - [lodestones/Chroma1-Base](https://huggingface.co/lodestones/Chroma1-Base) (Chroma) @@ -39,6 +40,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t - [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo) (Krea 2 Turbo) - [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) ### 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/__init__.py b/extensions_built_in/diffusion_models/__init__.py index 22baca3..7145136 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/extensions_built_in/diffusion_models/__init__.py @@ -24,6 +24,7 @@ from .boogu_image import BooguImageModel, BooguImageEditModel from .mageflow import MageFlowModel, MageFlowEditModel from .minimax_h3 import MinimaxH3Model, MinimaxH3Ref2VAModel, MinimaxH3FastModel from .flux3 import Flux3Model +from .moss_image import MossImageModel AI_TOOLKIT_MODELS = [ # put a list of models here @@ -67,4 +68,5 @@ AI_TOOLKIT_MODELS = [ MinimaxH3Ref2VAModel, MinimaxH3FastModel, Flux3Model, + MossImageModel, ] diff --git a/extensions_built_in/diffusion_models/moss_image/__init__.py b/extensions_built_in/diffusion_models/moss_image/__init__.py new file mode 100644 index 0000000..0fe3c71 --- /dev/null +++ b/extensions_built_in/diffusion_models/moss_image/__init__.py @@ -0,0 +1,3 @@ +from .moss_image import MossImageModel + +__all__ = ["MossImageModel"] diff --git a/extensions_built_in/diffusion_models/moss_image/moss_image.py b/extensions_built_in/diffusion_models/moss_image/moss_image.py new file mode 100644 index 0000000..2e9a444 --- /dev/null +++ b/extensions_built_in/diffusion_models/moss_image/moss_image.py @@ -0,0 +1,400 @@ +"""MOSS-Image (OpenMOSS-Team/MOSS-Image-SFT) for ai-toolkit. + + - DiT: 42-block single-stream transformer, 3072 wide, 128 latent channels + (``src/transformer.py``, Comfy-Org/MOSS-Image key layout) + - VAE: FLUX.2, 2x2 packed and BatchNorm normalized (the toolkit's BFL + ``AutoEncoder``; same packing as MOSS). The shared comfy flux2-vae file + (full decoder) replaces the vendor's small-decoder copy by default + - text: Qwen3.5-9B final-normed last hidden state of the chat-templated + prompt (thinking on, generation prompt appended), one unpadded forward per + prompt; stored per item as (L, 4096) in an ``AdvancedPromptEmbeds`` + +``model.name_or_path`` defaults to the vendor repo; ``Comfy-Org/MOSS-Image`` +(the repack), a local diffusers dir or a ``.safetensors`` path (local or +``org/repo/file``) also load. Comfy files on disk take precedence for the +vendor ids; see ``src/checkpoints.py`` for sourcing. The text encoder is +stock Qwen3.5-9B from a comfy file. + +Flow convention matches ai-toolkit (sigma 1 = noise, target = noise - clean). +""" + +import os +from typing import TYPE_CHECKING, List, Optional + +import torch +from safetensors import safe_open +from safetensors.torch import load_file +from transformers import Qwen2Tokenizer + +from toolkit.accelerator import unwrap_model +from toolkit.advanced_prompt_embeds import AdvancedPromptEmbeds +from toolkit.basic import flush +from toolkit.config_modules import GenerateImageConfig, ModelConfig +from toolkit.metadata import get_meta_for_safetensors +from toolkit.models.base_model import BaseModel +from toolkit.models.v2.text_encoders.qwen35 import Qwen35TextOnlyEncoder +from toolkit.models.v2.vae.flux2_kl import AutoEncoder, convert_diffusers_state_dict +from toolkit.samplers.custom_flowmatch_sampler import ( + CustomFlowMatchEulerDiscreteScheduler, +) + +from .src.checkpoints import ( + COMFY_REPO, + QWEN35_9B_REPO, + SFT_REPO, + VENDOR_VAE_FILE, + comfy_text_encoder_file, + comfy_transformer_file, + comfy_vae_file, + resolve_hub_file, +) +from .src.pipeline import ( + DEFAULT_SHIFT, + DEFAULT_UNCOND_SKIP_LAYERS, + MossImagePipeline, + stack_text_embeds, +) +from .src.transformer import MossImageTransformer + +if TYPE_CHECKING: + from toolkit.data_transfer_object.data_loader import DataLoaderBatchDTO + +HF_TOKEN = os.getenv("HF_TOKEN", None) +MAX_TEXT_TOKENS = 4096 + +scheduler_config = { + "num_train_timesteps": 1000, + "use_dynamic_shifting": False, +} + + +def _is_hub_repo(path: str) -> bool: + return ( + not os.path.exists(path) + and not path.endswith(".safetensors") + and path.count("/") == 1 + ) + + +class MossImageModel(BaseModel): + arch = "moss_image" + + def __init__( + self, + device, + model_config: ModelConfig, + dtype="bf16", + custom_pipeline=None, + noise_scheduler=None, + **kwargs, + ): + super().__init__( + device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs + ) + self.is_flow_matching = True + self.is_transformer = True + self.target_lora_modules = ["MossImageTransformer"] + self.supports_model_paths = True + self.use_old_lokr_format = False + self.lora_keys_use_comfy_prefix = True + self.latent_space_version = "moss_image" + + kw = self.model_config.model_kwargs + self.max_text_length = int(kw.get("max_text_length", MAX_TEXT_TOKENS)) + self.sample_shift = float(kw.get("sample_shift", DEFAULT_SHIFT)) + self.uncond_skip_layers = tuple( + kw.get("uncond_skip_layers", DEFAULT_UNCOND_SKIP_LAYERS) or () + ) + + @staticmethod + def get_train_scheduler(): + return CustomFlowMatchEulerDiscreteScheduler( + shift=DEFAULT_SHIFT, **scheduler_config + ) + + def get_bucket_divisibility(self): + return 16 # f8 conv + 2x2 packing + + # ------------------------------------------------------------------ + # Loading + # ------------------------------------------------------------------ + def _load_transformer(self) -> MossImageTransformer: + self.print_and_status_update("Loading transformer (MossImageTransformer)") + base = self.model_config.name_or_path + qtype = self.model_config.qtype if self.model_config.quantize else None + if os.path.isdir(base): + folder = os.path.join(base, "transformer") + if not os.path.isdir(folder): + folder = base + state_dict = {} + for name in sorted(os.listdir(folder)): + if name.endswith(".safetensors"): + with safe_open(os.path.join(folder, name), "pt") as f: + for key in f.keys(): + state_dict[key] = f.get_tensor(key) + if not state_dict: + raise FileNotFoundError(f"no MOSS-Image transformer weights in {folder}") + return MossImageTransformer.load_from_state_dict(state_dict, self.torch_dtype) + if base.endswith(".safetensors"): + path = resolve_hub_file( + base, hf_token=HF_TOKEN, status_fn=self.print_and_status_update + ) + else: + path = comfy_transformer_file( + base, + qtype=qtype, + hf_token=HF_TOKEN, + status_fn=self.print_and_status_update, + ) + self.print_and_status_update(f" - transformer weights: {path}") + return MossImageTransformer.load_model( + path, dtype=self.torch_dtype, qtype=qtype, quantize_on_load=False + ) + + def _load_text_encoder(self): + te_path = self.model_config.te_name_or_path + if te_path is None: + te_path = self.model_config.model_kwargs.get( + "text_encoder_path", QWEN35_9B_REPO + ) + te_kwargs = self.component_load_kwargs("te") + config_path = None + if te_path == QWEN35_9B_REPO and te_kwargs["use_comfy_weights"]: + te_path = comfy_text_encoder_file( + qtype=te_kwargs["qtype"], + hf_token=HF_TOKEN, + status_fn=self.print_and_status_update, + ) + elif te_path.endswith(".safetensors"): + te_path = resolve_hub_file( + te_path, hf_token=HF_TOKEN, status_fn=self.print_and_status_update + ) + if te_path.endswith(".safetensors"): + config_path = QWEN35_9B_REPO + self.print_and_status_update(f"Loading Qwen3.5 text encoder from {te_path}") + text_encoder = Qwen35TextOnlyEncoder.load_model( + te_path, + dtype=te_kwargs["dtype"], + config_path=config_path, + subfolder="", + qtype=te_kwargs["qtype"], # comfy candidate ranking only + quantize_on_load=False, + use_comfy_weights=te_kwargs["use_comfy_weights"], + ) + text_encoder.eval() + text_encoder.requires_grad_(False) + return text_encoder + + def _load_tokenizer(self): + base = self.model_config.name_or_path + source = SFT_REPO + if os.path.isdir(os.path.join(base, "tokenizer")) or ( + _is_hub_repo(base) and base != COMFY_REPO + ): + source = base + # the tokenizer's declared class: AutoTokenizer also resolves a + # config.json the folder lacks, which fails offline + return Qwen35TextOnlyEncoder._local_first( + Qwen2Tokenizer.from_pretrained, source, subfolder="tokenizer" + ) + + def _load_vae(self) -> AutoEncoder: + self.print_and_status_update("Loading VAE") + base = self.model_config.name_or_path + path = self.model_config.vae_path + if path is None and os.path.isdir(base): + local = os.path.join(base, VENDOR_VAE_FILE) + if os.path.exists(local): + path = local + if path is None: + path = comfy_vae_file( + hf_token=HF_TOKEN, status_fn=self.print_and_status_update + ) + elif not os.path.exists(path): + path = resolve_hub_file( + path, hf_token=HF_TOKEN, status_fn=self.print_and_status_update + ) + state_dict = load_file(path) + if any(k.startswith("encoder.down_blocks.") for k in state_dict): + state_dict = convert_diffusers_state_dict(state_dict) + # config sniffed from the checkpoint (small decoder) + vae = AutoEncoder.load_from_state_dict(state_dict, self.vae_torch_dtype) + return vae.eval().requires_grad_(False) + + def load_model(self): + self.print_and_status_update("Loading MOSS-Image model") + + transformer = self._load_transformer() + # quantize + offload + placement, all driven by model_config + transformer.aitk_post_load(**self.component_load_kwargs("transformer")) + flush() + + text_encoder = self._load_text_encoder() + text_encoder.aitk_post_load(**self.component_load_kwargs("te")) + flush() + tokenizer = self._load_tokenizer() + + vae = self._load_vae() + vae.to(self.vae_device_torch) + + self.noise_scheduler = MossImageModel.get_train_scheduler() + + self.vae = vae + self.text_encoder = text_encoder + self.tokenizer = tokenizer + self.model = transformer + self.pipeline = MossImagePipeline(self) + self.print_and_status_update("Model Loaded") + + # ------------------------------------------------------------------ + # Generation (training previews) + # ------------------------------------------------------------------ + def get_generation_pipeline(self): + return MossImagePipeline(self) + + def generate_single_image( + self, + pipeline: MossImagePipeline, + gen_config: GenerateImageConfig, + conditional_embeds: AdvancedPromptEmbeds, + unconditional_embeds: AdvancedPromptEmbeds, + generator: torch.Generator, + extra: dict, + ): + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + + sc = self.get_bucket_divisibility() + gen_config.width = int(gen_config.width // sc * sc) + gen_config.height = int(gen_config.height // sc * sc) + + result = pipeline( + conditional_embeds=conditional_embeds, + unconditional_embeds=unconditional_embeds, + height=gen_config.height, + width=gen_config.width, + num_inference_steps=gen_config.num_inference_steps, + guidance_scale=gen_config.guidance_scale, + latents=gen_config.latents, + generator=generator, + shift=self.sample_shift, + uncond_skip_layers=self.uncond_skip_layers, + ) + return result[0] + + # ------------------------------------------------------------------ + # Training hooks + # ------------------------------------------------------------------ + def get_noise_prediction( + self, + latent_model_input: torch.Tensor, # (B, 128, h, w) + timestep: torch.Tensor, # 0..1000 scale + text_embeddings: AdvancedPromptEmbeds, + batch: "DataLoaderBatchDTO" = None, + **kwargs, + ): + device = self.device_torch + dtype = self.torch_dtype + if self.model.device == torch.device("cpu"): + self.model.to(device) + + batch_size = latent_model_input.shape[0] + embeds_list = text_embeddings.text_embeds + if len(embeds_list) == 1 and batch_size > 1: + embeds_list = embeds_list * batch_size + with torch.no_grad(): + embeds, mask = stack_text_embeds(embeds_list, device, dtype) + t = timestep.to(device, dtype=torch.float32).reshape(-1) + if t.shape[0] != batch_size: + t = t.expand(batch_size) + + return self.transformer( + latent_model_input.to(device, dtype), t, embeds, mask + ) + + @torch.no_grad() + def get_prompt_embeds(self, prompt) -> AdvancedPromptEmbeds: + if isinstance(prompt, str): + prompt = [prompt] + if self.text_encoder.device == torch.device("cpu"): + self.text_encoder.to(self.device_torch) + + embeds = [] + # one unpadded forward per prompt, like the vendor (padding would leak + # into the gated-delta recurrence) + for text in prompt: + formatted = self.tokenizer.apply_chat_template( + [{"role": "user", "content": text}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=True, + ) + input_ids = self.tokenizer( + formatted, + return_tensors="pt", + truncation=True, + max_length=self.max_text_length, + )["input_ids"].to(self.device_torch) + out = self.text_encoder(input_ids=input_ids, use_cache=False) + embeds.append(out.last_hidden_state[0].to(self.torch_dtype)) + return AdvancedPromptEmbeds(text_embeds=embeds) + + def get_loss_target(self, *args, **kwargs): + noise = kwargs.get("noise") + batch = kwargs.get("batch") + return (noise - batch.latents).detach() + + def get_model_has_grad(self): + return False + + def get_te_has_grad(self): + return False + + # ------------------------------------------------------------------ + # VAE encode / decode + # ------------------------------------------------------------------ + @torch.no_grad() + def encode_images(self, image_list, device=None, dtype=None): + """Images (C, H, W) in [-1, 1] -> normalized latents (B, 128, H/16, W/16).""" + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + if self.vae.device == torch.device("cpu"): + self.vae.to(device) + batch = torch.stack([img.to(device, dtype=self.vae.dtype) for img in image_list]) + return self.vae.encode(batch).to(device, dtype=dtype) + + @torch.no_grad() + def decode_latents(self, latents: torch.Tensor, device=None, dtype=None): + if device is None: + device = self.vae_device_torch + if dtype is None: + dtype = self.vae_torch_dtype + if self.vae.device == torch.device("cpu"): + self.vae.to(device) + images = self.vae.decode(latents.to(device, self.vae.dtype)) + return images.to(device, dtype) + + # ------------------------------------------------------------------ + # Saving / bookkeeping + # ------------------------------------------------------------------ + def save_model(self, output_path, meta, save_dtype): + if not output_path.endswith(".safetensors"): + output_path = output_path + ".safetensors" + transformer: MossImageTransformer = unwrap_model(self.model) + transformer.save_model( + output_path, + dtype=save_dtype, + metadata=get_meta_for_safetensors(meta, name=self.arch), + quantized=False, + ) + + def get_base_model_version(self): + return "moss_image" + + def get_transformer_block_names(self) -> Optional[List[str]]: + return MossImageTransformer.get_transformer_block_names() + + def get_quantization_exclude_modules(self): + return MossImageTransformer.get_quantization_exclude_modules() diff --git a/extensions_built_in/diffusion_models/moss_image/src/__init__.py b/extensions_built_in/diffusion_models/moss_image/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/extensions_built_in/diffusion_models/moss_image/src/checkpoints.py b/extensions_built_in/diffusion_models/moss_image/src/checkpoints.py new file mode 100644 index 0000000..0abd564 --- /dev/null +++ b/extensions_built_in/diffusion_models/moss_image/src/checkpoints.py @@ -0,0 +1,274 @@ +"""Where MOSS-Image weights come from. + +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 + 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. +- ``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. + +The vendor text_encoder/ is stock Qwen3.5-9B (A_log / gated norm stored in +bf16), so the text encoder is a comfy file: the repack's int8 convrot or +Comfy-Org/Qwen3.5's bf16, whichever is local (ranked for the requested qtype), +else the one matching the qtype is downloaded. The VAE is the shared FLUX.2 +file: same encoder + BatchNorm stats as the vendor vae/ (which is the BFL +weights rounded to bf16), full decoder in place of the small one. The +tokenizer comes from the vendor repo. +""" + +import json +import math +import os +import shutil +import struct +from typing import Callable, List, Optional + +import torch +from safetensors import safe_open + +from toolkit.models.v2.resolver import ( + comfy_precision_rank, + resolve_comfy_candidates, + resolve_comfy_file, +) +from toolkit.paths import MODELS_PATH + +SFT_REPO = "OpenMOSS-Team/MOSS-Image-SFT" +BASE_REPO = "OpenMOSS-Team/MOSS-Image-Base" +COMFY_REPO = "Comfy-Org/MOSS-Image" +QWEN35_9B_REPO = "Qwen/Qwen3.5-9B" + +VENDOR_VAE_FILE = "vae/diffusion_pytorch_model.safetensors" +VENDOR_TRANSFORMER_INDEX = "transformer/diffusion_pytorch_model.safetensors.index.json" +COMFY_VAE_FILE = "vae/flux2-vae.safetensors" + +# variant -> (files in the repack, extra local-only candidates). The merged +# vendor SFT is the bf16 SFT the repack does not ship. +COMFY_TRANSFORMER_FILES = { + "sft": ( + ["diffusion_models/moss_image_sft_int8_convrot_attn.safetensors"], + ["diffusion_models/moss_image_sft_bf16.safetensors"], + ), + "base": ( + [ + "diffusion_models/moss_image_base_int8_convrot_attn.safetensors", + "diffusion_models/moss_image_base_bf16.safetensors", + ], + [], + ), +} +REPO_VARIANT = {SFT_REPO: "sft", BASE_REPO: "base", COMFY_REPO: "sft"} + +# (repo, file) candidates for the Qwen3.5-9B text encoder +COMFY_TEXT_ENCODER_FILES = [ + (COMFY_REPO, "text_encoders/qwen3.5_9b_int8_convrot.safetensors"), + ("Comfy-Org/Qwen3.5", "text_encoders/qwen3.5_9b_bf16.safetensors"), +] + + +def comfy_transformer_rel(repo_id: str) -> str: + name = repo_id.split("/")[-1].lower().replace("-", "_").replace(".", "_") + return f"diffusion_models/{name}_bf16.safetensors" + + +def _staging_dir(category: str) -> str: + # same filesystem as the target so the final rename is a move; removed + # afterwards together with the hub's .cache bookkeeping + return os.path.join(MODELS_PATH, category, ".moss_image_download") + + +def resolve_hub_file( + ref: str, + hf_token: Optional[str] = None, + status_fn: Optional[Callable[[str], None]] = None, +) -> str: + """A local path, or 'org/repo/rel/file.safetensors' found at / downloaded + to its comfy spot under MODELS_PATH.""" + if os.path.isfile(ref): + return ref + parts = ref.split("/") + if len(parts) < 3: + raise FileNotFoundError( + f"{ref} is not a local file or an 'org/repo/path/file.safetensors' reference" + ) + return resolve_comfy_candidates( + ["/".join(parts[2:])], + repo_id="/".join(parts[:2]), + hf_token=hf_token, + status_fn=status_fn, + ) + + +def comfy_transformer_file( + repo_id: str, + qtype: Optional[str] = None, + hf_token: Optional[str] = None, + status_fn: Optional[Callable[[str], None]] = None, +) -> str: + variant = REPO_VARIANT.get(repo_id) + if variant is not None: + repo_files, local_extra = COMFY_TRANSFORMER_FILES[variant] + candidates = repo_files + local_extra + if repo_id != COMFY_REPO and not (qtype or "").lower().startswith("convrot"): + # int8 only stands in for the vendor bf16 when it attaches as-is + candidates = [c for c in candidates if "int8" not in c] + found = resolve_comfy_candidates( + candidates, repo_id=COMFY_REPO, local_only=True, qtype=qtype + ) + if found is not None: + return found + if repo_id == COMFY_REPO: + return resolve_comfy_candidates( + repo_files, + repo_id=COMFY_REPO, + hf_token=hf_token, + status_fn=status_fn, + qtype=qtype, + ) + + rel_path = comfy_transformer_rel(repo_id) + found = resolve_comfy_file(rel_path, repo_id, local_only=True) + if found is not None: + return found + from huggingface_hub import hf_hub_download + + target = os.path.join(MODELS_PATH, rel_path) + staging = _staging_dir("diffusion_models") + if status_fn is not None: + status_fn(f"Downloading the {repo_id} transformer to {target}") + index_path = hf_hub_download( + repo_id, VENDOR_TRANSFORMER_INDEX, token=hf_token, local_dir=staging + ) + with open(index_path) as f: + shard_names = sorted(set(json.load(f)["weight_map"].values())) + shards = [ + hf_hub_download( + repo_id, f"transformer/{name}", token=hf_token, local_dir=staging + ) + for name in shard_names + ] + merge_safetensors(shards, target, fuse_qkv=True) + shutil.rmtree(staging, ignore_errors=True) + return target + + +def comfy_text_encoder_file( + qtype: Optional[str] = None, + hf_token: Optional[str] = None, + status_fn: Optional[Callable[[str], None]] = None, +) -> str: + rels = [rel for _, rel in COMFY_TEXT_ENCODER_FILES] + found = resolve_comfy_candidates(rels, repo_id=COMFY_REPO, local_only=True, qtype=qtype) + if found is not None: + return found + repo_id, rel = min( + COMFY_TEXT_ENCODER_FILES, key=lambda c: comfy_precision_rank(c[1], qtype=qtype) + ) + return resolve_comfy_candidates( + [rel], repo_id=repo_id, hf_token=hf_token, status_fn=status_fn + ) + + +def comfy_vae_file( + hf_token: Optional[str] = None, + status_fn: Optional[Callable[[str], None]] = None, +) -> str: + return resolve_comfy_candidates( + [COMFY_VAE_FILE], repo_id=COMFY_REPO, hf_token=hf_token, status_fn=status_fn + ) + + +_CAST_TO_BF16 = ("F32", "F16", "F64") + + +def merge_safetensors( + shards: List[str], + out_path: str, + fuse_qkv: bool = False, + chunk_bytes: int = 64 << 20, +): + """Concatenate safetensors shards into one bf16 file. bf16 (and non-float) + tensors are copied as raw bytes; other float tensors are cast to bf16. + ``fuse_qkv`` writes each ``*.proj.{q,k,v}.weight`` triple as one + ``*.proj.qkv.weight`` (row concat of contiguous tensors = byte concat). + Written to ``.part`` and renamed, so ``out_path`` may be one of the shards.""" + tensors = {} + for path in shards: + with open(path, "rb") as f: + header_len = struct.unpack(" (dtype, shape, [(src name, path, start, size)]) + entries = {} + for name, (dtype, shape, path, start, size) in tensors.items(): + if fuse_qkv and name.endswith((".proj.k.weight", ".proj.v.weight")): + continue + if fuse_qkv and name.endswith(".proj.q.weight"): + base = name[: -len("q.weight")] + parts = [base + f"{p}.weight" for p in ("q", "k", "v")] + srcs = [tensors[p] for p in parts] + if any(s[0] != dtype or s[1][1:] != shape[1:] for s in srcs): + raise ValueError(f"{base}q/k/v parts do not fuse") + fused_shape = [sum(s[1][0] for s in srcs)] + list(shape[1:]) + entries[base + "qkv.weight"] = ( + dtype, + fused_shape, + [(p, s[2], s[3], s[4]) for p, s in zip(parts, srcs)], + ) + continue + entries[name] = (dtype, shape, [(name, path, start, size)]) + + out_header, offset = {}, 0 + for name in sorted(entries): + dtype, shape, parts = entries[name] + size = sum(p[3] for p in parts) + if dtype in _CAST_TO_BF16: + dtype, size = "BF16", 2 * math.prod(shape) + out_header[name] = { + "dtype": dtype, + "shape": shape, + "data_offsets": [offset, offset + size], + } + offset += size + out_header["__metadata__"] = {"format": "pt"} + blob = json.dumps(out_header, separators=(",", ":")).encode("utf-8") + blob += b" " * (-len(blob) % 8) + + os.makedirs(os.path.dirname(out_path), exist_ok=True) + part_path = out_path + ".part" + handles = {} + try: + with open(part_path, "wb") as out: + out.write(struct.pack(" torch.Tensor: + """The vendor scheduler's sigmas (MOSSImageScheduler.set_timesteps) with + the terminal 0 appended: a shifted grid snapped to the training grid, + consecutive duplicates removed (so fewer than num_steps are possible).""" + + def grid(count, s): + lower = s / (total + s - 1) + values = torch.linspace(1, lower, count) + return s * values / (1 + (s - 1) * values) + + sigmas = grid(num_steps, shift) + training = grid(total, training_shift).flip(0) + upper = torch.searchsorted(training, sigmas).clamp(max=total - 1) + lower = (upper - 1).clamp(min=0) + nearest = torch.where( + (sigmas - training[lower]).abs() <= (training[upper] - sigmas).abs(), + lower, + upper, + ) + sigmas = torch.unique_consecutive(training[nearest]) + return torch.cat([sigmas, sigmas.new_zeros(1)]) + + +def stack_text_embeds( + embeds: List[torch.Tensor], device: torch.device, dtype: torch.dtype +): + """Per-item (L_i, D) prompt states -> right-padded (B, S, D) + (B, S) mask.""" + length = max(e.shape[0] for e in embeds) + stacked = torch.stack( + [F.pad(e.to(device, dtype), (0, 0, 0, length - e.shape[0])) for e in embeds] + ) + mask = torch.stack( + [ + torch.arange(length, device=device) < e.shape[0] + for e in embeds + ] + ) + return stacked, mask + + +class MossImagePipeline: + """Lightweight sampler used by ai-toolkit's preview generation.""" + + def __init__(self, model): + # the MossImageModel, for its transformer + encode/decode + self.model = model + + @property + def device(self): + return self.model.device_torch + + def to(self, *args, **kwargs): + return self + + def set_progress_bar_config(self, **kwargs): + pass + + @torch.no_grad() + def __call__( + self, + conditional_embeds, + unconditional_embeds, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = DEFAULT_STEPS, + guidance_scale: float = DEFAULT_GUIDANCE, + latents: Optional[torch.Tensor] = None, + generator: Optional[torch.Generator] = None, + shift: float = DEFAULT_SHIFT, + uncond_skip_layers: Sequence[int] = DEFAULT_UNCOND_SKIP_LAYERS, + **kwargs, + ): + model = self.model + device = model.device_torch + dtype = model.torch_dtype + transformer = model.transformer + + do_cfg = unconditional_embeds is not None and guidance_scale != 1.0 + cond, cond_mask = stack_text_embeds(conditional_embeds.text_embeds, device, dtype) + if do_cfg: + uncond, uncond_mask = stack_text_embeds( + unconditional_embeds.text_embeds, device, dtype + ) + + shape = (1, transformer.params.latent_channels, height // 16, width // 16) + if latents is None: + # fp32 noise: bf16 noise promoted afterwards cannot regain precision + latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32) + latents = latents.to(device, torch.float32) + + sigmas = moss_sigmas(num_inference_steps, shift=shift).to(device) + num_steps = len(sigmas) - 1 + for step_index in range(num_steps): + sigma, sigma_next = sigmas[step_index], sigmas[step_index + 1] + timestep = (sigma * NUM_TRAIN_TIMESTEPS).reshape(1) + # CFG combines fp32 predictions even though the DiT runs in bf16 + pred = transformer(latents, timestep, cond, cond_mask).float() + if do_cfg: + pred_uncond = transformer( + latents, + timestep, + uncond, + uncond_mask, + skip_layers=uncond_skip_layers, + ).float() + pred = pred_uncond + guidance_scale * (pred - pred_uncond) + if getattr(model, "sample_step_hook", None) is not None: + # inference engine preview: x0 estimate + model._emit_sample_step(latents - sigma * pred, step_index, num_steps) + latents = latents + (sigma_next - sigma) * pred + + images = model.decode_latents(latents, device=device, dtype=torch.float32) + images = images.float().clamp(-1.0, 1.0) + images = ((images + 1.0) * 127.5).round().to(torch.uint8) + images = images.permute(0, 2, 3, 1).cpu().numpy() + return [Image.fromarray(arr) for arr in images] diff --git a/extensions_built_in/diffusion_models/moss_image/src/transformer.py b/extensions_built_in/diffusion_models/moss_image/src/transformer.py new file mode 100644 index 0000000..26c033b --- /dev/null +++ b/extensions_built_in/diffusion_models/moss_image/src/transformer.py @@ -0,0 +1,340 @@ +"""MOSS-Image DiT (OpenMOSS-Team/MOSS-Image-SFT / -Base). + +Port of ``MOSSImageTransformer`` from the vendor's ``pipeline_moss_image.py`` +in the Comfy-Org/MOSS-Image key layout: the vendor's names with each block's +q/k/v fused into ``proj.qkv`` (vendor shards and split community repacks are +fused on load), so Comfy files load as-is and LoRA keys match ComfyUI. One stream: image tokens followed by text tokens +through 42 blocks that share one AdaLN modulation; the text first goes through +a 2-block refiner. The vendor evaluates every item alone at its true text +length; here a batch runs together with padded text masked out of every +attention (identical results) and the unmasked path when lengths match. +Activations follow the weight dtype exactly like the vendor's bf16 inference: +fp32 timestep sinusoid and RoPE phases, everything else in the compute dtype. +""" + +from dataclasses import dataclass +from typing import Iterable, Optional + +import torch +import torch.nn.functional as F +from torch import nn +from torch.utils.checkpoint import checkpoint + +from toolkit.models.v2._mixin import OstrisModelMixin + + +@dataclass +class MossImageParams: + dim: int = 3072 + ffn_dim: int = 8192 + num_heads: int = 24 + num_layers: int = 42 + text_dim: int = 4096 + freq_dim: int = 256 + latent_channels: int = 128 + rope_base: float = 10000.0 + + +class _RMSNorm(nn.Module): + def __init__(self, dim): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + y = F.rms_norm(x.float(), (x.shape[-1],), eps=1e-6) + return y.to(x.dtype) * self.weight + + +def _situ(hidden, gate): + return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * ( + 25.0 * torch.tanh(hidden / 25.0) + ) + + +def _rotate(x, rope): + sin, cos = rope.chunk(2, -1) + sin = torch.cat((sin, sin), -1).unsqueeze(-2) + cos = torch.cat((cos, cos), -1).unsqueeze(-2) + a, b = x.chunk(2, -1) + return x * cos + torch.cat((-b, a), -1) * sin + + +def _attention(q, k, v, rope, heads, mask): + shape = (*q.shape[:2], heads, q.shape[-1] // heads) + q = _rotate(q.reshape(shape), rope).transpose(1, 2) + k = _rotate(k.reshape(shape), rope).transpose(1, 2) + v = v.reshape(shape).transpose(1, 2) + out = F.scaled_dot_product_attention(q.to(v.dtype), k.to(v.dtype), v, attn_mask=mask) + return out.transpose(1, 2).flatten(2) + + +class _TextBlock(nn.Module): + def __init__(self, dim, ffn_dim, heads): + super().__init__() + self.heads = heads + self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim) + self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads) + self.qkv = nn.Linear(dim, 3 * dim) + self.proj = nn.Linear(dim, dim) + self.ffn_gate, self.ffn_up = nn.Linear(dim, ffn_dim), nn.Linear(dim, ffn_dim) + self.ffn_down = nn.Linear(ffn_dim, dim) + + def forward(self, x, rope, mask): + q, k, v = self.qkv(self.norm1(x)).chunk(3, -1) + q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2) + k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2) + x = x + self.proj(_attention(q, k, v, rope, self.heads, mask)) + y = self.norm2(x) + return x + self.ffn_down(_situ(self.ffn_up(y), self.ffn_gate(y))) + + +class _TextProjection(nn.Module): + def __init__(self, text_dim, dim, ffn_dim, heads): + super().__init__() + self.proj = nn.Linear(text_dim, dim) + self.blocks = nn.ModuleList([_TextBlock(dim, ffn_dim, heads) for _ in range(2)]) + + +class _Block(nn.Module): + def __init__(self, dim, ffn_dim, heads): + super().__init__() + self.heads = heads + self.norm1, self.norm2 = _RMSNorm(dim), _RMSNorm(dim) + self.q_norm, self.k_norm = _RMSNorm(dim // heads), _RMSNorm(dim // heads) + # q/k/v fused (rows q, k, v) like the ComfyUI files; the vendor keeps + # them split (fused on load) + self.proj = nn.ModuleDict( + { + "qkv": nn.Linear(dim, 3 * dim, bias=False), + "o": nn.Linear(dim, dim, bias=False), + "ffn_up_gate": nn.Linear(dim, 2 * ffn_dim, bias=False), + "ffn_down": nn.Linear(ffn_dim, dim, bias=False), + } + ) + + def forward(self, x, rope, modulation, mask): + shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = modulation.unbind(1) + y = self.norm1(x) * (1 + scale_a[:, None]) + shift_a[:, None] + q, k, v = self.proj["qkv"](y).chunk(3, -1) + q = self.q_norm(q.unflatten(-1, (self.heads, -1))).flatten(2) + k = self.k_norm(k.unflatten(-1, (self.heads, -1))).flatten(2) + attn = _attention(q, k, v, rope, self.heads, mask) + x = x + self.proj["o"](attn) * gate_a[:, None] + y = self.norm2(x) * (1 + scale_f[:, None]) + shift_f[:, None] + hidden, gate = self.proj["ffn_up_gate"](y).chunk(2, -1) + return x + self.proj["ffn_down"](_situ(hidden, gate)) * gate_f[:, None] + + +class MossImageTransformer(nn.Module, OstrisModelMixin): + @classmethod + def aitk_config_from_state_dict(cls, state_dict): + # sniffed from unquantized tensors only (block linears may be int8) + video = state_dict["video_embedder.weight"] + dim = video.shape[0] + return MossImageParams( + dim=dim, + ffn_dim=state_dict["text_embedder.blocks.0.ffn_up.weight"].shape[0], + num_heads=dim // state_dict["layers.0.q_norm.weight"].shape[0], + num_layers=sum( + 1 + for k in state_dict + if k.startswith("layers.") and k.endswith(".norm1.weight") + ), + text_dim=state_dict["text_embedder.proj.weight"].shape[1], + freq_dim=state_dict["time_embedding.0.weight"].shape[1], + latent_channels=video.shape[1], + ) + + @classmethod + def convert_state_dict_on_load(cls, state_dict): + # derived constants; recomputed in fp32 every forward + state_dict = {k: v for k, v in state_dict.items() if k != "rope.bands"} + fuse_split_qkv(state_dict) + return state_dict + + @classmethod + def convert_state_dict_on_save(cls, state_dict): + # the vendor's strict loader expects the (unused) frequency table + state_dict = dict(state_dict) + head_dim = state_dict["layers.0.q_norm.weight"].shape[0] + state_dict["rope.bands"] = _rope_bands("cpu", head_dim, MossImageParams.rope_base) + return state_dict + + @classmethod + def get_transformer_block_names(cls): + return ["layers"] + + @classmethod + def get_quantization_exclude_modules(cls): + # embedders, modulation, text refiner and output head stay full precision + return [ + "video_embedder", + "text_embedder*", + "time_embedding*", + "shared_block_modulation", + "video_head_modulation", + "final_linear_video", + ] + + def __init__(self, params: Optional[MossImageParams] = None): + super().__init__() + if params is None: + params = MossImageParams() + if params.dim % params.num_heads or (params.dim // params.num_heads) % 8: + raise ValueError("dim must divide into heads with head_dim divisible by 8") + self.params = params + dim, ffn_dim, heads = params.dim, params.ffn_dim, params.num_heads + self.video_embedder = nn.Linear(params.latent_channels, dim) + self.text_embedder = _TextProjection(params.text_dim, dim, ffn_dim, heads) + self.time_embedding = nn.Sequential( + nn.Linear(params.freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim) + ) + self.shared_block_modulation = nn.Linear(dim, 6 * dim) + self.layers = nn.ModuleList( + [_Block(dim, ffn_dim, heads) for _ in range(params.num_layers)] + ) + self.final_norm_video = _RMSNorm(dim) + self.final_linear_video = nn.Linear(dim, params.latent_channels, bias=False) + self.video_head_modulation = nn.Linear(dim, 2 * dim) + self.gradient_checkpointing = False + self._bands = {} + + @property + def device(self): + return self.final_norm_video.weight.device + + @property + def dtype(self): + # norms are never quantized; the block linears may be + return self.final_norm_video.weight.dtype + + def enable_gradient_checkpointing(self, enable: bool = True): + self.gradient_checkpointing = enable + + def disable_gradient_checkpointing(self): + self.gradient_checkpointing = False + + def _ckpt(self, module, *args): + if self.gradient_checkpointing and torch.is_grad_enabled(): + return checkpoint(module, *args, use_reentrant=False) + return module(*args) + + def _rope(self, lengths, height, width, text_len, device): + """(B, H*W + S, head_dim) fp32 rope; coordinates (r, f, h, w): image + tokens at (0, item text length, centered h, centered w), text token i + at (0, i, 0, 0).""" + bands = self._bands.get(device) + if bands is None: + bands = _rope_bands( + device, self.params.dim // self.params.num_heads, self.params.rope_base + ) + self._bands[device] = bands + batch = lengths.shape[0] + h, w = torch.meshgrid( + torch.arange(height, device=device, dtype=torch.float32), + torch.arange(width, device=device, dtype=torch.float32), + indexing="ij", + ) + image = torch.zeros(batch, height * width, 4, device=device) + image[..., 1] = lengths.float()[:, None] + image[..., 2] = (h - (height - 1) / 2).flatten() + image[..., 3] = (w - (width - 1) / 2).flatten() + text = torch.zeros(batch, text_len, 4, device=device) + text[..., 1] = torch.arange(text_len, device=device, dtype=torch.float32) + phase = torch.cat((image, text), 1).unsqueeze(-1) * bands + return torch.cat((phase.sin(), phase.cos()), -2).flatten(-2) + + def forward( + self, + hidden_states: torch.Tensor, # (B, latent_channels, H, W) + timestep: torch.Tensor, # 0..1000, 1000 = pure noise + encoder_hidden_states: torch.Tensor, # (B, S, text_dim) + encoder_attention_mask: Optional[torch.Tensor] = None, # (B, S) bool + skip_layers: Optional[Iterable[int]] = None, + ): + skip = frozenset(skip_layers or ()) + batch, _, height, width = hidden_states.shape + num_image = height * width + device = hidden_states.device + dtype = self.dtype + + context = encoder_hidden_states + mask = encoder_attention_mask + if mask is None: + mask = torch.ones(context.shape[:2], dtype=torch.bool, device=device) + mask = mask.to(device=device, dtype=torch.bool) + lengths = mask.sum(1) + if bool((lengths == 0).any()): + raise ValueError("each sample needs at least one valid text token") + text_len = int(lengths.max()) + padded = bool((lengths < text_len).any()) + if padded or not bool(mask[:, :text_len].all()): + # valid tokens first, in order (the vendor selects them by mask) + order = torch.argsort((~mask).to(torch.int8), dim=1, stable=True) + context = torch.gather( + context, 1, order[..., None].expand(-1, -1, context.shape[-1]) + ) + context = context[:, :text_len] + valid = torch.arange(text_len, device=device)[None] < lengths[:, None] + + times = torch.as_tensor(timestep, device=device, dtype=torch.float32) + times = times.reshape(-1).expand(batch) + half = self.params.freq_dim // 2 + freq = 10000 ** ( + -torch.arange(half, device=device, dtype=torch.float64) / half + ) + phase = times.double()[:, None] * freq + embedding = torch.cat((phase.cos(), phase.sin()), -1).float() + time = self.time_embedding(embedding.to(dtype)).float() + mod = self.shared_block_modulation(time.to(dtype)).view(batch, 6, -1) + head = self.video_head_modulation(time.to(dtype)).view(batch, 2, -1) + + rope = self._rope(lengths, height, width, text_len, device).to(dtype) + + text_mask = joint_mask = None + if padded: + text_mask = valid[:, None, None, :] + joint_mask = torch.cat( + (torch.ones(batch, num_image, dtype=torch.bool, device=device), valid), 1 + )[:, None, None, :] + + video = hidden_states.flatten(2).transpose(1, 2).to(dtype) + video = self.video_embedder(video.contiguous()) + text = self.text_embedder.proj(context.to(dtype)) + text_rope = rope[:, num_image:] + for block in self.text_embedder.blocks: + text = self._ckpt(block, text, text_rope, text_mask) + + x = torch.cat((video, text), 1) + for index, layer in enumerate(self.layers): + if index not in skip: + x = self._ckpt(layer, x, rope, mod, joint_mask) + + x = self.final_norm_video(x[:, :num_image]) + x = x * (1 + head[:, 1, None]) + head[:, 0, None] + x = self.final_linear_video(x) + return x.transpose(1, 2).reshape(batch, -1, height, width) + + +def fuse_split_qkv(state_dict): + """Vendor / community layout (``layers.N.proj.{q,k,v}``, plain or comfy + quantized) -> the fused ComfyUI ``layers.N.proj.qkv``. In place.""" + from toolkit.util.comfy_quant_import import fuse_split_quantized_keys + + for key in [k for k in state_dict if k.endswith(".proj.q.weight")]: + base = key[: -len(".q.weight")] + parts = [f"{base}.{p}" for p in ("q", "k", "v")] + if f"{base}.q.comfy_quant" in state_dict: + fuse_split_quantized_keys(state_dict, parts, f"{base}.qkv") + else: + state_dict[f"{base}.qkv.weight"] = torch.cat( + [state_dict.pop(f"{p}.weight") for p in parts], dim=0 + ) + return state_dict + + +def _rope_bands(device, head_dim: int, base: float): + # built on cpu in fp32 like the vendor's training, then moved + num_bands = head_dim // 8 + exponent = torch.arange(num_bands, dtype=torch.float32) / num_bands + return (1 / (base**exponent)).to(device) diff --git a/extensions_built_in/diffusion_models/ui.tsx b/extensions_built_in/diffusion_models/ui.tsx index b818927..62141a7 100644 --- a/extensions_built_in/diffusion_models/ui.tsx +++ b/extensions_built_in/diffusion_models/ui.tsx @@ -1754,6 +1754,49 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ "model.qie.match_target_res", ], }, + { + name: "moss_image", + label: "MOSS-Image SFT", + group: "image", + defaults: { + // default updates when [selected, unselected] in the UI + // the vendor repo: its transformer is merged once into the comfy models + // folder; a .safetensors path (fine-tune / int8 convrot repack) loads as named + "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 + "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/extensions_built_in/inference_engine/latent_preview.py b/extensions_built_in/inference_engine/latent_preview.py index 642b3aa..b6c064f 100644 --- a/extensions_built_in/inference_engine/latent_preview.py +++ b/extensions_built_in/inference_engine/latent_preview.py @@ -51,6 +51,7 @@ ARCH_LATENT_FORMAT = { "cogview4": "Flux", "prx_pixel": None, # pixel space "flux3": "Flux3Image", + "moss_image": "Flux2", # FLUX.2 encoder, same packing "iris3b": None, # pixel space } diff --git a/testing/test_model_loading.py b/testing/test_model_loading.py index 4d6e585..673387c 100644 --- a/testing/test_model_loading.py +++ b/testing/test_model_loading.py @@ -215,6 +215,30 @@ MODEL_TESTS = { "sample": {**IMG, "num_inference_steps": 12, "guidance_scale": 1.0}, "needs_control_image": True, }, + "moss_image": { + "model": { + "name_or_path": "OpenMOSS-Team/MOSS-Image-SFT", + "quantize": True, + "quantize_te": True, + "qtype": "convrot8", + "qtype_te": "convrot8", + }, + # vendor defaults are 50 steps, CFG 3 (negative branch skips block 22) + "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 + "model": { + "name_or_path": "Comfy-Org/MOSS-Image/diffusion_models/moss_image_sft_int8_convrot_attn.safetensors", + "te_name_or_path": "Comfy-Org/MOSS-Image/text_encoders/qwen3.5_9b_int8_convrot.safetensors", + "quantize": True, + "quantize_te": True, + "qtype": "convrot8", + "qtype_te": "convrot8", + }, + "sample": {**IMG, "num_inference_steps": 20, "guidance_scale": 3.0}, + }, # ---- legacy monolith archs (components adopted into v2 on load) ---- "sd1": { "model": {"name_or_path": "stable-diffusion-v1-5/stable-diffusion-v1-5"}, diff --git a/toolkit/models/registry.py b/toolkit/models/registry.py index 4a89167..aecef16 100644 --- a/toolkit/models/registry.py +++ b/toolkit/models/registry.py @@ -56,6 +56,7 @@ ARCH_REGISTRY: Dict[str, dict] = { "ltx2.3": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}}, "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}}, "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}}, diff --git a/toolkit/models/v2/text_encoders/qwen35.py b/toolkit/models/v2/text_encoders/qwen35.py new file mode 100644 index 0000000..cf0b5d0 --- /dev/null +++ b/toolkit/models/v2/text_encoders/qwen35.py @@ -0,0 +1,56 @@ +from transformers import Qwen3_5TextModel + +from .._mixin import OstrisTransformersMixin + +QWEN35_COMFY_REPO = "Comfy-Org/Qwen3.5" + + +class Qwen35TextOnlyEncoder(Qwen3_5TextModel, OstrisTransformersMixin): + """The Qwen3.5 language tower alone (moss_image). Text-only inputs, so the + default arange positions match the multimodal model's text path. A stock + Qwen repo id resolves to the Comfy-Org single-file repack under + MODELS_PATH/text_encoders/ (config from the Qwen repo); its lm_head, mtp + head and vision tower are dropped on load.""" + + aitk_subfolder = "text_encoder" + aitk_tokenizer_subfolder = "tokenizer" + # the int8 repacks keep A_log / gated-norm weights in fp32; one dtype + # throughout, like the bf16 path + aitk_cast_quantized_load = True + + aitk_comfy_repo = QWEN35_COMFY_REPO + aitk_comfy_weight_names = { + "Qwen/Qwen3.5-2B": ["text_encoders/qwen3.5_2b_bf16.safetensors"], + "Qwen/Qwen3.5-4B": ["text_encoders/qwen3.5_4b_bf16.safetensors"], + "Qwen/Qwen3.5-9B": ["text_encoders/qwen3.5_9b_bf16.safetensors"], + } + + _PREFIXES = ("model.language_model.", "language_model.", "model.") + _DROPPED = ("lm_head.", "mtp.", "model.visual.", "visual.") + + @classmethod + def get_transformer_block_names(cls): + return ["layers"] + + @classmethod + def aitk_load_config(cls, path, subfolder=None): + config = super().aitk_load_config(path, subfolder=subfolder) + # the Qwen repos carry the multimodal Qwen3_5Config + return getattr(config, "text_config", config) + + @classmethod + def convert_state_dict_on_load(cls, state_dict): + out = {} + for key, value in state_dict.items(): + if key.startswith(cls._DROPPED): + continue + for prefix in cls._PREFIXES: + if key.startswith(prefix): + key = key[len(prefix) :] + break + out[key] = value + return out + + @classmethod + def convert_state_dict_on_save(cls, state_dict): + return {f"model.language_model.{k}": v for k, v in state_dict.items()} diff --git a/version.py b/version.py index 95d5614..1d12fa4 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.13.26" +VERSION = "0.13.27"