MOSS-Image FP16-fixed (Base + SFT)

KohakuMoss MOSS-Image-Base MOSS-Image-SFT Comfy-Org

The MOSS-Image-Base and MOSS-Image-SFT transformers with their residual stream scaled down by c = 1/16 inside the weights, so the model runs in fp16 with 16x headroom instead of 5%. Same outputs as the official model up to rounding, closer to the fp32 reference than the official bf16 weights, in diffusers, ComfyUI and KohakuMoss.

Only the transformer changes. Text encoder, tokenizer, VAE, scheduler and pipeline code are the official ones.

file for
base/transformer/, sft/transformer/ diffusers MOSSImageTransformer (official config, fp16)
comfyui/moss_image_base_fp16_fixed.safetensors ComfyUI single file (fp16)
comfyui/moss_image_sft_fp16_fixed.safetensors ComfyUI single file (fp16)
reports/{base,sft}_fp16_fix.json every number on this page, per image
fp16 without the fix | fp16 with the fix | bf16 | fp32, ComfyUI

fp16 without the fix | fp16 with the fix | bf16 | fp32 (ComfyUI, 1024^2, seed 11).

Why: the stream sits at 95% of the fp16 limit

MOSS-Image adds every block's output into one residual stream. On the official weights its largest value is 62,089 (Base) / 61,828 (SFT) at block 26 of 42, measured in fp32 over 8 prompts + the empty prompt, 7 sigmas and 512^2, 1024^2 and 768x1344 (189 forwards), and 62,098 / 62,184 during real 50-step sampling. fp16 ends at 65,504.

stream per block, Base

Plain fp16 on the official weights did not overflow on any of these 189 forwards, and at 1024^2 its images look right (below). The margin is the problem: at block 26 it is 1.05x, and the stream grows with the pixel count. Anything that raises it by 5% (a larger image, a LoRA, a finetune) turns it into inf / NaN. The fixed weights keep 16.9x at the same block.

plain fp16 and the margin per block, Base

fp16 without the fix | fp16 with the fix | bf16 | fp32, 1024^2

fp16 without the fix | fp16 with the fix | bf16 | fp32 (KohakuMoss, 1024^2, prompt 0).

Where plain fp16 breaks. ComfyUI with --fp16-unet on the official Base file, 50 steps, CFG 3, euler / simple: all 8 prompts at 1024^2 finish (stream max 62,080), and so does 2048^2 (64,096, 2.2% below the limit). 2560^2 is the first size that gives a black image: on step 29 of 50 the stream reaches 65,504 at block 26, every later forward is NaN, and the image is black. KohakuMoss in plain fp16 goes non-finite at the same step and block. The same run in fp32 peaks at 67,048 (a single fp32 forward at sigma 0.7: 67,681), 2.4-3.3% over the fp16 limit; bf16 peaks at 67,072 and stays finite. The fixed weights peak at 4,188 (x 16 = 67,008, 15.6x headroom):

2560x2560: fp16 without the fix (NaN) | fp16 with the fix | bf16 | fp32

fp16 without the fix | fp16 with the fix | bf16 | fp32 (ComfyUI, 2560x2560, prompt 0 "tabby cat", seed 11; LPIPS* is the distance to fp32). In single forwards over 7 sigmas, 2560^2, 2048x3072, 3072x2048 and 3072^2 overflow; 1536^2, 2048^2, 768x2048 and 1024x4096 do not. A 226-token prompt moves the stream by under 1%. At 1024^2 no setting tried breaks: Base and SFT x CFG 3 / 5 / 7 / 10 x the 8 card prompts x seeds 11-14, plus 4 extreme prompts (226 tokens, dense text, saturated high-contrast, white high-key) at CFG 3 / 10, 288 ComfyUI runs, peak 62,208 (block 26), none non-finite; CFG and content move the peak by under 0.5%, the pixel count moves it (reports/plain_fp16_1024_hunt.json).

SFT charts

stream per block, SFT plain fp16 and the margin per block, SFT

How: the fold

RMSNorm does not see the scale of its input: RMSNorm(c x; eps c^2) = RMSNorm(x; eps). So the whole residual stream can be carried as c x: every writer into the stream is multiplied by c, every norm reading it gets eps x c^2, and every norm output (and so every attention / FFN input and the final prediction) is unchanged.

what is scaled

parameter change
video_embedder, text_embedder.proj, refiner proj / ffn_down (W, b) x c
shared_block_modulation rows of gate_attn and gate_ffn (W, b) x c
norm eps of the joint blocks, the text refiner and final_norm_video x c^2 (exact) or kept (stock)
q / k / v / o, ffn_up_gate, ffn_down, every other weight unchanged

