Spaces:
Running on Zero
Running on Zero
| """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 UI defaults): a | |
| repack DiT file already on disk wins (best precision for the requested | |
| 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. | |
| 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"} | |
| # 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 = [ | |
| (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, | |
| ) | |
| 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) | |
| 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("<Q", f.read(8))[0] | |
| header = json.loads(f.read(header_len)) | |
| header.pop("__metadata__", None) | |
| for name, meta in header.items(): | |
| start, end = meta["data_offsets"] | |
| tensors[name] = (meta["dtype"], meta["shape"], path, 8 + header_len + start, end - start) | |
| # out name -> (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("<Q", len(blob))) | |
| out.write(blob) | |
| for name in sorted(entries): | |
| dtype, _, parts = entries[name] | |
| for src_name, path, start, size in parts: | |
| if dtype in _CAST_TO_BF16: | |
| with safe_open(path, "pt") as f: | |
| tensor = f.get_tensor(src_name).to(torch.bfloat16) | |
| out.write(tensor.contiguous().view(torch.uint8).numpy().tobytes()) | |
| continue | |
| src = handles.get(path) | |
| if src is None: | |
| src = handles[path] = open(path, "rb") | |
| src.seek(start) | |
| remaining = size | |
| while remaining: | |
| chunk = src.read(min(chunk_bytes, remaining)) | |
| out.write(chunk) | |
| remaining -= len(chunk) | |
| finally: | |
| for src in handles.values(): | |
| src.close() | |
| os.replace(part_path, out_path) | |