KBlueLeaf's picture
Report the 1024^2 plain-fp16 search: no CFG, seed or prompt breaks it
ba63601 verified
|
Raw History Blame Contribute Delete
22.7 kB
---
license: apache-2.0
base_model:
- OpenMOSS-Team/MOSS-Image-Base
- OpenMOSS-Team/MOSS-Image-SFT
pipeline_tag: text-to-image
library_name: diffusers
tags:
- moss-image
- fp16
- comfyui
---
<style>
.mf-badges a { margin-right: 4px; }
.mf-gal { margin: 1em 0; }
.mf-gal > input { display: none; }
.mf-gal > label {
display: inline-block; padding: 4px 10px; margin: 0 4px 6px 0; cursor: pointer;
border: 1px solid #8a8a8a; border-radius: 6px; font-size: 0.85em; color: #8a8a8a;
}
.mf-gal > .mf-panel { display: none; }
.mf-gal > .mf-panel img { width: 100%; border-radius: 6px; }
.mf-gal > .mf-panel p { font-size: 0.85em; color: #8a8a8a; margin: 4px 0 0 0; }
#b0:checked ~ #pb0, #b1:checked ~ #pb1, #b2:checked ~ #pb2, #b3:checked ~ #pb3,
#b4:checked ~ #pb4, #b5:checked ~ #pb5, #b6:checked ~ #pb6, #b7:checked ~ #pb7,
#s0:checked ~ #ps0, #s1:checked ~ #ps1, #s2:checked ~ #ps2, #s3:checked ~ #ps3,
#s4:checked ~ #ps4, #s5:checked ~ #ps5, #s6:checked ~ #ps6, #s7:checked ~ #ps7,
#l0:checked ~ #pl0, #l1:checked ~ #pl1, #l2:checked ~ #pl2, #l3:checked ~ #pl3 { display: block; }
#b0:checked + label, #b1:checked + label, #b2:checked + label, #b3:checked + label,
#b4:checked + label, #b5:checked + label, #b6:checked + label, #b7:checked + label,
#s0:checked + label, #s1:checked + label, #s2:checked + label, #s3:checked + label,
#s4:checked + label, #s5:checked + label, #s6:checked + label, #s7:checked + label,
#l0:checked + label, #l1:checked + label, #l2:checked + label, #l3:checked + label {
background: #2e86ab; border-color: #2e86ab; color: #ffffff;
}
</style>
# MOSS-Image FP16-fixed (Base + SFT)
<div class="mf-badges">
<a href="https://github.com/KohakuBlueleaf/KohakuMossTrainer"><img src="https://img.shields.io/badge/made%20with-KohakuMoss-2e86ab" alt="KohakuMoss"></a>
<a href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base"><img src="https://img.shields.io/badge/official-MOSS--Image--Base-e4572e" alt="MOSS-Image-Base"></a>
<a href="https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT"><img src="https://img.shields.io/badge/official-MOSS--Image--SFT-e4572e" alt="MOSS-Image-SFT"></a>
<a href="https://huggingface.co/Comfy-Org"><img src="https://img.shields.io/badge/ComfyUI-text%20encoder%20%2F%20VAE-29bf12" alt="Comfy-Org"></a>
</div>
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 |
<img src="images/gallery/base_comfyui_p1_s11.webp" alt="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](images/base_stream_per_block.webp)
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](images/base_plain_fp16.webp)
<img src="images/base_plain_fp16_p0.webp" alt="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):
<img src="images/base_fp16_break_2560.webp" alt="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](reports/plain_fp16_1024_hunt.json)).
<details><summary>SFT charts</summary>
![stream per block, SFT](images/sft_stream_per_block.webp)
![plain fp16 and the margin per block, SFT](images/sft_plain_fp16.webp)
</details>
## 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](images/fold_block.svg)
| 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](images/base_c_sweep.webp)
| 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](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](reports/base_stream_fp16_error.json)).
![stream fp16 rounding error per block, plain vs fixed](images/base_stream_fp16_error.webp)
"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
<div class="mf-gal">
<input type="radio" name="gb" id="b0" checked><label for="b0">cat</label>
<input type="radio" name="gb" id="b1"><label for="b1">island</label>
<input type="radio" name="gb" id="b2"><label for="b2">sign</label>
<input type="radio" name="gb" id="b3"><label for="b3">library</label>
<input type="radio" name="gb" id="b4"><label for="b4">teapot</label>
<input type="radio" name="gb" id="b5"><label for="b5">night market</label>
<input type="radio" name="gb" id="b6"><label for="b6">lake</label>
<input type="radio" name="gb" id="b7"><label for="b7">locomotive</label>
<div class="mf-panel" id="pb0"><img src="images/gallery/base_comfyui_p0_s11.webp" alt=""><p>A photograph of a tabby cat sitting on a sunlit windowsill, shallow depth of field.</p></div>
<div class="mf-panel" id="pb1"><img src="images/gallery/base_comfyui_p1_s11.webp" alt=""><p>An isometric illustration of a small floating island with a lighthouse, waterfalls pouring off the edges into clouds, soft pastel colors, highly detailed.</p></div>
<div class="mf-panel" id="pb2"><img src="images/gallery/base_comfyui_p2_s11.webp" alt=""><p>A wooden sign on a forest trail that reads 'KOHAKU' in carved letters, morning fog.</p></div>
<div class="mf-panel" id="pb3"><img src="images/gallery/base_comfyui_p3_s11.webp" alt=""><p>An anime-style illustration of a girl reading a book in a quiet library, warm light.</p></div>
<div class="mf-panel" id="pb4"><img src="images/gallery/base_comfyui_p4_s11.webp" alt=""><p>A studio product photo of a ceramic teapot on a marble table, soft shadows.</p></div>
<div class="mf-panel" id="pb5"><img src="images/gallery/base_comfyui_p5_s11.webp" alt=""><p>A busy night market street in the rain, neon signs reflected on wet pavement.</p></div>
<div class="mf-panel" id="pb6"><img src="images/gallery/base_comfyui_p6_s11.webp" alt=""><p>An oil painting of a mountain lake at sunset with a small red boat.</p></div>
<div class="mf-panel" id="pb7"><img src="images/gallery/base_comfyui_p7_s11.webp" alt=""><p>A detailed pencil sketch of an old steam locomotive crossing a stone bridge.</p></div>
</div>
### SFT
<div class="mf-gal">
<input type="radio" name="gs" id="s0" checked><label for="s0">cat</label>
<input type="radio" name="gs" id="s1"><label for="s1">island</label>
<input type="radio" name="gs" id="s2"><label for="s2">sign</label>
<input type="radio" name="gs" id="s3"><label for="s3">library</label>
<input type="radio" name="gs" id="s4"><label for="s4">teapot</label>
<input type="radio" name="gs" id="s5"><label for="s5">night market</label>
<input type="radio" name="gs" id="s6"><label for="s6">lake</label>
<input type="radio" name="gs" id="s7"><label for="s7">locomotive</label>
<div class="mf-panel" id="ps0"><img src="images/gallery/sft_comfyui_p0_s11.webp" alt=""><p>A photograph of a tabby cat sitting on a sunlit windowsill, shallow depth of field.</p></div>
<div class="mf-panel" id="ps1"><img src="images/gallery/sft_comfyui_p1_s11.webp" alt=""><p>An isometric illustration of a small floating island with a lighthouse, waterfalls pouring off the edges into clouds, soft pastel colors, highly detailed.</p></div>
<div class="mf-panel" id="ps2"><img src="images/gallery/sft_comfyui_p2_s11.webp" alt=""><p>A wooden sign on a forest trail that reads 'KOHAKU' in carved letters, morning fog.</p></div>
<div class="mf-panel" id="ps3"><img src="images/gallery/sft_comfyui_p3_s11.webp" alt=""><p>An anime-style illustration of a girl reading a book in a quiet library, warm light.</p></div>
<div class="mf-panel" id="ps4"><img src="images/gallery/sft_comfyui_p4_s11.webp" alt=""><p>A studio product photo of a ceramic teapot on a marble table, soft shadows.</p></div>
<div class="mf-panel" id="ps5"><img src="images/gallery/sft_comfyui_p5_s11.webp" alt=""><p>A busy night market street in the rain, neon signs reflected on wet pavement.</p></div>
<div class="mf-panel" id="ps6"><img src="images/gallery/sft_comfyui_p6_s11.webp" alt=""><p>An oil painting of a mountain lake at sunset with a small red boat.</p></div>
<div class="mf-panel" id="ps7"><img src="images/gallery/sft_comfyui_p7_s11.webp" alt=""><p>A detailed pencil sketch of an old steam locomotive crossing a stone bridge.</p></div>
</div>
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.
<div class="mf-gal">
<input type="radio" name="gl" id="l0" checked><label for="l0">Base LoRA</label>
<input type="radio" name="gl" id="l1"><label for="l1">Base LoKr</label>
<input type="radio" name="gl" id="l2"><label for="l2">SFT LoRA</label>
<input type="radio" name="gl" id="l3"><label for="l3">SFT LoKr</label>
<div class="mf-panel" id="pl0"><img src="images/gallery/base_comfyui_lora0_p3_s11.webp" alt=""><p>fp16 without the fix + adapter | fp16 with the fix + adapter | bf16 + adapter | fp32 + adapter (ComfyUI)</p></div>
<div class="mf-panel" id="pl1"><img src="images/gallery/base_comfyui_lora1_p3_s11.webp" alt=""><p>fp16 without the fix + adapter | fp16 with the fix + adapter | bf16 + adapter | fp32 + adapter (ComfyUI)</p></div>
<div class="mf-panel" id="pl2"><img src="images/gallery/sft_comfyui_lora0_p3_s11.webp" alt=""><p>fp16 without the fix + adapter | fp16 with the fix + adapter | bf16 + adapter | fp32 + adapter (ComfyUI)</p></div>
<div class="mf-panel" id="pl3"><img src="images/gallery/sft_comfyui_lora1_p3_s11.webp" alt=""><p>fp16 without the fix + adapter | fp16 with the fix + adapter | bf16 + adapter | fp32 + adapter (ComfyUI)</p></div>
</div>
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
```python
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](https://github.com/KohakuBlueleaf/KohakuMossTrainer).
## 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.