c = 2^-4 is a power of two: the scaling is exact in fp32 and bf16. In fp16, 3.2% of the 101.08M scaled values become subnormal (3.25M) and 1,620 (Base) / 1,651 (SFT) flush to zero; with the bf16 -> fp16 cast of the other weights, 5,352 / 5,478 of the 5.08 billion weights are zero where the official weight is not. The o / ffn_down outputs peak at 1.3k / 0.7k before their gate, so they need no pre-gate scale (s = 1).

eps. The files keep the stock norm eps (1e-6), so diffusers and ComfyUI run them unmodified. Against the exact eps x c^2 that changes the output by a relative 1.3e-4 (Base) / 1.8e-4 (SFT) at c = 1/16, below fp16's own rounding step (4.9e-4). KohakuMoss reads the fold record in the file and can run either.

Choosing c

Each c below is a full run: 189 forwards for overflow, 4 sampled images (1024^2, 50 steps) against fp32.

c sweep, Base

c fp16 stream max headroom LPIPS vs fp32 scaled weights subnormal or flushed stock-eps rel. error
official bf16 61,184 1.07 0.135
1/2 31,056 2.1 0.033 0.42% 7.1e-6
1/4 15,528 4.2 0.044 0.83% 7.5e-6
1/8 7,764 8.4 0.033 1.65% 3.2e-5
1/16 3,882 16.9 0.040 3.22% 1.3e-4
1/32 1,941 33.7 0.042 6.12% 5.2e-4
1/64 970 67.5 0.034 11.49% 2.2e-3
1/256 243 270.0 0.157 38.14% 2.6e-2

(Base. SFT: same headroom and underflow, LPIPS 0.008-0.019 from 1/2 to 1/64, 0.068 at 1/256, official bf16 0.064; reports/sft_tables.md.)

  • Not larger. c is the largest power of two that puts the stream under 4,096, a 16x (4-bit) margin. 1/8 leaves 8.4x.
  • Not smaller. Quality is flat from 1/2 to 1/64 (no trend over 4 images) and breaks at 1/256. On the way down the subnormal share of the scaled weights doubles per halving, and the stock-eps error grows 4x per halving: at 1/32 it reaches fp16's rounding step, at 1/64 it is 4.5x past it. 1/16 is the smallest c whose stock-eps error stays below fp16 rounding.

Comparisons

1024^2, 50 steps, CFG 3, 8 prompts x seeds 11 / 12 (16 images), each runtime against its own fp32 run. LPIPS (AlexNet) / SSIM / PSNR on the saved 8-bit images; lower LPIPS, higher SSIM and PSNR is closer.

runtime variant fp16, no fix fp16-fixed official bf16 fixed vs bf16
ComfyUI Base 0.0062 / 0.986 / 44.7 dB 0.0062 / 0.986 / 44.8 dB 0.0501 / 0.914 / 29.3 dB 0.047
ComfyUI SFT 0.0035 / 0.990 / 42.6 dB 0.0042 / 0.990 / 43.2 dB 0.0413 / 0.919 / 29.3 dB 0.042
diffusers Base 0.0202 / 0.959 / 36.9 dB 0.0113 / 0.974 / 42.1 dB 0.0901 / 0.878 / 26.8 dB 0.089
diffusers SFT 0.0151 / 0.966 / 37.0 dB 0.0096 / 0.979 / 41.4 dB 0.0656 / 0.878 / 28.1 dB 0.065
KohakuMoss Base 0.0171 / 0.971 / 42.0 dB 0.0147 / 0.971 / 42.2 dB 0.0935 / 0.871 / 26.7 dB 0.090
KohakuMoss SFT 0.0053 / 0.985 / 42.2 dB 0.0122 / 0.975 / 42.0 dB 0.0683 / 0.873 / 27.4 dB 0.072

fp16 carries 3 more mantissa bits than bf16; with the stream in range that is what the fixed weights give back. "fixed vs bf16" is the size of bf16's own rounding drift, not an error of the fix. No 1024^2 run produced a non-finite latent or image: at this size the official weights still fit fp16 and plain fp16 is about as close to fp32 as the fixed weights. The difference is headroom (1.05x vs 16.9x), which is what the 2560^2 example above uses up.

