Add support for MOSS Image
This commit is contained in:
Jaret Burkett committed
1 parent
12395a103f
commit
e05a3a4f2c
14 files changed
+1290
-1
No files matched your search
@@ -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,
|
||||
]
|
||||
@@ -0,0 +1,3 @@
|
||||
from .moss_image import MossImageModel
|
||||
|
||||
__all__ = ["MossImageModel"]
|
||||
@@ -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()
|
||||
Whitespace-only changes.
@@ -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("<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)
|
||||
@@ -0,0 +1,143 @@
|
||||
"""MOSS-Image sampling (training previews / inference engine).
|
||||
|
||||
Matches the vendor pipeline: shift-5 sigma grid snapped to the nearest of the
|
||||
1000 training sigmas, fp32 Euler on fp32 noise, CFG 3 whose negative branch
|
||||
skips DiT block 22 (two forwards per step), empty negative prompt. Flow
|
||||
convention is the toolkit's: sigma 1 = noise, the DiT predicts noise - clean.
|
||||
"""
|
||||
|
||||
from typing import List, Optional, Sequence
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
|
||||
DEFAULT_SHIFT = 5.0
|
||||
DEFAULT_GUIDANCE = 3.0
|
||||
DEFAULT_STEPS = 30
|
||||
DEFAULT_UNCOND_SKIP_LAYERS = (22,)
|
||||
NUM_TRAIN_TIMESTEPS = 1000
|
||||
|
||||
|
||||
def moss_sigmas(
|
||||
num_steps: int,
|
||||
shift: float = DEFAULT_SHIFT,
|
||||
training_shift: float = 1.0,
|
||||
total: int = NUM_TRAIN_TIMESTEPS,
|
||||
) -> 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]
|
||||
@@ -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)
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user