Range, not precision. At 62k the fp16 step is 32, at the fixed stream's 3.9k it is 2; relative to the value both are 2^-11, because c = 2^-4 only shifts exponents. Measured on the same inputs (1024^2, prompt 0, sigma 0.9 / 0.7 / 0.5 / 0.3), the stream's fp16 rounding error against fp32 of the same weights is the same per block for both: at block 26 a relative 1.23e-3 (plain) vs 1.23e-3 (fixed), max error 96 vs 6.0 = 96 / 16; median ratio over the 42 blocks 1.003 (reports/base_stream_fp16_error.json).

stream fp16 rounding error per block, plain vs fixed

"fp16, no fix" ran on a V100, the other columns on an RTX 5090, with each seed's initial noise as the RTX 5090 draws it. The same plain-fp16 input on the two GPUs differs by LPIPS 0.039 / 0.012 (ComfyUI, Base / SFT), 0.015 / 0.025 (diffusers), 0.040 / 0.008 (KohakuMoss) over prompt 0 and both seeds; that is part of this column's distance. In fp32 the two GPUs agree to 0.0003.

Each panel: fp16 without the fix | fp16 with the fix | bf16 | fp32 (ComfyUI, seed 11). LPIPS* is the distance to the fp32 panel.

Base

A photograph of a tabby cat sitting on a sunlit windowsill, shallow depth of field.

An isometric illustration of a small floating island with a lighthouse, waterfalls pouring off the edges into clouds, soft pastel colors, highly detailed.

A wooden sign on a forest trail that reads 'KOHAKU' in carved letters, morning fog.

An anime-style illustration of a girl reading a book in a quiet library, warm light.

A studio product photo of a ceramic teapot on a marble table, soft shadows.

A busy night market street in the rain, neon signs reflected on wet pavement.

An oil painting of a mountain lake at sunset with a small red boat.

A detailed pencil sketch of an old steam locomotive crossing a stone bridge.

SFT

A photograph of a tabby cat sitting on a sunlit windowsill, shallow depth of field.

An isometric illustration of a small floating island with a lighthouse, waterfalls pouring off the edges into clouds, soft pastel colors, highly detailed.

A wooden sign on a forest trail that reads 'KOHAKU' in carved letters, morning fog.

An anime-style illustration of a girl reading a book in a quiet library, warm light.

A studio product photo of a ceramic teapot on a marble table, soft shadows.

A busy night market street in the rain, neon signs reflected on wet pavement.

An oil painting of a mountain lake at sunset with a small red boat.

A detailed pencil sketch of an old steam locomotive crossing a stone bridge.

Every runtime, both seeds and the adapter panels are in images/gallery/ (<variant>_<runtime>_p<prompt>_s<seed>.webp, <variant>_<runtime>_lora<i>_p<prompt>_s11.webp), all in the same four columns.

Whole pipeline in fp16

part bf16 vs fp32 fp16 vs fp32
Qwen3.5 text encoder, embedding cosine (mean / min, 9 prompts) 0.99979 / 0.99936 0.999996 / 0.999992
Flux2 VAE decode, PSNR (mean, 8 images) 57.9 dB 77.3 dB

Both run in plain fp16 with no change. The whole pipeline (fp16 text encoder + fp16-fixed transformer + fp16 VAE) is what the diffusers snippet below runs.

Adapters

A LoRA (rank 8) and a LoKr (factor 16), trained on the official Base weights in KohakuMoss, applied as-is to the fixed weights. Distance of each base + adapter to fp32 + adapter, then between fixed + adapter and official bf16 + adapter, next to the adapter's own effect:

runtime variant adapter fp16, no fix fp16-fixed bf16 fixed + adapter vs bf16 + adapter without adapter adapter effect
ComfyUI Base LoRA 0.003 / 0.993 / 44.7 0.006 / 0.990 / 43.8 0.068 / 0.897 / 26.6 0.056 0.047 0.287
ComfyUI Base LoKr 0.011 / 0.979 / 44.0 0.005 / 0.988 / 44.4 0.049 / 0.914 / 29.5 0.064 0.047 0.389
ComfyUI SFT LoRA 0.001 / 0.996 / 46.4 0.001 / 0.996 / 47.4 0.034 / 0.925 / 31.3 0.039 0.042 0.284
ComfyUI SFT LoKr 0.001 / 0.997 / 49.3 0.001 / 0.996 / 47.8 0.024 / 0.946 / 32.1 0.039 0.042 0.395
KohakuMoss Base LoRA 0.006 / 0.987 / 43.4 0.008 / 0.981 / 41.6 0.120 / 0.799 / 24.4 0.105 0.090 0.316
KohakuMoss Base LoKr 0.004 / 0.992 / 47.2 0.004 / 0.991 / 47.4 0.046 / 0.915 / 29.8 0.053 0.090 0.406
KohakuMoss SFT LoRA 0.003 / 0.992 / 43.0 0.002 / 0.994 / 45.4 0.038 / 0.922 / 29.7 0.053 0.072 0.288
KohakuMoss SFT LoKr 0.002 / 0.995 / 47.5 0.001 / 0.997 / 49.6 0.037 / 0.927 / 28.9 0.042 0.072 0.417

(The first three columns: LPIPS / SSIM / PSNR (dB), 8 images each, seed 11, against fp32 + adapter run on a V100, with plain fp16 + adapter also on the V100; the last three: LPIPS, 16 images each, RTX 5090.) The adapter changes the image 3-10x more than the fix does, fixed + adapter is as close to fp32 + adapter as plain fp16 is at 1024^2, and stays about as close to bf16 + adapter as the two bases are without it: it transfers.

fp16 without the fix + adapter | fp16 with the fix + adapter | bf16 + adapter | fp32 + adapter (ComfyUI)

fp16 without the fix + adapter | fp16 with the fix + adapter | bf16 + adapter | fp32 + adapter (ComfyUI)

fp16 without the fix + adapter | fp16 with the fix + adapter | bf16 + adapter | fp32 + adapter (ComfyUI)

fp16 without the fix + adapter | fp16 with the fix + adapter | bf16 + adapter | fp32 + adapter (ComfyUI)

Adapters on the joint blocks' q / k / v / o / ffn_up_gate / ffn_down (the usual targets) apply unchanged, since the fold leaves those weights alone. An adapter on video_embedder, text_embedder or shared_block_modulation needs its delta scaled by the same c; KohakuMoss's scripts/export/merge.py does that when merging into these files.

Usage

diffusers

import importlib

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "OpenMOSS-Team/MOSS-Image-Base",  # or MOSS-Image-SFT
    custom_pipeline="pipeline_moss_image",
    trust_remote_code=True,
    transformer=None,
    torch_dtype=torch.float16,
)
pipe.scheduler = type(pipe.scheduler).from_config(pipe.scheduler.config)
MOSSImageTransformer = importlib.import_module(type(pipe).__module__).MOSSImageTransformer
pipe.transformer = MOSSImageTransformer.from_pretrained(
    "KBlueLeaf/MOSS-Image-FP16-fixed",
    subfolder="base/transformer",  # or sft/transformer
    torch_dtype=torch.float16,
)
pipe.to("cuda")
image = pipe(
    "A photograph of a tabby cat sitting on a sunlit windowsill, shallow depth of field.",
    height=1024,
    width=1024,
    num_inference_steps=50,
    guidance_scale=3.0,
).images[0]

ComfyUI

  1. comfyui/moss_image_base_fp16_fixed.safetensors (or _sft_) into models/diffusion_models/; the Comfy-Org qwen3.5_9b_bf16.safetensors text encoder and flux2-vae_small.safetensors VAE as for the official model.
  2. Start ComfyUI with --fp16-unet. ComfyUI's MOSS-Image entry lists bf16 and fp32 only, so without the flag it runs this file in bf16 (or fp32 on GPUs without bf16).
  3. Load it with Load Diffusion Model in the official MOSS-Image workflow (euler / simple, 50 steps, CFG 3).

KohakuMoss

Training or sampling with PRECISION=fp16 and WEIGHTS (training) / TRANSFORMER (sampling) set to either file; the fold is read from the file's metadata and adapters are exported for the official weights. KohakuMoss can also fold the official weights on load (FP16_FIX=online); with the official weights and no fix it refuses PRECISION=fp16. See KohakuMoss.

Notes

  • Measured on an RTX 5090 (fp16 and bf16 matmuls with fp32 accumulation); the fp32 reference is the official weights in fp32. The 2560^2 example, the "fp16, no fix" column and the fp32 + adapter references ran on a V100 (bf16 there is emulated by torch; ComfyUI's own default on that GPU casts bf16 weights to fp32, so the 2560^2 bf16 run disables that cast). KohakuMoss and diffusers draw their initial noise with torch's CUDA generator, which gives different numbers on different GPUs; the V100 runs use the RTX 5090's noise for each seed.
  • The headroom figures cover the prompts, resolutions and sigmas above, and 2560^2 for the fixed Base weights (15.6x). Larger images, very long prompts and finetuned weights were not measured.
  • Other runtimes than the three above must keep the stock norm eps or use eps x c^2; either works (see eps).

License

Apache-2.0, as the official MOSS-Image weights.

Downloads last month
-
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for KBlueLeaf/MOSS-Image-FP16-fixed

Finetuned
(22)
this model