diff --git a/CHANGELOG.md b/CHANGELOG.md index bfc03fbaa..9dab8da62 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,12 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Added +- `gmlx pull`, `image generate`, `serve` and the ComfyUI pack run MOSS-Image + and MOSS-Image Base GGUFs, and their LoRAs. +- The ComfyUI pack runs MOSS-Image without the OpenMOSS plugin, and takes + `SkipLayerGuidanceDiTSimple` on it to leave blocks out of the negative pass. +- MOSS-Image takes an `upsampler`, on a caption prompt gmlx writes from the + OpenMOSS prompt-enhancer skill. - FLUX.3 klein rewrites each prompt with a chat model of your choice before encoding it, set as the entry's `upsampler` or by `--upsampler`. - FLUX.2 klein, Krea 2, Z-Image and Qwen-Image 2.1 take an `upsampler` too, @@ -53,6 +59,8 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). `--chat-template` as invalid Jinja. - `--thinking off` was silently ignored on a model with no thinking switch when its template has the word thinking in its text. It now warns. +- A Krea 2 GGUF whose stored projector shape is not a list of integers + crashed the load. It is refused with the reason. - A streamed image request refused while it ran, such as one too large for memory, got a server error frame where the same request unstreamed got a 400. - The `gmlx pull` plan of an image model could say its parts fit in stages diff --git a/README.md b/README.md index 9406dc1cb..d85df9dcb 100644 --- a/README.md +++ b/README.md @@ -98,8 +98,8 @@ The recording runs at true speed, with a 27B model in a local server. - Embeddings, rerank, speech-to-text and text-to-speech on the same server, for a local RAG and voice stack. See [Speech, embeddings and rerank](https://asher.github.io/gmlx/services.html). -- Image generation from Z-Image, Qwen-Image 2.1, Krea 2 Turbo, FLUX.2 klein - and FLUX.3 klein GGUFs, on the command line, on the OpenAI image routes +- Image generation from Z-Image, Qwen-Image 2.1, Krea 2 Turbo, MOSS-Image, + FLUX.2 klein and FLUX.3 klein GGUFs, on the command line, on the OpenAI image routes and in ComfyUI. It needs an M3 chip or newer. See [Image generation](https://asher.github.io/gmlx/images.html). - A decision API at `/v1/systemone` that returns a probability for each diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md index 5bcb89515..0969cfc8e 100644 --- a/THIRD_PARTY_NOTICES.md +++ b/THIRD_PARTY_NOTICES.md @@ -25,6 +25,9 @@ in [`licenses/`](licenses/). | [flux2](https://github.com/black-forest-labs/flux2) | Apache-2.0 | BFL's FLUX.2 reference at commit `50fe516`, ported to MLX and modified for gmlx: the DiT constants of `Klein4BParams`, `Klein9BParams` and `Flux2Params` and the tensor names and shapes of `Flux2` in `src/flux2/model.py`, the sampling defaults of `FLUX2_MODEL_INFO` in `src/flux2/util.py`, the prompt rules of `Qwen3Embedder` in `src/flux2/text_encoder.py`, the rounding of `denoise` and `denoise_cfg` and the reference image rules of `encode_image_refs` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein/family.py`, the constants of `Klein9BParams`, the 9B rows of `FLUX2_MODEL_INFO` and the Qwen3-8B encoder of `load_qwen3_embedder` in `src/flux2/text_encoder.py` in `gmlx/image/families/flux2_klein_9b/family.py`, the `flux.2-klein-9b-kv` row of `FLUX2_MODEL_INFO` in `gmlx/image/families/flux2_klein_9b_kv/family.py` and the step of `denoise_cached` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein_9b_kv/step.py`, the DiT forward of `src/flux2/model.py`, with `forward_kv_extract`, `forward_kv_cached` and `causal_attn_fn`, and the token ids of `prc_img`, `prc_txt` and `encode_image_refs` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein/model.py`, and the schedule of `compute_empirical_mu` and `get_schedule` in `src/flux2/sampling.py` in `flux2_mu` and `flux2_sigmas` of `gmlx/image/schedule.py`, and the reference image preparation of `encode_image_refs` and `default_prep` in `src/flux2/sampling.py` in `gmlx/image/families/flux2_klein/edit.py`; and the upsampler's text to image system message `SYSTEM_MESSAGE_UPSAMPLING_T2I` of `src/flux2/system_messages.py`, shipped unmodified as `gmlx/image/families/flux2_klein/system_message_upsampling_t2i.txt`, with the chat turn of `scripts/cli.py` in `gmlx/image/families/flux2_klein/family.py`. | | [Z-Image-Turbo Space](https://huggingface.co/spaces/Tongyi-MAI/Z-Image-Turbo) | apache-2.0, from the Space's card; `pe.py` carries no license header | At revision `768cb50`, the prompt enhancer's system prompt `prompt_template` of `pe.py`, shipped as its string in `gmlx/image/families/z_image/pe_prompt_template.txt`, with its use in `app.py` in `gmlx/image/families/z_image/family.py`. | | [krea-2](https://github.com/krea-ai/krea-2) | Apache-2.0 | At commit `db3984f`, the prompt expansion system prompt `docs/expansion.txt`, shipped unmodified in `gmlx/image/families/krea2/`, with its use of `docs/prompting.md` in `gmlx/image/families/krea2/family.py`. The Krea 2 weights are under their own license and do not ship with gmlx. | +| [MOSS-Image](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT) | Apache-2.0 | OpenMOSS-Team/MOSS-Image-SFT at revision `fc5de0d`, ported to MLX and modified for gmlx: the tensor names and shapes of `MOSSImageTransformer` in `pipeline_moss_image.py` and the constants of that file and of `transformer/config.json`, and the prompt template of `tokenizer/chat_template.jinja` with the token cap of the pipeline, in `gmlx/image/families/moss_image/family.py`, and the DiT forward of `pipeline_moss_image.py` in `gmlx/image/families/moss_image/model.py` (license text in [`licenses/MOSS-Image-LICENSE`](licenses/MOSS-Image-LICENSE)). | +| [MOSS-Image ComfyUI plugin](https://github.com/OpenMOSS/MOSS-Image/tree/main/comfyui) | Apache-2.0 | The OpenMOSS ComfyUI plugin at commit `c1355e6`: the model settings of `comfyui/moss-image/loaders.py` (a flow model on the FLUX.2 latent format, sampling shift 1.0, timestep multiplier 1000, bfloat16 and float32 inference, memory usage factor 4.0), which the `MOSSImage` model class in `gmlx/comfyui_pack/moss_image.py` takes. | +| [MOSS-Image prompt-enhancer skill](https://github.com/OpenMOSS/MOSS-Image/tree/main/skills/moss-image-prompt-enhancer) | Apache-2.0 | At commit `c1355e6`, the caption rules of `skills/moss-image-prompt-enhancer/SKILL.md`, restated and modified for gmlx as the system prompt `gmlx/image/families/moss_image/caption_prompt.txt`, with the skill's `{"caption": ...}` reply read in `gmlx/image/families/moss_image/family.py`. The skill's layout drafts and its tool are not used. | | [ComfyUI-GGUF](https://github.com/city96/ComfyUI-GGUF) | Apache-2.0 | The model detection keys of `convert.py` at commit `6ea2651`, which `gmlx/image/fingerprint.py` uses to refuse other diffusion model layouts by name. | | [misaki](https://github.com/hexgrad/misaki) | Apache-2.0 | Whole package vendored at `gmlx/_vendor/misaki` (commit `fba1236`): Kokoro's grapheme-to-phoneme front-end, used by the TTS service until upstream's next PyPI release. | | [musl](https://musl.libc.org) | MIT | Statically linked into the guest entry binary `gmlx/container/guest/gmlx-entry` as its C library ([`licenses/musl-COPYRIGHT`](licenses/musl-COPYRIGHT)). | diff --git a/docs/cli.md b/docs/cli.md index 827eb9080..fcb24cf60 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -640,8 +640,8 @@ gmlx pull qwen-image --quant Q5_K_M An image model name, `z-image-turbo`, `qwen-image`, `qwen-image-turbo`, `krea-2-turbo`, `flux-2-klein`, `flux-2-klein-base`, `flux-2-klein-9b`, -`flux-2-klein-base-9b`, `flux-2-klein-9b-kv`, `flux-3-klein` or -`flux-3-klein-base`, pulls its DiT, text encoder and VAE, checks the VAE +`flux-2-klein-base-9b`, `flux-2-klein-9b-kv`, `flux-3-klein`, +`flux-3-klein-base`, `moss-image` or `moss-image-base`, pulls its DiT, text encoder and VAE, checks the VAE against its published sha256, packs it to GGUF, and adds an `image_models` entry. The quants default to the ones for this Mac's memory, and an encoder already on disk is used @@ -971,6 +971,7 @@ the weights load. | Z-Image | Qwen3-4B | The FLUX.1 autoencoder | The variant's steps and guidance. | | Qwen-Image 2.1 | Qwen3-VL-8B, its text GGUF | The Qwen-Image 2.1 autoencoder | 40 steps and guidance 1. Turbo: 8 steps and guidance 1. | | Krea 2 | Qwen3-VL-4B-Instruct, its text GGUF | The Qwen-Image autoencoder | Turbo: 8 steps and guidance 1. | +| MOSS-Image | Qwen3.5-9B | The MOSS-Image autoencoder | 50 steps and guidance 3, with block 22 left out of the negative pass. Base: the same. | | FLUX.3 klein | Qwen3-VL-4B-Instruct, its text GGUF | The FLUX.2 autoencoder | 16 steps and guidance 1. Base: 50 steps and guidance 3.5. | | FLUX.2 klein | Qwen3-4B | The FLUX.2 autoencoder | 4 steps and guidance 1, which is the only guidance it takes. Base: 50 steps and guidance 4. | | FLUX.2 klein 9B | Qwen3-8B | The FLUX.2 autoencoder | The same as FLUX.2 klein's. Its license allows non-commercial use only. | @@ -1008,8 +1009,8 @@ loads from the encoder's mmproj GGUF, and the VAE encodes them. `--size` sets the size of the output, which defaults to the entry's `defaults.size` or the variant's, and not to a reference's. The fit check counts the references, and a request that does not fit is refused before any model -loads. Z-Image and Krea 2 take no reference images, and `--ref` exits 2 on -them. +loads. Z-Image, Krea 2 and MOSS-Image take no reference images, and +`--ref` exits 2 on them. `--ref` gives FLUX.3 klein up to 10 reference images too. Its text encoder reads the prompt alone, so it needs no `--encoder-mmproj`, and the VAE @@ -1150,7 +1151,8 @@ and compares every tensor bit for bit before the file replaces `--out`. It packs the VAE of each image family: the FLUX.1 autoencoder for `z_image`, the Qwen-Image 2.1 autoencoder for `qwen_image21`, whose official file is F32 and stays F32, the Qwen-Image autoencoder for -`krea2`, whose file is BF16, and the FLUX.2 autoencoder for `flux3_klein`, +`krea2`, whose file is BF16, the MOSS-Image autoencoder for `moss_image`, +whose file is F32, and the FLUX.2 autoencoder for `flux3_klein`, `flux2_klein` and `flux2_klein_9b`. BFL's FLUX.2 klein 4B file of that autoencoder is BF16 with its config, and the Comfy-Org `flux2-vae.safetensors` is F32 and needs no config. It packs from that diff --git a/docs/comfyui-reference.md b/docs/comfyui-reference.md index eb424525f..186c932b0 100644 --- a/docs/comfyui-reference.md +++ b/docs/comfyui-reference.md @@ -74,8 +74,10 @@ ComfyUI sees each gmlx model as its own model class, so its stock nodes for that model work too, such as `ModelSamplingAuraFlow` and `TextEncodeQwenImage21`. The prompt nodes of the example workflows are ComfyUI's `CLIP Text Encode (Prompt)`, and the latent node is its -`EmptySD3LatentImage`, or its `EmptyFlux2LatentImage` for FLUX.2 klein and -FLUX.3 klein. +`EmptySD3LatentImage`, or its `EmptyFlux2LatentImage` for FLUX.2 klein, +FLUX.3 klein and MOSS-Image. ComfyUI has no model class for MOSS-Image, +so the gmlx nodes register their own, with the settings of the OpenMOSS +ComfyUI plugin, and stock nodes such as `ModelSamplingSD3` work on it. The sampler, the noise, the guidance and the schedule run in ComfyUI. Each call to the model is one request to the gmlx server, which runs the DiT, @@ -105,6 +107,8 @@ that names it: - Hook keyframes and scheduled CLIP hooks. - A quantized Qwen-Image 2.1 prefix cache: `QwenImage21Cache` with a dtype other than the default. +- A MOSS-Image conditioning that masks some of its text rows, such as one + from a text encoder other than `gmlx CLIP Loader`. - A FLUX.2 klein reference method other than `index`, such as `offset` from `FluxKontextMultiReferenceLatentMethod`, but for the `index_timestep_zero` method that `Flux KV Cache` sets. @@ -118,5 +122,11 @@ keys and values of the reference images of each row of the batch from the first step of a run, the later steps of the run reuse those of each row, and the end of the run drops them. +`SkipLayerGuidanceDiTSimple` on MOSS-Image is not refused when it names +`single_layers` alone, with `start_percent` 0 and `end_percent` 1. The +server leaves those blocks out of the negative pass of each step. The +node is refused with `double_layers`, another window, a linked widget or +a block MOSS-Image lacks, and on any other model. + A refusal from a node's own setup shows on that node. A refusal during sampling shows on the sampler and names the node upstream that caused it. diff --git a/docs/comfyui.md b/docs/comfyui.md index 933569690..153250ca1 100644 --- a/docs/comfyui.md +++ b/docs/comfyui.md @@ -213,6 +213,17 @@ on with a `CFGGuider` and a `KSamplerSelect`. Krea 2 Turbo samples at `cfg` Krea 2 quant fits a 1024x1024 image on a Mac with 16 GB, so set the latent to 512x512 there. +MOSS-Image reads its prompt with the Qwen3.5-9B GGUF on the server, and +the gmlx nodes run it on the Mac's GPU without the OpenMOSS ComfyUI plugin +or the ComfyUI-MOSS-GGUF bridge. Its workflow, `moss-image-t2i`, takes the +three gmlx loaders with `moss-image` or `moss-image-base` in `model_name`, +`Empty Flux 2 Latent` for the image's size, `ModelSamplingSD3` at shift 5, +and `SkipLayerGuidanceDiTSimple` with `single_layers` set to 22, which +leaves block 22 out of the negative pass, as the OpenMOSS workflow does. +`KSampler` samples 50 steps at `cfg` 3 with `euler` and `sgm_uniform`. +That is ComfyUI's schedule, not the release's schedule that +`gmlx image generate` and the server run. + There are two more gmlx nodes. `gmlx Sigmas` gives the fixed schedule of a distilled model such as Qwen-Image 2.1 Turbo, for `SamplerCustom` or `SamplerCustomAdvanced`, though `KSampler` works on every model. The @@ -226,7 +237,7 @@ stock nodes work with a gmlx model and what runs where. ## The example workflows -ComfyUI's template browser lists ten workflows under `comfyui-gmlx`. Their +ComfyUI's template browser lists eleven workflows under `comfyui-gmlx`. Their files are in `custom_nodes/comfyui-gmlx/example_workflows`. Each has a note with the commands that install the gmlx nodes and pull its model. @@ -236,6 +247,7 @@ with the commands that install the gmlx nodes and pull its model. | `z-image-turbo-lora-t2i` | The same through a stock `LoraLoader` | A Z-Image LoRA named `z-image-turbo-style.safetensors` in `models/loras`, or another picked in the node | | `qwen-image-2.1-edit-two-references` | A Qwen-Image 2.1 edit from two reference images, through `TextEncodeQwenImage21` | `gmlx pull qwen-image`, and `reference-1.png` and `reference-2.png` in ComfyUI's `input` folder | | `krea-2-turbo-t2i` | Krea 2 Turbo text to image at 1024x1024 in 8 steps, through `BasicScheduler` and `SamplerCustomAdvanced` | `gmlx pull krea-2-turbo`, and on a Mac with 16 GB a 512x512 latent | +| `moss-image-t2i` | MOSS-Image text to image at 1024x1024 in 50 steps at `cfg` 3, with block 22 left out of the negative pass through `SkipLayerGuidanceDiTSimple` | `gmlx pull moss-image` | | `flux-2-klein-t2i` | FLUX.2 klein text to image at 1024x1024 in 4 steps, through `Flux2Scheduler` and `SamplerCustomAdvanced` | `gmlx pull flux-2-klein` | | `flux-2-klein-edit-two-references` | A FLUX.2 klein edit from two reference images, through `ReferenceLatent` | `gmlx pull flux-2-klein`, and `reference-1.png` and `reference-2.png` in ComfyUI's `input` folder | | `flux-2-klein-9b-t2i` | FLUX.2 klein 9B text to image at 1024x1024 in 4 steps, the same graph on the 9B model | `gmlx pull flux-2-klein-9b`, and a Mac with 24 GB or more | diff --git a/docs/image-api.md b/docs/image-api.md index 800f1558f..4e8053b51 100644 --- a/docs/image-api.md +++ b/docs/image-api.md @@ -90,7 +90,7 @@ GIF. | Model | Edits | Variations | |-------|-------|------------| -| Z-Image, Krea 2 | One image, which the result starts from at `strength`, 0.6 by default. More images get 400. | The image with an empty prompt, at the default `strength` | +| Z-Image, Krea 2, MOSS-Image | One image, which the result starts from at `strength`, 0.6 by default. More images get 400. | The image with an empty prompt, at the default `strength` | | Qwen-Image 2.1 | Up to 10 reference images, when the entry has `encoder_mmproj` | The image as a reference, with an instruction to make a variation of it | | FLUX.2 klein | Up to 10 reference images, with no vision tower | The image as a reference, with an instruction to make a variation of it | | FLUX.3 klein | Up to 10 reference images, with no vision tower | The image as a reference, with an instruction to make a variation of it | diff --git a/docs/images.md b/docs/images.md index 5d882d73c..f2f515116 100644 --- a/docs/images.md +++ b/docs/images.md @@ -1,8 +1,8 @@ # Image generation gmlx makes images from a text prompt, and edits images you give it, with -the Z-Image-Turbo, Qwen-Image 2.1, Krea 2 Turbo, FLUX.2 klein and FLUX.3 -klein models. You can run them from the command line, call them from an +the Z-Image-Turbo, Qwen-Image 2.1, Krea 2 Turbo, MOSS-Image, FLUX.2 klein +and FLUX.3 klein models. You can run them from the command line, call them from an app through the OpenAI image API, or use them in [ComfyUI](comfyui.md). Image generation came with gmlx 0.5 and needs a Mac with an M3 chip or newer and 16 GB of memory or more. @@ -22,6 +22,8 @@ get when you ask for none. | `qwen-image-turbo` | 14.4 GB | 17.5 GB | 1024x1024 | Research and evaluation only | | `qwen-image` | 14.4 GB | 17.6 GB | 1024x1024 | Research and evaluation only | | `krea-2-turbo` | 9.8 GB | 14.1 GB | 512x512 on 16 GB, else 1024x1024 | Free commercial use under $1M of yearly revenue | +| `moss-image` | 10.5 GB | 15.5 GB | 1024x1024 | Apache 2.0, any use | +| `moss-image-base` | 10.5 GB | 15.5 GB | 1024x1024 | Apache 2.0, any use | | `flux-2-klein` | 6.6 GB | 8.1 GB | 1024x1024 | Apache 2.0, any use | | `flux-2-klein-base` | 6.6 GB | 8.1 GB | 1024x1024 | Apache 2.0, any use | | `flux-2-klein-9b` | Needs 24 GB | 13.8 GB | 1024x1024 | Non-commercial only | @@ -180,8 +182,14 @@ most with short prompts and with text in the image. | FLUX.2 klein, all sizes and variants | Ships with gmlx, from Black Forest Labs | | Krea 2 | Ships with gmlx, from Krea | | Z-Image, Z-Image-Turbo | Ships with gmlx, from the Z-Image-Turbo demo of Tongyi-MAI | +| MOSS-Image, MOSS-Image Base | Ships with gmlx, written by gmlx from the prompt-enhancer skill of OpenMOSS | | Qwen-Image 2.1, Qwen-Image 2.1 Turbo | Your own file, as [Qwen-Image 2.1](#qwen-image-21-rewriter-prompt) describes | +OpenMOSS publishes its prompt enhancer as an agent skill, not a system +prompt, so gmlx's MOSS-Image prompt follows the skill's rules for a +caption and writes prose only. It writes no layout boxes, so run a layout +JSON prompt with `--no-upsample`. + Qwen3.5-9B writes captions close to those of BFL's own FLUX.3 upsampler, and its Q6_K GGUF is a good choice for every family. Any chat model of your configuration can write them. Add it under `models`, and name it as the @@ -287,6 +295,8 @@ gmlx node pack runs these image models in it, on the gmlx server. | `qwen-image` | Qwen-Image 2.1 | 40 | From up to 10 reference images | Qwen3-VL-8B-Instruct, with its mmproj | | `qwen-image-turbo` | Qwen-Image 2.1 Turbo | 8 | From up to 10 reference images | Qwen3-VL-8B-Instruct, with its mmproj | | `krea-2-turbo` | Krea 2 Turbo | 8 | From one image | Qwen3-VL-4B-Instruct | +| `moss-image` | MOSS-Image | 50 | From one image | Qwen3.5-9B | +| `moss-image-base` | MOSS-Image Base | 50 | From one image | Qwen3.5-9B | | `flux-2-klein` | FLUX.2 klein 4B | 4, with guidance 1 only | From up to 10 reference images | Qwen3-4B | | `flux-2-klein-base` | FLUX.2 klein Base 4B | 50 | From up to 10 reference images | Qwen3-4B | | `flux-2-klein-9b` | FLUX.2 klein 9B, non-commercial | 4, with guidance 1 only | From up to 10 reference images | Qwen3-8B | @@ -322,6 +332,11 @@ the later steps reuse them and run the image and the prompt alone. A takes on FLUX.2 klein 9B. Without reference images it runs as FLUX.2 klein 9B does. +MOSS-Image is OpenMOSS's fine-tuned model, and MOSS-Image Base the model +it is tuned from. Both run 50 steps at guidance 3, and leave block 22 out +of the negative pass, as OpenMOSS runs them. They read the prompt with +Qwen3.5-9B and edit from an init image, not from reference images. + `pull` takes the DiT at the quant for your Mac's memory, and the text encoder at Q8_0. A Mac takes the column of the largest class its memory holds. `--quant` and `--encoder-quant` pick others, and @@ -333,6 +348,8 @@ holds. `--quant` and `--encoder-quant` pick others, and | `qwen-image` | Q4_K_M | Q8_0 | Q8_0 | F16 | | `qwen-image-turbo` | AD-Q4_K | Q8_0 | BF16 | BF16 | | `krea-2-turbo` | Q3_K_M | Q6_K | Q8_0 | Q8_0 | +| `moss-image` | Q6_K | Q8_0 | Q8_0 | Q8_0 | +| `moss-image-base` | Q6_K | Q8_0 | Q8_0 | Q8_0 | | `flux-2-klein` | Q4_K_M | Q8_0 | Q8_0 | Q8_0 | | `flux-2-klein-base` | Q4_K_M | Q8_0 | Q8_0 | Q8_0 | | `flux-2-klein-9b` | None fits | Q4_K_M | Q8_0 | Q8_0 | @@ -346,6 +363,9 @@ default sets the entry's default size to 512x512, and a pull with another `--quant` sets none. Of Krea 2, gmlx runs Turbo, and it refuses a Krea 2 Raw GGUF by name. +On 16 GB a MOSS-Image pull takes the encoder at Q4_K_M, since the Q8_0 +encoder leaves too little memory for the DiT there. + No FLUX.2 klein 9B DiT that `pull` takes makes an image on 16 GB, so there a pull of a 9B model downloads nothing and names 24 GB, the smallest Mac that runs one. With `--quant` it takes that DiT, and @@ -364,11 +384,12 @@ references and its Q8_0 DiT none, and on 32 GB its Q8_0 DiT takes two. | Qwen-Image 2.1 | [Qwen Research License](https://huggingface.co/Qwen/Qwen-Image-2.1/blob/main/LICENSE) | Research and evaluation only. Commercial use needs a license from Qwen. | | Qwen-Image 2.1 Turbo | [Qwen Research License](https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo/blob/main/LICENSE) | Research and evaluation only. Commercial use needs a license from Qwen. | | Krea 2 Turbo | [Krea 2 Community License](https://cdn.jsdelivr.net/gh/krea-ai/krea-2@db3984fbc6e13b34c0064990fc2d95ac64d00058/assets/hf_samples/LICENSE.pdf) | Free commercial use for a company, affiliates counted, under $1M of revenue in the last twelve months. A deployment must filter what it makes. | +| MOSS-Image, MOSS-Image Base | [Apache 2.0](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT/blob/main/LICENSE) | Any, under the license's terms | | FLUX.2 klein 4B, FLUX.2 klein Base 4B | [Apache 2.0](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B/blob/main/LICENSE.md) | Any, under the license's terms | | FLUX.2 klein 9B, FLUX.2 klein Base 9B | [FLUX Non-Commercial License](https://huggingface.co/black-forest-labs/FLUX.2-klein-9B/blob/main/LICENSE.md) | Non-commercial use only, on the terms below | | FLUX.2 klein 9B KV | [FLUX Non-Commercial License](https://huggingface.co/black-forest-labs/FLUX.2-klein-9b-kv/blob/main/LICENSE) | Non-commercial use only, on the terms below | | FLUX.3 klein, FLUX.3 klein Base | [FLUX Kommunity License](https://huggingface.co/black-forest-labs/flux-3-image-klein/blob/main/LICENSE.md) | Non-commercial, and commercial use of the images under $5M of yearly revenue, on the terms below | -| Qwen3-4B, Qwen3-8B, Qwen3-VL-8B-Instruct, Qwen3-VL-4B-Instruct | Apache 2.0 | Any, under the license's terms | +| Qwen3-4B, Qwen3-8B, Qwen3-VL-8B-Instruct, Qwen3-VL-4B-Instruct, Qwen3.5-9B | Apache 2.0 | Any, under the license's terms | Under the [FLUX Kommunity License](https://huggingface.co/black-forest-labs/flux-3-image-klein/blob/main/LICENSE.md), a user whose yearly revenue, affiliates counted, is under $5M may use @@ -389,7 +410,9 @@ from Black Forest Labs. The system prompts gmlx ships for [prompt upsampling](#prompt-upsampling) are under Apache 2.0, from Black Forest Labs' flux-3-image-klein and flux2 repositories, Krea's krea-2 repository and Tongyi-MAI's Z-Image-Turbo -Space. Qwen's rewriter prompt, which +Space. The MOSS-Image caption prompt is gmlx's own, written from the +prompt-enhancer skill of OpenMOSS's MOSS-Image repository, which is under +Apache 2.0. Qwen's rewriter prompt, which `gmlx pull qwen-image` downloads, is under the Qwen Research License, for non-commercial use only. @@ -398,8 +421,10 @@ asks to filter adds its own. The GGUF files are quantized by unsloth, the Qwen-Image 2.1 Turbo files by AtomicChat, the Krea 2 Turbo files by Abiray, the FLUX.2 klein 9B KV files -by QuantStack and the FLUX.3 klein files by realrebelai, under the same -licenses. The Krea 2 VAE comes from Comfy-Org's Krea 2 repository. The +by QuantStack and the FLUX.3 klein and MOSS-Image files by realrebelai, +under the same licenses. The Krea 2 VAE comes from Comfy-Org's Krea 2 +repository, and the MOSS-Image VAE from OpenMOSS-Team's MOSS-Image-SFT +repository, under Apache 2.0. The FLUX.2 VAE, which FLUX.2 klein at both sizes and FLUX.3 klein decode with, comes from Black Forest Labs' FLUX.2 klein 4B repository, under Apache 2.0. The 9B repositories hold the same file. diff --git a/docs/internals/image-steps.md b/docs/internals/image-steps.md index 9fc19b1de..6a3acb564 100644 --- a/docs/internals/image-steps.md +++ b/docs/internals/image-steps.md @@ -120,6 +120,16 @@ the keys and values of that row. The header answers `hit` on those steps, and the end of the run drops the keys and values. A step with `ref_cache` and no run gets 400, as does one on a DiT that keeps no such cache. +A step with `skip`, a list of block indices from zero, on a server whose +features list `skip_blocks`, runs each of its rows without those blocks, +as ComfyUI's `SkipLayerGuidanceDiTSimple` runs the negative pass of +MOSS-Image. A client sends the negative rows in a step of their own with +that skip. A model's status gives the blocks its variant leaves out of +the negative pass under `sampling.neg_skip`, `[22]` on MOSS-Image and an +empty list on any other model. A skip on a DiT that skips no blocks, a +skip that names a block the DiT lacks, and a skip with `ref_cache` get +400. + A run that sits idle for 120 seconds ends within a few seconds after, and lets the model go. The server holds at most 32 runs, and a begin past that gets 429. An interrupt names its run and takes effect at the next step: @@ -160,6 +170,8 @@ steps in the bench with its sampler. Its steps take the `reference` convention, which embeds the time as `sample` does. The ComfyUI pack's steps take `comfy`, which embeds the time as ComfyUI's forward does, and for FLUX.2 klein that is the fp32 time. +Under both, a step's timestep is its sigma, except on MOSS-Image, where +it is the sigma times 1000, as in the OpenMOSS plugin's ComfyUI model. The record's transport is `steps`. It holds the median per step of the request's wall time less the forward time the server reports, and diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 24a01652a..dbf632162 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -262,8 +262,8 @@ default, or unload chat models with [`POST /unload`](api.md#endpoints). `gmlx pull` refuses an image model on an M1 or M2 Mac before it downloads anything, and `--force` downloads the files anyway. The models run in bfloat16, which M1 and M2 chips lack. On those Macs `gmlx image generate` -and the server refuse Qwen-Image 2.1, Krea 2, FLUX.2 klein and FLUX.3 -klein with `needs bfloat16, which the M1 chip lacks`. Z-Image-Turbo can +and the server refuse Qwen-Image 2.1, Krea 2, MOSS-Image, FLUX.2 klein +and FLUX.3 klein with `needs bfloat16, which the M1 chip lacks`. Z-Image-Turbo can run in float16, as the next entry describes. ### `runs in float16 on the M1 chip, which misses the reference checks` diff --git a/gmlx/comfyui_pack/README.md b/gmlx/comfyui_pack/README.md index 2867eee4c..9637a534c 100644 --- a/gmlx/comfyui_pack/README.md +++ b/gmlx/comfyui_pack/README.md @@ -10,6 +10,7 @@ Supported model families: - Z-Image (Z-Image-Turbo), text to image. - Qwen-Image 2.1 and Qwen-Image 2.1 Turbo, text to image and editing with up to 10 reference images. - Krea 2 Turbo, text to image. +- MOSS-Image and MOSS-Image Base, text to image. The pack runs them on Apple GPUs without the OpenMOSS ComfyUI plugin or the ComfyUI-MOSS-GGUF bridge. - FLUX.2 klein 4B and FLUX.2 klein Base 4B, text to image and editing with up to 10 reference images. - FLUX.2 klein 9B and FLUX.2 klein Base 9B, text to image and editing with up to 10 reference images. Their license allows non-commercial use only. - FLUX.2 klein 9B KV, text to image and editing with up to 10 reference images, where the stock `Flux KV Cache` node makes the steps after the first reuse the keys and values of the references. Its license allows non-commercial use only. @@ -41,12 +42,13 @@ At startup it checks the ComfyUI internals it uses. When one is missing, the nod 2. Add `gmlx Model Loader`, `gmlx CLIP Loader` and `gmlx VAE Loader`, and pick the model in `model_name` on each. The list holds the image models of the gmlx server, the default first, or only the default model's name when it has none. A model the server does not have stops the run with the `gmlx pull` command that adds it. 3. Connect them to stock nodes as you would any checkpoint: `CLIPTextEncode`, `KSampler`, `VAEDecode`, `LoraLoader`, `ModelSamplingAuraFlow` and so on. -The `example_workflows` folder has ten graphs, which ComfyUI's template browser lists under the pack: +The `example_workflows` folder has eleven graphs, which ComfyUI's template browser lists under the pack: - `z-image-turbo-t2i.json`: Z-Image-Turbo text to image. - `z-image-turbo-lora-t2i.json`: Z-Image-Turbo text to image through a stock `LoraLoader`. Put a Z-Image LoRA named `z-image-turbo-style.safetensors` in `models/loras`, or pick another one in the node. - `qwen-image-2.1-edit-two-references.json`: Qwen-Image 2.1 editing with two reference images through the stock `TextEncodeQwenImage21` node. Put `reference-1.png` and `reference-2.png` in ComfyUI's `input` folder, or pick other images. - `krea-2-turbo-t2i.json`: Krea 2 Turbo text to image through the stock `BasicScheduler`, `CFGGuider` and `SamplerCustomAdvanced`. It makes 1024x1024, and on a Mac with 16 GB set the latent to 512x512. +- `moss-image-t2i.json`: MOSS-Image text to image in 50 steps at cfg 3, through the stock `ModelSamplingSD3` at shift 5 and `SkipLayerGuidanceDiTSimple`, which leaves block 22 out of the negative pass, as the OpenMOSS workflow does. `KSampler` samples on ComfyUI's `sgm_uniform` schedule, not on the release's schedule that gmlx samples on. - `flux-2-klein-t2i.json`: FLUX.2 klein text to image through the stock `Flux2Scheduler`, `CFGGuider` and `SamplerCustomAdvanced`. `Flux2Scheduler` gives the schedule FLUX.2 klein samples on, which follows the image size, so keep its width and height those of the latent. - `flux-2-klein-edit-two-references.json`: FLUX.2 klein editing with two reference images, each encoded by `VAEEncode` and put on the prompt and its zeroed negative by the stock `ReferenceLatent`. The first reference sets the image size. Put `reference-1.png` and `reference-2.png` in ComfyUI's `input` folder, or pick other images. - `flux-2-klein-9b-t2i.json`: FLUX.2 klein 9B text to image, the same graph on the 9B model. It needs a Mac with 24 GB or more. @@ -98,9 +100,11 @@ Anything that would run inside the DiT forward in ComfyUI is refused with an err - attention overrides, RoPE options (`ScaleROPE`), weight wrappers and diffusion model wrappers; - ControlNet; - hook keyframes and scheduled CLIP hooks; -- a quantized Qwen-Image 2.1 prefix cache (`QwenImage21Cache` with a dtype other than default). +- a quantized Qwen-Image 2.1 prefix cache (`QwenImage21Cache` with a dtype other than default); +- a MOSS-Image conditioning that masks some of its text rows, such as one from a text encoder other than `gmlx CLIP Loader`. The stock `Flux KV Cache` node on FLUX.2 klein 9B KV is not refused: the server keeps the keys and values of the references of each row of the batch from the first step of a run, the later steps of the run reuse those of each row, and the end of the run drops them. +The stock `SkipLayerGuidanceDiTSimple` node on MOSS-Image is not refused when it names `single_layers` alone, with `start_percent` 0 and `end_percent` 1: the server leaves those blocks out of the negative pass of each step. It is refused with `double_layers`, another window, a linked widget or a block MOSS-Image lacks, and on any other model. A refusal from a node's own setup is reported on that node. A refusal during sampling is reported on the sampler and names the node upstream that caused it. Saving a gmlx MODEL, CLIP or VAE with the checkpoint save nodes is refused, since ComfyUI holds no weights for them. @@ -121,7 +125,7 @@ That runs the wire, client and in-process ComfyUI suites on the cpu device, then ### Parity with the stock loaders -Seven checks hold a gmlx MODEL against ComfyUI's own model on the same graph, noise, sigmas and conditioning: ComfyUI-GGUF for Z-Image Turbo, Krea 2 Turbo, FLUX.2 klein 4B, FLUX.2 klein 9B, FLUX.2 klein 9B KV and FLUX.3 klein, and ComfyUI's Qwen-Image 2.1 model on a GGUF with its arch key set. They need the model files and both loaders, so they are not part of `run_tests.py`. +Eight checks hold a gmlx MODEL against ComfyUI's own model on the same graph, noise, sigmas and conditioning: ComfyUI-GGUF for Z-Image Turbo, Krea 2 Turbo, FLUX.2 klein 4B, FLUX.2 klein 9B, FLUX.2 klein 9B KV and FLUX.3 klein, ComfyUI's Qwen-Image 2.1 model on a GGUF with its arch key set, and for MOSS-Image the OpenMOSS ComfyUI plugin with the ComfyUI-MOSS-GGUF bridge on ComfyUI's cpu device. They need the model files and both loaders, so they are not part of `run_tests.py`. The step check is the gate. The stock model in bf16 samples and records, at each step, the latent it is fed and the denoised output it returns. The gmlx MODEL and the stock model in fp32 are then fed the same latents and sigmas. At every step the relative L2 of the gmlx output from the stock bf16 output must be at most 1.5 times that of the stock fp32 output, or at most 2e-2 where that is larger. A fixed 2e-2 cannot hold on its own: the stock fp32 output is 0.11 from the stock bf16 output at the first step of Z-Image Turbo. @@ -140,7 +144,14 @@ The step check is the gate. The stock model in bf16 samples and records, at each | FLUX.3 klein | 512x512, 16 steps, cfg 1, stock LoRA loader | 0.029 | 0.081 | 1.24 | pass | | FLUX.3 klein | 512x512, 16 steps, cfg 1, one reference | 0.017 | 0.018 | 1.26 | pass | -On Z-Image Turbo the gmlx output is closer to the stock fp32 output than the stock bf16 output is, at every step. The Qwen-Image 2.1, Krea 2, FLUX.2 klein and FLUX.3 klein checks run ComfyUI with `--use-pytorch-cross-attention`. Without it, ComfyUI's attention on the Mac adds the text mask as a bias, so the stock model departs from the reference implementation at high sigma, and the ratio at the first step is about 3. +On Z-Image Turbo the gmlx output is closer to the stock fp32 output than the stock bf16 output is, at every step. The Qwen-Image 2.1, Krea 2, FLUX.2 klein, FLUX.3 klein and MOSS-Image checks run ComfyUI with `--use-pytorch-cross-attention`. Without it, ComfyUI's attention on the Mac adds the text mask as a bias, so the stock model departs from the reference implementation at high sigma, and the ratio at the first step is about 3. + +The MOSS-Image check runs at 256x256 with the stock model on ComfyUI's cpu device, so its stock bf16 output comes from CPU torch arithmetic, a bf16 path the pack never takes. Its rows, in the table below, read from the stock fp32 output instead: the gmlx distance from it must be at most 1.5 times the larger of the stock bf16 distance and the distance of gmlx sampling in process, without ComfyUI, on the same latents and sigmas, or at most 2e-2. The gmlx MODEL equals the in-process run bit for bit at every step, and the in-process distance is the larger at every step, at the median 1.55 times the stock bf16 one in text to image and 1.58 times with the LoRA. With block 22 kept in the negative pass, the worst step reads 3.7 and 3.6 times its gate. + +| Model | Graph | gmlx from stock fp32, worst step | gmlx in process, same step | stock bf16, same step | Worst ratio to in process | Worst ratio to stock bf16 | Result | +|---|---|---|---|---|---|---|---| +| MOSS-Image | 256x256, 50 steps, cfg 3, block 22 out of the negative pass | 0.099 | 0.099 | 0.037 | 1.00 | 2.69 | pass | +| MOSS-Image | 256x256, 50 steps, cfg 3, block 22 out of the negative pass, stock LoRA loader | 0.064 | 0.064 | 0.036 | 1.00 | 2.16 | pass | The final image check is reported beside it: the gmlx image must be at least as close to the stock fp32 image as the stock bf16 image is, within 6 dB. @@ -158,6 +169,8 @@ The final image check is reported beside it: the gmlx image must be at least as | FLUX.3 klein, 512x512, 16 steps | 33.50 dB | 32.96 dB | pass | | FLUX.3 klein, 512x512, 16 steps, stock LoRA loader | 32.57 dB | 32.29 dB | pass | | FLUX.3 klein, 512x512, 16 steps, one reference | 50.41 dB | 49.28 dB | pass | +| MOSS-Image, 256x256, 50 steps, cfg 3 | 29.75 dB | 27.44 dB | pass | +| MOSS-Image, 256x256, 50 steps, cfg 3, stock LoRA loader | 36.86 dB | 35.15 dB | pass | A fault planted in the Z-Image step convention shows why the step check is the gate. With the velocity sign dropped, the worst step ratio is 188 and the final image reads 8.35 dB. With the timestep read 3% low, the worst step ratio is 7.1, while the final image reads 13.49 dB against a bound of 14.30 dB and fails by only 0.8 dB, so a smaller timestep fault would pass the final image check. diff --git a/gmlx/comfyui_pack/__init__.py b/gmlx/comfyui_pack/__init__.py index a10480151..524fcaca7 100644 --- a/gmlx/comfyui_pack/__init__.py +++ b/gmlx/comfyui_pack/__init__.py @@ -14,10 +14,7 @@ compat.check() - from .model import register_prompt_handler from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS - - register_prompt_handler() except Exception as error: # noqa: BLE001: any failure gets the stand-in nodes if importlib.util.find_spec("comfy") is None: raise diff --git a/gmlx/comfyui_pack/clip.py b/gmlx/comfyui_pack/clip.py index cd13321e3..00f06a396 100644 --- a/gmlx/comfyui_pack/clip.py +++ b/gmlx/comfyui_pack/clip.py @@ -28,8 +28,9 @@ # The families whose stock ComfyUI tokenizer encodes an empty prompt as a # space under prevent_empty_text: QwenImageTokenizer, and Qwen3VLTokenizer # under Qwen-Image 2.1, Krea 2 and FLUX.3. ZImageTokenizer ignores the -# option, and FLUX.2's KleinTokenizer applies its chat template first, so -# the prompt it encodes is never empty. +# option, and FLUX.2's KleinTokenizer and the MOSSImageTokenizer of the +# OpenMOSS plugin apply their chat template first, so the prompt they +# encode is never empty. _SPACE_FOR_EMPTY = frozenset({"qwen_image", "qwen_image21", "krea2", "flux3_klein"}) diff --git a/gmlx/comfyui_pack/compat.py b/gmlx/comfyui_pack/compat.py index 861d5723d..299fe28f3 100644 --- a/gmlx/comfyui_pack/compat.py +++ b/gmlx/comfyui_pack/compat.py @@ -45,11 +45,16 @@ ("comfy.patcher_extension", "WrappersMP.DIFFUSION_MODEL"), ("comfy.patcher_extension", "get_all_wrappers"), ("comfy.model_base", "BaseModel._apply_model"), + ("comfy.model_base", "BaseModel.extra_conds"), + ("comfy.model_base", "ModelType.FLOW"), + ("comfy.conds", "CONDRegular"), + ("comfy.latent_formats", "Flux2"), ("comfy.model_detection", "model_config_from_unet"), ("comfy.model_detection", "unet_prefix_from_state_dict"), ("comfy.supported_models", "ZImage"), ("comfy.supported_models", "QwenImage21"), ("comfy.supported_models_base", "BASE.set_inference_dtype"), + ("comfy.supported_models_base", "BASE.get_model"), ("comfy.model_management", "load_models_gpu"), ("comfy.model_management", "processing_interrupted"), ("comfy.model_management", "throw_exception_if_processing_interrupted"), @@ -62,6 +67,7 @@ ("comfy.utils", "state_dict_prefix_replace"), ("comfy.utils", "set_attr"), ("comfy.weight_adapter", "LoRAAdapter"), + ("execution", "PromptQueue.get_current_queue_volatile"), ) # Arguments the pack passes by name to a ComfyUI signature. diff --git a/gmlx/comfyui_pack/discovery.py b/gmlx/comfyui_pack/discovery.py index 52abe843d..6a9e5ab99 100644 --- a/gmlx/comfyui_pack/discovery.py +++ b/gmlx/comfyui_pack/discovery.py @@ -44,6 +44,8 @@ "qwen-image": ("qwen-image-2.1",), "qwen-image-turbo": ("qwen-image-2.1-turbo",), "krea-2-turbo": (), + "moss-image": ("moss-image-sft",), + "moss-image-base": (), "flux-2-klein": (), "flux-2-klein-base": (), "flux-2-klein-9b": (), @@ -57,6 +59,8 @@ "qwen-image": ("qwen_image21", "base"), "qwen-image-turbo": ("qwen_image21", "turbo"), "krea-2-turbo": ("krea2", "turbo"), + "moss-image": ("moss_image", "sft"), + "moss-image-base": ("moss_image", "base"), "flux-2-klein": ("flux2_klein", "distilled"), "flux-2-klein-base": ("flux2_klein", "base"), "flux-2-klein-9b": ("flux2_klein_9b", "distilled"), diff --git a/gmlx/comfyui_pack/example_workflows/moss-image-t2i.json b/gmlx/comfyui_pack/example_workflows/moss-image-t2i.json new file mode 100644 index 000000000..3fb340aed --- /dev/null +++ b/gmlx/comfyui_pack/example_workflows/moss-image-t2i.json @@ -0,0 +1,556 @@ +{ + "id": "940a04a3-b546-59f1-bf45-a075d12783dc", + "revision": 0, + "last_node_id": 12, + "last_link_id": 11, + "nodes": [ + { + "id": 12, + "type": "Note", + "pos": [ + -440, + 80 + ], + "size": [ + 400, + 480 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "gmlx nodes: install", + "properties": {}, + "widgets_values": [ + "These gmlx nodes run the image model on a gmlx server on your Mac (Apple Silicon, M3 or newer). MOSS-Image samples 50 steps at cfg 3, two passes a step, so an image takes minutes. Its license is Apache-2.0.\n\nThe schedule is ComfyUI's: sgm_uniform after ModelSamplingSD3 at shift 5, with block 22 left out of the negative pass, as the OpenMOSS workflow has it. gmlx image and gmlx serve sample on the release's own schedule, so their images differ a little from this one.\n\n1. Install gmlx: https://asher.github.io/gmlx/installation.html\n2. Put the nodes into ComfyUI, then restart ComfyUI:\n\n gmlx comfyui install\n\n3. Download the model:\n\n gmlx pull moss-image\n\nGuide: https://asher.github.io/gmlx/comfyui.html" + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 1, + "type": "GmlxModelLoader", + "pos": [ + 0, + 80 + ], + "size": [ + 300, + 130 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 3 + ] + } + ], + "title": "gmlx Model Loader", + "properties": { + "Node name for S&R": "GmlxModelLoader" + }, + "widgets_values": [ + "moss-image", + "http://127.0.0.1:8080", + true + ] + }, + { + "id": 2, + "type": "GmlxCLIPLoader", + "pos": [ + 0, + 250 + ], + "size": [ + 300, + 130 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 1, + 2 + ] + } + ], + "title": "gmlx CLIP Loader", + "properties": { + "Node name for S&R": "GmlxCLIPLoader" + }, + "widgets_values": [ + "moss-image", + "http://127.0.0.1:8080", + true + ] + }, + { + "id": 3, + "type": "GmlxVAELoader", + "pos": [ + 0, + 420 + ], + "size": [ + 300, + 130 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 10 + ] + } + ], + "title": "gmlx VAE Loader", + "properties": { + "Node name for S&R": "GmlxVAELoader" + }, + "widgets_values": [ + "moss-image", + "http://127.0.0.1:8080", + true + ] + }, + { + "id": 4, + "type": "CLIPTextEncode", + "pos": [ + 340, + 80 + ], + "size": [ + 400, + 160 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 1 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 6 + ] + } + ], + "title": "Positive prompt", + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "A wide view of a refrigerator beside a kitchen counter with a fruit bowl. A small sticky note on the refrigerator door is handwritten: \"Remember to buy milk and eggs.\"" + ] + }, + { + "id": 5, + "type": "CLIPTextEncode", + "pos": [ + 340, + 280 + ], + "size": [ + 400, + 100 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 2 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 7 + ] + } + ], + "title": "Negative prompt", + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "" + ] + }, + { + "id": 6, + "type": "EmptyFlux2LatentImage", + "pos": [ + 340, + 420 + ], + "size": [ + 300, + 110 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 8 + ] + } + ], + "title": "Empty latent", + "properties": { + "Node name for S&R": "EmptyFlux2LatentImage" + }, + "widgets_values": [ + 1024, + 1024, + 1 + ] + }, + { + "id": 7, + "type": "ModelSamplingSD3", + "pos": [ + 340, + 570 + ], + "size": [ + 300, + 60 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 3 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 4 + ] + } + ], + "title": "Shift", + "properties": { + "Node name for S&R": "ModelSamplingSD3" + }, + "widgets_values": [ + 5.0 + ] + }, + { + "id": 8, + "type": "SkipLayerGuidanceDiTSimple", + "pos": [ + 340, + 670 + ], + "size": [ + 300, + 150 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 4 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 5 + ] + } + ], + "title": "Skip block 22 on the negative pass", + "properties": { + "Node name for S&R": "SkipLayerGuidanceDiTSimple" + }, + "widgets_values": [ + "", + "22", + 0.0, + 1.0 + ] + }, + { + "id": 9, + "type": "KSampler", + "pos": [ + 780, + 80 + ], + "size": [ + 300, + 270 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 5 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 6 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 7 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 8 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 9 + ] + } + ], + "title": "KSampler", + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 42, + "fixed", + 50, + 3.0, + "euler", + "sgm_uniform", + 1.0 + ] + }, + { + "id": 10, + "type": "VAEDecode", + "pos": [ + 1120, + 80 + ], + "size": [ + 210, + 50 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 9 + }, + { + "name": "vae", + "type": "VAE", + "link": 10 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 11 + ] + } + ], + "title": "VAE Decode", + "properties": { + "Node name for S&R": "VAEDecode" + }, + "widgets_values": [] + }, + { + "id": 11, + "type": "SaveImage", + "pos": [ + 1120, + 170 + ], + "size": [ + 420, + 470 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 11 + } + ], + "outputs": [], + "title": "Save Image", + "properties": { + "Node name for S&R": "SaveImage" + }, + "widgets_values": [ + "gmlx/moss-image" + ] + } + ], + "links": [ + [ + 1, + 2, + 0, + 4, + 0, + "CLIP" + ], + [ + 2, + 2, + 0, + 5, + 0, + "CLIP" + ], + [ + 3, + 1, + 0, + 7, + 0, + "MODEL" + ], + [ + 4, + 7, + 0, + 8, + 0, + "MODEL" + ], + [ + 5, + 8, + 0, + 9, + 0, + "MODEL" + ], + [ + 6, + 4, + 0, + 9, + 1, + "CONDITIONING" + ], + [ + 7, + 5, + 0, + 9, + 2, + "CONDITIONING" + ], + [ + 8, + 6, + 0, + 9, + 3, + "LATENT" + ], + [ + 9, + 9, + 0, + 10, + 0, + "LATENT" + ], + [ + 10, + 3, + 0, + 10, + 1, + "VAE" + ], + [ + 11, + 10, + 0, + 11, + 0, + "IMAGE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.75, + "offset": [ + 480, + 40 + ] + } + }, + "version": 0.4 +} diff --git a/gmlx/comfyui_pack/model.py b/gmlx/comfyui_pack/model.py index 241f303f9..26336753d 100644 --- a/gmlx/comfyui_pack/model.py +++ b/gmlx/comfyui_pack/model.py @@ -19,6 +19,7 @@ import collections import functools import logging +import re import time import numpy as np @@ -34,6 +35,7 @@ set_interrupt_check, tensor_hash) from .clip import TYPED, TypedPrompt from .compat import comfyui_version +from .moss_image import PACK_MODELS from .patcher import CPU, KV_METHOD, KV_PATCHES, GmlxBinding, GmlxModelPatcher log = logging.getLogger("gmlx") @@ -61,8 +63,13 @@ "TomePatchModel", "FreeU", "FreeU_V2", "ScaleROPE", "FluxKontextMultiReferenceLatentMethod"} -LAST_PROMPT: dict = {} STRICT_KEYS = False +# The node whose skips of single blocks on the negative pass a step takes +# as its skip set, on the families whose DiT skips blocks. +SLG_NODE = "SkipLayerGuidanceDiTSimple" +SKIP_FAMILIES = frozenset({"moss_image"}) +# The block replacement that node installs, which returns its input. +SLG_SKIP = f"{SLG_NODE}.execute..skip" def manifest_state_dict(manifest: dict, inline: bool = True) -> dict: @@ -81,7 +88,11 @@ def manifest_state_dict(manifest: dict, inline: bool = True) -> dict: def build_weightless(manifest: dict, model_class: str): """``(model, model_config)`` built on the meta device from a manifest, - after checking that ComfyUI detects ``model_class``.""" + after checking that ComfyUI detects ``model_class``, or by the pack's + own class of that name.""" + if model_class in PACK_MODELS: + cfg = PACK_MODELS[model_class](manifest) + return _built(cfg, {}) if not hasattr(comfy.supported_models, model_class): raise GmlxError(f"gmlx: ComfyUI {comfyui_version()} has no {model_class} model, which " f"the server serves {manifest.get('model')!r} as; update ComfyUI") @@ -103,6 +114,10 @@ def build_weightless(manifest: dict, model_class: str): raise GmlxRefusal(f"gmlx: the server's {model_class} DiT holds the fused qkv layout " "(single_blocks.0.qkv.weight), which a gmlx model does not take; " "serve a file that stores q_proj, k_proj and v_proj apart") + return _built(cfg, sd) + + +def _built(cfg, sd): dtypes = list(cfg.supported_inference_dtypes) dtype = torch.float32 if torch.float32 in dtypes else dtypes[0] cfg.set_inference_dtype(dtype, None, device=CPU) @@ -173,38 +188,54 @@ def walk_back(prompt: dict, node_id: str, classes: set) -> list: return sorted(hits, key=lambda s: (len(s), s)) -def register_prompt_handler() -> bool: - """Keep the last submitted graph, so a forward-time refusal can name the - node upstream of the sampler that caused it.""" +def _running() -> tuple: + """``(graph, node id)`` of the prompt that runs: the queue's running + item of the server's last prompt id, else its only running item, and + the node the server says runs. ``({}, None)`` when none is known. A + prompt queued behind it is not read: ComfyUI hands the on_prompt + handlers each prompt when it is queued, not when it runs.""" try: import server - inst = getattr(server.PromptServer, "instance", None) - if inst is None: - return False - - def on_prompt(json_data): - LAST_PROMPT.clear() - LAST_PROMPT.update(json_data if isinstance(json_data, dict) else {}) - return json_data - - inst.add_on_prompt_handler(on_prompt) - return True + inst = server.PromptServer.instance + running, _ = inst.prompt_queue.get_current_queue_volatile() + pid = getattr(inst, "last_prompt_id", None) + sid = getattr(inst, "last_node_id", None) except Exception: # noqa: BLE001: best effort; the refusal still names the cause - return False + return {}, None + items = [it for it in running if it[1] == pid] or (running if len(running) == 1 else []) + graph = items[0][2] if items else None + if not isinstance(graph, dict) or sid is None: + return {}, None + return graph, sid def _upstream_names() -> list: - try: - import server - sid = getattr(server.PromptServer.instance, "last_node_id", None) - except Exception: # noqa: BLE001 - return [] - prompt = LAST_PROMPT.get("prompt") or {} - if sid is None or not prompt: + prompt, sid = _running() + if not prompt: return [] return [f"{n} ({prompt[n]['class_type']})" for n in walk_back(prompt, sid, IN_FORWARD_CLASSES)] +def _layers(text) -> list: + """The block numbers of a layers widget: every run of digits.""" + return [int(v) for v in re.findall(r"\d+", str(text))] + + +def _slg_inputs() -> tuple: + """``(node id, inputs)`` of the SkipLayerGuidanceDiTSimple node upstream + of the running sampler. Refuses none and two that differ.""" + prompt, sid = _running() + hits = walk_back(prompt, sid, {SLG_NODE}) if prompt else [] + if not hits: + _refuse("transformer_options['patches_replace'] on single blocks, from a node other " + f"than {SLG_NODE},") + inputs = [prompt[n].get("inputs") or {} for n in hits] + if any(i != inputs[0] for i in inputs[1:]): + raise GmlxRefusal(f"gmlx: the {SLG_NODE} nodes {', '.join(hits)} skip other blocks; " + "a gmlx model takes one skip set a step") + return hits[0], inputs[0] + + def _refuse(what: str): hits = _upstream_names() where = f"; upstream node(s): {', '.join(hits)}" if hits else "" @@ -273,6 +304,57 @@ def _check(self, control, to: dict, kwargs: dict, kv: bool = False): "gmlx model; use dtype default") return cache.get("device") != "off" + def _skip(self, to: dict) -> tuple: + """The blocks the rows leave out, from the single-block skips that + SkipLayerGuidanceDiTSimple puts on the negative pass, checked + against the node's widgets. Empty when the call carries no block + replacements, or on a family whose DiT skips none, which the check + refuses.""" + pr = to.get("patches_replace") + if not pr or self.b.family not in SKIP_FAMILIES: + return () + dit = (pr.get("dit") or {}) if isinstance(pr, dict) else {} + if set(pr) - {"dit"} or not dit or any( + not (isinstance(k, tuple) and len(k) == 2 + and k[0] in ("single_block", "double_block")) + or getattr(f, "__qualname__", None) != SLG_SKIP for k, f in dit.items()): + _refuse("a transformer_options['patches_replace'] entry other than the " + f"block skips of {SLG_NODE}") + blocks = sorted(int(k[1]) for k in dit if k[0] == "single_block") + node, inputs = _slg_inputs() + where = f"the {SLG_NODE} node {node}" + widgets = ("double_layers", "single_layers", "start_percent", "end_percent") + linked = [w for w in widgets if isinstance(inputs.get(w), list)] + if linked: + raise GmlxRefusal(f"gmlx: {where} takes {', '.join(linked)} from another node; " + "a gmlx model reads them as widgets") + double = sorted({*_layers(inputs.get("double_layers", "")), + *(int(k[1]) for k in dit if k[0] == "double_block")}) + if double: + raise GmlxRefusal(f"gmlx: {where} skips double_layers {double}; the MOSS-Image " + "DiT has single blocks only, so leave double_layers empty") + start = float(inputs.get("start_percent", 0.0)) + end = float(inputs.get("end_percent", 1.0)) + if (start, end) != (0.0, 1.0): + raise GmlxRefusal(f"gmlx: {where} skips its blocks from {start} to {end} of the " + "run; a gmlx model skips them over the whole run, at " + "start_percent 0 and end_percent 1") + if sorted(set(_layers(inputs.get("single_layers", "")))) != blocks: + raise GmlxRefusal(f"gmlx: this call skips single blocks {blocks}, which are not " + f"the single_layers of {where}") + if any(c != 1 for c in to.get("cond_or_uncond") or [0]): + raise GmlxRefusal(f"gmlx: {where} skips blocks of a positive row; a gmlx model " + "skips them on the negative pass alone") + caps = getattr(self.b.client, "capabilities", None) or {} + if "skip_blocks" not in (caps.get("features") or ()): + raise GmlxRefusal("gmlx: this gmlx server skips no blocks in a step; update it") + n = int(getattr(self.base_model.diffusion_model, "blocks", 0) or 0) + bad = [b for b in blocks if not 0 <= b < n] + if bad: + raise GmlxRefusal(f"gmlx: {where} skips blocks {bad}, and the DiT has blocks 0 to " + f"{n - 1}") + return tuple(blocks) + def _hash(self, key, t: torch.Tensor, held: dict) -> str: h = self.memo.get(key) if key is not None else None if h is None: @@ -352,7 +434,9 @@ def __call__(self, x, timesteps, context=None, control=None, transformer_options kv = getattr(patcher, "gmlx_ref_cache", frozenset()) if kv and kv != KV_PATCHES: _refuse("part of the Flux KV Cache node's patches") - prefix_cache = self._check(control, to, kwargs, bool(kv)) + skip = self._skip(to) + checked = {k: v for k, v in to.items() if k != "patches_replace"} if skip else to + prefix_cache = self._check(control, checked, kwargs, bool(kv)) sigmas = self._schedule(to) # A latent format with a frame axis, such as Krea 2's, gives one frame. framed = x.ndim == 5 and x.shape[2] == 1 @@ -373,6 +457,8 @@ def __call__(self, x, timesteps, context=None, control=None, transformer_options if kv and kwargs.get("ref_latents"): # The server keeps the references' K and V of the run's first step. fields["ref_cache"] = True + if skip: + fields["skip"] = list(skip) xs = x.float() ts = timesteps.float().reshape(-1) t1 = time.perf_counter() diff --git a/gmlx/comfyui_pack/moss_image.py b/gmlx/comfyui_pack/moss_image.py new file mode 100644 index 000000000..b60add45f --- /dev/null +++ b/gmlx/comfyui_pack/moss_image.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MOSS-Image as a ComfyUI model, which ComfyUI has no class of its own for. + +The pack registers ``MOSSImage`` with the settings of the OpenMOSS plugin: +a flow model on the FLUX.2 latent format, sampling shift 1.0 and a +timestep of sigma times 1000. Its DiT holds the manifest's tensors on the +meta device under their own names, so ComfyUI's generic LoRA key map +reads them, as it reads a model of the plugin. A conditioning keeps its +own length: rows of other lengths never join one call. +""" + +from __future__ import annotations + +import torch + +import comfy.conds +import comfy.latent_formats +import comfy.model_base +import comfy.supported_models_base + +from .client import GmlxError, GmlxRefusal + +# Tensors the release model and the GGUF both hold, checked before a build. +REQUIRED = ("video_embedder.weight", "text_embedder.proj.weight", + "shared_block_modulation.weight", "final_linear_video.weight", + "layers.0.proj.ffn_up_gate.weight", "text_embedder.blocks.0.qkv.weight") + + +class WeightlessDiT(torch.nn.Module): + """The DiT's tensors on the meta device under the manifest's names. + The pack replaces its forward with the step call.""" + + def __init__(self, tensors=(), dtype=None, device=None, operations=None, **_): + super().__init__() + self.dtype = dtype or torch.float32 + for name, shape in tensors: + *path, leaf = name.split(".") + mod = self + for part in path: + child = mod._modules.get(part) + if child is None: + child = torch.nn.Module() + mod.add_module(part, child) + mod = child + mod.register_parameter(leaf, torch.nn.Parameter( + torch.empty(tuple(shape), dtype=self.dtype, device="meta"), + requires_grad=False)) + self.blocks = sum(1 for n, _ in tensors + if n.startswith("layers.") and n.endswith(".proj.q.weight")) + + def forward(self, *args, **kwargs): + raise GmlxError("gmlx: the MOSS-Image DiT runs on the server") + + +class MOSSImage(comfy.model_base.BaseModel): + def __init__(self, model_config, device=None): + super().__init__(model_config, comfy.model_base.ModelType.FLOW, device=device, + unet_model=WeightlessDiT) + + def extra_conds(self, **kwargs): + mask = kwargs.get("attention_mask") + if mask is not None and not bool(torch.all(mask != 0)): + raise GmlxRefusal("gmlx: this conditioning masks some of its text rows, which a " + "gmlx MOSS-Image model reads all of; encode it with the gmlx CLIP") + out = super().extra_conds(**kwargs) + cross_attn = kwargs.get("cross_attn") + if cross_attn is not None: + out["c_crossattn"] = comfy.conds.CONDRegular(cross_attn) + return out + + +class MOSSImageConfig(comfy.supported_models_base.BASE): + sampling_settings = {"shift": 1.0, "multiplier": 1000} + latent_format = comfy.latent_formats.Flux2 + supported_inference_dtypes = [torch.bfloat16, torch.float32] + memory_usage_factor = 4.0 + + def get_model(self, state_dict, prefix="", device=None): + return MOSSImage(self, device=device) + + +def model_config(manifest: dict) -> MOSSImageConfig: + """The config of the weightless model of a MOSS-Image manifest.""" + names = {t["name"] for t in manifest["tensors"]} + missing = [n for n in REQUIRED if n not in names] + if missing: + raise GmlxError(f"gmlx: the server's manifest of {manifest.get('model')!r} is not " + f"a MOSS-Image DiT: it has no {', '.join(missing)}") + return MOSSImageConfig({"tensors": tuple((t["name"], tuple(t["shape"])) + for t in manifest["tensors"])}) + + +# The models the pack builds itself, by the class name the server gives. +PACK_MODELS = {"MOSSImage": model_config} diff --git a/gmlx/commands/image.py b/gmlx/commands/image.py index ee19f927b..79cc7c6c2 100644 --- a/gmlx/commands/image.py +++ b/gmlx/commands/image.py @@ -1141,6 +1141,7 @@ def bench_run_steps(pipe, client, model: str, prompt: str, *, height: int, from gmlx.image import api from gmlx.image.samplers import make_sampler, noise_from_seed + from gmlx.serve.image_steps import CONVENTIONS from gmlx.serve.image_steps_loopback import work def start(): @@ -1195,6 +1196,11 @@ def combine(pos, neg, g): return work(lambda: np.array(api.guide(mx.array(pos), mx.array(neg), g, act))) + # The negative pass leaves out the variant's blocks, as sample does, + # and each timestep is in the reference convention of the family. + row = api.sampler_defaults(pipe.family.name, pipe.variant) + neg_skip = row.neg_skip if row is not None and cfg != 1.0 else () + conv = CONVENTIONS["reference"][pipe.family.name] marks = [t_enc] first = len(client.records) run = client.begin(model) @@ -1202,7 +1208,8 @@ def combine(pos, neg, g): x = client.sample(model, noise, sigmas, rows[0], cfg=cfg, neg_rows=rows[1] if len(rows) > 1 else None, run=run, update=update, combine=combine, - on_step=lambda i, _x: marks.append(clock())) + on_step=lambda i, _x: marks.append(clock()), + neg_skip=neg_skip, timestep_scale=1.0 / conv.scale) finally: client.end(run) t_dec = clock() diff --git a/gmlx/commands/image_pull.py b/gmlx/commands/image_pull.py index 61c0ee6be..90ced557f 100644 --- a/gmlx/commands/image_pull.py +++ b/gmlx/commands/image_pull.py @@ -805,9 +805,28 @@ def pull_recipe(a, r: catalog.Recipe) -> int: encoder_quant_given=bool(a.encoder_quant), only=only, roots=roots, base_dir=base_dir, nest=nest, cfg=cfg, ram=ram) + plan.size = catalog.default_size(r, ram, plan.dit_quant) + enc = plan.part("encoder") + if (enc is not None and not enc.download and not a.encoder_quant + and plan.encoder_quant != enc_q and ram + and catalog.makes_image(r, cls)): + # Another quant on disk serves, unless no image fits with it and + # one fits with the class's own. + image = _image_plan(plan, ram) + if image is not None and not image[1].fits: + own = plan_pull(r, dit_quant=dit_q, encoder_quant=enc_q, + encoder_quant_given=True, only=only, + roots=roots, base_dir=base_dir, nest=nest, + cfg=cfg, ram=ram) + own.size = plan.size + fits = _image_plan(own, ram) + if fits is not None and fits[1].fits: + own.part("encoder").note = ( + f"{enc.name} on disk does not fit an image on a {cls} " + f"GB Mac; the pull takes {enc_q}") + plan = own plan.default = (plan.dit_quant, plan.encoder_quant) == \ catalog.default_quants(r, ram) and catalog.makes_image(r, cls) - plan.size = catalog.default_size(r, ram, plan.dit_quant) need = _need(plan) free = manage._disk_free(base_dir) if a.json: diff --git a/gmlx/commands/manage.py b/gmlx/commands/manage.py index d8ff39c8e..9eb441c3a 100644 --- a/gmlx/commands/manage.py +++ b/gmlx/commands/manage.py @@ -279,13 +279,14 @@ def _image_refusal(report: remote.HeaderReport, return word if labeled else None -def _image_companion_bytes(family: str | None) -> int | None: - """Bytes of the default encoder (with its vision tower) and VAE of an - image family, for the combined figure.""" - from gmlx.image.catalog import RECIPES +def _image_companion_bytes(family: str | None, + ram: int | None) -> int | None: + """Bytes of the encoder (with its vision tower) and VAE a pull of an + image family takes on this Mac, for the combined figure.""" + from gmlx.image.catalog import RECIPES, companion_bytes for r in RECIPES.values(): if r.family == family: - return r.encoder_bytes + r.vae_bytes + return companion_bytes(r, ram) return None @@ -315,9 +316,11 @@ def _verdict(ref: remote.Ref, report: remote.HeaderReport, *, from gmlx.image.dtypes import chip_refusal from gmlx.image.registry import get_family chip_error = chip_refusal(image_family, get_family(image_family).title) + from gmlx.load.memfit import classify_fit, total_ram_bytes + ram = total_ram_bytes() combined = None if image == "dit" and size_bytes: - extra = _image_companion_bytes(image_family) + extra = _image_companion_bytes(image_family, ram) # The load gate counts the 16-bit tensors the DiT casts too. cast = 0 if scans and image_family: @@ -335,8 +338,6 @@ def _verdict(ref: remote.Ref, report: remote.HeaderReport, *, from gmlx.load.preflight import hadamard_fold_error fold_err = hadamard_fold_error(report.arch, report.hadamard) runs = report.loadable_codecs and arch_ok and fold_err is None - from gmlx.load.memfit import classify_fit, total_ram_bytes - ram = total_ram_bytes() return { "ref": ref.raw, "kind": ref.kind, @@ -1107,9 +1108,10 @@ def _planned_download_bytes(ref: remote.Ref, dest_dir: str, An image model name (z-image-turbo, qwen-image, qwen-image-turbo, krea-2-turbo, flux-2-klein, flux-2-klein-base, flux-2-klein-9b, -flux-2-klein-base-9b, flux-2-klein-9b-kv, flux-3-klein, flux-3-klein-base) -pulls the DiT, its text encoder and its VAE, uses the parts your model -folders already hold, and adds the model under image_models: in the config. +flux-2-klein-base-9b, flux-2-klein-9b-kv, flux-3-klein, flux-3-klein-base, +moss-image, moss-image-base) pulls the DiT, its text encoder and its VAE, +uses the parts your model folders already hold, and adds the model under +image_models: in the config. It prints the model's license before it downloads. The license of flux-2-klein-9b, flux-2-klein-base-9b and flux-2-klein-9b-kv allows non-commercial use only. @@ -1162,7 +1164,8 @@ def cmd_pull(argv: list | None = None, prog: str = "gmlx pull") -> int: "qwen-image, qwen-image-turbo, krea-2-turbo, " "flux-2-klein, flux-2-klein-base, flux-2-klein-9b, " "flux-2-klein-base-9b, flux-2-klein-9b-kv, " - "flux-3-klein, flux-3-klein-base) " + "flux-3-klein, flux-3-klein-base, moss-image, " + "moss-image-base) " "pulls the model with its text encoder and VAE. See " "the examples below.") ap.add_argument("--to", "--out", default=None, metavar="DIR", diff --git a/gmlx/image/api.py b/gmlx/image/api.py index 351272f5c..f69b7c826 100644 --- a/gmlx/image/api.py +++ b/gmlx/image/api.py @@ -70,6 +70,7 @@ class SamplerDefaults: shift: float size: tuple = (1024, 1024) # width, height strength: float = 0.6 # img2img denoise, as diffusers sets it + neg_skip: tuple = () # blocks the negative pass skips def sampler_defaults(family: str, variant: str) -> SamplerDefaults | None: @@ -80,7 +81,8 @@ def sampler_defaults(family: str, variant: str) -> SamplerDefaults | None: row = get_family(family).sampling_row(variant) if row is None: return None - return SamplerDefaults(steps=row.steps, cfg=row.cfg, shift=row.shift) + return SamplerDefaults(steps=row.steps, cfg=row.cfg, shift=row.shift, + neg_skip=tuple(row.neg_skip)) def fixed_steps(family, variant: str) -> int | None: @@ -222,7 +224,7 @@ def sigmas(self, steps: int, tokens: int | None = None) -> np.ndarray: raise ValueError(f"no schedule for {self.family.title}; pass sigmas") -def _borrowed(encoder_model, spec): +def _borrowed(encoder_model, spec, gguf: str | None = None): from .text_encoder import borrow_encoder if len(encoder_model) == 3: @@ -231,7 +233,7 @@ def _borrowed(encoder_model, spec): model, tokenizer = encoder_model args = getattr(model, "args", None) config = dict(vars(args)) if args is not None else {} - return borrow_encoder(model, tokenizer, config, spec) + return borrow_encoder(model, tokenizer, config, spec, gguf=gguf) def _header_bytes(path, kind: str) -> int: @@ -253,13 +255,16 @@ def _header_bytes(path, kind: str) -> int: def load(dit, *, text_encoder=None, vae=None, mmproj=None, - encoder_model=None, variant=None, memory: str = "auto", - log=print) -> ImagePipeline: + encoder_model=None, encoder_gguf=None, variant=None, + memory: str = "auto", log=print) -> ImagePipeline: """Load an image pipeline from a DiT GGUF, a text encoder GGUF or a resident chat model given as ``encoder_model=(model, tokenizer)``, and - a VAE. ``mmproj`` is the mmproj GGUF of the text encoder: with it the - encoder loads with its vision tower, which a request with reference - images needs and a text request does not. With ``memory="auto"`` the + a VAE. ``encoder_gguf`` is the GGUF ``encoder_model`` was loaded from: + a family whose encode runs on F32 norms (MOSS-Image) reads them from + it, and without it runs on the model's own. ``mmproj`` is the mmproj + GGUF of the text encoder: with it the encoder loads with its vision + tower, which a request with reference images needs and a text request + does not. With ``memory="auto"`` the weights of all are gated on their header sizes before anything loads; ``memory="off"`` skips that gate, and each load still gates its own weights.""" @@ -272,6 +277,9 @@ def load(dit, *, text_encoder=None, vae=None, mmproj=None, raise ValueError(f"memory must be 'auto' or 'off', not {memory!r}") if text_encoder is not None and encoder_model is not None: raise ValueError("pass text_encoder or encoder_model, not both") + if encoder_gguf is not None and encoder_model is None: + raise ValueError("encoder_gguf is the file of encoder_model; pass " + "encoder_model") if mmproj is not None and text_encoder is None: raise ValueError("an mmproj loads with its text encoder GGUF; pass " "text_encoder") @@ -307,7 +315,9 @@ def load(dit, *, text_encoder=None, vae=None, mmproj=None, # What the tower holds as loaded, which the plan counts. weights["mmproj"] = handle.vision_bytes() elif encoder_model is not None: - handle = _borrowed(encoder_model, spec) + handle = _borrowed(encoder_model, spec, + gguf=None if encoder_gguf is None + else str(encoder_gguf)) seconds["encoder"] = time.perf_counter() - t0 ae = None t0 = time.perf_counter() @@ -878,7 +888,8 @@ def _table_id(pipe: ImagePipeline, patch_table) -> str: def denoise_step(pipe: ImagePipeline, latents, sigma, conds, *, patch_table=None, target=None, time: str = "reference", - ref_cache: dict | None = None) -> mx.array: + ref_cache: dict | None = None, + skip=()) -> mx.array: """The velocity of every row of ``latents`` at ``sigma``, a scalar or one value per row, under ``conds``, one Conditioning or one per row. ``target`` is the sigma the step goes to, likewise, which a flow map @@ -898,7 +909,8 @@ def denoise_step(pipe: ImagePipeline, latents, sigma, conds, *, under its conditioning and its place in the call. That is how ComfyUI's Flux KV Cache node samples. A family whose step runs in parts takes ``ref_cache`` under another time rule than the reference's - only.""" + only. ``skip`` holds blocks, zero-based, that every row of the call + leaves out, for a DiT that skips blocks.""" if time not in getattr(pipe.dit, "time_rules", ("reference",)): raise ValueError(f"the {pipe.family.title} DiT has no time rule " f"{time!r}") @@ -910,6 +922,13 @@ def denoise_step(pipe: ImagePipeline, latents, sigma, conds, *, raise ValueError(f"the {pipe.family.title} DiT caches its " f"references by itself under the reference's " f"time rule") + skip = tuple(skip) + if any(isinstance(i, bool) or not isinstance(i, numbers.Integral) + for i in skip): + raise ValueError(f"skip holds block indices, not {list(skip)}") + skip = tuple(sorted({int(i) for i in skip})) + if skip and not getattr(pipe.dit, "skips_blocks", False): + raise ValueError(f"the {pipe.family.title} DiT skips no blocks") table = _table_id(pipe, patch_table) x = mx.array(latents).astype(mx.float32) if x.ndim != 4: @@ -929,7 +948,8 @@ def per_row(v, what): raise ValueError(f"{len(rows)} conditionings for {b} rows") with state_scope(pipe, rows): out = [_row(pipe, x[i:i + 1], sig[i], rows[i], table, - None if tgt is None else tgt[i], time, ref_cache, i) + None if tgt is None else tgt[i], time, ref_cache, i, + skip) for i in range(b)] return out[0] if b == 1 else mx.concatenate(out, axis=0) @@ -1109,15 +1129,17 @@ def _parts_row(pipe: ImagePipeline, parts, x: mx.array, sigma: np.float32, def _row(pipe: ImagePipeline, x: mx.array, sigma: np.float32, cond: Conditioning, table, target=None, time: str = "reference", ref_cache: dict | None = None, - row: int = 0) -> mx.array: + row: int = 0, skip: tuple = ()) -> mx.array: """One row. A DiT that takes reference images gets the tokens of the conditioning's reference latents from its ``references``, and a flow map the ``target`` its modulation reads. A time rule other than the reference's reaches the DiT through its modulation, and a family whose step runs in parts runs it whole then. ``ref_cache`` is that of - ``denoise_step``, and ``row`` the row's place in its call.""" + ``denoise_step``, and ``row`` the row's place in its call. ``skip`` + reaches the DiT only when it holds a block.""" dit = pipe.dit modulated = hasattr(dit, "modulation") and hasattr(dit, "caption") + kw = {"skip": skip} if skip else {} if time != "reference" and not modulated: raise ValueError(f"the {pipe.family.title} DiT takes the time rule " f"{time!r} through no modulation") @@ -1125,15 +1147,19 @@ def _row(pipe: ImagePipeline, x: mx.array, sigma: np.float32, if target is not None: raise ValueError(f"the {pipe.family.title} DiT takes no target " f"time") + if skip: + raise ValueError(f"the {pipe.family.title} step parts skip no " + f"blocks") return _parts_row(pipe, pipe.family.step, x, sigma, cond, table) _need_latents(pipe, cond) - kw = {} if cond.ref_latents: if not hasattr(dit, "references"): raise ValueError(f"the {pipe.family.title} DiT takes no " f"reference images") refs = dit.references([mx.array(z)[None] for z in cond.ref_latents]) if ref_cache is not None: + if skip: + raise ValueError("a cached reference row skips no blocks") return _cached_row(pipe, x, sigma, cond, table, time, refs, ref_cache, _ref_key(cond, row)) kw["refs"] = refs @@ -1269,7 +1295,7 @@ def sample(pipe: ImagePipeline, cond: Conditioning, *, seed: int | None = None, noise=None, sampler: str = "euler", sigmas=None, init_latents=None, denoise: float = 1.0, mask=None, patch_table=None, on_step=None, should_stop=None, - plan=None) -> mx.array: + plan=None, neg_skip=None) -> mx.array: """Sample latents for ``cond``. With ``init_latents`` the run is img2img from strength ``denoise``, and a ``mask`` (1 where the image changes) makes it inpainting. ``noise`` is (B, C, H, W) float32, else @@ -1281,10 +1307,19 @@ def sample(pipe: ImagePipeline, cond: Conditioning, *, rounds as its reference loop does. ``plan`` is the memory plan of the request: where it runs one row at a time, a guided step evaluates the positive velocity before it runs the negative row, and the two give - the latents of a step that evaluates both together.""" + the latents of a step that evaluates both together. ``neg_skip`` holds + the blocks the negative pass of a guided step skips, the variant's + sampling row's when None.""" patch_table = _table_id(pipe, patch_table) if cfg is None: cfg = pipe.defaults.cfg + if neg_skip is None: + row = sampler_defaults(pipe.family.name, pipe.variant) + neg_skip = () if row is None or cfg == 1.0 else row.neg_skip + else: + neg_skip = tuple(neg_skip) + if neg_skip and cfg == 1.0: + raise ValueError("neg_skip needs guidance: cfg is 1.0") flow_map = bool(getattr(pipe.dit.cfg, "flow_map", False)) check_sampling(pipe.family, pipe.variant, flow_map, cfg, sampler) if cfg != 1.0 and neg is None: @@ -1360,7 +1395,8 @@ def sample(pipe: ImagePipeline, cond: Conditioning, *, mx.eval(v) v = guide(v, denoise_step(pipe, x, s[i], neg, patch_table=patch_table, - target=to), cfg, act) + target=to, skip=neg_skip), + cfg, act) x = smp.step(x, v, i) if m is not None: x = inpaint_blend(x, m, noised(s[i + 1], noise, init)) diff --git a/gmlx/image/catalog.py b/gmlx/image/catalog.py index 169746f92..b6311f2b2 100644 --- a/gmlx/image/catalog.py +++ b/gmlx/image/catalog.py @@ -152,6 +152,22 @@ def encoder_name(self, quant: str) -> str: prompt_license_url="https://huggingface.co/Qwen/Qwen-Image-2.1-PE-T2I/" "blob/f3ed7985c788ad75b3ab7223e0c4c51e2a43545b/LICENSE") +# MOSS-Image: realrebelai's GGUFs of OpenMOSS's SFT and Base checkpoints, +# conditioned on Qwen3.5-9B. Its VAE is the release's file, float32 in the +# diffusers layout with its config.json, in the SFT repository. +_MOSS = dict( + family="moss_image", dit_repo="realrebelai/MOSS-Image_GGUFs", + dit_quants=_K_QUANTS, encoder_repo="unsloth/Qwen3.5-9B-GGUF", + encoder_file="Qwen3.5-9B-{quant}.gguf", encoder_quants=("Q4_K_M", "Q8_0"), + vae_repo="OpenMOSS-Team/MOSS-Image-SFT", + vae_revision="fc5de0d56fc910cbca6f6625a70f5f582b26deb1", + vae_file="vae/diffusion_pytorch_model.safetensors", + vae_config="vae/config.json", + vae_sha256="c3bdf94dc67ec69b30ea9fe1fa8340085679cae64d019f7c7e4a9921f1a7f17a", + vae_pack="moss-image-vae-F32.gguf", + license="apache-2.0", encoder_license="apache-2.0", + encoder_bytes=9_527_502_048, vae_bytes=249_514_688) + RECIPES = {r.name: r for r in ( Recipe( name="z-image-turbo", title="Z-Image-Turbo", family="z_image", @@ -222,6 +238,17 @@ def encoder_name(self, quant: str) -> str: license="krea-2-community", license_url=_KREA2_LICENSE, license_note=_KREA2_NOTE, encoder_license="apache-2.0", encoder_bytes=4_280_406_816, vae_bytes=253_807_712), + Recipe( + name="moss-image", title="MOSS-Image", variant="sft", + dit_file="MOSS-Image-SFT-{quant}.gguf", + license_url="https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT/" + "blob/main/LICENSE", + aliases=("moss-image-sft",), **_MOSS), + Recipe( + name="moss-image-base", title="MOSS-Image Base", variant="base", + dit_file="MOSS-Image-Base-{quant}.gguf", + license_url="https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base/" + "blob/main/LICENSE", **_MOSS), Recipe( name="flux-2-klein", title="FLUX.2 klein", variant="distilled", dit_repo="unsloth/FLUX.2-klein-4B-GGUF", @@ -308,6 +335,16 @@ def encoder_name(self, quant: str) -> str: # DiT four. "flux-2-klein-9b-kv": {16: ("Q4_K_M", "Q8_0"), 24: ("Q4_K_M", "Q8_0"), 32: ("Q8_0", "Q8_0"), 64: ("Q8_0", "Q8_0")}, + # A MOSS-Image request at 1024 pixels with guidance on the longest + # prompt: on 16 GB the Q8_0 encoder alone needs 9.49 GiB of 8.55, so the + # class takes the Q4_K_M encoder, and a Q6_K DiT needs 8.10 GiB with the + # staged load, the encoder released and a tiled decode. On 24 GB a Q8_0 + # DiT with the Q8_0 encoder needs 14.92 GiB of 15.2 with the encoder + # released. + "moss-image": {16: ("Q6_K", "Q4_K_M"), 24: ("Q8_0", "Q8_0"), + 32: ("Q8_0", "Q8_0"), 64: ("Q8_0", "Q8_0")}, + "moss-image-base": {16: ("Q6_K", "Q4_K_M"), 24: ("Q8_0", "Q8_0"), + 32: ("Q8_0", "Q8_0"), 64: ("Q8_0", "Q8_0")}, # A Q8_0 DiT fits a 1024 pixel image from the 24 GB class up, with the # encoder released on 24 GB, where Base also runs its two guided rows # one at a time. @@ -334,9 +371,26 @@ def encoder_name(self, quant: str) -> str: "Qwen3-8B-Q8_0.gguf": 661_225_472, "Qwen3-VL-8B-Instruct-Q8_0.gguf": 661_225_472, "Qwen3-VL-8B-Instruct-Q4_K_M.gguf": 510_504_960, + "Qwen3.5-9B-Q8_0.gguf": 1_080_688_640, + "Qwen3.5-9B-Q4_K_M.gguf": 834_355_200, +} + + +# Encoder files a class default pulls other than the one encoder_bytes +# sizes, by name. +ENCODER_BYTES = { + "Qwen3.5-9B-Q4_K_M.gguf": 5_680_522_464, + "Qwen3.5-9B-Q8_0.gguf": 9_527_502_048, } +def companion_bytes(r: Recipe, ram_bytes: int | None) -> int: + """Bytes of the encoder (with its vision tower) and the VAE that a + pull of ``r`` takes on this Mac, at the class's encoder quant.""" + name = r.encoder_name(default_quants(r, ram_bytes)[1]) + return ENCODER_BYTES.get(name, r.encoder_bytes) + r.vae_bytes + + # The image size an entry pulled with the class default DiT gets in its # defaults, where that DiT makes images smaller than the family's size. DEFAULT_SIZES = {"krea-2-turbo": {16: "512x512"}} diff --git a/gmlx/image/families/krea2/family.py b/gmlx/image/families/krea2/family.py index dd22b1582..d0a2e42c6 100644 --- a/gmlx/image/families/krea2/family.py +++ b/gmlx/image/families/krea2/family.py @@ -18,6 +18,8 @@ from __future__ import annotations +import numbers + # The prompt template, the encoder layers, the timestep shift of each # variant and the scheduler values are those of diffusers' # src/diffusers/pipelines/krea2/pipeline_krea2.py and @@ -55,6 +57,7 @@ # The projector over the encoder layers: (layers, 1), or the layers alone # in a file that flattened it. PROJECTOR = "txtfusion.projector.weight" +ORIG_SHAPE = "comfy.gguf.orig_shape." + PROJECTOR # The two matrices of the final layer that one converter writes and the # reference model does not read. LAST_MLP = ("last.up.weight", "last.down.weight") @@ -160,14 +163,20 @@ def _projector_layers(view: HeaderView) -> int | None: def projector_shape(view: HeaderView) -> tuple | None: """The torch shape a flattened projector is stored with, from its ``comfy.gguf.orig_shape`` key, or None when the file keeps the - shape.""" - raw = view.kv.get("comfy.gguf.orig_shape." + PROJECTOR) + shape or the key is not a list of integers.""" + raw = view.kv.get(ORIG_SHAPE) s = view.shape(PROJECTOR) - if raw is None or s is None or len(s) != 1: + if raw is None or s is None or len(s) != 1 or not _integers(raw): return None return tuple(int(v) for v in raw) +def _integers(raw) -> bool: + return isinstance(raw, (list, tuple)) and all( + isinstance(v, numbers.Integral) and not isinstance(v, bool) + for v in raw) + + def layout(view: HeaderView) -> str: """``last_mlp`` when the file has the final layer's two extra matrices, else ``plain``.""" @@ -184,10 +193,14 @@ def fingerprint(view: HeaderView) -> Match | Refusal | None: if n != 12: diffs.append(f"{PROJECTOR} is {view.shape(PROJECTOR)}, Krea 2 has " "(12, 1)") - orig = projector_shape(view) - if len(view.shape(PROJECTOR) or ()) == 1 and orig != (1, 12): + flat = len(view.shape(PROJECTOR) or ()) == 1 + raw = view.kv.get(ORIG_SHAPE) + if flat and raw is not None and not _integers(raw): + diffs.append(f"{PROJECTOR} is flattened, and its stored shape " + f"{raw!r} is not a list of integers") + elif flat and projector_shape(view) != (1, 12): diffs.append(f"{PROJECTOR} is flattened, and its stored shape is " - f"{orig}, not (1, 12)") + f"{projector_shape(view)}, not (1, 12)") diffs += [f"{g} has {view.count(g)} blocks, Krea 2 has {c}" for g, c in _COUNTS if view.count(g) != c] if sum(view.has(t) for t in LAST_MLP) == 1: diff --git a/gmlx/image/families/moss_image/__init__.py b/gmlx/image/families/moss_image/__init__.py new file mode 100644 index 000000000..b4c9ca8ff --- /dev/null +++ b/gmlx/image/families/moss_image/__init__.py @@ -0,0 +1 @@ +"""MOSS-Image: the OpenMOSS single-stream DiT, SFT and Base.""" diff --git a/gmlx/image/families/moss_image/caption_prompt.txt b/gmlx/image/families/moss_image/caption_prompt.txt new file mode 100644 index 000000000..d8257abf3 --- /dev/null +++ b/gmlx/image/families/moss_image/caption_prompt.txt @@ -0,0 +1,15 @@ +You write captions for MOSS-Image, a text to image model that follows detailed descriptions of images. The user turn is a request for an image. Rewrite it as one caption of the image to make. + +Keep everything the request asks for: its subjects and how many of each, what they do and how they relate, the medium and style, anything it excludes, and any words that must appear in the image. Settle ordinary gaps yourself. Do not invent facts, names, dates or specifications the request does not give. + +Describe what is visible: the subject and its action, the appearance that matters, the setting, the camera and framing, the medium, the lighting, the materials, the colors and the mood. Add detail where it makes the image more specific. Give spatial relations in words, such as left of, behind or in the lower corner. Leave out quality tags such as masterpiece or 8k, styles the request does not ask for, word counts and notes about how the caption was written. + +Keep a minimal request minimal. Develop a rich request in full, without turning it into a stock layout of a title, a main picture and a footer. + +For a designed image such as a poster, a sign or a page, describe the text and the imagery together: the reading order, which parts are largest, how they group, and text large enough to read. Write the actual words, never placeholders such as "headline here". Do not add names, watermarks or text the request does not ask for. + +Write the caption in English unless the request asks for another language. Quote every word that appears in the image exactly as the request gives it, in its own language, with its punctuation, capitals and line breaks. Only the words the request asks to show appear as text. A thing the request names to depict, such as the dog of a birthday card with a dog, is drawn and never written out, and a quoted sentence in the request is not always text to draw: tell the words to show apart from instructions about the image. + +The caption is cohesive prose and may have more than one paragraph. Give no coordinates, boxes or layout fields. + +Reply with one JSON object and nothing else: {"caption": "..."} diff --git a/gmlx/image/families/moss_image/family.py b/gmlx/image/families/moss_image/family.py new file mode 100644 index 000000000..a4165b842 --- /dev/null +++ b/gmlx/image/families/moss_image/family.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 AND MIT +# SPDX-FileCopyrightText: Copyright 2026 OpenMOSS Team +# Ported from OpenMOSS-Team/MOSS-Image-SFT at revision fc5de0d5: the +# tensor names and shapes of MOSSImageTransformer in +# pipeline_moss_image.py, the constants of that file and of +# transformer/config.json, and the prompt template of +# tokenizer/chat_template.jinja with the token cap of the pipeline, and +# modified for gmlx (Apache-2.0; see licenses/MOSS-Image-LICENSE). +"""MOSS-Image fingerprint, constants, config synthesis and text encoder. + +Shapes are GGUF order, input dim first. The file's tensor names are the +training names of the release's MOSSImageTransformer, and the rule takes +a file whose names and shapes are exactly that model's. + +SFT and Base share every name and shape, so the variant comes from the +hash of three small float tensors, which every published quant keeps in +F16. A file of another image family given as MOSS-Image is refused by +that family's name. +""" + +from __future__ import annotations + +import re + +from gmlx.image.fingerprint import HeaderView, Match, Refusal +from gmlx.image import upsample +from gmlx.image.registry import (ConstantsRow, Family, SamplingRow, + UpsampleParts, VariantRow) + +TITLE = "MOSS-Image" + +# The release config: width, feed-forward width, text width, latent +# channels, sinusoid width, head dim and RoPE bands per axis. +_D, _F, _T, _C, _FQ, _HD, _BANDS = 3072, 8192, 4096, 128, 256, 128, 16 +_DEPTH, _TEXT_DEPTH = 42, 2 + +_TOP = { + "video_embedder.weight": (_C, _D), + "video_embedder.bias": (_D,), + "text_embedder.proj.weight": (_T, _D), + "text_embedder.proj.bias": (_D,), + "rope.bands": (_BANDS,), + "time_embedding.0.weight": (_FQ, _D), + "time_embedding.0.bias": (_D,), + "time_embedding.2.weight": (_D, _D), + "time_embedding.2.bias": (_D,), + "shared_block_modulation.weight": (_D, 6 * _D), + "shared_block_modulation.bias": (6 * _D,), + "final_norm_video.weight": (_D,), + "final_linear_video.weight": (_D, _C), + "video_head_modulation.weight": (_D, 2 * _D), + "video_head_modulation.bias": (2 * _D,), +} +_TEXT_BLOCK = { + "norm1.weight": (_D,), + "norm2.weight": (_D,), + "q_norm.weight": (_HD,), + "k_norm.weight": (_HD,), + "qkv.weight": (_D, 3 * _D), + "qkv.bias": (3 * _D,), + "proj.weight": (_D, _D), + "proj.bias": (_D,), + "ffn_gate.weight": (_D, _F), + "ffn_gate.bias": (_F,), + "ffn_up.weight": (_D, _F), + "ffn_up.bias": (_F,), + "ffn_down.weight": (_F, _D), + "ffn_down.bias": (_D,), +} +_BLOCK = { + "norm1.weight": (_D,), + "norm2.weight": (_D,), + "q_norm.weight": (_HD,), + "k_norm.weight": (_HD,), + "proj.q.weight": (_D, _D), + "proj.k.weight": (_D, _D), + "proj.v.weight": (_D, _D), + "proj.o.weight": (_D, _D), + "proj.ffn_up_gate.weight": (_D, 2 * _F), + "proj.ffn_down.weight": (_F, _D), +} +BLOCK_GROUPS = (("layers", _DEPTH, _BLOCK), + ("text_embedder.blocks", _TEXT_DEPTH, _TEXT_BLOCK)) +_GROUP = re.compile(r"(layers|text_embedder\.blocks)\.(\d+)\.") + +CONSTANTS = ConstantsRow( + values={ + "rope_base": 10000.0, + "norm_eps": 1e-06, + # The base of the timestep sinusoid's frequencies. + "time_period": 10000.0, + "num_train_timesteps": 1000, + # scheduler/scheduler_config.json: the shift of the sampling + # schedule; training_shift 1.0 leaves the snap grid unshifted. + "scheduler_shift": 5.0, + # The pipeline's image processor: 8 pixels a VAE latent, 2 a patch. + "vae_scale_factor": 16, + "size_multiple": 16, + }, + repo="OpenMOSS-Team/MOSS-Image-SFT", + revision="fc5de0d56fc910cbca6f6625a70f5f582b26deb1", +) + +IDENTITY_TENSORS = ("video_embedder.bias", "final_norm_video.weight", + "time_embedding.0.bias") +_QUANTS = ("Q2_K", "Q3_K_L", "Q3_K_M", "Q3_K_S", "Q4_K_M", "Q4_K_S", + "Q5_K_M", "Q5_K_S", "Q6_K", "Q8_0") +_REPO = "realrebelai/MOSS-Image_GGUFs" +_REVISION = "4f21f583c9213eb2230a09e9b33bfa86605aecf2" +_SFT = VariantRow("sft", _REPO, _REVISION, + tuple(f"MOSS-Image-SFT-{q}.gguf" for q in _QUANTS)) +_BASE = VariantRow("base", _REPO, _REVISION, + tuple(f"MOSS-Image-Base-{q}.gguf" for q in _QUANTS)) +# The release's sampling for both variants: 50 steps at guidance 3, +# which is gmlx's cfg, as MOSSImagePipeline guides by u + g (c - u), with +# block 22 left out of the negative pass. SFT from the pipeline's defaults +# and README at the constants' revision, Base from the user guide. +SAMPLING = { + "sft": SamplingRow(50, 3.0, 5.0, "OpenMOSS-Team/MOSS-Image-SFT", + "fc5de0d56fc910cbca6f6625a70f5f582b26deb1", + neg_skip=(22,)), + "base": SamplingRow(50, 3.0, 5.0, "github.com/OpenMOSS/MOSS-Image", + "c1355e6b1986dfe8fba3c0455783a05899212342", + neg_skip=(22,)), +} +VARIANTS = { + "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60": _SFT, + "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626": _BASE, +} + + +def expected() -> dict: + """Every tensor name of the file and its shape.""" + out = dict(_TOP) + for group, count, block in BLOCK_GROUPS: + for i in range(count): + out.update({f"{group}.{i}.{n}": s for n, s in block.items()}) + return out + + +def _past_count(name: str, view: HeaderView) -> bool: + """Whether a block tensor sits in a block the counts already name.""" + m = _GROUP.match(name) + if m is None: + return False + want = {g: c for g, c, _ in BLOCK_GROUPS}[m[1]] + return int(m[2]) >= min(want, view.count(m[1])) + + +def _other_family(view: HeaderView) -> Refusal | None: + """The refusal of a file that another family takes or refuses, by + that family's name, for a file given as MOSS-Image.""" + from gmlx.image.fingerprint import strip_prefix + from gmlx.image.registry import families + + for fam in families(): + if fam.name == FAMILY.name: + continue + got = fam.fingerprint(strip_prefix(view, fam.strip_prefixes)) + if isinstance(got, Refusal): + return got + if got is not None: + return Refusal(f"this is a {fam.title} GGUF, not a {TITLE} one") + return None + + +def fingerprint(view: HeaderView) -> Match | Refusal | None: + if not (view.has("shared_block_modulation.weight") + and view.has("video_embedder.weight")): + return _other_family(view) + want = expected() + diffs = [f"{g} has {view.count(g)} blocks, {TITLE} has {c}" + for g, c, _ in BLOCK_GROUPS if view.count(g) != c] + diffs += [f"{n} is {view.shape(n)}, {TITLE} has {s}" + for n, s in want.items() + if view.has(n) and view.shape(n) != s] + diffs += [f"{n} is missing" for n in want + if not view.has(n) and not _past_count(n, view)] + diffs += [f"{n} is not a {TITLE} tensor" for n in view.tensors + if n not in want and not _past_count(n, view)] + if diffs: + shown = ", ".join(diffs[:8]) + more = f" (+{len(diffs) - 8} more)" if len(diffs) > 8 else "" + return Refusal(f"a {TITLE} relative that gmlx image does not " + f"support: {shown}{more}") + return Match() + + +def synth_config(view: HeaderView) -> dict: + """The keys of the official transformer/config.json, from shapes and + constants, plus the derived values the module tree needs.""" + c = CONSTANTS.values + s = view.dims + latent, dim = s("video_embedder.weight") + head_dim = s("layers.0.q_norm.weight")[0] + return { + "dim": dim, + "ffn_dim": s("layers.0.proj.ffn_down.weight")[0], + "freq_dim": s("time_embedding.0.weight")[0], + "latent_channels": latent, + "num_heads": dim // head_dim, + "num_layers": view.count("layers"), + "rope_base": c["rope_base"], + "text_dim": s("text_embedder.proj.weight")[0], + "head_dim": head_dim, + "num_bands": s("rope.bands")[0], + "num_text_blocks": view.count("text_embedder.blocks"), + "norm_eps": c["norm_eps"], + "time_period": c["time_period"], + "num_train_timesteps": c["num_train_timesteps"], + } + + +# The prompt template is the release's chat_template.jinja rendered for +# one user turn with the generation prompt and thinking on, and the cap is +# the release pipeline's max_sequence_length. +ENCODER_TEMPLATE = ("<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + "\n") +ENCODER_MAX_TOKENS = 4096 +ENCODER_IDENTITY = ( + ("blk.0.attn_norm.weight", "model.layers.0.input_layernorm.weight", + "622011779892c5d8a3fceb52c4b9dccdb4aee84d1c49099080b2b06e215e7421"), + ("blk.17.post_attention_norm.weight", + "model.layers.17.post_attention_layernorm.weight", + "18afbd19c59b5fcc7e6b7c483a53bcb6ddf8d7d34cba0986a28f0a5a96944b1d"), + ("output_norm.weight", "model.norm.weight", + "451aeaff886b5580361ec7a79cd90278494f0e4d2c4645bfb5482416018785ba"), +) +ENCODER_IDENTITY_SOURCE = ( + "unsloth/Qwen3.5-9B-GGUF at 3885219b6810b007914f3a7950a8d1b469d598a5") + + +def encoder(): + """The text encoder of the family: Qwen3.5-9B read after its final + norm, as the release reads last_hidden_state, on the chat template of + one user turn with thinking on, which trims the prompt. Each prompt + runs alone at its own length, with no pad rows. Its RMSNorms run on + their F32 weights from the GGUF, which hold the offset 1 + w, so the + encode does not round them to bfloat16.""" + from gmlx.image.text_encoder import EncoderSpec, IdentityTensor + + return EncoderSpec( + name="qwen3.5-9b", + arch="qwen35", + model_type="qwen3_5", + hidden=4096, + layers=32, + tap_layers=(32,), + template=ENCODER_TEMPLATE, + max_tokens=ENCODER_MAX_TOKENS, + identity=tuple(IdentityTensor(*row) for row in ENCODER_IDENTITY), + identity_source=ENCODER_IDENTITY_SOURCE, + final_norm=True, + trim=True, + fp32_norms=True, + ) + + +def build(config: dict): + from .model import MossImageDiT + + return MossImageDiT(config) + + +def lora_map(config: dict): + from .lora_map import lora_map as build_map + + return build_map(config) + + +def schedule(steps: int, *, tokens: int | None, variant: str | None = None): + """The release's snapped schedule of ``steps``, the same for every + image size and variant.""" + from gmlx.image.schedule import moss_sigmas + + return moss_sigmas(steps, CONSTANTS.values["scheduler_shift"], + CONSTANTS.values["num_train_timesteps"]) + + +# The optional upsampler, on a chat model of the user's choice. OpenMOSS +# ships an agent skill, not a system prompt: skills/moss-image-prompt- +# enhancer/SKILL.md of github.com/OpenMOSS/MOSS-Image at commit c1355e6 +# (Apache-2.0). caption_prompt.txt is gmlx's own caption-only system +# prompt, written from the skill's guidance for a caption, with no layout +# boxes and no tools. The skill names no sampling: gmlx decodes greedily, +# with no presence penalty and thinking off, as for Krea 2. The reply is +# the skill's {"caption": ...} object, or the caption as written. +UPSAMPLE = UpsampleParts( + prompts="caption_prompt.txt", + sha256="389ce41e0d9f4e11df7f0c078cef7a4af772b7312b19fc7d637e5808345cbba1", + key=None, + max_tokens=2048, + temperature=0.0, + presence_penalty=0.0, + thinking="off", + repo="github.com/OpenMOSS/MOSS-Image", + revision="c1355e6b1986dfe8fba3c0455783a05899212342", + package=__package__, + suggested="Qwen3.5-9B", + parse=upsample.json_field("caption"), +) + + +FAMILY = Family( + name="moss_image", + title=TITLE, + order=40, + arch_strings=frozenset({"moss_image", None}), + strip_prefixes=("model.diffusion_model.", "diffusion_model."), + fingerprint=fingerprint, + synth_config=synth_config, + constants=CONSTANTS, + build=build, + lora_map=lora_map, + encoder=encoder, + # The release builds the bands from constants and ignores the stored + # ones; the file's copy keeps its float32 values. + fp32_keep=("rope.bands",), + identity_tensors=IDENTITY_TENSORS, + variants=VARIANTS, + # A file of neither release, as a fine-tune, samples as SFT: both + # releases guide alike, block 22 left out of the negative pass. + unknown_variant="sft", + schedule=schedule, + sampling=SAMPLING, + upsample=UPSAMPLE, +) diff --git a/gmlx/image/families/moss_image/lora_map.py b/gmlx/image/families/moss_image/lora_map.py new file mode 100644 index 000000000..cd2cb5df6 --- /dev/null +++ b/gmlx/image/families/moss_image/lora_map.py @@ -0,0 +1,51 @@ +"""MOSS-Image LoRA key map. + +The GGUF keeps the names of the release's MOSSImageTransformer, which +the release trains under, the OpenMOSS ComfyUI plugin loads and PEFT +names under ``transformer.``. In each block those are ``proj.q``, +``proj.k``, ``proj.v``, ``proj.o``, ``proj.ffn_up_gate`` and +``proj.ffn_down``; in each text refiner block of ``text_embedder`` they +are ``qkv``, ``proj``, ``ffn_gate``, ``ffn_up`` and ``ffn_down``; and +outside the blocks they are ``video_embedder``, ``text_embedder.proj``, +``time_embedding.0``, ``time_embedding.2``, ``shared_block_modulation``, +``video_head_modulation`` and ``final_linear_video``. Each name maps to +itself. The fused weights take adapters under their own names: every +known layout holds ``proj.ffn_up_gate`` and the refiner ``qkv`` as one +module, so no target is a band. No text encoder key resolves: the map +has no encoder targets. + +ComfyUI reads the adapter of a model it has no class of its own for +under ``diffusion_model.`` and ``lora_unet_``, in these names only. +""" + +from __future__ import annotations + +from gmlx.image.lora.reader import LoraKeyMap + +_BLOCK = ("proj.q", "proj.k", "proj.v", "proj.o", "proj.ffn_up_gate", + "proj.ffn_down") +_REFINER = ("qkv", "proj", "ffn_gate", "ffn_up", "ffn_down") +_TOP = ("video_embedder", "text_embedder.proj", "time_embedding.0", + "time_embedding.2", "shared_block_modulation", + "video_head_modulation", "final_linear_video") +# The prefixes under which ComfyUI's key map finds a MOSS-Image adapter. +COMFY_PREFIXES = ("lora_unet_", "diffusion_model.") + + +def _counts(config: dict) -> tuple[int, int]: + """Blocks and text refiner blocks, from the config synthesized from a + header or the fields of a built model's config.""" + if "depth" in config: + return int(config["depth"]), int(config["text_depth"]) + return int(config["num_layers"]), int(config.get("num_text_blocks", 2)) + + +def lora_map(config: dict) -> LoraKeyMap: + blocks, refiners = _counts(config) + names = [f"layers.{i}.{leaf}" for i in range(blocks) for leaf in _BLOCK] + names += [f"text_embedder.blocks.{i}.{leaf}" for i in range(refiners) + for leaf in _REFINER] + names += list(_TOP) + return LoraKeyMap(family="moss_image", title="MOSS-Image", + targets={n: (n, None) for n in names}, + comfy_prefixes=COMFY_PREFIXES) diff --git a/gmlx/image/families/moss_image/model.py b/gmlx/image/families/moss_image/model.py new file mode 100644 index 000000000..e21b3034e --- /dev/null +++ b/gmlx/image/families/moss_image/model.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: Apache-2.0 AND MIT +# SPDX-FileCopyrightText: Copyright 2026 OpenMOSS Team +# Ported from MOSSImageTransformer in OpenMOSS-Team/MOSS-Image-SFT +# pipeline_moss_image.py (its blocks, text projection, RoPE bands and +# coordinates, timestep sinusoid, modulations and output head) at revision +# fc5de0d5 and modified for gmlx (Apache-2.0; see +# licenses/MOSS-Image-LICENSE). +"""MOSS-Image DiT module tree and forward. + +Attribute paths equal the GGUF tensor names, which are the release's +training names, so numbered elements are plain lists and ``proj`` of a +block is a dict. + +Latents arrive patched, 128 channels a token, and each latent pixel is +one token in row-major order. The text features pass ``text_embedder``: +a projection to the model width and two refiner blocks with biases and no +modulation, over the text alone. The image tokens and then the text rows +run through every block as one sequence, and the velocity is the final +head on the image tokens, with no sign flip. Each sample runs at its own +text length, so no row is padding. + +One timestep embedding drives every block: ``shared_block_modulation`` +maps it to six vectors (shift, scale and gate of the attention, then of +the feed-forward) that every block reads. ``video_head_modulation`` maps +it to the shift and scale of the head. The sinusoid of the timestep is +formed in float64, cosines first, and rounded to float32 once, as the +release forms it. The time vectors then run in float32 and round once. + +The feed-forward is ``ffn_down(situ(hidden, gate))`` over the two halves +of ``ffn_up_gate``, where situ is a soft-capped gated SiLU, computed in +float32 and rounded once. + +RoPE has four axes (r, f, h, w) of 16 bands at base 10000. The bands, +the phases and their sines and cosines are float32 as the release builds +them. An image token at (h, w) of an H by W grid sits at (0, L, h - (H - +1) / 2, w - (W - 1) / 2) for a text of L rows, and text row i at (0, i, +0, 0). The rotation pairs dim j with dim j + 64; ``interleave`` puts each +pair side by side, in q, k and their norm weights alike, so the +interleaved rotation of ``ops.qk_norm_rope`` computes the release's and +the attention scores are unchanged. + +The forward keeps no caches, but it takes the two parts of a step that do +not depend on the latents from its caller. ``modulation`` derives from +sigma alone the time vectors, and ``caption`` derives the text rows from +the conditioning alone. A step given either computes the same values as +one without. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from gmlx.image import ops, pacing +from gmlx.image.dtypes import ImageDtypeError +from gmlx.image.ops.products import own, precise + + +@dataclass(frozen=True) +class MossConfig: + dim: int + ffn_dim: int + heads: int + head_dim: int + depth: int + text_dim: int + text_depth: int + freq_dim: int + latent: int + bands: int + rope_base: float + eps: float + time_period: float + train_timesteps: int + # A latent pixel is one token: the DiT has no patch of its own. + patch: int = 1 + + @property + def in_channels(self) -> int: + """Channels of a latent token, the noise the sampler draws.""" + return self.latent + + @classmethod + def from_dict(cls, c: dict) -> "MossConfig": + dim, heads = int(c["dim"]), int(c["num_heads"]) + head_dim = int(c.get("head_dim", dim // max(heads, 1))) + bands = int(c.get("num_bands", head_dim // 8)) + if heads * head_dim != dim or head_dim % 8: + raise ValueError(f"MOSS-Image takes {heads} heads that divide " + f"its width {dim} into a head dim that is a " + "multiple of 8") + if 8 * bands != head_dim: + raise ValueError(f"MOSS-Image takes {head_dim // 8} RoPE bands " + f"a head dim of {head_dim}, not {bands}") + if int(c["freq_dim"]) % 2: + raise ValueError("the MOSS-Image timestep sinusoid takes an " + "even width") + return cls( + dim=dim, ffn_dim=int(c["ffn_dim"]), heads=heads, + head_dim=head_dim, depth=int(c["num_layers"]), + text_dim=int(c["text_dim"]), + text_depth=int(c.get("num_text_blocks", 2)), + freq_dim=int(c["freq_dim"]), latent=int(c["latent_channels"]), + bands=bands, rope_base=float(c["rope_base"]), + eps=float(c.get("norm_eps", 1e-6)), + time_period=float(c.get("time_period", 10000.0)), + train_timesteps=int(c.get("num_train_timesteps", 1000))) + + +def rope_bands(n: int, base: float) -> np.ndarray: + """The release's bands ``1 / base ** (k / n)``, k = 0 to n - 1, each + op in float32.""" + e = np.arange(n, dtype=np.float32) / np.float32(n) + return (np.float32(1) / (np.float32(base) ** e)).astype(np.float32) + + +def positions(h: int, w: int, n_txt: int) -> np.ndarray: + """The (r, f, h, w) coordinates (h * w + n_txt, 4), float32, of an h + by w image grid in row-major order and then of ``n_txt`` text rows.""" + rows, cols = np.meshgrid(np.arange(h), np.arange(w), indexing="ij") + img = np.zeros((h * w, 4), np.float32) + img[:, 1] = n_txt + img[:, 2] = rows.reshape(-1) - np.float32((h - 1) / 2) + img[:, 3] = cols.reshape(-1) - np.float32((w - 1) / 2) + txt = np.zeros((n_txt, 4), np.float32) + txt[:, 1] = np.arange(n_txt) + return np.concatenate([img, txt]) + + +def rope_tables(coords: np.ndarray, bands: np.ndarray) -> tuple: + """cos and sin tables (L, axes * bands), fp32, of coordinates (L, + axes): the angle of column j is axis j // bands at band j % bands. + The phase is a float32 product, as the release forms it, and its + cosine and sine round to float32 once.""" + phase = (coords[:, :, None] * bands[None, None, :]).reshape( + coords.shape[0], -1).astype(np.float64) + return (mx.array(np.cos(phase).astype(np.float32)), + mx.array(np.sin(phase).astype(np.float32))) + + +def time_sinusoid(t, dim: int, period: float) -> mx.array: + """The embedding (B, dim) of timesteps ``t``: cosines then sines of t + times ``period ** (-k / (dim / 2))``, in float64, rounded to float32 + once.""" + half = dim // 2 + freq = period ** (-np.arange(half, dtype=np.float64) / half) + phase = np.asarray(t, np.float64).reshape(-1, 1) * freq[None] + emb = np.concatenate([np.cos(phase), np.sin(phase)], axis=-1) + return mx.array(emb.astype(np.float32)) + + +def interleave(x: mx.array) -> mx.array: + """The last dim of ``x`` with dims j and j + D/2 side by side, the + pairs the release rotates together.""" + *lead, d = x.shape + return x.reshape(*lead, 2, d // 2).swapaxes(-1, -2).reshape(*lead, d) + + +def situ(hidden: mx.array, gate: mx.array) -> mx.array: + """``(4 tanh(gate / 4) sigmoid(gate)) (25 tanh(hidden / 25))`` in + float32, rounded once to the input dtype.""" + g = gate.astype(mx.float32) + h = hidden.astype(mx.float32) + y = (4 * mx.tanh(g / 4) * mx.sigmoid(g)) * (25 * mx.tanh(h / 25)) + return y.astype(hidden.dtype) + + +def qk_rope(q: mx.array, k: mx.array, q_w: mx.array, k_w: mx.array, + cos: mx.array, sin: mx.array, eps: float) -> tuple: + """q and k (B, L, H, D) normed over the head dim and rotated by the + release's pairs, as (B, H, L, D) in the dims of ``interleave``.""" + return ops.qk_norm_rope(interleave(q), interleave(k), interleave(q_w), + interleave(k_w), cos, sin, eps) + + +def gated_projection(x: mx.array, gate: mx.array, lin, h: mx.array, + branch) -> mx.array: + """``x + gate * lin(h)``: on mlx-kquant's residual epilogue when + ``ops.resid_leaf`` takes the leaf, else ``x + gate * branch()`` on the + stock ops, where ``branch`` is the module call that gives ``lin(h)``.""" + leaf = ops.resid_leaf(lin, h, x, gate) + if leaf is not None: + return ops.resid_fused(h, leaf, x, gate) + return ops.gated_residual(x, gate, branch()) + + +class Rope(nn.Module): + """The file's copy of the bands. The forward builds its own, as the + release does.""" + + def __init__(self, bands: int): + super().__init__() + self.bands = mx.zeros((bands,)) + + +class TextBlock(nn.Module): + """A refiner block of the text projection: pre-norm attention and + feed-forward with biases, no modulation.""" + + def __init__(self, cfg: MossConfig): + super().__init__() + d, f = cfg.dim, cfg.ffn_dim + self.heads, self.head_dim, self.eps = cfg.heads, cfg.head_dim, cfg.eps + self.norm1 = ops.RMSNorm(d, cfg.eps) + self.norm2 = ops.RMSNorm(d, cfg.eps) + self.q_norm = ops.RMSNorm(cfg.head_dim, cfg.eps) + self.k_norm = ops.RMSNorm(cfg.head_dim, cfg.eps) + self.qkv = nn.Linear(d, 3 * d) + self.proj = nn.Linear(d, d) + self.ffn_gate = nn.Linear(d, f) + self.ffn_up = nn.Linear(d, f) + self.ffn_down = nn.Linear(f, d) + + def attention(self, h: mx.array, cos: mx.array, + sin: mx.array) -> mx.array: + """The attention rows (B, L, dim) that ``proj`` takes.""" + b, n, _ = h.shape + q, k, v = mx.split(self.qkv(h), 3, axis=-1) + shape = (b, n, self.heads, self.head_dim) + q, k = qk_rope(q.reshape(shape), k.reshape(shape), self.q_norm.weight, + self.k_norm.weight, cos, sin, self.eps) + o = ops.sdpa(q, k, v.reshape(shape).transpose(0, 2, 1, 3)) + return o.transpose(0, 2, 1, 3).reshape(b, n, -1) + + def hidden(self, y: mx.array) -> mx.array: + """The rows that ``ffn_down`` takes.""" + return situ(self.ffn_up(y), self.ffn_gate(y)) + + def __call__(self, x: mx.array, cos: mx.array, + sin: mx.array) -> mx.array: + x = x + self.proj(self.attention(self.norm1(x), cos, sin)) + return x + self.ffn_down(self.hidden(self.norm2(x))) + + +class TextProjection(nn.Module): + def __init__(self, cfg: MossConfig): + super().__init__() + self.proj = nn.Linear(cfg.text_dim, cfg.dim) + self.blocks = [TextBlock(cfg) for _ in range(cfg.text_depth)] + + +class Block(nn.Module): + """A joint block: modulated pre-norm attention and feed-forward over + the image and text rows, each added under its gate.""" + + def __init__(self, cfg: MossConfig): + super().__init__() + d = cfg.dim + self.heads, self.head_dim, self.eps = cfg.heads, cfg.head_dim, cfg.eps + self.norm1 = ops.RMSNorm(d, cfg.eps) + self.norm2 = ops.RMSNorm(d, cfg.eps) + self.q_norm = ops.RMSNorm(cfg.head_dim, cfg.eps) + self.k_norm = ops.RMSNorm(cfg.head_dim, cfg.eps) + self.proj = {n: nn.Linear(d, d, bias=False) + for n in ("q", "k", "v", "o")} + self.proj["ffn_up_gate"] = nn.Linear(d, 2 * cfg.ffn_dim, bias=False) + self.proj["ffn_down"] = nn.Linear(cfg.ffn_dim, d, bias=False) + + def attention(self, h: mx.array, cos: mx.array, + sin: mx.array) -> mx.array: + """The attention rows (B, L, dim) that ``proj.o`` takes.""" + b, n, _ = h.shape + shape = (b, n, self.heads, self.head_dim) + p = self.proj + q, k = qk_rope(p["q"](h).reshape(shape), p["k"](h).reshape(shape), + self.q_norm.weight, self.k_norm.weight, cos, sin, + self.eps) + o = ops.sdpa(q, k, p["v"](h).reshape(shape).transpose(0, 2, 1, 3)) + return o.transpose(0, 2, 1, 3).reshape(b, n, -1) + + def hidden(self, h: mx.array) -> mx.array: + """``situ`` of the two halves of ``ffn_up_gate(h)``, the hidden + half first: the rows that ``proj.ffn_down`` takes.""" + hid, gate = mx.split(self.proj["ffn_up_gate"](h), 2, axis=-1) + return situ(hid, gate) + + def __call__(self, x: mx.array, mod: tuple, cos: mx.array, + sin: mx.array) -> mx.array: + shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod + p = self.proj + h = ops.modulate(self.norm1(x), scale_a) + shift_a + rows = self.attention(h, cos, sin) + x = gated_projection(x, gate_a, p["o"], rows, lambda: p["o"](rows)) + h = ops.modulate(self.norm2(x), scale_f) + shift_f + hid = self.hidden(h) + return gated_projection(x, gate_f, p["ffn_down"], hid, + lambda: p["ffn_down"](hid)) + + +@dataclass(frozen=True) +class StepModulation: + """What one step derives from sigma: the six vectors every block + reads, then the head's shift and scale, each (B, 1, dim) in the + activation dtype.""" + blocks: tuple + shift: mx.array + scale: mx.array + + def arrays(self) -> list: + return [*self.blocks, self.shift, self.scale] + + +@dataclass(frozen=True) +class Caption: + """What one step derives from the conditioning: the text rows after + ``text_embedder``, (B, L, dim).""" + rows: mx.array + + def arrays(self) -> list: + return [self.rows] + + +class MossImageDiT(nn.Module): + # A guided step can leave blocks out of its negative pass. + skips_blocks = True + + def __init__(self, config: dict): + super().__init__() + cfg = MossConfig.from_dict(config) + self.cfg = cfg + self.act_dtype = mx.bfloat16 + d = cfg.dim + self.video_embedder = nn.Linear(cfg.latent, d) + self.text_embedder = TextProjection(cfg) + self.rope = Rope(cfg.bands) + self.time_embedding = [nn.Linear(cfg.freq_dim, d), nn.SiLU(), + nn.Linear(d, d)] + self.shared_block_modulation = nn.Linear(d, 6 * d) + self.layers = [Block(cfg) for _ in range(cfg.depth)] + self.final_norm_video = ops.RMSNorm(d, cfg.eps) + self.final_linear_video = nn.Linear(d, cfg.latent, bias=False) + self.video_head_modulation = nn.Linear(d, 2 * d) + self._bands = rope_bands(cfg.bands, cfg.rope_base) + + @property + def act_dtype(self) -> mx.Dtype: + return self._act_dtype + + @act_dtype.setter + def act_dtype(self, dtype: mx.Dtype) -> None: + if dtype == mx.float16: + raise ImageDtypeError("MOSS-Image never runs in float16") + self._act_dtype = dtype + + def timesteps(self, sigma, batch: int) -> np.ndarray: + """The timesteps the release's scheduler passes for sigmas: sigma + times its training timesteps, in float32.""" + s = np.broadcast_to(np.asarray(sigma, np.float32), (batch,)) + return s * np.float32(self.cfg.train_timesteps) + + def time_vectors(self, sigma, batch: int) -> tuple: + """The timestep embedding (B, dim) and the modulation rows of the + blocks (B, 6 dim) and of the head (B, 2 dim), in float32. Under + bf16 activations each leaf runs as a float32 product by + ``precise``; under fp32 each runs its own call.""" + lin = precise if self.act_dtype != mx.float32 else own + t = time_sinusoid(self.timesteps(sigma, batch), self.cfg.freq_dim, + self.cfg.time_period) + te = self.time_embedding + temb = lin(te[2], nn.silu(lin(te[0], t))) + return (temb, lin(self.shared_block_modulation, temb), + lin(self.video_head_modulation, temb)) + + def modulation(self, sigma, batch: int) -> StepModulation: + _, blocks, head = self.time_vectors(sigma, batch) + act = self.act_dtype + + def chunks(rows, n): + return tuple(mx.split(rows.astype(act)[:, None, :], n, axis=-1)) + + shift, scale = chunks(head, 2) + return StepModulation(blocks=chunks(blocks, 6), shift=shift, + scale=scale) + + def text_rope(self, n_txt: int) -> tuple: + """cos and sin tables of ``n_txt`` text rows alone.""" + return rope_tables(positions(0, 0, n_txt), self._bands) + + def caption(self, feats: mx.array, taps: dict | None = None) -> Caption: + """The text rows of features (B, L, text dim).""" + te = self.text_embedder + x = te.proj(feats.astype(self.act_dtype)) + if taps is not None: + taps["text_embedder.proj"] = x + cos, sin = self.text_rope(x.shape[1]) + for i, blk in enumerate(te.blocks): + x = blk(x, cos, sin) + if taps is not None: + taps[f"text_embedder.blocks.{i}"] = x + return Caption(x) + + def rope_tables(self, h: int, w: int, n_txt: int) -> tuple: + """cos and sin tables of the joint rows: the image grid, then the + text.""" + return rope_tables(positions(h, w, n_txt), self._bands) + + def head(self, x: mx.array, shift: mx.array, + scale: mx.array) -> mx.array: + """The output head on image rows (B, N, dim).""" + y = ops.modulate(self.final_norm_video(x), scale) + shift + return self.final_linear_video(y) + + def __call__(self, latents: mx.array, sigma, feats: mx.array | None, + taps: dict | None = None, *, + mod: StepModulation | None = None, + caption: Caption | None = None, skip=()) -> mx.array: + """Velocity (B, C, H, W) fp32 for latents (B, C, H, W) at + ``sigma``, with text features (B, L, text dim). ``mod`` and + ``caption``, from ``modulation`` and ``caption``, stand in for + the work they hold. The blocks in ``skip``, zero-based, are left + out. ``taps`` collects the output of every block by module + path.""" + skip = frozenset(int(i) for i in skip) + if any(not 0 <= i < len(self.layers) for i in skip): + raise ValueError(f"MOSS-Image skips blocks 0 to " + f"{len(self.layers) - 1}, not " + f"{sorted(skip)}") + b, c, h, w = latents.shape + if mod is None: + mod = self.modulation(sigma, b) + if caption is None: + if feats is None: + raise ValueError("pass text features or their caption") + caption = self.caption(feats, taps) + txt = caption.rows + if txt.shape[0] != b: + txt = mx.broadcast_to(txt, (b, *txt.shape[1:])) + tokens = latents.reshape(b, c, h * w).transpose(0, 2, 1) + img = self.video_embedder(tokens.astype(self.act_dtype)) + if taps is not None: + taps["video_embedder"] = img + x = mx.concatenate([img, txt], axis=1) + cos, sin = self.rope_tables(h, w, txt.shape[1]) + for i, blk in enumerate(self.layers): + if i in skip: + continue + x = blk(x, mod.blocks, cos, sin) + pacing.point(x) + if taps is not None: + taps[f"layers.{i}"] = x + out = self.head(x[:, :h * w], mod.shift, mod.scale) + if taps is not None: + taps["final_linear_video"] = out + return out.transpose(0, 2, 1).reshape(b, c, h, w).astype(mx.float32) diff --git a/gmlx/image/memfit.py b/gmlx/image/memfit.py index 8b87d466d..accb0d7f1 100644 --- a/gmlx/image/memfit.py +++ b/gmlx/image/memfit.py @@ -263,6 +263,24 @@ def bytes(self, formula: float, rows: int = 1, read: float = 0) -> int: # second row of a guided step takes the 9B's slope. FLUX2_KLEIN_9B_KV_JOINED = StepFit(base=3.1 * GiB, k=2.9, row_k=0.5) +# One MOSS-Image SFT Q4_K_M step through ``api.sample`` under bounded_cache +# on an M5 Max, in fresh processes, over the settled process. The cells +# are in tests/fixtures/image/memfit/moss_image_cells.json: 256 to 2048 +# pixels with 80 text rows, 256 to 2048 with 4096, and 1024 with 640. A +# text row reads more than an image row of the formula: at nearly the +# same formula a 256 pixel step with 4096 text rows reads 2.50 GiB and a +# 1024 pixel step with 80 reads 1.57. A line through the 1536 and 2048 +# pixel cells of 80 text rows has a slope of 5.1, raised 10 percent to +# 5.6. The base is 0.25 GiB above every cell, the 256 pixel cell with +# 4096 text rows the closest, and the 2048 pixel cell with 80 is priced +# 1.46 GiB above its reading. A guided step, at 512 to 2048 pixels with +# 80 text rows and at 1024 and 2048 with 4096, holds both passes: at 1024 +# pixels with 4096 rows it reads 4.18 GiB against 3.28 unguided, and a +# second row of slope 1.0 prices it under its reading. Its slope of 2.0 +# leaves 0.31 GiB above that cell, the closest guided one, and 0.38 and +# 0.50 above the 2048 pixel cells. +MOSS_IMAGE_STEP = StepFit(base=1.64 * GiB, k=5.6, row_k=2.0) + # Measured footprint of one 512-token Qwen3-4B Q4_K_M encode on an M5 Max: # 0.35 GiB above the weights, 12.5 times the per-token formula, rounded up. ENCODE_LIVE = 16.0 @@ -514,6 +532,21 @@ def _qwen_image_vae_crop(op: str, config=None) -> int: return reach(m.OFFICIAL if config is None else config, op) +# MOSS-Image autoencoder decode footprints under bounded_cache on the +# official file, over the loaded weights, per pixel at its 96 decode +# channels. Whole square passes at 256, 512, 768 and 1024 pixels on an M5 +# Max, in fresh processes with a cold and a warm pass, two runs each, fit +# a line in pixels by least squares, as the Qwen-Image 2.1 fits do. The +# slope is raised 10 percent and the base covers every whole pass plus +# 0.25 GiB. The kq and exact paths are fit together. The FLUX.1 fits at +# 96 channels price a whole fp32 decode at 512 pixels at 3.66 GiB, and it +# added 4.08. +MOSS_VAE_DECODE = {"exact": {2: VaeFit(base=1.44 * GiB, k=19.7), + 4: VaeFit(base=3.24 * GiB, k=11.7)}, + "winograd": {2: VaeFit(base=2.19 * GiB, k=24.3), + 4: VaeFit(base=2.60 * GiB, k=22.5)}} + + VAE_ROWS = { # The FLUX.1 autoencoder. "z_image": VaeMemoryRow( @@ -559,6 +592,22 @@ def _qwen_image_vae_crop(op: str, config=None) -> int: measured=True, channels={"decode": 128, "encode": 128}, whole={"decode": VAE_DECODE, "encode": VAE_ENCODE}, tiled=VAE_TILED, overlap=128, scale=8), + # The MOSS-Image autoencoder: the FLUX.2 one with its decoder at 3/4 of + # the encoder's widths. Its whole decode has its own fits, at 96 + # channels. The encoder is the FLUX.2 one, and the encode and the + # tiled passes keep the FLUX.1 fits. Measured on the official file on + # an M5 Max, every pass grew the footprint by less than its price: at + # 2048 pixels a whole bf16 decode by 10.16 GiB against 16.21, in fp32 + # by 18.13 against 20.79, a tiled bf16 decode by 6.54 against 7.47, in + # fp32 by 11.34 against 12.58, a whole bf16 encode by 8.70 against + # 12.23 and a tiled fp32 one by 10.39 against 11.54. Past the passes + # the decode fits were made from, the closest was the tiled fp32 + # decode at 1024 pixels, 4.33 against 4.56. Its tiles overlap by 128 + # pixels, as the FLUX.2 decoder's do. + "moss_image": VaeMemoryRow( + measured=True, channels={"decode": 96, "encode": 128}, + whole={"decode": MOSS_VAE_DECODE, "encode": VAE_ENCODE}, + tiled=VAE_TILED, overlap=128, scale=8), } # FLUX.2 klein at both sizes and KV read the file the FLUX.3 klein row # was measured on. @@ -764,6 +813,20 @@ def step_peak(self, step: int, prefill: int, rows: int = 1, joined_step=FLUX2_KLEIN_9B_KV_JOINED, ref_build=FLUX2_KLEIN_9B_KV_REF_BUILD, ref_read=FLUX2_KLEIN_9B_KV_REF_READ), + # MOSS-Image joins up to 4096 text rows to the image rows in every + # block. The MLP's fused gate and up product and its gated half are the + # widest outputs: 3072 wide, an 8192 MLP. A plan that knows the prompt + # prices its text rows, and one that does not the most. An encode of + # the encoder alone, Qwen3.5-9B with its final hidden row of 4096 in + # float32, reads over the settled process 0.37 GiB at 80 rows on the + # Q4_K_M file and 0.25 on the Q8_0 one, and 1.32 GiB at 4096 rows on + # either: a base of 0.62 GiB and the float32 row times ENCODE_LIVE + # leave 0.25 GiB above every cell, the Q4_K_M one of 80 rows the + # closest. + "moss_image": PeakRow(token_px=16, row_multiple=1, caption_rows=4096, + width=2 * 8192 + 8192, step=MOSS_IMAGE_STEP, + encode_row_bytes=4096 * 4, + encode_base=int(0.62 * GiB), text_rows=4096), } def joined_row(family: str) -> PeakRow: diff --git a/gmlx/image/registry.py b/gmlx/image/registry.py index c9ca6ed67..1f1479a90 100644 --- a/gmlx/image/registry.py +++ b/gmlx/image/registry.py @@ -31,7 +31,8 @@ class SamplingRow: steps, as (level, steps) pairs; the variant's own steps are high. ``sigmas`` is a fixed schedule without its final zero, for a variant distilled to it: such a variant runs that many steps and no other - count.""" + count. ``neg_skip`` holds the blocks, zero-based, that the negative + pass of a guided step skips.""" steps: int cfg: float shift: float @@ -39,6 +40,7 @@ class SamplingRow: revision: str quality: tuple = () sigmas: tuple = () + neg_skip: tuple = () @dataclass(frozen=True) diff --git a/gmlx/image/schedule.py b/gmlx/image/schedule.py index 1c481533b..f9e71d293 100644 --- a/gmlx/image/schedule.py +++ b/gmlx/image/schedule.py @@ -16,6 +16,10 @@ # flux2_mu and flux2_sigmas follow compute_empirical_mu and get_schedule of # src/flux2/sampling.py of BFL's FLUX.2 reference at commit 50fe516 # (Apache-2.0; see licenses/flux2-LICENSE). +# SPDX-FileCopyrightText: Copyright 2026 OpenMOSS Team +# moss_sigmas follows MOSSImageScheduler.set_timesteps of +# pipeline_moss_image.py in OpenMOSS-Team/MOSS-Image-SFT at revision +# fc5de0d5 (Apache-2.0; see licenses/MOSS-Image-LICENSE). """Sigma schedules, computed on the host in numpy with float32 results. A schedule is part of a family's numerics, so it must equal the reference @@ -43,6 +47,13 @@ float32, which ends at zero, and applies the static shift with its alpha, the same for every image size. +MOSS-Image's reference shifts a float32 linspace statically and snaps +each sigma to the training grid ``k / 1000``, which can leave fewer steps +than asked for. It forms that grid with a float32 linspace too, from 1 +to ``1 / 1000``, and ``moss_sigmas`` forms the same one. ``k / 1000`` +rounded once differs from it by one ulp at 386 points, and would break a +tie between two points the other way. + FLUX.2 klein's reference forms the same float32 linspace and shifts it exponentially by a mu of the image token count and the step count. It divides Python floats by float32 tensors, which torch forms as the @@ -224,6 +235,28 @@ def flux2_sigmas(steps: int, tokens: int) -> np.ndarray: return ((_F32(1) / (e + d)) * e).astype(_F32) +def moss_sigmas(steps: int, shift: float = 5.0, + train_steps: int = 1000) -> np.ndarray: + """MOSS-Image sigmas: the static shift of a linspace from 1 to + ``shift / (train_steps + shift - 1)``, each snapped to the nearest + point of the grid ``k / train_steps``, k = 1 to ``train_steps``, the + lower one on a tie, with consecutive repeats dropped and a final zero. + A schedule can have fewer than ``steps`` steps. The grid is the + reference's float32 linspace from 1 to ``1 / train_steps``, reversed. + """ + if steps < 1: + raise RequestRefusal(f"steps must be at least 1, not {steps}") + v = linspace_f32(1.0, shift / (train_steps + shift - 1), steps) + s = shift_static(v, shift) + grid = linspace_f32(1.0, 1.0 / train_steps, train_steps)[::-1] + hi = np.minimum(np.searchsorted(grid, s, side="left"), train_steps - 1) + lo = np.maximum(hi - 1, 0) + pick = np.where(np.abs(s - grid[lo]) <= np.abs(grid[hi] - s), lo, hi) + snapped = grid[pick] + keep = np.concatenate([[True], snapped[1:] != snapped[:-1]]) + return np.append(snapped[keep], _F32(0)).astype(_F32) + + def check_sigmas(sigmas) -> np.ndarray: """An external schedule as float32, refused unless it is one dimension with at least two values, finite and not increasing.""" diff --git a/gmlx/image/text_encoder.py b/gmlx/image/text_encoder.py index 8d817debf..24204cdf8 100644 --- a/gmlx/image/text_encoder.py +++ b/gmlx/image/text_encoder.py @@ -6,6 +6,12 @@ walks the loaded model's own layer list and stops there. It changes and subclasses no class of the model. +A family whose reference reads the last hidden state after the final +norm, as MOSS-Image reads Qwen3.5, sets ``final_norm``. Its features come +from the decoder stack's own call on a fresh cache, which builds the +masks of its attention and recurrent layers, and not from ``tap``, whose +block calls pass the one causal mask that only attention layers take. + Each prompt is rendered into the family's fixed template, which the family owns and never reads from the GGUF chat template, and truncated to the family's token cap. It is encoded alone at batch 1 with no padding, so the @@ -43,14 +49,17 @@ class of the model. A text prompt taps that class by the same walk, and import contextlib import dataclasses +import functools import hashlib import importlib import json import os +import threading import unicodedata from dataclasses import dataclass import mlx.core as mx +import mlx.nn as nn import numpy as np from .fingerprint import ImageModelError, canonical_bf16 @@ -65,6 +74,17 @@ class IdentityTensor: sha256: str +# GGUF norm roles of a block and their modules under it, by encoder arch. +# output_norm is the stack's final norm. +_NORMS = { + "qwen35": (("attn_norm", "input_layernorm"), + ("post_attention_norm", "post_attention_layernorm"), + ("attn_q_norm", "self_attn.q_norm"), + ("attn_k_norm", "self_attn.k_norm"), + ("ssm_norm", "linear_attn.norm")), +} + + @dataclass(frozen=True) class EncoderSpec: name: str # "qwen3-4b" @@ -95,6 +115,32 @@ class EncoderSpec: fp32_qkv: bool = False # q, k and v cast to float32 after # their products, for the qk norm, # the rotary and the attention + final_norm: bool = False # the features are the text model's + # output after its final norm, from + # its own call; tap_layers is (layers,) + trim: bool = False # the prompt loses its leading and + # trailing whitespace, as a chat + # template's trim filter takes it + fp32_norms: bool = False # an encode runs the RMSNorms of the + # decoder stack on their F32 GGUF + # weights, each norm's output in its + # input dtype; the loaded model keeps + # its own weights + + def __post_init__(self): + if self.fp32_norms and self.arch not in _NORMS: + raise ValueError(f"{self.name}: no F32 norm map for arch " + f"{self.arch!r}") + if not self.final_norm: + return + if tuple(self.tap_layers) != (self.layers,): + raise ValueError(f"{self.name} reads after the final norm, so " + f"its tap is after layer {self.layers} alone") + if (self.drop or self.stack or self.pad_multiple or self.fp32_qkv + or self.activation is not None): + raise ValueError(f"{self.name} reads after the final norm, " + "which takes no dropped rows, stacked taps, " + "pads, float32 qkv or activation dtype") @property def template_id(self) -> str: @@ -111,6 +157,12 @@ def template_id(self) -> str: where = ("each" if self.pad_position is None else self.pad_position) text += f"\x00pad={self.pad_multiple},{self.pad_id},{where}" + if self.final_norm: + text += "\x00final_norm" + if self.trim: + text += "\x00trim" + if self.fp32_norms: + text += "\x00fp32_norms" return hashlib.sha256(text.encode()).hexdigest() @@ -174,11 +226,14 @@ def encoder_spec(family) -> EncoderSpec: def render(spec: EncoderSpec, prompt: str) -> str: - """The prompt in the family template. An empty prompt is rendered as - the spec's ``empty`` text.""" + """The prompt in the family template, without its leading and + trailing whitespace for a spec that trims. An empty prompt is rendered + as the spec's ``empty`` text.""" head, sep, tail = spec.template.partition("{}") if not sep: raise ValueError(f"template of {spec.name} has no prompt slot") + if spec.trim: + prompt = prompt.strip() return head + (prompt or spec.empty) + tail @@ -226,12 +281,19 @@ def rows_at_most(spec: EncoderSpec, prompt: str) -> int: return max(0, n - spec.drop) + pad_rows(spec, n) -def _text_stack(model): - """The embedding and the decoder layer list of a loaded text model, - or of the language model of a multimodal one.""" +def _inner(model): + """The decoder stack of a loaded text model, or of the language model + of a multimodal one.""" inner = getattr(model, "model", None) if inner is None: inner = getattr(getattr(model, "language_model", None), "model", None) + return inner + + +def _text_stack(model): + """The embedding and the decoder layer list of a loaded text model, + or of the language model of a multimodal one.""" + inner = _inner(model) embed = getattr(inner, "embed_tokens", None) layers = getattr(inner, "layers", None) if embed is None or layers is None: @@ -298,6 +360,28 @@ def tap(model, ids, *, layers: tuple, activation: str | None = None, return [out[k] for k in layers] +def final_hidden(model, ids) -> mx.array: + """The output (1, L, hidden) of the decoder stack after its final + norm, for one prompt's token ids, from the stack's own call at batch 1 + on a fresh cache, as a chat prefill starts. The call builds its + attention and recurrent masks, and the positions of a text prompt + start at 0. The call of the wrapping language model is not used: it + applies the output head, and a multimodal one keeps the positions of + its last call, which a resident chat model needs for its own next + one.""" + inner = _inner(model) + if inner is None or not callable(inner): + raise TypeError(f"{type(model).__name__} has no decoder stack to " + "call") + owner = model if hasattr(model, "make_cache") else getattr( + model, "language_model", None) + make_cache = getattr(owner, "make_cache", None) + if make_cache is None: + raise TypeError(f"{type(model).__name__} makes no cache") + x = mx.array(np.asarray(ids, np.int32).reshape(1, -1)) + return inner(x, cache=make_cache()) + + class _Float32Out: """A product whose output is cast to float32.""" @@ -335,6 +419,140 @@ def float32_qkv(model): setattr(attn, name, p) +@functools.lru_cache(maxsize=4) +def _gguf_norms(path: str, arch: str, stamp: tuple) -> tuple: + """``(module path, F32 weight)`` of each norm _NORMS maps for ``arch`` + in the GGUF at ``path``, the path under the decoder stack. ``stamp`` + is the file's size and mtime, so a changed file reads again.""" + roles = dict(_NORMS[arch]) + out = [] + for name, t in _header_view(path).tensors.items(): + parts = name.split(".") + if name == "output_norm.weight": + where = "norm" + elif (len(parts) == 4 and parts[0] == "blk" and parts[3] == "weight" + and parts[2] in roles): + where = f"layers.{parts[1]}.{roles[parts[2]]}" + else: + continue + if t.codec != "F32" or t.path is None: + raise EncoderMismatchError( + f"{path}: {name} is {t.codec}, not the F32 norm weight the " + "encoder reads") + fd = os.open(t.path, os.O_RDONLY) + try: + raw = os.pread(fd, t.nbytes, t.offset) + finally: + os.close(fd) + out.append((where, np.frombuffer(raw, np.float32).copy())) + return tuple(out) + + +def _module(inner, where: str): + obj = inner + for part in where.split("."): + if obj is None: + break + obj = obj[int(part)] if part.isdigit() else getattr(obj, part, None) + return obj + + +def read_norms(path: str, model, spec: EncoderSpec) -> dict: + """The F32 norm weights of the encoder GGUF at ``path`` by module path + under the decoder stack of ``model``. Each must round to the weight + the model holds, so they are the model's own norms.""" + st = os.stat(path) + inner = _inner(model) + out = {} + for where, a in _gguf_norms(os.path.abspath(path), spec.arch, + (st.st_size, st.st_mtime_ns)): + mod = _module(inner, where) + held = getattr(mod, "weight", None) + w = mx.array(a) + if (not isinstance(held, mx.array) or held.shape != w.shape + or not mx.array_equal(w.astype(held.dtype), held).item()): + raise EncoderMismatchError( + f"{path}: the norm at {where} is not the loaded encoder's") + out[where] = w + if not out: + raise EncoderMismatchError(f"{path}: no norm weights to read") + return out + + +class _Float32RMSNorm(nn.RMSNorm): + """An RMSNorm on an F32 weight whose output keeps its input dtype.""" + + def __init__(self, weight: mx.array, eps: float): + super().__init__(weight.shape[0], eps=eps) + self.weight = weight + + def __call__(self, x: mx.array) -> mx.array: + return mx.fast.rms_norm(x, self.weight, self.eps).astype(x.dtype) + + +# Models whose norms float32_norms has swapped: id of the decoder stack +# to [entries, norms, swaps]. The first entry swaps, the last exit +# restores, so nested or interleaved entries leave the model's own. +_NORM_SWAPS: dict = {} +_NORM_LOCK = threading.Lock() + + +def _swap_norms(inner, norms: dict) -> list: + swaps = [] + for where, w in norms.items(): + head, _, attr = where.rpartition(".") + parent = _module(inner, head) if head else inner + mod = getattr(parent, attr) + if isinstance(mod, nn.RMSNorm): + swaps.append((parent, attr, mod, _Float32RMSNorm(w, mod.eps))) + else: + swaps.append((mod, "weight", mod.weight, w)) + try: + for obj, attr, _, new in swaps: + setattr(obj, attr, new) + except BaseException: + _restore_norms(swaps) + raise + return swaps + + +def _restore_norms(swaps: list) -> None: + for obj, attr, old, _ in reversed(swaps): + setattr(obj, attr, old) + + +@contextlib.contextmanager +def float32_norms(model, norms: dict): + """Run an encode with each norm of ``norms`` on its F32 weight, and + restore the model's own norms after. An RMSNorm is swapped for one + whose output keeps its input dtype; another norm, such as a gated + one whose output takes the dtype of its input, keeps its module and + gets the weight. An entry while the model is swapped shares the swap, + which the last exit restores; one with other norms is refused.""" + inner = _inner(model) + key = id(inner) + with _NORM_LOCK: + held = _NORM_SWAPS.get(key) + if held is None: + _NORM_SWAPS[key] = [1, norms, _swap_norms(inner, norms)] + elif held[1] is not norms and ( + held[1].keys() != norms.keys() + or any(held[1][k] is not v for k, v in norms.items())): + raise RuntimeError("the encoder model already runs on other F32 " + "norms") + else: + held[0] += 1 + try: + yield + finally: + with _NORM_LOCK: + held = _NORM_SWAPS[key] + held[0] -= 1 + if not held[0]: + del _NORM_SWAPS[key] + _restore_norms(held[2]) + + class _PadRope: """A rotary embedding that rotates the last of ``n + 1`` rows at the pad position and the others as the wrapped one does.""" @@ -429,11 +647,15 @@ def vision_rope(model, spec: EncoderSpec): def _model_param(model, path: str): - obj = model - for part in path.split("."): - obj = obj[int(part)] if part.isdigit() else getattr(obj, part, None) + """The tensor at ``path`` in a loaded model. A path under ``model.`` + starts at the decoder stack, which a multimodal or wrapped text model + holds under its language model.""" + head, _, rest = path.partition(".") + obj = _inner(model) if head == "model" else getattr(model, head, None) + for part in rest.split(".") if rest else (): if obj is None: - return None + break + obj = obj[int(part)] if part.isdigit() else getattr(obj, part, None) return obj @@ -610,6 +832,9 @@ def __init__(self, model, tokenizer, spec: EncoderSpec, *, owned: bool, self.weights_owner = weights_owner self.dtype_name = spec.activation or _float_dtype(model) self.adapters = "" # text adapter state, in the cache key + # F32 norm weights an encode runs on, by module path; see + # EncoderSpec.fp32_norms. + self.norms: dict = {} @property def released(self) -> bool: @@ -636,11 +861,17 @@ def encode(self, prompt: str) -> TextTap: spec = self.spec ids = token_ids(spec, self.tokenizer, prompt) pads = pad_rows(spec, len(ids)) - with vision_rope(self.model, spec): - hidden = tap(self.model, ids, layers=spec.tap_layers, - activation=spec.activation, pad=pads, - pad_id=spec.pad_id, pad_position=spec.pad_position, - fp32_qkv=spec.fp32_qkv) + norms = (float32_norms(self.model, self.norms) if self.norms + else contextlib.nullcontext()) + with vision_rope(self.model, spec), norms: + if spec.final_norm: + hidden = [final_hidden(self.model, ids)] + else: + hidden = tap(self.model, ids, layers=spec.tap_layers, + activation=spec.activation, pad=pads, + pad_id=spec.pad_id, + pad_position=spec.pad_position, + fp32_qkv=spec.fp32_qkv) mx.eval(hidden) return TextTap(ids=tuple(ids), hidden=tuple(hidden), template_id=spec.template_id, @@ -767,20 +998,35 @@ def load_encoder(path: str, spec: EncoderSpec, *, override: bool = False, check_encoder(model, config, spec, override=True) _drop_head(model) tower = " with its vision tower" if mmproj is not None else "" - log(f"[image] text encoder {spec.name}{tower}, tap after layer " - f"{', '.join(str(n) for n in spec.tap_layers)} of {spec.layers}") + read = ("read after the final norm" if spec.final_norm else + "tap after layer " + ", ".join(str(n) for n in spec.tap_layers) + + f" of {spec.layers}") + log(f"[image] text encoder {spec.name}{tower}, {read}") shards = find_split_shards(path) key = (weights_source_key(*shards) if mmproj is None else weights_source_key(*shards, mmproj)) - return EncoderHandle(model, tokenizer, spec, owned=True, source_key=key, - weights_owner=owner, processor=processor) + handle = EncoderHandle(model, tokenizer, spec, owned=True, + source_key=key, weights_owner=owner, + processor=processor) + if spec.fp32_norms: + handle.norms = read_norms(path, model, spec) + return handle def borrow_encoder(model, tokenizer, config: dict, spec: EncoderSpec, *, - override: bool = False) -> EncoderHandle: + override: bool = False, + gguf: str | None = None) -> EncoderHandle: """Wrap a resident chat model as the encoder. It is checked against the - spec and never released.""" + spec and never released. ``gguf`` is the chat model's file, which a + spec with F32 norms reads them from; without it an encode runs on the + model's own norms.""" check_encoder(model, config, spec, override=override) digests = json.dumps(model_identity(model, spec), sort_keys=True) - key = "borrowed:" + hashlib.sha256(digests.encode()).hexdigest() - return EncoderHandle(model, tokenizer, spec, owned=False, source_key=key) + norms = (read_norms(gguf, model, spec) if spec.fp32_norms and gguf + else {}) + key = "borrowed:" + hashlib.sha256( + (digests + ("\x00fp32_norms" if norms else "")).encode()).hexdigest() + handle = EncoderHandle(model, tokenizer, spec, owned=False, + source_key=key) + handle.norms = norms + return handle diff --git a/gmlx/image/vae/flux1_ae.py b/gmlx/image/vae/flux1_ae.py index c23d7c8bf..a9434c83b 100644 --- a/gmlx/image/vae/flux1_ae.py +++ b/gmlx/image/vae/flux1_ae.py @@ -13,7 +13,9 @@ channels last (B, H, W, C). Channels per level are [128, 256, 512, 512], each down block has two resnets and each up block three, the mid block is a resnet, one attention head of 512 and a resnet, and there are no quant -convs. GroupNorm is the owned contiguous-group wrapper. +convs. GroupNorm is the owned contiguous-group wrapper. ``tree_config`` +reads the decoder's widths from its shapes, for a decoder narrower than +the encoder. ``original_name`` maps the original FLUX.1 layout, as in ``ae.safetensors``, to diffusers names, and the FLUX.2 one, which keeps its quant convs inside @@ -58,6 +60,13 @@ class Flux1AEConfig: eps: float scaling_factor: float = 1.0 # model space is (z - shift) * scaling shift_factor: float = 0.0 + decoder_block_out_channels: tuple = () # () when the encoder's + + @property + def decoder_channels(self) -> tuple: + """Channels per level of the decoder, from the full resolution + side.""" + return self.decoder_block_out_channels or self.block_out_channels @property def scale(self) -> int: @@ -69,10 +78,12 @@ def config_from_shapes(shapes: dict) -> Flux1AEConfig | None: """The config from diffusers-named torch shapes, or None when the names are not this autoencoder's. A tree with both quant convs and the BatchNorm statistics is the FLUX.2 autoencoder's; with some of them it - is this one's with tensors the load does not expect.""" + is this one's with tensors the load does not expect. The decoder is as + wide as the encoder.""" if all(n in shapes for n in FLUX2_MARKS): return None - return tree_config(shapes) + cfg = tree_config(shapes) + return None if cfg is None or cfg.decoder_block_out_channels else cfg # Tensors the FLUX.1 autoencoder lacks and the FLUX.2 one has. @@ -98,6 +109,15 @@ def tree_config(shapes: dict) -> Flux1AEConfig | None: if s is None or len(s) != 4: return None chans.append(int(s[0])) + # Up block i of diffusers counts from the latent side. + if _count(shapes, r"decoder\.up_blocks\.(\d+)\.") != levels: + return None + dec = [] + for i in reversed(range(levels)): + s = shapes.get(f"decoder.up_blocks.{i}.resnets.0.conv1.weight") + if s is None or len(s) != 4: + return None + dec.append(int(s[0])) return Flux1AEConfig( in_channels=int(shapes["encoder.conv_in.weight"][1]), out_channels=int(shapes["decoder.conv_out.weight"][0]), @@ -106,6 +126,7 @@ def tree_config(shapes: dict) -> Flux1AEConfig | None: layers_per_block=per_block, norm_num_groups=CONSTANTS["norm_num_groups"], eps=CONSTANTS["eps"], + decoder_block_out_channels=() if dec == chans else tuple(dec), ) @@ -344,7 +365,7 @@ def tail(self, x: mx.array) -> mx.array: class Decoder(nn.Module): def __init__(self, cfg: Flux1AEConfig): super().__init__() - ch = tuple(reversed(cfg.block_out_channels)) + ch = tuple(reversed(cfg.decoder_channels)) g, eps = cfg.norm_num_groups, cfg.eps self.conv_in = Conv2d(cfg.latent_channels, ch[0], 3, padding=1) self.mid_block = MidBlock(ch[0], g, eps) @@ -398,6 +419,8 @@ def dtype(self) -> mx.Dtype: def channels(self, op: str) -> int: """Channels at full resolution in ``op``, "decode" or "encode".""" + if op == "decode": + return self.config.decoder_channels[0] return self.config.block_out_channels[0] def decode(self, z: mx.array) -> mx.array: diff --git a/gmlx/image/vae/flux2_ae.py b/gmlx/image/vae/flux2_ae.py index d0f93684d..00abd2d18 100644 --- a/gmlx/image/vae/flux2_ae.py +++ b/gmlx/image/vae/flux2_ae.py @@ -19,6 +19,10 @@ ``(z - mean) / sqrt(var + eps)``. Both maps run in float32, and the statistics stay float32 whatever the decode dtype. +The decoder may run narrower than the encoder, as the MOSS-Image one +does: ``tree_config`` reads its widths from the decoder's shapes, and +``config_from_shapes`` takes only a decoder as wide as the encoder. + Module paths equal the diffusers tensor names. The original layout, as in BFL's ``ae.safetensors``, keeps the quant convs inside the encoder and the decoder, and ``flux1_ae.original_name`` maps it. @@ -57,10 +61,10 @@ def model_channels(self) -> int: return self.patch * self.patch * self.latent_channels -def config_from_shapes(shapes: dict) -> Flux2AEConfig | None: - """The config from diffusers-named torch shapes, or None when the names - are not this autoencoder's: the FLUX.1 tree with both quant convs and - the BatchNorm statistics.""" +def tree_config(shapes: dict) -> Flux2AEConfig | None: + """The config of the FLUX.1 tree with both quant convs and the + BatchNorm statistics, at the decoder widths its shapes give, or None + when the names are not that tree's.""" if any(n not in shapes for n in flux1_ae.FLUX2_MARKS): return None tree = flux1_ae.tree_config(shapes) @@ -70,11 +74,20 @@ def config_from_shapes(shapes: dict) -> Flux2AEConfig | None: return Flux2AEConfig( **{k: getattr(tree, k) for k in ( "in_channels", "out_channels", "latent_channels", - "block_out_channels", "layers_per_block")}, + "block_out_channels", "layers_per_block", + "decoder_block_out_channels")}, norm_num_groups=CONSTANTS["norm_num_groups"], eps=CONSTANTS["eps"], patch=patch[0], bn_eps=CONSTANTS["batch_norm_eps"]) +def config_from_shapes(shapes: dict) -> Flux2AEConfig | None: + """The config from diffusers-named torch shapes, or None when the names + are not this autoencoder's: the tree of ``tree_config`` with a decoder + as wide as its encoder.""" + cfg = tree_config(shapes) + return None if cfg is None or cfg.decoder_block_out_channels else cfg + + def expected_unused(cfg: Flux2AEConfig) -> dict: """The tensor a file may hold that stays unread: the batch count of the BatchNorm, an int64 scalar.""" @@ -82,10 +95,20 @@ def expected_unused(cfg: Flux2AEConfig) -> dict: def config_disagreements(config: dict) -> list: - """The fields of a vae/config.json that differ from ``CONSTANTS``. A - field the config leaves out is none.""" + """The fields of a vae/config.json that differ from ``CONSTANTS``, and + decoder widths other than the encoder's. A field the config leaves + out is none, and so are null decoder widths.""" + out = constant_disagreements(config, CONSTANTS) + dec = config.get("decoder_block_out_channels") + if dec is not None and dec != config.get("block_out_channels", dec): + out.append("decoder_block_out_channels") + return out + + +def constant_disagreements(config: dict, constants: dict) -> list: + """The fields of ``config`` that differ from ``constants``.""" out = [] - for k, v in CONSTANTS.items(): + for k, v in constants.items(): if k not in config: continue got = config[k] diff --git a/gmlx/image/vae/moss_image_ae.py b/gmlx/image/vae/moss_image_ae.py new file mode 100644 index 000000000..97ff25f4a --- /dev/null +++ b/gmlx/image/vae/moss_image_ae.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 AND MIT +# SPDX-FileCopyrightText: Copyright 2026 OpenMOSS Team +# The constants of vae/config.json in OpenMOSS-Team/MOSS-Image-SFT at +# revision fc5de0d5 (Apache-2.0; see licenses/MOSS-Image-LICENSE). +"""The MOSS-Image autoencoder: the FLUX.2 autoencoder with a narrower +decoder. + +The encoder, the quant convs, the BatchNorm statistics and the 2x2 +packing of model space are the FLUX.2 autoencoder's. The decoder runs at +[96, 192, 384, 384] channels per level from the full resolution side, +where the encoder runs at [128, 256, 512, 512], so the module is +``flux2_ae.Flux2AE`` at those decoder widths. +""" + +from __future__ import annotations + +from . import flux2_ae + +# Decoder widths from the official vae/config.json, where the other config +# values equal the FLUX.2 autoencoder's constants. +DECODER = (96, 192, 384, 384) +CONSTANTS = {**flux2_ae.CONSTANTS, "decoder_block_out_channels": DECODER} +CONSTANTS_SOURCE = ("OpenMOSS-Team/MOSS-Image-SFT", + "fc5de0d56fc910cbca6f6625a70f5f582b26deb1") + + +def config_from_shapes(shapes: dict) -> flux2_ae.Flux2AEConfig | None: + """The config from diffusers-named torch shapes, or None when the names + are not this autoencoder's: the FLUX.2 tree with the decoder widths of + ``DECODER``.""" + cfg = flux2_ae.tree_config(shapes) + if cfg is None or cfg.decoder_block_out_channels != DECODER: + return None + return cfg + + +def relative(shapes: dict) -> str | None: + """What a FLUX.2 tree holds whose decoder is neither as wide as its + encoder nor at ``DECODER``, or None.""" + cfg = flux2_ae.tree_config(shapes) + if cfg is None or cfg.decoder_block_out_channels in ((), DECODER): + return None + return (f"a FLUX.2 autoencoder with decoder widths " + f"{list(cfg.decoder_block_out_channels)}") + + +def config_disagreements(config: dict) -> list: + """The fields of a vae/config.json that differ from ``CONSTANTS``. A + field the config leaves out is none.""" + return flux2_ae.constant_disagreements(config, CONSTANTS) diff --git a/gmlx/image/vae/reader.py b/gmlx/image/vae/reader.py index fd6a65b12..19118634c 100644 --- a/gmlx/image/vae/reader.py +++ b/gmlx/image/vae/reader.py @@ -35,7 +35,8 @@ from mlx.utils import tree_flatten from ..fingerprint import ImageModelError -from . import flux1_ae, flux2_ae, qwen_image21_vae, qwen_image_vae +from . import (flux1_ae, flux2_ae, moss_image_ae, qwen_image21_vae, + qwen_image_vae) from .conv import conv_path, set_conv_path STRIP_PREFIXES = ("vae.", "first_stage_model.") @@ -112,6 +113,14 @@ def _own_space(c: dict) -> dict: "flux2_klein": _FLUX2_AE, "flux2_klein_9b": _FLUX2_AE, "flux2_klein_9b_kv": _FLUX2_AE, + # The FLUX.2 autoencoder with a narrower decoder; a FLUX.2 tree with + # other decoder widths is its relative. + "moss_image": VaeRow( + "the MOSS-Image autoencoder", moss_image_ae.config_from_shapes, + flux2_ae.Flux2AE, 32, _own_space, flux2_ae.expected_unused, + relative=moss_image_ae.relative, + config_disagreements=moss_image_ae.config_disagreements, + fp32_keep=flux2_ae.BN_STATS), } diff --git a/gmlx/serve/image_steps.py b/gmlx/serve/image_steps.py index b85ad4f30..73d8a09a0 100644 --- a/gmlx/serve/image_steps.py +++ b/gmlx/serve/image_steps.py @@ -63,7 +63,7 @@ MAX_TENSOR_BYTES = 1 << 30 MAX_ADAPTER_BYTES = 1 << 30 FEATURES = ("prefix_cache", "noise_mask", "rgba_vae", "tiled_vae", - "ref_cache") + "ref_cache", "skip_blocks") @dataclass(frozen=True) @@ -82,7 +82,8 @@ class Convention: "krea2": Convention(), "flux2_klein": Convention(time="fp32"), "flux2_klein_9b": Convention(time="fp32"), "flux2_klein_9b_kv": Convention(time="fp32"), - "flux3_klein": Convention()}, + "flux3_klein": Convention(), + "moss_image": Convention(scale=0.001)}, } # Under "reference" a step reads the timestep and returns its output as # under "comfy", and embeds the time by the DiT's reference rule, as a @@ -124,6 +125,12 @@ class Convention: "latent": {"channels": 128, "stride": 16, "patch": 1}, "vae": {"channels_out": 3, "stride": 16, "space": "model"}}, + # ComfyUI has no MOSS-Image model of its own: the pack registers + # MOSSImage, on the FLUX.2 latent format, as the OpenMOSS plugin does. + "moss_image": {"model_class": "MOSSImage", "image_model": "moss_image", + "latent": {"channels": 128, "stride": 16, "patch": 1}, + "vae": {"channels_out": 3, "stride": 16, + "space": "model"}}, } @@ -520,7 +527,8 @@ def sampling(fam, variant: str) -> dict | None: rule = SIZE_SCHEDULES.get(fam.name) return {"steps": row.steps, "cfg": row.cfg, "shift": row.shift, "sigmas": [*map(float, row.sigmas), 0.0] if row.sigmas else None, - "schedule": rule(fam) if rule and not row.sigmas else None} + "schedule": rule(fam) if rule and not row.sigmas else None, + "neg_skip": [int(i) for i in row.neg_skip]} def _torch_shape(kv: dict, name: str, ne) -> list: @@ -787,11 +795,12 @@ def fn(arrays): def run_step(served, x, sigmas, conds, table, run, gen, targets=None, time_rule: str = "reference", - ref_cache: bool = False) -> tuple: + ref_cache: bool = False, skip=()) -> tuple: """One ``denoise_step`` over the rows on the image worker, to the - ``targets`` of a flow map DiT, on the DiT's ``time_rule``, and with - ``ref_cache`` on the run's cache of reference K and V. Returns - (velocity f32 evaluated, forward seconds, prefix cache answer).""" + ``targets`` of a flow map DiT, on the DiT's ``time_rule``, with + ``ref_cache`` on the run's cache of reference K and V, and with every + row leaving out the blocks of ``skip``. Returns (velocity f32 + evaluated, forward seconds, prefix cache answer).""" import mlx.core as mx from gmlx.image import api @@ -806,7 +815,8 @@ def run_step(served, x, sigmas, conds, table, run, gen, t0 = time.perf_counter() with paced(step_pacer(run, gen)): v = api.denoise_step(pipe, x, sigmas, conds, patch_table=table, - target=targets, time=time_rule, ref_cache=held) + target=targets, time=time_rule, ref_cache=held, + skip=skip) v = v.astype(mx.float32) mx.eval(v) # Without chat the step is one evaluation, which an interrupt or an end diff --git a/gmlx/serve/image_steps_client.py b/gmlx/serve/image_steps_client.py index 559a543b6..ca09c726b 100644 --- a/gmlx/serve/image_steps_client.py +++ b/gmlx/serve/image_steps_client.py @@ -202,18 +202,22 @@ def define_table(self, model: str, body: dict) -> str: def step(self, model: str, run: str | None, x, timestep, rows: list, patch_table: str | None = None, - sample_sigmas=None, ref_cache: bool = False) -> np.ndarray: + sample_sigmas=None, ref_cache: bool = False, + skip=()) -> np.ndarray: """The model output for latents ``x`` (B, C, h, w) at ``timestep`` under the client's convention, one row dict per batch row. ``sample_sigmas`` is the schedule the sampler runs on, which a flow map DiT reads its target times from. ``ref_cache`` steps on the run's cache of reference K and V, as ComfyUI's Flux KV Cache node - samples. Records the times of the call.""" + samples. Every row leaves out the blocks of ``skip``. Records the + times of the call.""" t = np.asarray(timestep, np.float32).reshape(-1) fields = {"model": model, "run": run, "convention": self.convention, "patch_table": patch_table, "rows": rows} if ref_cache: fields["ref_cache"] = True + if skip: + fields["skip"] = [int(i) for i in skip] tensors = {"x": np.ascontiguousarray(x, np.float32), "timestep": t} if sample_sigmas is not None: tensors["sample_sigmas"] = np.asarray(sample_sigmas, @@ -230,11 +234,16 @@ def step(self, model: str, run: str | None, x, timestep, rows: list, def sample(self, model: str, noise, sigmas, rows: list, *, cfg: float = 1.0, neg_rows: list | None = None, patch_table: str | None = None, run: str | None = None, - update=None, combine=None, on_step=None) -> np.ndarray: + update=None, combine=None, on_step=None, + neg_skip=(), timestep_scale: float = 1.0) -> np.ndarray: """Sample from ``noise`` (B, C, h, w) over ``sigmas`` with one step request per sigma. With ``cfg`` other than 1 each request carries the negative rows after the positive ones and the client guides - with ``combine(pos, neg, cfg)``, by default ``guide``. + with ``combine(pos, neg, cfg)``, by default ``guide``. With + ``neg_skip``, the blocks the negative pass leaves out, the + negative rows go in a request of their own. Each timestep is its + sigma times ``timestep_scale`` in float32, the convention's + timestep for a family whose convention scales it. ``update(x, v, i)`` moves ``x`` from ``sigmas[i]`` to ``sigmas[i + 1]``; the default is Euler in float32, as ``gmlx.image.samplers.Euler`` computes it. Each request carries the @@ -247,11 +256,22 @@ def sample(self, model: str, noise, sigmas, rows: list, *, if guided and (neg_rows is None or len(neg_rows) != len(rows)): raise ValueError("guidance needs one negative row per row") step_rows = rows + (neg_rows if guided else []) + if neg_skip and not guided: + raise ValueError("neg_skip needs guidance: cfg is 1.0") + t = (s * np.float32(timestep_scale)).astype(np.float32) for i in range(len(s) - 1): - xin = np.concatenate([x, x]) if guided else x - ts = np.full((len(step_rows),), s[i], np.float32) - out = self.step(model, run, xin, ts, step_rows, patch_table, - sample_sigmas=s) + if neg_skip: + ts = np.full((b,), t[i], np.float32) + pos = self.step(model, run, x, ts, rows, patch_table, + sample_sigmas=s) + neg = self.step(model, run, x, ts, neg_rows, patch_table, + sample_sigmas=s, skip=neg_skip) + out = np.concatenate([pos, neg]) + else: + xin = np.concatenate([x, x]) if guided else x + ts = np.full((len(step_rows),), t[i], np.float32) + out = self.step(model, run, xin, ts, step_rows, patch_table, + sample_sigmas=s) v = (combine or guide)(out[:b], out[b:], cfg) if guided else out x = update(x, v, i) if update is not None else euler(x, v, s, i) if on_step is not None: diff --git a/gmlx/serve/patches/image_steps.py b/gmlx/serve/patches/image_steps.py index a133537dc..4cb691d54 100644 --- a/gmlx/serve/patches/image_steps.py +++ b/gmlx/serve/patches/image_steps.py @@ -477,6 +477,28 @@ def _table(served, tid): return core +def _skip(served, raw) -> tuple: + """The blocks a step's rows leave out, zero-based, from its ``skip`` + field.""" + from ..image_steps import StepError + + if not isinstance(raw, list) or any( + isinstance(i, bool) or not isinstance(i, int) for i in raw): + raise StepError(400, "skip is a list of block indices") + if not raw: + return () + dit = served.pipe.dit + title = served.pipe.family.title + if not getattr(dit, "skips_blocks", False): + raise StepError(400, f"the {title} DiT skips no blocks") + n = len(dit.layers) + bad = sorted({i for i in raw if not 0 <= i < n}) + if bad: + raise StepError(400, f"the {title} DiT has blocks 0 to {n - 1}, " + f"and skip names {bad}") + return tuple(sorted(set(raw))) + + async def step(request: Request): from ..image_steps import (CONVENTIONS, MAX_ROWS, STATE, Missing, StepError, flow_targets, row_hashes, run_step) @@ -527,6 +549,10 @@ async def step(request: Request): if ref_cache and parts and conv.time == "reference": raise StepError(400, f"the {title} DiT caches its references " "by itself under the reference convention") + skip = _skip(served, fields.get("skip", [])) + if skip and ref_cache: + raise StepError(400, "a step on a cache of reference K and V " + "skips no blocks") table = None try: table = _table(served, fields.get("patch_table")) @@ -538,7 +564,8 @@ async def step(request: Request): ts = np.broadcast_to(np.asarray(t_w.numpy(), np.float32).reshape(-1), (len(rows),)) - sig = [float(conv.scale * t + conv.shift) for t in ts] + # In float64, so a scaled timestep gives back its time in float32. + sig = [conv.scale * float(t) + conv.shift for t in ts] targets = None if getattr(served.pipe.dit.cfg, "flow_map", False): # The schedule is in the client's timestep convention. @@ -548,7 +575,7 @@ async def step(request: Request): "a flow map: a step body holds " "sample_sigmas, the schedule the client " "samples on") - targets = [float(conv.scale * t + conv.shift) + targets = [conv.scale * float(t) + conv.shift for t in flow_targets(s_w.numpy(), ts)] gen = run.gen if run is not None else 0 @@ -603,7 +630,7 @@ def work(): with api.state_scope(served.pipe, scope): v, fwd_s, cache = run_step(served, x, sig, conds, table, run, gen, targets, conv.time, - ref_cache) + ref_cache, skip) return _host(v), fwd_s, cache v, fwd_s, cache = await _on_worker(work) diff --git a/gmlx/serve/patches/images.py b/gmlx/serve/patches/images.py index dc0d9df2b..ab23594ac 100644 --- a/gmlx/serve/patches/images.py +++ b/gmlx/serve/patches/images.py @@ -307,6 +307,12 @@ def _borrow_for(spec): app_mod = importlib.import_module("gmlx._vendor.mlx_vlm.server.app") runtime = importlib.import_module("gmlx._vendor.mlx_vlm.server.runtime").runtime app_mod.get_cached_model(spec.encoder_model) + # The chat model's file, which an encoder with F32 norms reads them + # from. + from gmlx.serve import bridge_vlm + + rm = bridge_vlm._RESOLVED_MODELS.get(spec.encoder_model) + gguf = getattr(rm, "path", None) guard = runtime.response_generator rg = getattr(guard, "_rg", guard) hold = getattr(guard, "_hold", None) @@ -325,7 +331,7 @@ def job(engine_rg, should_stop): # steps; the encode is one step. served_pipe.encoder = api._borrowed( (_text_model(engine_rg.model), engine_rg.tokenizer), - encoder_spec(served_pipe.family)) + encoder_spec(served_pipe.family), gguf=gguf) try: return fn() finally: diff --git a/licenses/MOSS-Image-LICENSE b/licenses/MOSS-Image-LICENSE new file mode 100644 index 000000000..7baaab6ce --- /dev/null +++ b/licenses/MOSS-Image-LICENSE @@ -0,0 +1,219 @@ +MOSS-Image - https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT + +gmlx ports code from OpenMOSS-Team/MOSS-Image-SFT at revision +fc5de0d56fc910cbca6f6625a70f5f582b26deb1, modified for gmlx and MLX. Each +ported file carries the Apache-2.0 header and the copyright line of its +source. + +- gmlx/image/families/moss_image/family.py: the tensor names and shapes + of MOSSImageTransformer in pipeline_moss_image.py, its RMSNorm eps and + timestep constants, the config of transformer/config.json, and the + prompt template of tokenizer/chat_template.jinja for one user turn with + thinking on and the token cap of MOSSImagePipeline.encode_prompt. +- gmlx/image/families/moss_image/model.py: the DiT of + pipeline_moss_image.py (its joint and text blocks, the situ + feed-forward, the text projection, the RoPE bands and coordinates, the + timestep sinusoid, the shared block modulation and the output head). + + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/pyproject.toml b/pyproject.toml index ed5b15068..42c3b102e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -182,6 +182,7 @@ include = ["gmlx*"] "gmlx.image.families.flux2_klein" = ["system_message_upsampling_t2i.txt"] "gmlx.image.families.krea2" = ["expansion.txt"] "gmlx.image.families.z_image" = ["pe_prompt_template.txt"] +"gmlx.image.families.moss_image" = ["caption_prompt.txt"] "gmlx.container" = ["files/Containerfile", "guest/gmlx-entry"] gmlx = ["_vendor/misaki/LICENSE", "_vendor/misaki/data/*", "_vendor/misaki/g2pkc/*.csv", "_vendor/misaki/g2pkc/*.txt", diff --git a/scripts/image_oracle/check2_moss_image.py b/scripts/image_oracle/check2_moss_image.py new file mode 100644 index 000000000..dd4c48ad2 --- /dev/null +++ b/scripts/image_oracle/check2_moss_image.py @@ -0,0 +1,293 @@ +#!/usr/bin/env python3 +"""Check 2 oracle for MOSS-Image: the Qwen3.5-9B text encoder read after +its final norm. + +MOSSImagePipeline.encode_prompt renders each prompt with the tokenizer's +chat template for one user turn, with the generation prompt and thinking +on, tokenizes it with truncation at 4096 tokens, runs it alone through +the text model and keeps last_hidden_state: (L, 4096) per prompt, with no +pad rows. This script runs the same steps on the model that transformers +builds from the encoder GGUF itself (arch qwen35), on the CPU in float32, +with every quantized tensor dequantized at load by gguf-py in float32, +which is ggml's arithmetic (``exact_reads``). Before the run, one matrix +of each codec must agree bit for bit with the parameter that holds it, so +no transform of the load has changed it. + +The bf16 run rounds that model to bfloat16 once, as the release loads its +encoder in bfloat16, and each linear product then runs as a float32 GEMM +over the bfloat16 operands with one rounding (oracle.emulate_bf16_matmul). +The tokenizer is the release's, from --tokenizer, loaded as +Qwen2Tokenizer, the class the release's model_index.json names. The +manifest records the class and its pre-tokenizer. + +The prompts are check 2's common set, plus a Devanagari and a Thai prompt +whose combining marks split differently under a pre-tokenizer that +keeps marks with their letters and one that does not. + +For each prompt the fixture holds the token ids and the fp32 and bf16 +features. The manifest records the bf16 against fp32 gap of each prompt +and sets the tolerances: relative L2 at twice that gap, and a per-token +cosine floor of 1 - 4 (1 - c), c the oracle's own smallest per-token +cosine of the prompt. + + python scripts/image_oracle/run.py --python --threads 16 \\ + check2_moss_image.py --encoder Qwen3.5-9B-Q8_0.gguf \\ + --tokenizer /tokenizer --name check2_moss_image_q8_0 + +--layers N keeps the first N layers, for a memory ladder. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +import fixtures +import oracle +from check2_z_image import PROMPTS, token_cosine + +NAME = "check2_moss_image" +FAMILY = "moss_image" +SCRIPT = "scripts/image_oracle/check2_moss_image.py" +RULE = ("relative L2 = 2 x the oracle's own bf16 against fp32 relative L2 " + "of the same prompt; per-token cosine at least 1 - 4 (1 - c), c the " + "oracle's own smallest per-token cosine of the same prompt") +MAX_LEN = 4096 +# Combining marks: a vowel sign and a virama in Devanagari, two vowel signs +# in Thai. +MARKS = {"devanagari": "a temple in \u0928\u092e\u0938\u094d\u0924\u0947 style", + "thai": "a shop sign that reads \u0e2a\u0e27\u0e31\u0e2a\u0e14\u0e35"} + + +def exact_reads() -> None: + """Every block tensor of the GGUF dequantized by gguf-py in float32, + which is ggml's arithmetic, as transformers reads it. transformers + 5.18 without the ``kernels`` package falls back to a dense load whose + unpack step skips this file's tensors and leaves their bytes in the + parameters; with the values unpacked here it passes them through.""" + import gguf + import torch + from transformers.integrations.gguf import gguf_conversion_mapping as cm + from transformers.integrations.gguf import reader as rd + + native = rd.LazyGgufTensor.__getitem__ + unpack = cm.Dequantize.convert + + def getitem(self, key): + if self.ggml_type in rd._TORCH_DTYPE: + return native(self, key) + raw = np.copy(self.data).reshape(self.shape) + values = gguf.quants.dequantize( + raw, gguf.GGMLQuantizationType(self.ggml_type)) + return torch.from_numpy(np.ascontiguousarray(values, np.float32)) + + def convert(self, input_dict, *args, **kwargs): + if all((t[0] if isinstance(t, list) else t).dtype != torch.uint8 + for t in input_dict.values()): + return input_dict + return unpack(self, input_dict, *args, **kwargs) + + rd.LazyGgufTensor.__getitem__ = getitem + cm.Dequantize.convert = convert + + +def load_encoder(torch, path: Path, layers: int | None = None): + """The text model of the GGUF on the CPU in float32, every quantized + tensor dequantized at load. ``layers`` keeps the first N layers, for a + memory ladder.""" + from transformers import AutoConfig, AutoModel + + exact_reads() + kw: dict = dict(gguf_file=path.name, local_files_only=True) + if layers is not None: + cfg = AutoConfig.from_pretrained(str(path.parent), **kw) + kw.update(num_hidden_layers=layers, + layer_types=list(cfg.layer_types)[:layers]) + model = AutoModel.from_pretrained(str(path.parent), dtype=torch.float32, + device_map="cpu", **kw).eval() + bad = {(str(p.device), p.dtype) for p in model.parameters()} - { + ("cpu", torch.float32)} + if bad: + raise SystemExit(f"error: the encoder loaded as {sorted(map(str, bad))}" + ", not float32 on the CPU") + return model + + +# GGUF tensors the check tries, in order. A parameter holds the first three +# as they are; the others may come with their heads reordered. +_KINDS = ("ffn_down", "ffn_gate", "ffn_up", "attn_output", "ssm_out", + "ssm_alpha", "ssm_beta", "attn_gate", "attn_qkv") + + +def gguf_codecs(path: Path) -> dict: + import gguf + + out: dict = {} + for t in gguf.GGUFReader(str(path)).tensors: + out[t.tensor_type.name] = out.get(t.tensor_type.name, 0) + 1 + return out + + +def check_dequant(model, path: Path, layers: int) -> dict: + """For each quantized codec, one matrix of the first ``layers`` blocks + dequantized by gguf-py in float32 must be held by a parameter of the + same size: bit for bit, or, for a matrix the load reorders by head, + as the same multiset of float32 values. Raises when a codec of those + blocks has no such parameter. Returns the codec to the GGUF and + parameter names and how they matched.""" + import gguf + + params = [(n, p.detach().numpy()) for n, p in model.named_parameters()] + tensors = {t.name: t for t in gguf.GGUFReader(str(path)).tensors} + float_types = {"F32", "F16", "BF16"} + wanted = {t.tensor_type.name for n, t in tensors.items() + if n.startswith(tuple(f"blk.{i}." for i in range(layers))) + and t.tensor_type.name not in float_types} + done: dict = {} + for kind in _KINDS: + for i in range(layers): + t = tensors.get(f"blk.{i}.{kind}.weight") + if t is None or t.tensor_type.name not in wanted - set(done): + continue + want = np.asarray(gguf.quants.dequantize(t.data, t.tensor_type), + np.float32) + same = [n for n, a in params if a.shape == want.shape + and a.tobytes() == want.tobytes()] + if same: + done[t.tensor_type.name] = [t.name, same[0], "bitwise"] + continue + flat = np.sort(want, axis=None) + same = [n for n, a in params if a.size == want.size + and np.sort(a, axis=None).tobytes() == flat.tobytes()] + if same: + done[t.tensor_type.name] = [t.name, same[0], "reordered"] + missing = wanted - set(done) + if missing: + raise SystemExit(f"error: no parameter holds a {sorted(missing)} " + "matrix as gguf-py dequantizes it") + return done + + +def tokenizer(path: Path): + from transformers import Qwen2Tokenizer + + return Qwen2Tokenizer.from_pretrained(str(path), local_files_only=True) + + +def pre_tokenizer(tok) -> dict: + import json + + return json.loads(tok.backend_tokenizer.to_str())["pre_tokenizer"] + + +def token_ids(tok, prompt: str) -> np.ndarray: + """The pipeline's ids of one prompt.""" + text = tok.apply_chat_template( + [{"role": "user", "content": prompt}], tokenize=False, + add_generation_prompt=True, enable_thinking=True) + ids = tok(text, return_tensors="np", truncation=True, + max_length=MAX_LEN)["input_ids"][0] + return ids.astype(np.int64) + + +def features(torch, model, ids: np.ndarray) -> dict: + x = torch.from_numpy(ids[None]) + t0 = time.perf_counter() + out = model(input_ids=x, use_cache=False, return_dict=True) + secs = time.perf_counter() - t0 + return {"features": oracle.to_numpy(out.last_hidden_state[0]), + "seconds": secs} + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--encoder", required=True, type=Path, + help="Qwen3.5-9B text GGUF") + ap.add_argument("--tokenizer", required=True, type=Path, + help="the release's tokenizer folder") + ap.add_argument("--name", default=NAME) + ap.add_argument("--layers", type=int, + help="keep the first N layers, for a memory ladder; the " + "fixture name takes a _l suffix") + ap.add_argument("--prompts", nargs="*", + default=[*PROMPTS, *MARKS]) + ap.add_argument("--prompt", nargs=2, action="append", default=[], + metavar=("NAME", "TEXT"), + help="one more prompt by name, such as the prompt of " + "check 8's trajectory") + a = ap.parse_args(argv) + if a.layers is not None: + a.name = f"{a.name}_{a.layers}l" + known = {**PROMPTS, **MARKS} + prompts = {**{k: known[k] for k in a.prompts}, **dict(a.prompt)} + torch = oracle.setup() + tok = tokenizer(a.tokenizer) + ids = {k: token_ids(tok, v) for k, v in prompts.items()} + fx = fixtures.FixtureWriter(a.name, check=2, family=FAMILY, + script=SCRIPT) + fx.add_file("encoder", a.encoder) + for f in ("tokenizer.json", "tokenizer_config.json", + "chat_template.jinja"): + if (a.tokenizer / f).is_file(): + fx.add_file(f"tokenizer/{f}", a.tokenizer / f) + t0 = time.perf_counter() + model = load_encoder(torch, a.encoder, a.layers) + layers = model.config.num_hidden_layers + info = {"load_s": round(time.perf_counter() - t0, 1), "layers": layers, + "gguf_codecs": gguf_codecs(a.encoder), + "dequant_checked": check_dequant(model, a.encoder, layers)} + print(f"[check2] loaded {layers} layers in {info['load_s']}s, dequant " + f"checked {sorted(info['dequant_checked'])}", flush=True) + runs = {} + for mode in ("fp32", "bf16"): + if mode == "bf16": + # Parameters only: from_pretrained in bfloat16 keeps the + # float32 RoPE buffers. + for p in model.parameters(): + p.data = p.data.to(torch.bfloat16) + oracle.emulate_bf16_matmul(mode == "bf16") + runs[mode] = {k: features(torch, model, v) for k, v in ids.items()} + print(f"[check2] {mode}: " + f"{ {k: round(r['seconds'], 1) for k, r in runs[mode].items()} }", + flush=True) + oracle.emulate_bf16_matmul(False) + fx.inputs.update(prompts=prompts, max_sequence_length=MAX_LEN, + tokenizer_class=type(tok).__name__, + pre_tokenizer=pre_tokenizer(tok), + read="last_hidden_state, after the final norm", + template="the tokenizer's chat template, one user turn, " + "add_generation_prompt and enable_thinking", + matmul_bf16=("bf16 operands, float32 GEMM, one bf16 " + "rounding"), **info) + rel, cos, rows, secs = {}, {}, {}, {} + for k in prompts: + f32, b16 = runs["fp32"][k], runs["bf16"][k] + fx.add(f"ids_{k}", ids[k]) + fx.add(f"features_{k}", f32["features"]) + fx.add(f"features_{k}_bf16", b16["features"]) + gap = fixtures.rel_l2(b16["features"], f32["features"]) + c = token_cosine(b16["features"], f32["features"]) + rel[k] = 2 * gap + cos[k] = float(1 - 4 * (1 - c.min())) + rows[k] = int(f32["features"].shape[0]) + secs[k] = {m: round(runs[m][k]["seconds"], 1) for m in runs} + fx.results[f"bf16_vs_fp32_{k}"] = { + "rel_l2": gap, "token_cosine_min": float(c.min()), + "worst_token": int(c.argmin())} + print(f"[check2] {k}: {rows[k]} rows, bf16 vs fp32 rel L2 " + f"{gap:.3e}, min token cosine {c.min():.5f}", flush=True) + fx.inputs["rows"] = rows + fx.results["seconds"] = secs + fx.tolerances.update(rel_l2=rel, token_cosine_min=cos, rule=RULE) + fx.note("The bf16 model is the float32 model rounded to bfloat16 once, " + "as the release loads its encoder in bfloat16.") + print(f"[check2] wrote {fx.write()}", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/image_oracle/check3_moss_image.py b/scripts/image_oracle/check3_moss_image.py new file mode 100644 index 000000000..27ab4d800 --- /dev/null +++ b/scripts/image_oracle/check3_moss_image.py @@ -0,0 +1,472 @@ +#!/usr/bin/env python3 +"""Check 3 oracle for MOSS-Image: one denoise step of the release's +MOSSImageTransformer, block by block, on the weights of a DiT GGUF. + +The class lives in the release's pipeline_moss_image.py, not in +diffusers, so the script loads that file by path from --pipeline and +builds the model itself. The GGUF tensor names are the class's own +parameter names, so no name map is needed. Each quantized matrix goes in +a GGUFLinear over the file's own bytes, which dequantizes on every call +in float32 arithmetic (oracle ``dequant="exact"``), and each float tensor +is set by name in the mode's dtype: the release loads every weight in +bfloat16. The file's rope.bands is checked against the class's own +buffer, which the class rebuilds and never loads. --blocks N keeps the +first N blocks, for a memory ladder or a short fixture. + +One step is one call as MOSSImagePipeline makes it for one prompt: +latents (1, 128, H / 16, W / 16), the unpadded encoder features of the +prompt, and the scheduler's timestep, sigma times 1000 in float32 in both +modes. The velocity is the output, with no sign flip. + +For a size in --rows the fixture holds, in float32: the inputs, the +timestep embedding and the two modulation outputs, the text projection +and the output of each text block, the image embedding, the input of the +first block and the output of every block, the head and the velocity. +The bf16 model runs the same call, and each text block, each block and +the head also run teacher-forced on the fp32 model's input to them. The +manifest sets each row's tolerance to twice its bf16 against fp32 gap: +teacher-forced for those rows, free running for the others and for the +velocity. For the other sizes the fixture holds the inputs and the +velocities. + +The features are --features, an .npz with --features-key (L, 4096), as +check 2 writes them, or seeded normal features of std --feature-std. +The latents are seeded normal at --sigma, or with --latents and --step a +point on a sampling trajectory (an .npz with ``xs`` (steps, 1, 128, h, w) +and ``sigmas``). + + python scripts/image_oracle/run.py --python --threads 16 \\ + check3_moss_image.py --pipeline /pipeline_moss_image.py \\ + --dit MOSS-Image-SFT-Q4_K_M.gguf --sizes 512 --name check3_moss_image + +--self-test writes a tiny model of the class on random weights, so the +gmlx side can check its conventions on the CPU. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import re +import sys +import time +from pathlib import Path + +import numpy as np + +import fixtures +import oracle + +CLS = "MOSSImageTransformer" +FAMILY = "moss_image" +SCRIPT = "scripts/image_oracle/check3_moss_image.py" +RULE = ("tolerance = 2 x the oracle's own bf16 against fp32 relative L2: " + "teacher-forced for each text block, each block and the head, free " + "running for the other rows and for the velocity") +BANDS = "rope.bands" + +# The tiny model of --self-test: every width of the real one scaled down. +TINY = {"dim": 64, "ffn_dim": 48, "num_heads": 2, "num_layers": 2, + "text_dim": 24, "freq_dim": 16, "latent_channels": 8, + "rope_base": 10000.0} +TINY_CASE = {"latent": (3, 4), "text": 5, "sigma": 0.6875} + + +def release(path: Path): + """The release's pipeline module, loaded by path.""" + spec = importlib.util.spec_from_file_location("moss_release", path) + if spec is None or spec.loader is None: + raise SystemExit(f"error: cannot load {path}") + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + return mod + + +def gguf_state(torch, path: Path, blocks: int | None): + """The file's tensors by name, quantized ones as GGUFParameter, with + the codec census and the stored bands. Blocks from ``blocks`` on are + left out.""" + import gguf + from diffusers.quantizers.gguf.utils import GGUFParameter + + reader = gguf.GGUFReader(str(path)) + floats = (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16, + gguf.GGMLQuantizationType.BF16) + out, census, bands = {}, {}, None + for t in reader.tensors: + name = t.name + for p in ("model.diffusion_model.", "diffusion_model."): + name = name.removeprefix(p) + census[t.tensor_type.name] = census.get(t.tensor_type.name, 0) + 1 + if name == BANDS: + bands = np.array(t.data, np.float32) + continue + m = re.match(r"layers\.(\d+)\.", name) + if m and blocks is not None and int(m[1]) >= blocks: + continue + data = np.array(t.data) + if t.tensor_type == gguf.GGMLQuantizationType.BF16: + data = (data.view(np.uint16).astype(np.uint32) << 16).view( + np.float32) + if t.tensor_type in floats: + out[name] = torch.from_numpy(np.ascontiguousarray(data)) + else: + out[name] = GGUFParameter(torch.from_numpy(data), + quant_type=t.tensor_type) + return out, census, bands + + +_TYPED: dict = {} + + +def typed(dtype): + """A GGUFParameter whose ``dtype`` reads as the compute dtype. The + class casts each block's input to ``proj["q"].weight.dtype``, which on + the stored bytes would be uint8.""" + from diffusers.quantizers.gguf.utils import GGUFParameter + + if dtype not in _TYPED: + _TYPED[dtype] = type(f"GGUFParameter_{dtype}".replace(".", "_"), + (GGUFParameter,), + {"dtype": property(lambda self: dtype)}) + return _TYPED[dtype] + + +def build(torch, rel, state: dict, config: dict, mode: str): + """The class on ``state``, quantized leaves as GGUFLinear over a view + of the stored bytes typed for the mode.""" + from accelerate import init_empty_weights + from diffusers.quantizers.gguf.utils import (GGUFParameter, + _replace_with_gguf_linear) + + dtype = oracle.torch_dtype(mode) + with init_empty_weights(): + model = getattr(rel, CLS)(**config) + _replace_with_gguf_linear(model, dtype, state) + names = {n for n, _ in model.named_parameters()} + if names != set(state): + raise SystemExit(f"names: missing {sorted(names - set(state))[:4]}, " + f"extra {sorted(set(state) - names)[:4]}") + for name, v in state.items(): + owner, _, leaf = name.rpartition(".") + mod = model.get_submodule(owner) + if isinstance(v, GGUFParameter): + v = typed(dtype)(v.as_tensor(), quant_type=v.quant_type) + else: + v = torch.nn.Parameter(v.to(dtype), requires_grad=False) + mod._parameters[leaf] = v + assert not any(p.is_meta for p in model.parameters()) + bufs = dict(model.named_buffers()) + assert set(bufs) == {BANDS}, sorted(bufs) + assert bufs[BANDS].dtype == torch.float32 + return model.eval() + + +class Capture: + """Forward hooks that keep the rows of one call by gmlx module path.""" + + def __init__(self, model, blocks: bool): + self.rows: dict = {} + self.inputs: dict = {} + self.handles = [] + te = model.text_embedder + mods = [("time_embedding", model.time_embedding), + ("shared_block_modulation", model.shared_block_modulation), + ("video_head_modulation", model.video_head_modulation), + ("text_embedder.proj", te.proj), + ("video_embedder", model.video_embedder), + ("final_norm_video", model.final_norm_video), + ("final_linear_video", model.final_linear_video)] + mods += [(f"text_embedder.blocks.{i}", b) + for i, b in enumerate(te.blocks)] + if blocks: + mods += [(f"layers.{i}", b) for i, b in enumerate(model.layers)] + for name, mod in mods: + self.handles.append(mod.register_forward_hook(self._keep(name))) + self.handles.append( + mod.register_forward_pre_hook(self._feed(name))) + + def _keep(self, name): + def hook(_mod, _args, out): + self.rows[name] = out.detach() + return hook + + def _feed(self, name): + def hook(_mod, args): + self.inputs[name] = tuple(a.detach() if hasattr(a, "detach") + else a for a in args) + return hook + + def remove(self): + for h in self.handles: + h.remove() + + +def timestep(torch, sigma: float): + """The scheduler's timestep: sigma times 1000, in float32.""" + s = np.float32(sigma) * np.float32(1000) + return torch.tensor([s], dtype=torch.float32) + + +def call(torch, model, lat: np.ndarray, feats: np.ndarray, sigma: float, + mode: str, blocks: bool): + dtype = oracle.torch_dtype(mode) + hs = torch.from_numpy(lat) + ehs = torch.from_numpy(feats[None]).to(dtype) + cap = Capture(model, blocks) + try: + out = model(hs, timestep(torch, sigma), ehs, return_dict=False)[0] + finally: + cap.remove() + return oracle.to_numpy(out), cap + + +def head(torch, model, x, mod): + """The release's output head on image rows ``x`` with the head + modulation ``mod`` (1, 2 dim), as its forward runs it.""" + m = mod.view(1, 2, -1) + y = model.final_norm_video(x.to(model.final_norm_video.weight.dtype)) + y = y * (1 + m[:, 1, None]) + m[:, 0, None] + return model.final_linear_video(y.to(model.final_linear_video.weight.dtype)) + + +def teacher_forced(torch, model16, cap32: Capture) -> dict: + """The bf16 output of each text block, each block and the head fed + the fp32 model's inputs to them, rounded to bfloat16 as the bf16 run + holds them.""" + def bf(a): + if isinstance(a, torch.Tensor) and a.is_floating_point(): + return a.to(torch.bfloat16) + return a + + out = {} + for name, args in cap32.inputs.items(): + if name.startswith("layers."): + mod = model16.layers[int(name.split(".")[1])] + elif name.startswith("text_embedder.blocks."): + mod = model16.text_embedder.blocks[int(name.split(".")[2])] + else: + continue + out[name] = oracle.to_numpy(mod(*(bf(a) for a in args))) + x = cap32.inputs["final_norm_video"][0] + mod = cap32.rows["video_head_modulation"] + out["final_linear_video"] = oracle.to_numpy( + head(torch, model16, bf(x), bf(mod))) + return out + + +def run_case(torch, models: dict, lat, feats, sigma, rows: bool, name: str, + inputs: dict, files: dict, notes: list, *, check: int = 3, + script: str = SCRIPT, more=None) -> dict: + """One step in fp32 and bf16 written as fixture ``name``. ``more``, + when given, is called with the writer and the velocities before the + write.""" + fx = fixtures.FixtureWriter(name, check=check, family=FAMILY, + script=script) + fx.inputs.update(inputs) + for role, path in files.items(): + fx.add_file(role, path) + for n in notes: + fx.note(n) + fx.add("latents", lat) + fx.add("features", feats) + fx.add("sigma", np.float32(sigma)) + fx.add("timestep", oracle.to_numpy(timestep(torch, sigma))) + vel, caps, times = {}, {}, {} + for mode in ("fp32", "bf16"): + t0 = time.perf_counter() + oracle.emulate_bf16_matmul(mode == "bf16") + vel[mode], caps[mode] = call(torch, models[mode], lat, feats, sigma, + mode, rows) + times[f"{mode}_s"] = round(time.perf_counter() - t0, 1) + fx.add("velocity", vel["fp32"]) + fx.add("velocity_bf16", vel["bf16"]) + gap = fixtures.compare(vel["bf16"], vel["fp32"]) + fx.results["velocity_bf16_vs_fp32"] = gap + fx.tolerances["velocity"] = {"rel_l2": 2 * gap["rel_l2"], + "cosine_min": 1 - 4 * (1 - gap["cosine"])} + if rows: + r32 = {k: oracle.to_numpy(v) for k, v in caps["fp32"].rows.items() + if k != "final_norm_video"} + r16 = {k: oracle.to_numpy(v) for k, v in caps["bf16"].rows.items() + if k != "final_norm_video"} + oracle.emulate_bf16_matmul(True) + t0 = time.perf_counter() + tf = teacher_forced(torch, models["bf16"], caps["fp32"]) + times["bf16_teacher_forced_s"] = round(time.perf_counter() - t0, 1) + oracle.emulate_bf16_matmul(False) + fx.add("in.layers.0", + oracle.to_numpy(caps["fp32"].inputs["layers.0"][0])) + fx.add("in.final_norm_video", + oracle.to_numpy(caps["fp32"].inputs["final_norm_video"][0])) + gaps = {} + for k, a in r32.items(): + fx.add(k, a) + forced = k in tf + gaps[k] = fixtures.compare(tf[k] if forced else r16[k], a) + gaps[k]["teacher_forced"] = forced + fx.tolerances[k] = {"rel_l2": 2 * gaps[k]["rel_l2"], + "teacher_forced": forced} + fx.results["rows_bf16_vs_fp32"] = gaps + fx.results["free_bf16_vs_fp32"] = { + k: fixtures.compare(r16[k], r32[k])["rel_l2"] + for k in r32 if k.startswith("layers.")} + fx.results["times"] = times + fx.tolerances["rule"] = RULE + if more is not None: + more(fx, vel) + path = fx.write() + print(f"[check3] {name}: velocity bf16 vs fp32 rel L2 " + f"{gap['rel_l2']:.3e}, {times}, wrote {path}", flush=True) + return vel + + +def tiny_state(torch, rel, seed: int) -> tuple[dict, np.ndarray]: + """Random tiny weights by name, float32 numpy arrays in torch order, + norm weights near one, and the class's bands.""" + torch.manual_seed(seed) + model = getattr(rel, CLS)(**TINY) + rng = np.random.default_rng(seed) + out = {} + for name, p in model.named_parameters(): + shape = tuple(p.shape) + if "norm" in name.split(".")[-2]: + v = 1 + rng.standard_normal(shape) * 0.2 + else: + fan = shape[-1] if len(shape) > 1 else 16 + v = rng.standard_normal(shape) / np.sqrt(fan) + out[name] = v.astype(np.float32) + return out, oracle.to_numpy(model.rope.bands) + + +def self_test(torch, rel, a) -> int: + """The tiny model, fp32 and bf16, on random weights and inputs.""" + state, bands = tiny_state(torch, rel, a.seed) + rng = np.random.default_rng(a.seed + 1) + h, w = TINY_CASE["latent"] + lat = rng.standard_normal((1, TINY["latent_channels"], h, w)).astype( + np.float32) + feats = rng.standard_normal((TINY_CASE["text"], TINY["text_dim"])).astype( + np.float32) + torch_state = {g: torch.from_numpy(v) for g, v in state.items()} + models = {m: build(torch, rel, torch_state, TINY, m) + for m in ("fp32", "bf16")} + name = f"selftest_{FAMILY}_tiny" + fx_inputs = {"config": TINY, "case": TINY_CASE, "seed": a.seed} + run_case(torch, models, lat, feats, TINY_CASE["sigma"], True, name, + fx_inputs, {"pipeline": a.pipeline}, [ + f"the weights are in {name}_weights by name in torch " + "order, with the class's rope.bands buffer"]) + wx = fixtures.FixtureWriter(name + "_weights", check=3, family=FAMILY, + script=SCRIPT) + wx.inputs.update(fx_inputs) + for g, v in state.items(): + wx.add(g, v) + wx.add(BANDS, bands) + wx.write() + return 0 + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--pipeline", type=Path, required=True, + help="the release's pipeline_moss_image.py") + ap.add_argument("--dit", type=Path, help="MOSS-Image DiT GGUF") + ap.add_argument("--config", type=Path, + help="the release's transformer/config.json") + ap.add_argument("--blocks", type=int, help="keep the first N blocks") + ap.add_argument("--sizes", type=int, nargs="+", default=[512, 1024]) + ap.add_argument("--rows", type=int, nargs="*", default=[512], + help="sizes whose fixture holds every row") + ap.add_argument("--features", type=Path, + help=".npz with features (L, 4096)") + ap.add_argument("--features-key", default="features", + help="array of --features to read") + ap.add_argument("--features-source", default="", + help="what the features are, for the manifest") + ap.add_argument("--latents", type=Path, + help=".npz with xs and sigmas of a sampling trajectory") + ap.add_argument("--step", type=int, + help="step of --latents whose x and sigma the step takes") + ap.add_argument("--latents-source", default="", + help="what --latents is, for the manifest") + ap.add_argument("--text-len", type=int, default=48) + ap.add_argument("--feature-std", type=float, default=1.0) + ap.add_argument("--sigma", type=float, default=0.75) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--dequant", choices=oracle.DEQUANT, default="exact") + ap.add_argument("--name", default="check3_moss_image") + ap.add_argument("--self-test", action="store_true") + a = ap.parse_args(argv) + torch = oracle.setup() + oracle.set_dequant(a.dequant) + rel = release(a.pipeline) + if a.self_test: + return self_test(torch, rel, a) + if a.dit is None or a.config is None: + ap.error("--dit and --config are required") + import json + + cfg = {k: v for k, v in json.loads(a.config.read_text()).items() + if not k.startswith("_")} + t0 = time.perf_counter() + state, census, stored = gguf_state(torch, a.dit, a.blocks) + if a.blocks is not None: + cfg["num_layers"] = a.blocks + models = {m: build(torch, rel, state, cfg, m) for m in ("fp32", "bf16")} + own = oracle.to_numpy(models["fp32"].rope.bands) + if stored is None or stored.tobytes() != own.tobytes(): + raise SystemExit("error: the file's rope.bands differ from the " + "class's") + print(f"[check3] loaded {a.dit.name} in " + f"{time.perf_counter() - t0:.1f}s, codecs {census}", flush=True) + rng = np.random.default_rng(a.seed) + if a.features is not None: + feats = np.load(a.features)[a.features_key].astype(np.float32) + source = a.features_source or a.features.name + else: + feats = (rng.standard_normal((a.text_len, cfg["text_dim"])) + * a.feature_std).astype(np.float32) + source = f"seeded normal, std {a.feature_std}" + files = {"dit": a.dit, "pipeline": a.pipeline, "config": a.config} + if a.features is not None: + files["features"] = a.features + traj = None + if a.latents is not None: + if a.step is None or len(a.sizes) != 1: + ap.error("--latents takes --step and one size in --sizes") + files["latents"] = a.latents + with np.load(a.latents) as z: + x = z["xs"][a.step].astype(np.float32) + traj = (x.reshape(1, x.shape[1], *x.shape[-2:]), + float(z["sigmas"][a.step])) + if traj[0].shape[-1] * 16 != a.sizes[0]: + ap.error(f"--latents holds {traj[0].shape}, not size " + f"{a.sizes[0]}") + for size in a.sizes: + n = size // 16 + if traj is None: + lat = np.random.default_rng(a.seed + size).standard_normal( + (1, cfg["latent_channels"], n, n)).astype(np.float32) + sigma, origin = a.sigma, {"seed": a.seed + size} + else: + lat, sigma = traj + origin = {"latents": a.latents_source or a.latents.name, + "step": a.step} + inputs = {"size": size, "blocks": cfg["num_layers"], + "sigma": sigma, **origin, "features": source, + "text_rows": int(feats.shape[0]), + "dequant": oracle.dequant_mode(), + "matmul_bf16": "bf16 operands, float32 GEMM, one bf16 " + "rounding", + "timestep": "sigma x 1000 in float32 in both modes", + "gguf_codecs": census} + tag = f"_{a.blocks}b" if a.blocks is not None else "" + run_case(torch, models, lat, feats, sigma, size in a.rows, + f"{a.name}{tag}_{size}", inputs, files, []) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/image_oracle/check4_moss_image.py b/scripts/image_oracle/check4_moss_image.py new file mode 100644 index 000000000..2c91e6b48 --- /dev/null +++ b/scripts/image_oracle/check4_moss_image.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Check 4 oracle for MOSS-Image: sigmas and timesteps of the release's +MOSSImageScheduler, as MOSSImagePipeline sets them. + +The class lives in the release's pipeline_moss_image.py, which the script +loads by path from --pipeline, with the official scheduler config. For +each step count the fixture holds the scheduler's float32 sigmas with the +final zero and its timesteps, and once the training grid the sigmas snap +to. gmlx's sigmas must equal these bit for bit, with the same count. The +grid is a float32 linspace that differs from k / 1000 rounded once by one +ulp at 386 points, and on that other grid a tie can go to the other +point, a whole grid step off, with the count unchanged: at 822 steps one +sigma does. + + python scripts/image_oracle/run.py --python \\ + check4_moss_image.py --pipeline /pipeline_moss_image.py \\ + --config /scheduler/scheduler_config.json +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +from pathlib import Path + +import numpy as np + +import fixtures +import oracle + +NAME = "check4_moss_image" +# 250 and 1000 steps snap some sigmas to one grid point, so the release +# runs fewer steps than asked for. At 822 a sigma ties between two grid +# points, which a grid of k / 1000 rounded once breaks the other way. +STEPS = (50, 28, 8, 250, 1000, 822) + + +def release(path: Path): + spec = importlib.util.spec_from_file_location("moss_release", path) + if spec is None or spec.loader is None: + raise SystemExit(f"error: cannot load {path}") + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + return mod + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--pipeline", type=Path, required=True) + ap.add_argument("--config", type=Path, required=True, + help="the release's scheduler/scheduler_config.json") + ap.add_argument("--out", type=Path, + help="fixture directory (default $GMLX_IMAGE_FIXTURES)") + a = ap.parse_args(argv) + torch = oracle.setup() + rel = release(a.pipeline) + cfg = {k: v for k, v in json.loads(a.config.read_text()).items() + if not k.startswith("_")} + fx = fixtures.FixtureWriter( + NAME, check=4, family="moss_image", + script=f"scripts/image_oracle/{Path(__file__).name}", + argv=[Path(__file__).name]) + fx.add_file("pipeline", a.pipeline) + fx.add_file("scheduler_config", a.config) + counts = {} + for n in STEPS: + sched = rel.MOSSImageScheduler( + num_train_timesteps=cfg["num_train_timesteps"], + shift=cfg["shift"], training_shift=cfg["training_shift"]) + sched.set_timesteps(n) + sig = sched.sigmas.numpy().astype(np.float32) + fx.add(f"sigmas_{n}", sig) + fx.add(f"timesteps_{n}", sched.timesteps.numpy().astype(np.float32)) + counts[n] = int(sig.size - 1) + total = cfg["num_train_timesteps"] + grid = (torch.linspace(1, 1 / total, total).flip(0)).numpy() + fx.add("training_grid", grid.astype(np.float32)) + exact = (np.arange(1, total + 1) / total).astype(np.float32) + off = int(np.count_nonzero(grid != exact)) + fx.inputs.update(steps=list(STEPS), scheduler=cfg) + fx.results.update(steps_run=counts, grid_points_off_k_over_total=off, + grid_max_ulps=int(np.max(np.abs( + grid.view(np.int32) - exact.view(np.int32))))) + fx.tolerances.update( + sigmas="bitwise", count="exact", + timesteps="sigma x 1000 in float32", + rule="each sigma equal to the release's bit for bit and the same " + "number of steps") + print(f"wrote {fx.write(a.out)}; steps {counts}; grid points off " + f"k / {total}: {off}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/image_oracle/check5_flux3_klein_vae.py b/scripts/image_oracle/check5_flux3_klein_vae.py index 09b961a7f..86227b55f 100644 --- a/scripts/image_oracle/check5_flux3_klein_vae.py +++ b/scripts/image_oracle/check5_flux3_klein_vae.py @@ -32,6 +32,12 @@ so a run at 512 and 1024 pixels with float64 peaked at 32.8 GiB on an M5 Max, and the oracle runs no larger size. +MOSS-Image reads the FLUX.2 autoencoder with a narrower decoder, from +OpenMOSS-Team/MOSS-Image-SFT: ``--family moss_image``, ``--config`` a +copy of that repository and ``--name check5_moss_image_vae`` write its +fixture. Its pipeline decodes through the same BatchNorm inverse and 2x2 +unpack, which the script checks. + FLUX.2 klein reads the same autoencoder, which gmlx pulls from BFL's FLUX.2 klein 4B repository in BF16. ``--vae`` that file, ``--config`` a copy of that repository, ``--name check5_flux2_klein_vae`` and @@ -40,7 +46,8 @@ --self-test builds the class at a tiny config with seeded weights and BatchNorm statistics, which the fixture holds by diffusers name, and runs -one 48x64 image through it. +one 48x64 image through it; with ``--family moss_image`` the tiny decoder +is narrower than its encoder. python scripts/image_oracle/run.py --python --threads 16 \\ check5_flux3_klein_vae.py --config --vae flux2-vae.safetensors \\ @@ -79,6 +86,12 @@ "latent_channels": 32, "norm_num_groups": 32, "down_block_types": ["DownEncoderBlock2D"] * 3, "up_block_types": ["UpDecoderBlock2D"] * 3} +# MOSS-Image's autoencoder: the same tree with a decoder narrower than the +# encoder at every level. +TINY_MOSS = {**TINY, "block_out_channels": [64, 64, 96], + "decoder_block_out_channels": [32, 32, 64]} +TINY_NAMES = {"flux3_klein": TINY_NAME, + "moss_image": "selftest_moss_image_vae_tiny"} TINY_SEED = 29 # The BatchNorm's batch count, an int64 training counter no forward reads. COUNT = "bn.num_batches_tracked" @@ -103,6 +116,18 @@ def pipeline_maps(): return P._patchify_latents, P._unpatchify_latents +def same_maps_as_moss(torch) -> None: + """MOSSImagePipeline decodes the inverse of the BatchNorm and then + unpacks channel c * 4 + dy * 2 + dx to pixel (2y + dy, 2x + dx), which + is a pixel shuffle by 2. The FLUX.2 klein pipeline's unpatchify must + be that map, so the fixture holds for MOSS-Image.""" + _, unpatchify = pipeline_maps() + z = torch.randn(2, 128, 3, 5, dtype=torch.float64) + if not torch.equal(unpatchify(z), torch.nn.functional.pixel_shuffle(z, 2)): + raise SystemExit("error: the FLUX.2 klein unpatchify is not " + "MOSS-Image's pixel shuffle") + + def _stats(torch, vae): """The pipeline's BatchNorm mean and std, (1, 4C, 1, 1).""" mean = vae.bn.running_mean.view(1, -1, 1, 1) @@ -303,7 +328,10 @@ def tiny_weights(torch, vae) -> dict: def self_test(torch, a) -> int: from diffusers import AutoencoderKLFlux2 - vae32 = AutoencoderKLFlux2(**TINY).eval() + if a.family not in TINY_NAMES: + raise SystemExit(f"error: no tiny model for {a.family}") + tiny = TINY_MOSS if a.family == "moss_image" else TINY + vae32 = AutoencoderKLFlux2(**tiny).eval() weights = tiny_weights(torch, vae32) state = vae32.state_dict() state.update({k: torch.from_numpy(v.astype(np.float32)) @@ -311,9 +339,10 @@ def self_test(torch, a) -> int: vae32.load_state_dict(state, strict=True) tracked = int(vae32.bn.num_batches_tracked) return write( - torch, a, TINY_NAME, vae32.float().eval(), [TINY_SIZE], {}, - [Path(__file__).name, "--self-test"], - {"config": TINY, "weights_seed": TINY_SEED, + torch, a, TINY_NAMES[a.family], vae32.float().eval(), [TINY_SIZE], + {}, [Path(__file__).name, "--self-test"] + + (["--family", a.family] if a.family != "flux3_klein" else []), + {"config": tiny, "weights_seed": TINY_SEED, "weights": "every float tensor by its diffusers name, float16 " "values the model holds in float32, under weight.", "num_batches_tracked": tracked}, @@ -414,7 +443,8 @@ def main(argv=None) -> int: ap.add_argument("--sizes", type=int, nargs="*", default=[512, 1024], help="square sizes in pixels, encoded and decoded") ap.add_argument("--name", default=NAME) - ap.add_argument("--family", choices=("flux3_klein", "flux2_klein"), + ap.add_argument("--family", + choices=("flux3_klein", "flux2_klein", "moss_image"), default="flux3_klein", help="the family the manifest records, which reads the " "autoencoder") @@ -424,6 +454,8 @@ def main(argv=None) -> int: help="fixture directory (default $GMLX_IMAGE_FIXTURES)") a = ap.parse_args(argv) torch = oracle.setup() + if a.family == "moss_image": + same_maps_as_moss(torch) if a.self_test: return self_test(torch, a) if None in (a.config, a.vae): diff --git a/scripts/image_oracle/check6_moss_image.py b/scripts/image_oracle/check6_moss_image.py new file mode 100644 index 000000000..abcfb668c --- /dev/null +++ b/scripts/image_oracle/check6_moss_image.py @@ -0,0 +1,261 @@ +#!/usr/bin/env python3 +"""Check 6 oracle for MOSS-Image: a whole guided sampling run on injected +noise, to the final latents. + +The run is MOSSImagePipeline.__call__ itself, from the release's +pipeline_moss_image.py loaded by path, with the DiT of --dit built as in +check3_moss_image and the release's MOSSImageScheduler on its official +config. The call takes the encoder features of --features as +prompt_embeds and negative_prompt_embeds, each with a mask that keeps +every row: the check 2 oracle's float32 features of the plain prompt and +of the empty prompt, which is the release's default negative prompt. The +two have different lengths, and the call runs each pass at its own. It +runs at the release's defaults: 50 steps, guidance 3, block 22 left out +of the negative pass, the latents in float32 and output_type latent. A +wrapper on the scheduler's step keeps the latents after every step. The +noise is numpy's PCG64 standard normal for the seed, NCHW, which is how +gmlx makes noise from a seed. + +The fp32 run has every op in float32. The bf16 run is the same call on +the DiT in bfloat16, the features cast to bfloat16, products as one +float32 GEMM over bf16 operands with one rounding. The release turns each +velocity to float32 before the guidance, so the guidance and the Euler +steps run in float32 in both modes. + +One fixture per size holds the noise, both features, the schedule, the +fp32 latents after the steps of --keep and after the last one, and with +a bf16 run the bf16 latents at the same steps and the bf16 against fp32 +gap after every step. At GATE_SIZE the tolerance of gmlx's bf16 latents +after a step is twice that gap, so that size needs a bf16 run. At every +other size gmlx's exact fp32 run is held to EXACT_REL_L2 after each kept +step and the last, and its bf16 run is read against the GATE_SIZE gates. + + python scripts/image_oracle/run.py --python --threads 16 \\ + check6_moss_image.py --pipeline /pipeline_moss_image.py \\ + --dit MOSS-Image-SFT-Q4_K_M.gguf --config /transformer/config.json \\ + --scheduler-config /scheduler/scheduler_config.json \\ + --features check2_moss_image_q8_0.npz --sizes 512 1024 --bf16 512 +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +import numpy as np + +import check3_moss_image as c3 +import fixtures +import oracle + +NAME = "check6_moss_image" +SCRIPT = "scripts/image_oracle/check6_moss_image.py" +SEED = 6 +STEPS = 50 +GUIDANCE = 3.0 +NEG_SKIP = (22,) +KEEP = (1, 2, 5, 10, 25) +GATE_SIZE = 512 +EXACT_REL_L2 = 1e-3 +RULE = (f"latents after a kept step and after the last. At {GATE_SIZE}: " + "gmlx's bf16 run at relative L2 twice the oracle's own bf16 " + "against fp32 relative L2 after that step. At any other size: " + "gmlx's exact fp32 run against the fp32 latents at exact_rel_l2, " + f"and its bf16 run read as a share of the {GATE_SIZE} gate of the " + "step, not gated") + + +def scheduler(rel, config: dict): + return rel.MOSSImageScheduler( + num_train_timesteps=config["num_train_timesteps"], + shift=config["shift"], training_shift=config["training_shift"]) + + +def run(torch, rel, model, sched_cfg: dict, noise: np.ndarray, pos, neg, + mode: str, steps: int): + """The release's call. Returns the latents after every step, the + sigmas and timesteps of the schedule and the seconds of every + step.""" + dtype = oracle.torch_dtype(mode) + sched = scheduler(rel, sched_cfg) + pipe = rel.MOSSImagePipeline(transformer=model, scheduler=sched) + pipe.set_progress_bar_config(disable=True) + after, marks = [], [time.perf_counter()] + step = sched.step + + def keep(*args, **kw): + out = step(*args, **kw) + after.append(oracle.to_numpy(out[0])) + marks.append(time.perf_counter()) + return out + + sched.step = keep + + def embeds(f): + return torch.from_numpy(f[None]).to(dtype) + + _, _, h, w = noise.shape + with torch.no_grad(): + (lat,) = pipe(prompt_embeds=embeds(pos), + negative_prompt_embeds=embeds(neg), + height=h * 16, width=w * 16, num_inference_steps=steps, + guidance_scale=GUIDANCE, + uncond_skip_layers=list(NEG_SKIP), + latents=torch.from_numpy(noise), output_type="latent", + return_dict=False) + final = oracle.to_numpy(lat) + if not after or not np.array_equal(after[-1], final): + raise SystemExit(f"error: {mode}: the kept steps end apart from " + "the call's latents") + return (after, oracle.to_numpy(sched.sigmas).astype(np.float32), + oracle.to_numpy(sched.timesteps).astype(np.float32), + np.diff(marks)) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--pipeline", type=Path, required=True, + help="the release's pipeline_moss_image.py") + ap.add_argument("--dit", type=Path, required=True, + help="MOSS-Image DiT GGUF") + ap.add_argument("--config", type=Path, required=True, + help="the release's transformer/config.json") + ap.add_argument("--scheduler-config", type=Path, required=True, + help="the release's scheduler/scheduler_config.json") + ap.add_argument("--features", type=Path, required=True, + help=".npz with the features of both prompts") + ap.add_argument("--positive-key", default="features_plain") + ap.add_argument("--negative-key", default="features_empty") + ap.add_argument("--features-source", + default="check2_moss_image_q8_0 fp32 features") + ap.add_argument("--blocks", type=int, help="keep the first N blocks") + ap.add_argument("--sizes", type=int, nargs="+", default=[512, 1024]) + ap.add_argument("--bf16", type=int, nargs="*", default=[GATE_SIZE], + help="sizes that also run in bf16") + ap.add_argument("--steps", type=int, default=STEPS) + ap.add_argument("--keep", type=int, nargs="*", default=list(KEEP), + help="steps whose latents the fixture keeps") + ap.add_argument("--seed", type=int, default=SEED) + ap.add_argument("--dequant", choices=oracle.DEQUANT, default="exact") + ap.add_argument("--name", default=NAME) + ap.add_argument("--out", type=Path, + help="fixture directory (default $GMLX_IMAGE_FIXTURES)") + a = ap.parse_args(argv) + if any(s % 16 for s in a.sizes) or set(a.bf16) - set(a.sizes): + ap.error("sizes must be multiples of 16, and --bf16 among --sizes") + if GATE_SIZE in a.sizes and GATE_SIZE not in a.bf16: + ap.error(f"size {GATE_SIZE} sets the gates and needs a bf16 run") + torch = oracle.setup() + oracle.set_dequant(a.dequant) + rel = c3.release(a.pipeline) + cfg = {k: v for k, v in json.loads(a.config.read_text()).items() + if not k.startswith("_")} + sched_cfg = {k: v for k, v in + json.loads(a.scheduler_config.read_text()).items() + if not k.startswith("_")} + t0 = time.perf_counter() + state, census, _ = c3.gguf_state(torch, a.dit, a.blocks) + if a.blocks is not None: + cfg["num_layers"] = a.blocks + if any(i >= cfg["num_layers"] for i in NEG_SKIP): + ap.error(f"--blocks must keep block {max(NEG_SKIP)}") + models = {m: c3.build(torch, rel, state, cfg, m) + for m in (("fp32", "bf16") if a.bf16 else ("fp32",))} + print(f"[check6] loaded {a.dit.name} in " + f"{time.perf_counter() - t0:.1f}s, codecs {census}", flush=True) + with np.load(a.features) as z: + pos = z[a.positive_key].astype(np.float32) + neg = z[a.negative_key].astype(np.float32) + print(f"[check6] text rows: positive {pos.shape[0]}, negative " + f"{neg.shape[0]}", flush=True) + + for size in a.sizes: + n = size // 16 + rng = np.random.Generator(np.random.PCG64(a.seed)) + noise = rng.standard_normal((1, cfg["latent_channels"], n, n), + dtype=np.float32) + name = f"{a.name}_{size}" + fx = fixtures.FixtureWriter( + name, check=6, family="moss_image", script=SCRIPT, + argv=[Path(__file__).name, "--dit", a.dit.name, "--features", + a.features.name, "--sizes", str(size), + *(["--bf16", str(size)] if size in a.bf16 else [])]) + for role, path in (("dit", a.dit), ("pipeline", a.pipeline), + ("config", a.config), + ("scheduler_config", a.scheduler_config), + ("features", a.features)): + fx.add_file(role, path) + after, sigmas, timesteps, secs = run( + torch, rel, models["fp32"], sched_cfg, noise, pos, neg, "fp32", + a.steps) + kept = [k for k in a.keep if k < len(after)] + fx.add("noise", noise) + fx.add("features_positive", pos) + fx.add("features_negative", neg) + fx.add("sigmas", sigmas) + fx.add("timesteps", timesteps) + fx.add("latents", after[-1]) + for k in kept: + fx.add(f"latents_step{k}", after[k - 1]) + row = {"steps_run": len(after), + "fp32_step_s": [round(float(s), 1) for s in secs], + "fp32_s": round(float(secs.sum()), 1)} + print(f"[check6] {size}: fp32 {row['fp32_s']} s, {len(after)} " + f"steps, median step {np.median(secs):.1f} s", flush=True) + tol = {} + if size in a.bf16: + oracle.emulate_bf16_matmul(True) + try: + after16, _, _, secs16 = run( + torch, rel, models["bf16"], sched_cfg, noise, pos, neg, + "bf16", a.steps) + finally: + oracle.emulate_bf16_matmul(False) + fx.add("latents_bf16", after16[-1]) + for k in kept: + fx.add(f"latents_bf16_step{k}", after16[k - 1]) + by_step = [fixtures.rel_l2(x, y) for x, y in zip(after16, after)] + gap = fixtures.compare(after16[-1], after[-1]) + row.update(bf16_s=round(float(secs16.sum()), 1), + latents_bf16_vs_fp32=gap, + bf16_vs_fp32_by_step=by_step) + tol = {"latents_rel_l2": 2 * gap["rel_l2"], + **{f"step{k}_rel_l2": 2 * by_step[k - 1] for k in kept}} + print(f"[check6] {size}: bf16 {row['bf16_s']} s, final latents " + f"rel L2 {gap['rel_l2']:.3e}, cosine {gap['cosine']:.6f}", + flush=True) + if size != GATE_SIZE: + tol = {**tol, "exact_rel_l2": EXACT_REL_L2} + fx.inputs.update( + size=size, steps=a.steps, steps_run=len(after), + blocks=cfg["num_layers"], guidance_scale=GUIDANCE, + uncond_skip_layers=list(NEG_SKIP), negative_prompt="", + scheduler=sched_cfg, seed=a.seed, features=a.features_source, + positive_key=a.positive_key, negative_key=a.negative_key, + text_rows={"positive": int(pos.shape[0]), + "negative": int(neg.shape[0])}, + kept_steps=kept, gate_size=GATE_SIZE, + noise="numpy Generator(PCG64(seed)).standard_normal(shape, " + "dtype=float32), NCHW, shape (1, 128, size / 16, " + "size / 16)", + latents="(1, 128, size / 16, size / 16) float32 in model " + "space, after the step", gguf_codecs=census) + fx.results.update( + oracle={"device": "cpu", "threads": torch.get_num_threads(), + "dequant": oracle.dequant_mode(), + "bf16_matmul": "bf16 operands, float32 GEMM, one bf16 " + "rounding", "timestep": "sigma x 1000 in float32 in " + "both modes", "guidance": "float32 velocities in both " + "modes, as the release casts them"}, + **row) + fx.tolerances.update(rule=RULE, **tol) + path = fx.write(a.out) + print(f"[check6] wrote {path}", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/image_oracle/check7_moss_image_lora.py b/scripts/image_oracle/check7_moss_image_lora.py new file mode 100644 index 000000000..1ba9469bc --- /dev/null +++ b/scripts/image_oracle/check7_moss_image_lora.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python3 +"""Check 7 oracle for MOSS-Image, step half: check 3's step at 512 with a +synthetic LoRA applied. + +No community MOSS-Image adapter is known, so the script writes one from +a seed: one pair on every linear layer of three blocks and of a text +refiner block, the fused ``proj.ffn_up_gate`` and ``qkv`` among them, +and on six layers outside the blocks, at ranks 4 to 16, each with an +alpha of half its rank and random factors scaled so that its delta has +--ratio of the Frobenius norm of its dequantized weight. The default +moves a leaf by about 8 percent, as a published LoRA moves its leaves by +5 to 15: one fixed spread of the factors would move the head, whose +weights are small, by most of its output. The release's MOSSImageTransformer +names its layers as the GGUF does, so the pairs carry those names. PEFT +is not in the oracle environment, so the script applies each pair as +PEFT does when it leaves the weights unmerged: a forward hook on the +target layer adds + + strength * (alpha / rank) * linear(linear(x, down), up) + +to the layer's output, in the mode's dtype. In bf16 the two products run +as the base layers do, as float32 products over bfloat16 operands with +one bfloat16 rounding. + +The step is check 3's at 512: the latents, the features and the sigma of +the --source fixture, every row in fp32 and the bf16 gaps that set the +tolerances, with check 3's rule, so the gmlx side reads the model under a +patch table with the rows it reads the base model with. The fixture +holds the pairs as ``down.`` and ``up.``, so the gmlx +side writes the same adapter in each dialect, and the relative L2 of +each fp32 row against the same row of the source fixture. + + python scripts/image_oracle/run.py --python --threads 16 \\ + check7_moss_image_lora.py --pipeline /pipeline_moss_image.py \\ + --dit MOSS-Image-SFT-Q4_K_M.gguf --config /transformer/config.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np + +import check3_moss_image as c3 +import fixtures +import oracle + +NAME = "check7_moss_image_synthetic_512" +SOURCE = "check3_moss_image_512" +SCRIPT = "scripts/image_oracle/check7_moss_image_lora.py" +# The layers of the synthetic adapter: every linear layer of these +# blocks and text refiner blocks, and these layers outside the blocks. +BLOCKS = (0, 21, 41) +REFINERS = (1,) +TOP = ("video_embedder", "text_embedder.proj", "time_embedding.2", + "shared_block_modulation", "video_head_modulation", + "final_linear_video") +_BLOCK = ("proj.q", "proj.k", "proj.v", "proj.o", "proj.ffn_up_gate", + "proj.ffn_down") +_REFINER = ("qkv", "proj", "ffn_gate", "ffn_up", "ffn_down") + + +def synthetic_pairs(torch, model, state: dict, seed: int, + ratio: float) -> dict: + """Module name to (down, up, alpha) of the synthetic adapter, in + float32, each delta at ``ratio`` of its weight's Frobenius norm.""" + from diffusers.quantizers.gguf.utils import (GGUFParameter, + dequantize_gguf_tensor) + + names = [f"layers.{i}.{leaf}" for i in BLOCKS for leaf in _BLOCK] + names += [f"text_embedder.blocks.{i}.{leaf}" for i in REFINERS + for leaf in _REFINER] + names += list(TOP) + have = dict(model.named_modules()) + rng = np.random.default_rng(seed) + out = {} + for j, name in enumerate(names): + lin = have[name] + rank = 4 * (1 + j % 4) + alpha = rank / 2 + w = state[f"{name}.weight"] + if isinstance(w, GGUFParameter): + w = dequantize_gguf_tensor(w) + norm = float(w.to(torch.float64).norm()) + down = rng.standard_normal((rank, lin.in_features)) \ + / np.sqrt(lin.in_features) + up = rng.standard_normal((lin.out_features, rank)) / np.sqrt(rank) + delta = np.linalg.norm(alpha / rank * (up @ down)) + up *= ratio * norm / delta + out[name] = (torch.from_numpy(down.astype(np.float32)), + torch.from_numpy(up.astype(np.float32)), alpha) + return out + + +def install(torch, model, pairs: dict, strength: float, mode: str) -> list: + """A forward hook per target that adds its pair's product.""" + import torch.nn.functional as F + + dtype = oracle.torch_dtype(mode) + handles = [] + for name, (down, up, alpha) in pairs.items(): + mod = model.get_submodule(name) + scale = strength * alpha / down.shape[0] + d, u = down.to(dtype), up.to(dtype) + + def hook(_mod, args, out, d=d, u=u, scale=scale): + x = args[0].to(dtype) + return out + (F.linear(F.linear(x, d), u) * scale).to(out.dtype) + + handles.append(mod.register_forward_hook(hook)) + return handles + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + ap.add_argument("--pipeline", type=Path, required=True, + help="the release's pipeline_moss_image.py") + ap.add_argument("--dit", type=Path, required=True, + help="MOSS-Image DiT GGUF of the source fixture") + ap.add_argument("--config", type=Path, required=True, + help="the release's transformer/config.json") + ap.add_argument("--seed", type=int, default=7, + help="seed of the synthetic adapter") + ap.add_argument("--ratio", type=float, default=0.08, + help="Frobenius norm of each delta over its weight's") + ap.add_argument("--strength", type=float, default=1.0) + ap.add_argument("--source", default=SOURCE, + help="check 3 fixture whose inputs the step reads") + ap.add_argument("--dequant", choices=oracle.DEQUANT, default="exact") + ap.add_argument("--name", default=NAME) + a = ap.parse_args(argv) + torch = oracle.setup() + oracle.set_dequant(a.dequant) + rel = c3.release(a.pipeline) + root = fixtures.fixture_root() + m3 = fixtures.read_manifest(a.source, root=root) + if m3["files"]["dit"]["name"] != a.dit.name: + raise SystemExit(f"error: {a.source} read {m3['files']['dit']}") + with np.load(root / m3["npz"], allow_pickle=False) as z: + src = {k: z[k] for k in z.files} + cfg = {k: v for k, v in json.loads(a.config.read_text()).items() + if not k.startswith("_")} + state, census, _ = c3.gguf_state(torch, a.dit, None) + models = {m: c3.build(torch, rel, state, cfg, m) + for m in ("fp32", "bf16")} + pairs = synthetic_pairs(torch, models["fp32"], state, a.seed, a.ratio) + for mode, model in models.items(): + install(torch, model, pairs, a.strength, mode) + inputs = dict(m3["inputs"]) + inputs.update( + source=a.source, strength=a.strength, lora_pairs=len(pairs), + lora_ranks=sorted({int(d.shape[0]) for d, _, _ in pairs.values()}), + lora_alphas=sorted({al for _, _, al in pairs.values()}), + lora_source=f"synthetic, seed {a.seed}, each delta at {a.ratio} of " + "its weight's Frobenius norm, ranks 4 to 16, alpha " + "half the rank", + targets=sorted(pairs), + lora_pair_alphas={n: al for n, (_, _, al) in pairs.items()}, + application="unmerged: strength * (alpha / rank) * up(down(x)) " + "added to each target's output in the mode's dtype") + + def more(fx, _vel): + for name, (down, up, _) in pairs.items(): + fx.add(f"down.{name}", down.numpy()) + fx.add(f"up.{name}", up.numpy()) + fx.results["adapter_effect"] = { + k: fixtures.rel_l2(fx.arrays[k], src[k]) + for k in fx.arrays if k in src and k not in ( + "latents", "features", "sigma", "timestep", + "velocity_bf16") and not k.startswith("in.")} + + c3.run_case(torch, models, src["latents"], src["features"], + float(src["sigma"]), True, a.name, inputs, + {"dit": a.dit, "pipeline": a.pipeline, "config": a.config}, + ["the adapter is applied unmerged as PEFT applies it; the " + "gmlx side runs it under a patch table"], + check=7, script=SCRIPT, more=more) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/comfyui_pack/fixtures/manifest-moss_image.json b/tests/comfyui_pack/fixtures/manifest-moss_image.json new file mode 100644 index 000000000..b36aa2e30 --- /dev/null +++ b/tests/comfyui_pack/fixtures/manifest-moss_image.json @@ -0,0 +1 @@ +{"arch":"moss_image","count":463,"family":"moss_image","inline":{},"model":"moss-image","naming":"gguf","tensors":[{"dtype":"f16","ggml_type":"F16","name":"final_linear_video.weight","shape":[128,3072]},{"dtype":"f16","ggml_type":"F16","name":"final_norm_video.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.0.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.0.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.0.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.0.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.0.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.0.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.0.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.0.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.0.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.0.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.1.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.1.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.1.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.1.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.1.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.1.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.1.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.1.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.1.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.1.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.10.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.10.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.10.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.10.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.10.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.10.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.10.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.10.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.10.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.10.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.11.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.11.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.11.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.11.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.11.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.11.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.11.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.11.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.11.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.11.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.12.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.12.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.12.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.12.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.12.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.12.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.12.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.12.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.12.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.12.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.13.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.13.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.13.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.13.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.13.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.13.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.13.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.13.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.13.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.13.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.14.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.14.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.14.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.14.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.14.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.14.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.14.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.14.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.14.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.14.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.15.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.15.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.15.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.15.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.15.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.15.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.15.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.15.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.15.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.15.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.16.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.16.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.16.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.16.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.16.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.16.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.16.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.16.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.16.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.16.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.17.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.17.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.17.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.17.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.17.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.17.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.17.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.17.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.17.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.17.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.18.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.18.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.18.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.18.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.18.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.18.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.18.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.18.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.18.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.18.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.19.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.19.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.19.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.19.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.19.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.19.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.19.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.19.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.19.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.19.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.2.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.2.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.2.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.2.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.2.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.2.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.2.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.2.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.2.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.2.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.20.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.20.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.20.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.20.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.20.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.20.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.20.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.20.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.20.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.20.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.21.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.21.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.21.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.21.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.21.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.21.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.21.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.21.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.21.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.21.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.22.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.22.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.22.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.22.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.22.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.22.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.22.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.22.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.22.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.22.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.23.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.23.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.23.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.23.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.23.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.23.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.23.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.23.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.23.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.23.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.24.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.24.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.24.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.24.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.24.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.24.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.24.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.24.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.24.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.24.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.25.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.25.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.25.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.25.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.25.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.25.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.25.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.25.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.25.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.25.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.26.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.26.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.26.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.26.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.26.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.26.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.26.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.26.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.26.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.26.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.27.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.27.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.27.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.27.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.27.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.27.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.27.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.27.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.27.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.27.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.28.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.28.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.28.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.28.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.28.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.28.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.28.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.28.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.28.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.28.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.29.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.29.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.29.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.29.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.29.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.29.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.29.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.29.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.29.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.29.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.3.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.3.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.3.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.3.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.3.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.3.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.3.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.3.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.3.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.3.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.30.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.30.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.30.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.30.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.30.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.30.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.30.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.30.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.30.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.30.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.31.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.31.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.31.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.31.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.31.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.31.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.31.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.31.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.31.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.31.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.32.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.32.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.32.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.32.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.32.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.32.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.32.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.32.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.32.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.32.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.33.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.33.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.33.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.33.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.33.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.33.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.33.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.33.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.33.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.33.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.34.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.34.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.34.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.34.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.34.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.34.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.34.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.34.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.34.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.34.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.35.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.35.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.35.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.35.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.35.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.35.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.35.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.35.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.35.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.35.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.36.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.36.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.36.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.36.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.36.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.36.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.36.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.36.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.36.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.36.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.37.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.37.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.37.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.37.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.37.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.37.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.37.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.37.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.37.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.37.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.38.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.38.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.38.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.38.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.38.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.38.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.38.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.38.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.38.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.38.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.39.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.39.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.39.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.39.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.39.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.39.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.39.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.39.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.39.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.39.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.4.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.4.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.4.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.4.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.4.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.4.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.4.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.4.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.4.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.4.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.40.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.40.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.40.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.40.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.40.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.40.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.40.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.40.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.40.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.40.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.41.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.41.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.41.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.41.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.41.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.41.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.41.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.41.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.41.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.41.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.5.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.5.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.5.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.5.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.5.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.5.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.5.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.5.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.5.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.5.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.6.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.6.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.6.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.6.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.6.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.6.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.6.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.6.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.6.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.6.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.7.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.7.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.7.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.7.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.7.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.7.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.7.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.7.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.7.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.7.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.8.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.8.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.8.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.8.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.8.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.8.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.8.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.8.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.8.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.8.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.9.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"layers.9.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.9.norm2.weight","shape":[3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.9.proj.ffn_down.weight","shape":[3072,8192]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.9.proj.ffn_up_gate.weight","shape":[16384,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.9.proj.k.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.9.proj.o.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q4_K","name":"layers.9.proj.q.weight","shape":[3072,3072]},{"dtype":"quant","ggml_type":"Q6_K","name":"layers.9.proj.v.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"layers.9.q_norm.weight","shape":[128]},{"dtype":"f32","ggml_type":"F32","name":"rope.bands","shape":[16]},{"dtype":"f16","ggml_type":"F16","name":"shared_block_modulation.bias","shape":[18432]},{"dtype":"f16","ggml_type":"F16","name":"shared_block_modulation.weight","shape":[18432,3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.ffn_down.bias","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.ffn_down.weight","shape":[3072,8192]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.ffn_gate.bias","shape":[8192]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.ffn_gate.weight","shape":[8192,3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.ffn_up.bias","shape":[8192]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.ffn_up.weight","shape":[8192,3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.norm2.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.proj.bias","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.proj.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.qkv.bias","shape":[9216]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.0.qkv.weight","shape":[9216,3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.ffn_down.bias","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.ffn_down.weight","shape":[3072,8192]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.ffn_gate.bias","shape":[8192]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.ffn_gate.weight","shape":[8192,3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.ffn_up.bias","shape":[8192]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.ffn_up.weight","shape":[8192,3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.k_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.norm1.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.norm2.weight","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.proj.bias","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.proj.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.q_norm.weight","shape":[128]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.qkv.bias","shape":[9216]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.blocks.1.qkv.weight","shape":[9216,3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.proj.bias","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"text_embedder.proj.weight","shape":[3072,4096]},{"dtype":"f16","ggml_type":"F16","name":"time_embedding.0.bias","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"time_embedding.0.weight","shape":[3072,256]},{"dtype":"f16","ggml_type":"F16","name":"time_embedding.2.bias","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"time_embedding.2.weight","shape":[3072,3072]},{"dtype":"f16","ggml_type":"F16","name":"video_embedder.bias","shape":[3072]},{"dtype":"f16","ggml_type":"F16","name":"video_embedder.weight","shape":[3072,128]},{"dtype":"f16","ggml_type":"F16","name":"video_head_modulation.bias","shape":[6144]},{"dtype":"f16","ggml_type":"F16","name":"video_head_modulation.weight","shape":[6144,3072]}]} \ No newline at end of file diff --git a/tests/comfyui_pack/headless.py b/tests/comfyui_pack/headless.py index 8140596cb..f2c1f9219 100755 --- a/tests/comfyui_pack/headless.py +++ b/tests/comfyui_pack/headless.py @@ -198,6 +198,32 @@ async def run(self, prompt: dict, on_progress=None, timeout=300.0) -> dict: res["messages"] = sorted(set(res["messages"])) return res + async def run_queued(self, prompts: list, timeout=300.0) -> list: + """Queue every prompt at once, then wait for each to end.""" + pids = [] + for p in prompts: + async with self.http.post(self.c.url + "/prompt", + json={"prompt": p, "client_id": self.client_id}) as r: + pids.append((await r.json())["prompt_id"]) + res = {pid: {"prompt_id": pid, "status": None} for pid in pids} + while any(r["status"] is None for r in res.values()): + msg = await self.ws.receive(timeout=timeout) + if msg.type != aiohttp.WSMsgType.TEXT: + if msg.type == aiohttp.WSMsgType.BINARY: + continue + raise RuntimeError(f"websocket closed: {msg.type}") + d = json.loads(msg.data) + data = d.get("data") or {} + r = res.get(data.get("prompt_id")) + if r is None: + continue + if d.get("type") == "execution_success": + r["status"] = "success" + elif d.get("type") == "execution_error": + r["status"] = "error" + r["error"] = data.get("exception_message") + return [res[pid] for pid in pids] + async def interrupt(self): async with self.http.post(self.c.url + "/interrupt", json={}) as r: await r.read() @@ -412,6 +438,50 @@ async def flux3_t2i(self, upsample=False): self.check("the steps sample on the caption", enc and conds == {enc[0]["cond"]}, [conds, enc[:1]]) + async def moss_t2i(self): + """The example workflow: block 22 left out of each negative pass, found + by the walk-back from the running KSampler, on ComfyUI's sgm_uniform + timesteps after ModelSamplingSD3 at shift 5.""" + wf = workflow("moss-image-t2i.json") + wf["6"]["inputs"].update(width=64, height=48) + wf["9"]["inputs"]["steps"] = 2 + r = await self.s.run(wf) + log = self.stub.log() + self.check("success", r["status"] == "success", r.get("error")) + enc = kinds(log, "encode") + self.check("the prompt and the empty negative", sorted(e["text"] for e in enc) + == ["", wf["4"]["inputs"]["text"]], [e["text"] for e in enc]) + steps = kinds(log, "step") + self.check("a positive and a negative call a step", len(steps) == 4 and + {e["rows"] for e in steps} == {1}, [e["rows"] for e in steps]) + self.check("block 22 left out of the negative pass alone", + [e["skip"] for e in steps] == [None, [22]] * 2, [e["skip"] for e in steps]) + self.check("model space latent", {tuple(e["x_shape"]) for e in steps} == {(1, 128, 3, 4)}, + [e["x_shape"] for e in steps][:1]) + t = np.linspace(1.0, 5 * 0.001 / (1 + 4 * 0.001), 3)[:-1] + want = [float(v) * 1000 for v in 5 * t / (1 + 4 * t) for _ in (0, 1)] + got = [e["timestep"][0] for e in steps] + self.check("timesteps of sgm_uniform at shift 5", np.allclose(got, want, atol=1e-3), + {"got": got, "want": want}) + self.check("decoded", len(kinds(log, "vae_decode")) == 1, len(kinds(log, "vae_decode"))) + + async def moss_queued(self): + """A prompt queued behind the running one, with other single_layers, + changes neither run: each takes the skips of its own graph.""" + a = workflow("moss-image-t2i.json") + a["6"]["inputs"].update(width=64, height=48) + a["9"]["inputs"].update(steps=2, seed=7) + b = copy.deepcopy(a) + b["9"]["inputs"]["seed"] = 8 + b["8"]["inputs"]["single_layers"] = "21" + ra, rb = await self.s.run_queued([a, b]) + self.check("the running prompt", ra["status"] == "success", ra.get("error")) + self.check("the queued prompt", rb["status"] == "success", rb.get("error")) + steps = kinds(self.stub.log(), "step") + self.check("each run skips its own blocks", + [e["skip"] for e in steps] == [None, [22]] * 2 + [None, [21]] * 2, + [e["skip"] for e in steps]) + async def flux3_edit(self, upsample=False): wf = workflow("flux-3-klein-edit-two-references.json") wf["8"]["inputs"]["resolution.megapixels"] = 0.25 @@ -604,6 +674,8 @@ async def amain(a) -> dict: await c.case("flux-2-klein edit with two references", c.flux2_edit) await c.case("FluxKontextMultiReferenceLatentMethod refused on flux-2-klein, naming it", c.flux2_reference_method) + await c.case("moss-image t2i under SkipLayerGuidanceDiTSimple", c.moss_t2i) + await c.case("moss-image: a prompt queued behind the running one", c.moss_queued) if comfy_version(a.comfy) >= (0, 41, 0): await c.case("flux-3-klein t2i", c.flux3_t2i) await c.case("flux-3-klein edit with two references", c.flux3_edit) diff --git a/tests/comfyui_pack/stub_server.py b/tests/comfyui_pack/stub_server.py index ac16c4aa0..3804eff89 100755 --- a/tests/comfyui_pack/stub_server.py +++ b/tests/comfyui_pack/stub_server.py @@ -63,7 +63,13 @@ "latent": {"channels": 128, "stride": 16, "patch": 1}, "vae": {"channels_out": 3, "stride": 16, "space": "model"}, "edit": True, "context": ("txt_in.weight", 1)}, + "moss_image": {"model_class": "MOSSImage", "image_model": "moss_image", + "latent": {"channels": 128, "stride": 16, "patch": 1}, + "vae": {"channels_out": 3, "stride": 16, "space": "model"}, "edit": False, + "context": ("text_embedder.proj.weight", 1)}, } +# The families whose DiT leaves out the blocks a step names. +SKIPS = {"moss_image"} TURBO_SIGMAS = [1.0, 0.978453, 0.95418, 0.926626, 0.89508, 0.845148, 0.704534, 0.414568, 0.0] # The schedule rules of the server's status, for the families whose # schedule follows the image size. @@ -95,6 +101,8 @@ "sigmas": None, "schedule": None}), "flux-3-klein-base": ("flux3_klein", "base", {"steps": 50, "cfg": 3.5, "shift": 6.93, "sigmas": None, "schedule": None}), + "moss-image": ("moss_image", "sft", {"steps": 50, "cfg": 3.0, "shift": 5.0, "sigmas": None, + "schedule": None, "neg_skip": [22]}), } # The manifest of a model whose file is not its family's fixture. FIXTURE = {"flux-3-klein-base": "flux3_klein-base"} @@ -284,7 +292,8 @@ def do_GET(self): **contract, "gmlx_version": st.gmlx_version, "wire": {"container": "safetensors", "dtypes": ["f32", "f16", "bf16", "u8"]}, "families": sorted(COMFY), "conventions": ["comfy"], "patch_kinds": ["lora"], - "features": ["prefix_cache", "noise_mask", "rgba_vae", "tiled_vae", "ref_cache"], + "features": ["prefix_cache", "noise_mask", "rgba_vae", "tiled_vae", "ref_cache", + "skip_blocks"], "limits": {"max_tensor_bytes": 1 << 30, "max_adapter_bytes": 1 << 30, "max_rows": 16}}) if sub == "/models": return self._json(200, {"models": [{"model": m, "family": f, "variant": v, "resident": True} @@ -498,6 +507,10 @@ def _step(self, body, t0): if missing: st.event(kind="step_missing", run=rid, missing=len(set(missing))) return self._err(409, "not held", missing=sorted(set(missing))) + skip = req.get("skip") + if skip is not None and (MODELS[req["model"]][0] not in SKIPS or not isinstance(skip, list) + or not all(isinstance(i, int) for i in skip)): + return self._err(400, f"{req['model']!r} skips no blocks") sched = tensors.get("sample_sigmas") targets = None if st.flow_map(req["model"]): @@ -531,7 +544,7 @@ def _step(self, body, t0): for r in rows], image_slots=[r.get("image_slots") for r in rows], prefix_cache=[r.get("prefix_cache", True) for r in rows], - ref_cache=req.get("ref_cache", False), + ref_cache=req.get("ref_cache", False), skip=skip, row_keys=sorted({k for r in rows for k in r}), patch_table=req.get("patch_table"), fields=sorted(req), tensors=sorted(tensors), targets=targets, sample_sigmas=None if sched is None else [float(v) for v in sched[3].reshape(-1)]) diff --git a/tests/comfyui_pack/support.py b/tests/comfyui_pack/support.py index 60c94ae5f..6cb6d2cd4 100644 --- a/tests/comfyui_pack/support.py +++ b/tests/comfyui_pack/support.py @@ -221,6 +221,35 @@ def write_flux2_lora(path: str, manifest: dict, rank: int = 4, alpha: float = 2. return [layer for _, layer, _ in FLUX2_PAIRS] +# The layers of write_moss_lora, inside the blocks and out: the fused +# feed-forward input and the text refiner's qkv among them. +MOSS_STEMS = ("layers.0.proj.q", "layers.0.proj.ffn_up_gate", "layers.41.proj.ffn_down", + "text_embedder.blocks.1.qkv", "text_embedder.proj", "time_embedding.2", + "shared_block_modulation", "video_head_modulation", "video_embedder", + "final_linear_video") + + +def write_moss_lora(path: str, manifest: dict, rank: int = 4, alpha: float = 2.0, + seed: int = 0) -> list: + """A small MOSS-Image LoRA on the GGUF's names, which are the training + names: the blocks in the kohya ``lora_unet_`` form, the rest in the + dotted ``diffusion_model`` form. Returns the target layer names.""" + import torch + from safetensors.torch import save_file + shapes = {t["name"]: t["shape"] for t in manifest["tensors"]} + g = torch.Generator().manual_seed(seed) + sd = {} + for stem in MOSS_STEMS: + dout, dn = shapes[f"{stem}.weight"] + k = ("lora_unet_" + stem.replace(".", "_")) if stem.startswith("layers.") \ + else "diffusion_model." + stem + sd[f"{k}.lora_down.weight"] = (torch.randn(rank, dn, generator=g) * 0.01).to(torch.bfloat16) + sd[f"{k}.lora_up.weight"] = (torch.randn(dout, rank, generator=g) * 0.01).to(torch.bfloat16) + sd[f"{k}.alpha"] = torch.tensor(alpha) + save_file(sd, path) + return list(MOSS_STEMS) + + def manifest(family: str) -> dict: with open(os.path.join(FIXTURES, f"manifest-{family}.json")) as f: return json.load(f) diff --git a/tests/comfyui_pack/test_client.py b/tests/comfyui_pack/test_client.py index 796ca525a..a6de2bfc3 100644 --- a/tests/comfyui_pack/test_client.py +++ b/tests/comfyui_pack/test_client.py @@ -43,7 +43,7 @@ def test_every_versioned_request_names_the_contract(self): s, _, d = self.c.request("GET", C.PREFIX + "/models", versioned=False) self.assertEqual(s, 400) self.assertIn(b"X-Gmlx-Steps-Contract", d) - self.assertEqual(len(self.c.models()), 11) + self.assertEqual(len(self.c.models()), 12) def test_status_and_manifest(self): st = self.c.status("z-image-turbo") diff --git a/tests/comfyui_pack/test_comfy.py b/tests/comfyui_pack/test_comfy.py index 28f559168..bdba4e661 100644 --- a/tests/comfyui_pack/test_comfy.py +++ b/tests/comfyui_pack/test_comfy.py @@ -17,6 +17,7 @@ import tempfile import threading import time +import types import unittest import uuid @@ -28,6 +29,8 @@ support.setup_comfy() P = support.load_pack() +import comfy.latent_formats # noqa: E402 +import comfy.model_base # noqa: E402 import comfy.model_management as mm # noqa: E402 import comfy.model_patcher # noqa: E402 import comfy.patcher_extension # noqa: E402 @@ -47,6 +50,7 @@ F29ID, F29BID = "flux-2-klein-9b", "flux-2-klein-base-9b" F29KID = "flux-2-klein-9b-kv" F3ID, F3BID = "flux-3-klein", "flux-3-klein-base" +MID = "moss-image" STEPS = 4 LORAS = tempfile.mkdtemp(prefix="gmlx-pack-loras-") folder_paths.add_model_folder_path("loras", LORAS) @@ -60,6 +64,9 @@ support.write_flux2_lora(os.path.join(LORAS, LORA_F2), support.manifest("flux2_klein"))} LORA_F29 = "flux2-9b-test-lora.safetensors" support.write_flux2_lora(os.path.join(LORAS, LORA_F29), support.manifest("flux2_klein_9b")) +LORA_M = "moss-test-lora.safetensors" +LORA_M_KEYS = {f"diffusion_model.{s}.weight" for s in + support.write_moss_lora(os.path.join(LORAS, LORA_M), support.manifest("moss_image"))} # FLUX.3 came to ComfyUI after the pinned checkout; there its models are # refused by name. HAS_FLUX3 = hasattr(comfy.supported_models, "Flux3") @@ -1185,6 +1192,240 @@ def test_a_typed_negative_samples_on_the_prompt_as_typed(self): self.assertEqual(got, {ep["cond"], en["cond_typed"]}) +def slg_graph(*nodes) -> dict: + """A submitted graph: the gmlx loader, ModelSamplingSD3 at shift 5, the + SkipLayerGuidanceDiTSimple nodes of ``nodes`` in a chain, each the + widgets it takes, and the KSampler, node 9.""" + g = {"1": {"class_type": "GmlxModelLoader", "inputs": {"model_name": MID}}, + "2": {"class_type": "ModelSamplingSD3", "inputs": {"model": ["1", 0], "shift": 5.0}}} + up = "2" + for i, widgets in enumerate(nodes): + nid = str(3 + i) + g[nid] = {"class_type": "SkipLayerGuidanceDiTSimple", + "inputs": {"model": [up, 0], "double_layers": "", "single_layers": "22", + "start_percent": 0.0, "end_percent": 1.0, **widgets}} + up = nid + g["9"] = {"class_type": "KSampler", "inputs": {"model": [up, 0], "cfg": 3.0}} + return {"prompt": g} + + +class MossImage(Base): + """MOSS-Image, which ComfyUI has no class for: the pack's weightless model + with the OpenMOSS plugin's settings, its LoRA key map, and block 22 left + out of the negative pass under SkipLayerGuidanceDiTSimple.""" + + def running(self, graph: dict): + """The pack sees ``graph`` as the prompt whose KSampler runs, as + ComfyUI's queue holds it.""" + from unittest import mock + item = (0, "p1", graph["prompt"], {}, ["9"], []) + queue = types.SimpleNamespace(get_current_queue_volatile=lambda: ([item], [])) + fake = types.SimpleNamespace(PromptServer=types.SimpleNamespace( + instance=types.SimpleNamespace(last_node_id="9", last_prompt_id="p1", + prompt_queue=queue))) + return mock.patch.dict(sys.modules, {"server": fake}) + + def slg(self, model, single="22", double="", start=0.0, end=1.0): + from comfy_extras.nodes_model_advanced import ModelSamplingSD3 + from comfy_extras.nodes_slg import SkipLayerGuidanceDiTSimple + m = ModelSamplingSD3().patch(model, 5.0)[0] + return SkipLayerGuidanceDiTSimple.execute(m, start, end, double, single)[0] + + def test_moss_image_builds_on_meta_with_the_plugins_settings(self): + gc.collect() + r0 = rss() + m = self.model(MID) + gc.collect() + grew = (rss() - r0) / 2 ** 20 + cfg = m.model.model_config + self.assertEqual(type(cfg).__name__, "MOSSImageConfig") + self.assertEqual(type(m.model).__name__, "MOSSImage") + self.assertEqual(m.model.model_type, comfy.model_base.ModelType.FLOW) + self.assertIsInstance(m.model.latent_format, comfy.latent_formats.Flux2) + ms = m.get_model_object("model_sampling") + self.assertEqual((float(ms.shift), float(ms.multiplier)), (1.0, 1000.0)) + self.assertEqual(m.model.diffusion_model.blocks, 42) + names = {t["name"] for t in self.client.manifest(MID)["tensors"]} + sd = m.model.diffusion_model.state_dict() + self.assertEqual(set(sd), names) + self.assertTrue(all(p.device.type == "meta" for p in m.model.parameters())) + self.assertLess(grew, 200.0) + self.assertEqual(m.model.get_dtype(), torch.float32) + + def test_a_manifest_that_is_no_moss_image_dit_is_refused_by_name(self): + man = self.client.manifest(MID) + cut = {**man, "tensors": [t for t in man["tensors"] + if t["name"] != "layers.0.proj.ffn_up_gate.weight"]} + with self.assertRaises(C.GmlxError) as e: + P.model.build_weightless(cut, "MOSSImage") + self.assertIn("not a MOSS-Image DiT: it has no layers.0.proj.ffn_up_gate.weight", + str(e.exception)) + + def test_a_moss_lora_on_its_training_names_becomes_a_table(self): + """ComfyUI's generic key map reads the DiT's own names under the kohya + and dotted forms; the fused feed-forward input and the refiner's qkv + take their pairs whole.""" + m = comfy_nodes.LoraLoaderModelOnly().load_lora_model_only(self.model(MID), LORA_M, 0.8)[0] + self.assertEqual(set(m.patches), LORA_M_KEYS) + self.assertEqual({t[3] for v in m.patches.values() for t in v}, {None}) + tid = m.sync_patches() + table, _ = P.patches.build_table(MID, m.patches, family="moss_image") + self.assertEqual(tid, C.table_id(table)[0]) + (put,) = self.stub.events("put_table") + self.assertEqual((put["id"], put["entries"], put["bands"], put["strengths"]), + (tid, len(LORA_M_KEYS), 0, [0.8])) + self.assertEqual(len(self.stub.events("put_adapter")), 2 * len(LORA_M_KEYS)) + + def test_a_guided_sample_leaves_block_22_out_of_the_negative_pass_alone(self): + m = self.slg(self.model(MID)) + with self.running(slg_graph({})): + out = self.sample(m, MID, cfg=3.0, steps=3) + steps = self.stub.events("step") + self.assertEqual(len(steps), 2 * 3) + self.assertEqual({s["rows"] for s in steps}, {1}) + self.assertEqual([s["skip"] for s in steps], [None, [22]] * 3) + self.assertEqual(["skip" in s["fields"] for s in steps], [False, True] * 3) + self.assertEqual({tuple(s["x_shape"]) for s in steps}, {(1, 128, 4, 4)}) + sig = comfy.samplers.calculate_sigmas(m.get_model_object("model_sampling"), "simple", 3) + want = [float(np.float32(v) * np.float32(1000.0)) for v in sig[:-1].tolist() for _ in (0, 1)] + np.testing.assert_allclose([s["timestep"][0] for s in steps], want, rtol=1e-6) + self.assertGreater(want[0] / 1000.0, 0.99) + self.assertTrue(bool(torch.isfinite(out).all())) + # Without the node both passes run whole, in one call. + self.stub.reset() + self.sample(self.model(MID), MID, cfg=3.0, steps=1) + (s,) = self.stub.events("step") + self.assertEqual((s["rows"], s["skip"]), (2, None)) + + def test_skip_layer_guidance_refusals_name_each_cause(self): + base = self.model(MID) + node = "the SkipLayerGuidanceDiTSimple node 3" + cases = { + "double layers": (dict(double="7"), slg_graph({"double_layers": "7"}), + f"{node} skips double_layers [7]"), + "window": (dict(start=0.0, end=0.5), slg_graph({"end_percent": 0.5}), + f"{node} skips its blocks from 0.0 to 0.5 of the run"), + "other blocks than the widget": ({}, slg_graph({"single_layers": "21"}), + "this call skips single blocks [22], which are not " + f"the single_layers of {node}"), + "linked widget": ({}, slg_graph({"single_layers": ["12", 0]}), + f"{node} takes single_layers from another node"), + "past the last block": (dict(single="42"), slg_graph({"single_layers": "42"}), + f"{node} skips blocks [42], and the DiT has blocks 0 to 41"), + "no node": ({}, {"prompt": {"9": {"class_type": "KSampler", "inputs": {}}}}, + "from a node other than SkipLayerGuidanceDiTSimple"), + "two nodes that differ": ({}, slg_graph({}, {"single_layers": "21"}), + "the SkipLayerGuidanceDiTSimple nodes 3, 4 skip other blocks"), + } + for name, (kw, graph, want) in cases.items(): + with self.subTest(name): + m = self.slg(base, **kw) + with self.running(graph), self.assertRaises(C.GmlxRefusal) as e: + self.sample(m, MID, cfg=3.0, steps=1) + self.assertIn(want, str(e.exception)) + self.assertEqual([s["skip"] for s in self.stub.events("step")], [None] * len(cases)) + + def test_a_skip_the_server_or_the_call_cannot_take_is_refused(self): + f = self.model(MID).model.diffusion_model.forward + cond = self.cond(MID, "x")[0][0] + x, t = torch.zeros(1, 128, 4, 4), torch.tensor([500.0]) + + def skip(args, extra): + return args + skip.__qualname__ = P.model.SLG_SKIP + neg = {"cond_or_uncond": [1], "patches_replace": {"dit": {("single_block", 22): skip}}} + with self.running(slg_graph({})): + f(x, t, context=cond, transformer_options=neg) + self.assertEqual(self.stub.events("step")[-1]["skip"], [22]) + cases = { + "a positive row": ({**neg, "cond_or_uncond": [0]}, + "skips blocks of a positive row"), + "another replacement": ({**neg, "patches_replace": {"dit": { + ("single_block", 22): lambda a, e: a}}}, + "entry other than the block skips of SkipLayerGuidanceDiTSimple"), + "another kind of patch": ({**neg, "patches_replace": { + "dit": {("single_block", 22): skip}, "attn2": {}}}, + "entry other than the block skips of SkipLayerGuidanceDiTSimple")} + for name, (to, want) in cases.items(): + with self.subTest(name), self.assertRaises(C.GmlxRefusal) as e: + f(x, t, context=cond, transformer_options=to) + self.assertIn(want, str(e.exception)) + caps = self.client.capabilities + self.client.capabilities = {**caps, "features": [ + v for v in caps["features"] if v != "skip_blocks"]} + try: + with self.assertRaises(C.GmlxRefusal) as e: + f(x, t, context=cond, transformer_options=neg) + self.assertIn("this gmlx server skips no blocks in a step; update it", + str(e.exception)) + finally: + self.client.capabilities = caps + # A family whose DiT skips no blocks keeps its refusal of the replacement. + z = self.model(ZID).model.diffusion_model.forward + with self.assertRaises(C.GmlxRefusal) as e: + z(torch.zeros(1, 16, 4, 4), torch.tensor([0.5]), context=self.cond(ZID, "x")[0][0], + transformer_options=neg) + self.assertIn("transformer_options['patches_replace']", str(e.exception)) + self.assertEqual(len(self.stub.events("step")), 1) + + def test_a_masked_conditioning_is_refused_and_a_whole_one_samples(self): + m = self.model(MID) + pos = self.cond(MID, "a red bicycle") + n = pos[0][0].shape[1] + whole = [[pos[0][0], {"attention_mask": torch.ones(1, n)}]] + self.sample(m, MID, cfg=1.0, steps=1, pos=whole) + self.assertEqual(len(self.stub.events("step")), 1) + mask = torch.ones(1, n) + mask[0, -2:] = 0 + with self.assertRaises(C.GmlxRefusal) as e: + self.sample(m, MID, cfg=1.0, steps=1, pos=[[pos[0][0], {"attention_mask": mask}]]) + self.assertIn("masks some of its text rows", str(e.exception)) + + def test_the_moss_clip_sends_an_empty_prompt_as_it_is(self): + # MOSSImageTokenizer applies its chat template first. + clip = P.nodes.GmlxCLIPLoader().load(MID, self.stub.url, False)[0] + (cond,) = comfy_nodes.CLIPTextEncode().encode(clip, "a lighthouse") + self.assertEqual(tuple(cond[0][0].shape), (1, 32, 4096)) + clip.encode_from_tokens_scheduled(clip.tokenize("", prevent_empty_text=True)) + self.assertEqual([e["text"] for e in self.stub.events("encode")], ["a lighthouse", ""]) + + def test_the_lora_key_map_equals_the_openmoss_plugins(self): + """The patches ComfyUI makes of one LoRA on the pack's model and on + the OpenMOSS plugin's own model of the same file are the same.""" + import importlib.util + root = os.environ.get("GMLX_TEST_MOSS_PLUGIN") + if not root: + self.skipTest("the plugin comparison needs GMLX_TEST_MOSS_PLUGIN, the OpenMOSS " + "plugin's moss-image folder") + spec = importlib.util.spec_from_file_location( + "moss_image_plugin", os.path.join(root, "__init__.py"), submodule_search_locations=[root]) + mod = sys.modules.get("moss_image_plugin") + if mod is None: + mod = importlib.util.module_from_spec(spec) + sys.modules["moss_image_plugin"] = mod + spec.loader.exec_module(mod) + loaders = importlib.import_module("moss_image_plugin.loaders") + man = self.client.manifest(MID) + state = {t["name"]: torch.empty(t["shape"], device="meta") for t in man["tensors"]} + cfg = loaders.model_config(state) + cfg.set_inference_dtype(torch.float32, None, device=torch.device("cpu")) + theirs = comfy.model_patcher.ModelPatcher(cfg.get_model(state, device=torch.device("meta")), + torch.device("cpu"), torch.device("cpu")) + ours = self.model(MID) + self.assertEqual(set(ours.model.diffusion_model.state_dict()), + set(theirs.model.diffusion_model.state_dict())) + sd = comfy.utils.load_torch_file(os.path.join(LORAS, LORA_M), safe_load=True) + mw, _ = comfy.sd.load_lora_for_models(ours, None, sd, 0.8, 0.0) + mp, _ = comfy.sd.load_lora_for_models(theirs, None, sd, 0.8, 0.0) + self.assertEqual(set(mw.patches), set(mp.patches)) + self.assertEqual(set(mw.patches), LORA_M_KEYS) + for k in mw.patches: + for ta, tb in zip(mw.patches[k], mp.patches[k]): + self.assertEqual((ta[0], ta[2], ta[3]), (tb[0], tb[2], tb[3])) + for x, y in zip(ta[1].weights, tb[1].weights): + self.assertTrue((x is None and y is None) or torch.equal(torch.as_tensor(x), torch.as_tensor(y))) + + class Widgets(unittest.TestCase): """The loaders list the server's image models, the default first, and show the address they use; a name the server lacks is refused with its @@ -1224,7 +1465,7 @@ def tearDown(self): os.environ[k] = v def test_the_loaders_list_the_server_models_default_first(self): - want = [QID, ZID, QTID, KID, F2ID, F2BID, F29ID, F29BID, F29KID, F3ID, F3BID] + want = [QID, ZID, QTID, KID, F2ID, F2BID, F29ID, F29BID, F29KID, F3ID, F3BID, MID] for name in ("GmlxModelLoader", "GmlxCLIPLoader", "GmlxVAELoader", "GmlxServerStatus"): node = P.NODE_CLASS_MAPPINGS[name] req = node.INPUT_TYPES()["required"] @@ -1244,7 +1485,7 @@ def test_the_server_default_image_model_comes_first(self): self.probed() names, opts = P.nodes.GmlxModelLoader.INPUT_TYPES()["required"]["model_name"] self.assertEqual(names, [KID, ZID, QID, QTID, F2ID, F2BID, F29ID, F29BID, F29KID, - F3ID, F3BID]) + F3ID, F3BID, MID]) self.assertEqual(opts["default"], KID) finally: stub.close() diff --git a/tests/comfyui_pack/workflows/moss-image-t2i.json b/tests/comfyui_pack/workflows/moss-image-t2i.json new file mode 100644 index 000000000..166c395c9 --- /dev/null +++ b/tests/comfyui_pack/workflows/moss-image-t2i.json @@ -0,0 +1,160 @@ +{ + "1": { + "inputs": { + "model_name": "moss-image", + "server_url": "http://127.0.0.1:8080", + "start_server": true + }, + "class_type": "GmlxModelLoader", + "_meta": { + "title": "gmlx Model Loader" + } + }, + "2": { + "inputs": { + "model_name": "moss-image", + "server_url": "http://127.0.0.1:8080", + "start_server": true + }, + "class_type": "GmlxCLIPLoader", + "_meta": { + "title": "gmlx CLIP Loader" + } + }, + "3": { + "inputs": { + "model_name": "moss-image", + "server_url": "http://127.0.0.1:8080", + "start_server": true + }, + "class_type": "GmlxVAELoader", + "_meta": { + "title": "gmlx VAE Loader" + } + }, + "4": { + "inputs": { + "clip": [ + "2", + 0 + ], + "text": "A wide view of a refrigerator beside a kitchen counter with a fruit bowl. A small sticky note on the refrigerator door is handwritten: \"Remember to buy milk and eggs.\"" + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "Positive prompt" + } + }, + "5": { + "inputs": { + "clip": [ + "2", + 0 + ], + "text": "" + }, + "class_type": "CLIPTextEncode", + "_meta": { + "title": "Negative prompt" + } + }, + "6": { + "inputs": { + "width": 1024, + "height": 1024, + "batch_size": 1 + }, + "class_type": "EmptyFlux2LatentImage", + "_meta": { + "title": "Empty latent" + } + }, + "7": { + "inputs": { + "model": [ + "1", + 0 + ], + "shift": 5.0 + }, + "class_type": "ModelSamplingSD3", + "_meta": { + "title": "Shift" + } + }, + "8": { + "inputs": { + "model": [ + "7", + 0 + ], + "double_layers": "", + "single_layers": "22", + "start_percent": 0.0, + "end_percent": 1.0 + }, + "class_type": "SkipLayerGuidanceDiTSimple", + "_meta": { + "title": "Skip block 22 on the negative pass" + } + }, + "9": { + "inputs": { + "model": [ + "8", + 0 + ], + "positive": [ + "4", + 0 + ], + "negative": [ + "5", + 0 + ], + "latent_image": [ + "6", + 0 + ], + "seed": 42, + "steps": 50, + "cfg": 3.0, + "sampler_name": "euler", + "scheduler": "sgm_uniform", + "denoise": 1.0 + }, + "class_type": "KSampler", + "_meta": { + "title": "KSampler" + } + }, + "10": { + "inputs": { + "samples": [ + "9", + 0 + ], + "vae": [ + "3", + 0 + ] + }, + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + } + }, + "11": { + "inputs": { + "images": [ + "10", + 0 + ], + "filename_prefix": "gmlx/moss-image" + }, + "class_type": "SaveImage", + "_meta": { + "title": "Save Image" + } + } +} diff --git a/tests/commands/test_image_pull.py b/tests/commands/test_image_pull.py index 7010c580d..634524734 100644 --- a/tests/commands/test_image_pull.py +++ b/tests/commands/test_image_pull.py @@ -742,6 +742,8 @@ def test_the_json_plan_names_the_size_the_entry_gets(lib, monkeypatch, F2 = catalog.RECIPES["flux-2-klein"] F2B = catalog.RECIPES["flux-2-klein-base"] K = catalog.RECIPES["krea-2-turbo"] +M = catalog.RECIPES["moss-image"] +MB = catalog.RECIPES["moss-image-base"] # The family each pack of the FLUX.2 autoencoder names: the pull's. _PACKED_FOR: dict = {} @@ -779,6 +781,17 @@ def _shared_scan(roots, *, recursive=True, families=None): "distilled")) elif f.startswith("Krea-2-Turbo-"): out.append(catalog.ImageFile(p, "dit", "krea2", "turbo")) + elif f.startswith("MOSS-Image-SFT-"): + out.append(catalog.ImageFile(p, "dit", "moss_image", + "sft")) + elif f.startswith("MOSS-Image-Base-"): + out.append(catalog.ImageFile(p, "dit", "moss_image", + "base")) + elif f.startswith("Qwen3.5-9B-"): + out.append(catalog.ImageFile(p, "encoder", "moss_image")) + elif f == M.vae_pack: + out.append(catalog.ImageFile(p, "vae", "moss_image", + sha256=M.vae_sha256)) elif f.startswith("Qwen3-VL-4B-Instruct-"): out += [catalog.ImageFile(p, "encoder", fam) for fam in ("krea2", "flux3_klein")] @@ -1056,6 +1069,66 @@ def test_a_flux3_pull_after_krea_2_shares_its_encoder_file(lib, monkeypatch): assert base.encoder == krea.encoder and base.vae == flux.vae +@pytest.mark.parametrize("r, ram_gb, quant, enc", [ + (M, 16, "Q6_K", "Q4_K_M"), (M, 32, "Q8_0", "Q8_0"), + (MB, 24, "Q8_0", "Q8_0")]) +def test_a_moss_image_pull_takes_its_class_quants_and_packs_the_release_vae( + lib, monkeypatch, capsys, r, ram_gb, quant, enc): + """The DiT from realrebelai's repo, Qwen3.5-9B from unsloth's at the + class's quant, and the release's float32 VAE with its config.json at + the pinned revision, packed in its own dtype.""" + root, cfg = lib + got, packed, m = _shared_pull(lib, monkeypatch, r.name, ram_gb) + stem = "MOSS-Image-Base" if r.variant == "base" else "MOSS-Image-SFT" + assert got == [ + ("realrebelai/MOSS-Image_GGUFs", f"{stem}-{quant}.gguf", "main"), + ("unsloth/Qwen3.5-9B-GGUF", f"Qwen3.5-9B-{enc}.gguf", "main"), + ("OpenMOSS-Team/MOSS-Image-SFT", + "vae/diffusion_pytorch_model.safetensors", r.vae_revision), + ("OpenMOSS-Team/MOSS-Image-SFT", "vae/config.json", r.vae_revision)] + ((_src, kw),) = packed + assert kw["family"] == "moss_image" and kw.get("dtype") is None + assert m.path == f"realrebelai__MOSS-Image_GGUFs/{stem}-{quant}.gguf" + assert m.encoder == f"unsloth__Qwen3.5-9B-GGUF/Qwen3.5-9B-{enc}.gguf" + assert m.vae == "OpenMOSS-Team__MOSS-Image-SFT/moss-image-vae-F32.gguf" + assert m.encoder_mmproj is None and not m.defaults.get("size") + assert (f"license: {r.title} apache-2.0 (https://huggingface.co/" + f"OpenMOSS-Team/{stem}/blob/main/LICENSE); encoder apache-2.0" + in capsys.readouterr().out) + + +def test_a_moss_image_base_pull_after_sft_takes_only_its_dit(lib, + monkeypatch): + _, _, sft = _shared_pull(lib, monkeypatch, "moss-image") + got, packed, base = _shared_pull(lib, monkeypatch, "moss-image-base") + assert got == [(MB.dit_repo, "MOSS-Image-Base-Q8_0.gguf", "main")] + assert packed == [] + assert base.encoder == sft.encoder and base.vae == sft.vae + + +@pytest.mark.parametrize("held, ram_gb, args, takes", [ + ("Q8_0", 16, (), "Q4_K_M"), ("Q8_0", 16, ("--encoder-quant", "Q8_0"), "Q8_0"), + ("Q4_K_M", 32, (), "Q4_K_M")]) +def test_a_held_encoder_serves_a_moss_image_pull_where_it_fits( + lib, monkeypatch, capsys, held, ram_gb, args, takes): + """Another quant on disk serves the pull when an image fits with it. + Qwen3.5-9B Q8_0 does not fit a 16 GB Mac, so the pull takes the class's + Q4_K_M and says so, unless --encoder-quant asks for the held one.""" + root, _ = lib + d = root / "unsloth__Qwen3.5-9B-GGUF" + d.mkdir() + with open(d / f"Qwen3.5-9B-{held}.gguf", "wb") as f: + f.truncate(catalog.ENCODER_BYTES[f"Qwen3.5-9B-{held}.gguf"]) + got, _, m = _shared_pull(lib, monkeypatch, "moss-image", ram_gb, args) + out = capsys.readouterr().out + assert m.encoder == f"unsloth__Qwen3.5-9B-GGUF/Qwen3.5-9B-{takes}.gguf" + pulled = (M.encoder_repo, f"Qwen3.5-9B-{takes}.gguf", "main") in got + assert pulled == (held != takes) + note = ("Qwen3.5-9B-Q8_0.gguf on disk does not fit an image on a 16 GB " + "Mac; the pull takes Q4_K_M") + assert (note in out) == (held != takes) + + def test_a_flux3_pull_takes_no_encoder_quant_but_q8_0(lib, capsys): root, cfg = lib assert manage.cmd_pull(["flux-3-klein", "--encoder-quant", "Q4_K_M", @@ -1119,6 +1192,26 @@ def test_an_image_dit_skips_the_chat_arch_gate_and_fits_with_its_parts(): assert v["usable"] and v["arch_error"] is None +@pytest.mark.parametrize("ram_gb, enc", [(16, 5_680_522_464), + (32, 9_527_502_048)]) +def test_a_moss_image_dit_fits_with_the_encoder_its_class_pulls( + monkeypatch, ram_gb, enc): + """The fit of a DiT counts the encoder quant the pull takes on this + Mac: Qwen3.5-9B Q4_K_M on 16 GB, Q8_0 above.""" + monkeypatch.setattr("gmlx.load.memfit.total_ram_bytes", + lambda: ram_gb << 30) + ref = remote.parse_ref( + "hf:realrebelai/MOSS-Image_GGUFs/MOSS-Image-SFT-Q6_K.gguf") + v = manage._verdict(ref, _report("moss_image"), size_bytes=4_553_650_496) + assert v["image_family"] == "moss_image" + assert v["combined_bytes"] == 4_553_650_496 + enc + M.vae_bytes + if ram_gb == 16: + assert v["fit"] == "fits" + assert manage._image_fit_note(v) == ( + "9.8 GB with its text encoder and VAE, fits on this Mac's " + "16 GB RAM") + + @pytest.mark.parametrize("gen, chip", [(13, "M1"), (14, "M2")]) def test_a_dit_this_chip_does_not_run_is_refused_by_name(monkeypatch, capsys, gen, chip): diff --git a/tests/fixtures/image/manifest.json b/tests/fixtures/image/manifest.json index 53d977301..2393fdd6f 100644 --- a/tests/fixtures/image/manifest.json +++ b/tests/fixtures/image/manifest.json @@ -279,5 +279,48 @@ "sha256": "72ac6e652196cd9f450d4d3256d197978d3fa690e22dc21361d3bd4fb4e8250e", "license": "apache-2.0", "sha256_source": "The hub's LFS record of the file whose header the fixture holds, which the downloaded file matches." + }, + "moss_image_sft_q4_k_m.header.json": { + "repo": "realrebelai/MOSS-Image_GGUFs", + "revision": "4f21f583c9213eb2230a09e9b33bfa86605aecf2", + "file": "MOSS-Image-SFT-Q4_K_M.gguf", + "size": 3702092096, + "sha256": "c6dda89c18c0c0897d09ed01a515ee31415c4baea02bb9c2dedbfca66eec50b6", + "kv_kept": "general.architecture", + "note": "MOSS-Image-Base-Q4_K_M.gguf (sha256 88e336e4429f81085b3a1595f01432a07de50837282e0ecace5c323d74a9d38f) has the same names, shapes and codecs" + }, + "moss_image_transformer_config.json": { + "repo": "OpenMOSS-Team/MOSS-Image-SFT", + "revision": "fc5de0d56fc910cbca6f6625a70f5f582b26deb1", + "file": "transformer/config.json", + "size": 240, + "sha256": "3291c3a1c7adc22b8182a0de25ca8ec97474d72715ce84adbab09cc0be24f666" + }, + "moss_image_identity.json": { + "repo": "realrebelai/MOSS-Image_GGUFs", + "revision": "4f21f583c9213eb2230a09e9b33bfa86605aecf2", + "file": "the identity tensors of every MOSS-Image .gguf, by HTTP range reads", + "note": "each file size is from the hub listing at the revision" + }, + "moss_image_chat_template.jinja": { + "repo": "OpenMOSS-Team/MOSS-Image-SFT", + "revision": "fc5de0d56fc910cbca6f6625a70f5f582b26deb1", + "file": "tokenizer/chat_template.jinja", + "size": 7756, + "sha256": "a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715" + }, + "moss_image_vae.header.json": { + "repo": "OpenMOSS-Team/MOSS-Image-SFT", + "revision": "fc5de0d56fc910cbca6f6625a70f5f582b26deb1", + "file": "vae/diffusion_pytorch_model.safetensors", + "size": 249521372, + "sha256": "c3bdf94dc67ec69b30ea9fe1fa8340085679cae64d019f7c7e4a9921f1a7f17a" + }, + "moss_image_vae_config.json": { + "repo": "OpenMOSS-Team/MOSS-Image-SFT", + "revision": "fc5de0d56fc910cbca6f6625a70f5f582b26deb1", + "file": "vae/config.json", + "size": 838, + "sha256": "cc77f3f0ce0357039c71018dc19edaa4fcceb0c2f7488fdc2b9c0580d6baafa1" } } diff --git a/tests/fixtures/image/memfit/moss_image_cells.json b/tests/fixtures/image/memfit/moss_image_cells.json new file mode 100644 index 000000000..8f725a833 --- /dev/null +++ b/tests/fixtures/image/memfit/moss_image_cells.json @@ -0,0 +1,237 @@ +{ + "what": "MOSS-Image SFT Q4_K_M (realrebelai MOSS-Image-SFT-Q4_K_M.gguf) on an M5 Max, one fresh process per cell, through api.sample under the bounded cache, with random features of text_rows rows; a guided cell runs cfg 3 with a negative of the same rows and block 22 left out of its pass. steps_gib is the footprint of each of three steps over the settled process, the DiT loaded. encode holds one prompt encode of the unsloth Qwen3.5-9B encoder alone, by file, at rows rows, the footprint over the settled process.", + "cells": [ + { + "size": 256, + "text_rows": 80, + "guided": false, + "steps_gib": [ + 0.8792, + 0.8793, + 0.8794 + ] + }, + { + "size": 512, + "text_rows": 80, + "guided": false, + "steps_gib": [ + 1.1987, + 1.1988, + 1.1988 + ] + }, + { + "size": 768, + "text_rows": 80, + "guided": false, + "steps_gib": [ + 1.4619, + 1.4621, + 1.4621 + ] + }, + { + "size": 1024, + "text_rows": 80, + "guided": false, + "steps_gib": [ + 1.5731, + 1.5734, + 1.5734 + ] + }, + { + "size": 1280, + "text_rows": 80, + "guided": false, + "steps_gib": [ + 2.0243, + 2.0245, + 2.0246 + ] + }, + { + "size": 1536, + "text_rows": 80, + "guided": false, + "steps_gib": [ + 2.7351, + 2.7351, + 2.6968 + ] + }, + { + "size": 2048, + "text_rows": 80, + "guided": false, + "steps_gib": [ + 4.4045, + 4.1951, + 4.1847 + ] + }, + { + "size": 1024, + "text_rows": 640, + "guided": false, + "steps_gib": [ + 1.7751, + 1.7753, + 1.7753 + ] + }, + { + "size": 256, + "text_rows": 4096, + "guided": false, + "steps_gib": [ + 2.5045, + 2.5045, + 2.3483 + ] + }, + { + "size": 512, + "text_rows": 4096, + "guided": false, + "steps_gib": [ + 2.6968, + 2.6968, + 2.6645 + ] + }, + { + "size": 1024, + "text_rows": 4096, + "guided": false, + "steps_gib": [ + 3.2757, + 3.0648, + 2.6594 + ] + }, + { + "size": 1536, + "text_rows": 4096, + "guided": false, + "steps_gib": [ + 4.5186, + 3.4133, + 3.41 + ] + }, + { + "size": 2048, + "text_rows": 4096, + "guided": false, + "steps_gib": [ + 5.7607, + 5.0914, + 5.0798 + ] + }, + { + "size": 512, + "text_rows": 80, + "guided": true, + "steps_gib": [ + 1.2351, + 1.2352, + 1.2353 + ] + }, + { + "size": 1024, + "text_rows": 80, + "guided": true, + "steps_gib": [ + 2.4877, + 2.4877, + 2.4509 + ] + }, + { + "size": 1536, + "text_rows": 80, + "guided": true, + "steps_gib": [ + 4.3422, + 4.0801, + 4.1182 + ] + }, + { + "size": 1024, + "text_rows": 4096, + "guided": true, + "steps_gib": [ + 4.1826, + 3.8712, + 3.8713 + ] + }, + { + "size": 2048, + "text_rows": 80, + "guided": true, + "steps_gib": [ + 6.983, + 6.991, + 6.9911 + ] + }, + { + "size": 2048, + "text_rows": 4096, + "guided": true, + "steps_gib": [ + 8.2565, + 8.2645, + 8.2645 + ] + } + ], + "encode": [ + { + "file": "Qwen3.5-9B-Q4_K_M.gguf", + "rows": 80, + "gib": 0.3696 + }, + { + "file": "Qwen3.5-9B-Q4_K_M.gguf", + "rows": 640, + "gib": 0.4378 + }, + { + "file": "Qwen3.5-9B-Q4_K_M.gguf", + "rows": 2048, + "gib": 0.8289 + }, + { + "file": "Qwen3.5-9B-Q4_K_M.gguf", + "rows": 4096, + "gib": 1.3174 + }, + { + "file": "Qwen3.5-9B-Q8_0.gguf", + "rows": 80, + "gib": 0.2453 + }, + { + "file": "Qwen3.5-9B-Q8_0.gguf", + "rows": 640, + "gib": 0.3551 + }, + { + "file": "Qwen3.5-9B-Q8_0.gguf", + "rows": 2048, + "gib": 0.7563 + }, + { + "file": "Qwen3.5-9B-Q8_0.gguf", + "rows": 4096, + "gib": 1.3249 + } + ] +} diff --git a/tests/fixtures/image/moss_image_chat_template.jinja b/tests/fixtures/image/moss_image_chat_template.jinja new file mode 100644 index 000000000..a585dec89 --- /dev/null +++ b/tests/fixtures/image/moss_image_chat_template.jinja @@ -0,0 +1,154 @@ +{%- set image_count = namespace(value=0) %} +{%- set video_count = namespace(value=0) %} +{%- macro render_content(content, do_vision_count, is_system_content=false) %} + {%- if content is string %} + {{- content }} + {%- elif content is iterable and content is not mapping %} + {%- for item in content %} + {%- if 'image' in item or 'image_url' in item or item.type == 'image' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain images.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set image_count.value = image_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Picture ' ~ image_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|image_pad|><|vision_end|>' }} + {%- elif 'video' in item or item.type == 'video' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain videos.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set video_count.value = video_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Video ' ~ video_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|video_pad|><|vision_end|>' }} + {%- elif 'text' in item %} + {{- item.text }} + {%- else %} + {{- raise_exception('Unexpected item type in content.') }} + {%- endif %} + {%- endfor %} + {%- elif content is none or content is undefined %} + {{- '' }} + {%- else %} + {{- raise_exception('Unexpected content type.') }} + {%- endif %} +{%- endmacro %} +{%- if not messages %} + {{- raise_exception('No messages provided.') }} +{%- endif %} +{%- if tools and tools is iterable and tools is not mapping %} + {{- '<|im_start|>system\n' }} + {{- "# Tools\n\nYou have access to the following functions:\n\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n" }} + {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n\n\n\nReminder:\n- Function calls MUST follow the specified format: an inner block must be nested within XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n' }} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {%- if content %} + {{- '\n\n' + content }} + {%- endif %} + {%- endif %} + {{- '<|im_end|>\n' }} +{%- else %} + {%- if messages[0].role == 'system' %} + {%- set content = render_content(messages[0].content, false, true)|trim %} + {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" %} + {%- set content = render_content(message.content, false)|trim %} + {%- if not(content.startswith('') and content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if ns.multi_step_tool %} + {{- raise_exception('No user query found in messages.') }} +{%- endif %} +{%- for message in messages %} + {%- set content = render_content(message.content, true)|trim %} + {%- if message.role == "system" %} + {%- if not loop.first %} + {{- raise_exception('System message must be at the beginning.') }} + {%- endif %} + {%- elif message.role == "user" %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- set reasoning_content = reasoning_content|trim %} + {%- if loop.index0 > ns.last_query_index %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content + '\n\n\n' + content }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.function is defined %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {%- if loop.first %} + {%- if content|trim %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n\n' }} + {%- endif %} + {%- else %} + {{- '\n\n\n' }} + {%- endif %} + {%- if tool_call.arguments is defined %} + {%- for args_name, args_value in tool_call.arguments|items %} + {{- '\n' }} + {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %} + {{- args_value }} + {{- '\n\n' }} + {%- endfor %} + {%- endif %} + {{- '\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.previtem and loop.previtem.role != "tool" %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if not loop.last and loop.nextitem.role != "tool" %} + {{- '<|im_end|>\n' }} + {%- elif loop.last %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- else %} + {{- raise_exception('Unexpected message role.') }} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n' }} + {%- endif %} +{%- endif %} \ No newline at end of file diff --git a/tests/fixtures/image/moss_image_identity.json b/tests/fixtures/image/moss_image_identity.json new file mode 100644 index 000000000..f34ef83ea --- /dev/null +++ b/tests/fixtures/image/moss_image_identity.json @@ -0,0 +1,228 @@ +{ + "_note": "The identity tensors video_embedder.bias, final_norm_video.weight and time_embedding.0.bias of every MOSS-Image GGUF of the repository, by HTTP range reads: their codecs and the sha256 of their bf16 bits, as gmlx.image.fingerprint.weight_identity forms it", + "repo": "realrebelai/MOSS-Image_GGUFs", + "revision": "4f21f583c9213eb2230a09e9b33bfa86605aecf2", + "files": { + "MOSS-Image-SFT-Q2_K.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60", + "n_tensors": 463, + "size": 2252894528 + }, + "MOSS-Image-SFT-Q3_K_L.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60", + "n_tensors": 463, + "size": 3070390592 + }, + "MOSS-Image-SFT-Q3_K_M.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60", + "n_tensors": 463, + "size": 2891084096 + }, + "MOSS-Image-SFT-Q3_K_S.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60", + "n_tensors": 463, + "size": 2695704896 + }, + "MOSS-Image-SFT-Q4_K_M.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60", + "n_tensors": 463, + "size": 3702092096 + }, + "MOSS-Image-SFT-Q4_K_S.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60", + "n_tensors": 463, + "size": 3464245568 + }, + "MOSS-Image-SFT-Q5_K_M.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60", + "n_tensors": 463, + "size": 4114968896 + }, + "MOSS-Image-SFT-Q5_K_S.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60", + "n_tensors": 463, + "size": 3921948992 + }, + "MOSS-Image-SFT-Q6_K.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60", + "n_tensors": 463, + "size": 4553650496 + }, + "MOSS-Image-SFT-Q8_0.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "96b09d61487460c3f321abd04a71a2ca989025150226ca1c30be7a57c0d86c60", + "n_tensors": 463, + "size": 5705576768 + }, + "MOSS-Image-Base-Q2_K.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626", + "n_tensors": 463, + "size": 2252894528 + }, + "MOSS-Image-Base-Q3_K_L.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626", + "n_tensors": 463, + "size": 3070390592 + }, + "MOSS-Image-Base-Q3_K_M.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626", + "n_tensors": 463, + "size": 2891084096 + }, + "MOSS-Image-Base-Q3_K_S.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626", + "n_tensors": 463, + "size": 2695704896 + }, + "MOSS-Image-Base-Q4_K_M.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626", + "n_tensors": 463, + "size": 3702092096 + }, + "MOSS-Image-Base-Q4_K_S.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626", + "n_tensors": 463, + "size": 3464245568 + }, + "MOSS-Image-Base-Q5_K_M.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626", + "n_tensors": 463, + "size": 4114968896 + }, + "MOSS-Image-Base-Q5_K_S.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626", + "n_tensors": 463, + "size": 3921948992 + }, + "MOSS-Image-Base-Q6_K.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626", + "n_tensors": 463, + "size": 4553650496 + }, + "MOSS-Image-Base-Q8_0.gguf": { + "identity_codecs": [ + "F16", + "F16", + "F16" + ], + "sha256": "22e89962212b91739f15827592163257e8717b26870e57d446e48846259c2626", + "n_tensors": 463, + "size": 5705576768 + } + }, + "rope_bands": { + "note": "rope.bands of MOSS-Image-SFT-Q4_K_M.gguf and MOSS-Image-Base-Q4_K_M.gguf, F32, as uint32 bits", + "bits": [ + 1065353216, + 1058010522, + 1050798235, + 1043732615, + 1036831949, + 1030116802, + 1023510243, + 1016180025, + 1008981770, + 1001931932, + 995049016, + 988353832, + 981668463, + 974350870, + 967166816, + 960132947 + ] + } +} diff --git a/tests/fixtures/image/moss_image_sft_q4_k_m.header.json b/tests/fixtures/image/moss_image_sft_q4_k_m.header.json new file mode 100644 index 000000000..4d6bdb362 --- /dev/null +++ b/tests/fixtures/image/moss_image_sft_q4_k_m.header.json @@ -0,0 +1,470 @@ +{ + "gguf_version": 3, + "kv": {"general.architecture": "moss_image"}, + "n_tensors": 463, + "tensors": [ + ["final_linear_video.weight", [3072, 128], "F16"], + ["final_norm_video.weight", [3072], "F16"], + ["layers.0.k_norm.weight", [128], "F16"], + ["layers.0.norm1.weight", [3072], "F16"], + ["layers.0.norm2.weight", [3072], "F16"], + ["layers.0.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.0.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.0.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.0.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.0.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.0.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.0.q_norm.weight", [128], "F16"], + ["layers.1.k_norm.weight", [128], "F16"], + ["layers.1.norm1.weight", [3072], "F16"], + ["layers.1.norm2.weight", [3072], "F16"], + ["layers.1.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.1.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.1.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.1.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.1.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.1.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.1.q_norm.weight", [128], "F16"], + ["layers.10.k_norm.weight", [128], "F16"], + ["layers.10.norm1.weight", [3072], "F16"], + ["layers.10.norm2.weight", [3072], "F16"], + ["layers.10.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.10.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.10.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.10.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.10.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.10.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.10.q_norm.weight", [128], "F16"], + ["layers.11.k_norm.weight", [128], "F16"], + ["layers.11.norm1.weight", [3072], "F16"], + ["layers.11.norm2.weight", [3072], "F16"], + ["layers.11.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.11.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.11.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.11.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.11.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.11.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.11.q_norm.weight", [128], "F16"], + ["layers.12.k_norm.weight", [128], "F16"], + ["layers.12.norm1.weight", [3072], "F16"], + ["layers.12.norm2.weight", [3072], "F16"], + ["layers.12.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.12.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.12.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.12.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.12.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.12.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.12.q_norm.weight", [128], "F16"], + ["layers.13.k_norm.weight", [128], "F16"], + ["layers.13.norm1.weight", [3072], "F16"], + ["layers.13.norm2.weight", [3072], "F16"], + ["layers.13.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.13.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.13.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.13.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.13.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.13.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.13.q_norm.weight", [128], "F16"], + ["layers.14.k_norm.weight", [128], "F16"], + ["layers.14.norm1.weight", [3072], "F16"], + ["layers.14.norm2.weight", [3072], "F16"], + ["layers.14.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.14.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.14.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.14.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.14.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.14.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.14.q_norm.weight", [128], "F16"], + ["layers.15.k_norm.weight", [128], "F16"], + ["layers.15.norm1.weight", [3072], "F16"], + ["layers.15.norm2.weight", [3072], "F16"], + ["layers.15.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.15.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.15.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.15.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.15.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.15.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.15.q_norm.weight", [128], "F16"], + ["layers.16.k_norm.weight", [128], "F16"], + ["layers.16.norm1.weight", [3072], "F16"], + ["layers.16.norm2.weight", [3072], "F16"], + ["layers.16.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.16.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.16.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.16.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.16.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.16.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.16.q_norm.weight", [128], "F16"], + ["layers.17.k_norm.weight", [128], "F16"], + ["layers.17.norm1.weight", [3072], "F16"], + ["layers.17.norm2.weight", [3072], "F16"], + ["layers.17.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.17.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.17.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.17.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.17.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.17.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.17.q_norm.weight", [128], "F16"], + ["layers.18.k_norm.weight", [128], "F16"], + ["layers.18.norm1.weight", [3072], "F16"], + ["layers.18.norm2.weight", [3072], "F16"], + ["layers.18.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.18.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.18.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.18.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.18.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.18.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.18.q_norm.weight", [128], "F16"], + ["layers.19.k_norm.weight", [128], "F16"], + ["layers.19.norm1.weight", [3072], "F16"], + ["layers.19.norm2.weight", [3072], "F16"], + ["layers.19.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.19.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.19.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.19.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.19.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.19.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.19.q_norm.weight", [128], "F16"], + ["layers.2.k_norm.weight", [128], "F16"], + ["layers.2.norm1.weight", [3072], "F16"], + ["layers.2.norm2.weight", [3072], "F16"], + ["layers.2.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.2.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.2.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.2.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.2.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.2.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.2.q_norm.weight", [128], "F16"], + ["layers.20.k_norm.weight", [128], "F16"], + ["layers.20.norm1.weight", [3072], "F16"], + ["layers.20.norm2.weight", [3072], "F16"], + ["layers.20.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.20.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.20.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.20.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.20.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.20.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.20.q_norm.weight", [128], "F16"], + ["layers.21.k_norm.weight", [128], "F16"], + ["layers.21.norm1.weight", [3072], "F16"], + ["layers.21.norm2.weight", [3072], "F16"], + ["layers.21.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.21.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.21.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.21.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.21.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.21.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.21.q_norm.weight", [128], "F16"], + ["layers.22.k_norm.weight", [128], "F16"], + ["layers.22.norm1.weight", [3072], "F16"], + ["layers.22.norm2.weight", [3072], "F16"], + ["layers.22.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.22.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.22.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.22.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.22.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.22.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.22.q_norm.weight", [128], "F16"], + ["layers.23.k_norm.weight", [128], "F16"], + ["layers.23.norm1.weight", [3072], "F16"], + ["layers.23.norm2.weight", [3072], "F16"], + ["layers.23.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.23.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.23.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.23.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.23.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.23.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.23.q_norm.weight", [128], "F16"], + ["layers.24.k_norm.weight", [128], "F16"], + ["layers.24.norm1.weight", [3072], "F16"], + ["layers.24.norm2.weight", [3072], "F16"], + ["layers.24.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.24.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.24.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.24.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.24.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.24.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.24.q_norm.weight", [128], "F16"], + ["layers.25.k_norm.weight", [128], "F16"], + ["layers.25.norm1.weight", [3072], "F16"], + ["layers.25.norm2.weight", [3072], "F16"], + ["layers.25.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.25.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.25.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.25.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.25.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.25.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.25.q_norm.weight", [128], "F16"], + ["layers.26.k_norm.weight", [128], "F16"], + ["layers.26.norm1.weight", [3072], "F16"], + ["layers.26.norm2.weight", [3072], "F16"], + ["layers.26.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.26.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.26.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.26.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.26.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.26.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.26.q_norm.weight", [128], "F16"], + ["layers.27.k_norm.weight", [128], "F16"], + ["layers.27.norm1.weight", [3072], "F16"], + ["layers.27.norm2.weight", [3072], "F16"], + ["layers.27.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.27.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.27.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.27.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.27.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.27.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.27.q_norm.weight", [128], "F16"], + ["layers.28.k_norm.weight", [128], "F16"], + ["layers.28.norm1.weight", [3072], "F16"], + ["layers.28.norm2.weight", [3072], "F16"], + ["layers.28.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.28.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.28.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.28.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.28.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.28.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.28.q_norm.weight", [128], "F16"], + ["layers.29.k_norm.weight", [128], "F16"], + ["layers.29.norm1.weight", [3072], "F16"], + ["layers.29.norm2.weight", [3072], "F16"], + ["layers.29.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.29.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.29.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.29.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.29.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.29.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.29.q_norm.weight", [128], "F16"], + ["layers.3.k_norm.weight", [128], "F16"], + ["layers.3.norm1.weight", [3072], "F16"], + ["layers.3.norm2.weight", [3072], "F16"], + ["layers.3.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.3.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.3.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.3.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.3.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.3.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.3.q_norm.weight", [128], "F16"], + ["layers.30.k_norm.weight", [128], "F16"], + ["layers.30.norm1.weight", [3072], "F16"], + ["layers.30.norm2.weight", [3072], "F16"], + ["layers.30.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.30.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.30.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.30.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.30.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.30.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.30.q_norm.weight", [128], "F16"], + ["layers.31.k_norm.weight", [128], "F16"], + ["layers.31.norm1.weight", [3072], "F16"], + ["layers.31.norm2.weight", [3072], "F16"], + ["layers.31.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.31.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.31.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.31.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.31.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.31.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.31.q_norm.weight", [128], "F16"], + ["layers.32.k_norm.weight", [128], "F16"], + ["layers.32.norm1.weight", [3072], "F16"], + ["layers.32.norm2.weight", [3072], "F16"], + ["layers.32.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.32.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.32.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.32.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.32.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.32.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.32.q_norm.weight", [128], "F16"], + ["layers.33.k_norm.weight", [128], "F16"], + ["layers.33.norm1.weight", [3072], "F16"], + ["layers.33.norm2.weight", [3072], "F16"], + ["layers.33.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.33.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.33.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.33.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.33.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.33.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.33.q_norm.weight", [128], "F16"], + ["layers.34.k_norm.weight", [128], "F16"], + ["layers.34.norm1.weight", [3072], "F16"], + ["layers.34.norm2.weight", [3072], "F16"], + ["layers.34.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.34.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.34.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.34.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.34.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.34.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.34.q_norm.weight", [128], "F16"], + ["layers.35.k_norm.weight", [128], "F16"], + ["layers.35.norm1.weight", [3072], "F16"], + ["layers.35.norm2.weight", [3072], "F16"], + ["layers.35.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.35.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.35.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.35.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.35.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.35.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.35.q_norm.weight", [128], "F16"], + ["layers.36.k_norm.weight", [128], "F16"], + ["layers.36.norm1.weight", [3072], "F16"], + ["layers.36.norm2.weight", [3072], "F16"], + ["layers.36.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.36.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.36.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.36.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.36.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.36.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.36.q_norm.weight", [128], "F16"], + ["layers.37.k_norm.weight", [128], "F16"], + ["layers.37.norm1.weight", [3072], "F16"], + ["layers.37.norm2.weight", [3072], "F16"], + ["layers.37.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.37.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.37.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.37.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.37.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.37.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.37.q_norm.weight", [128], "F16"], + ["layers.38.k_norm.weight", [128], "F16"], + ["layers.38.norm1.weight", [3072], "F16"], + ["layers.38.norm2.weight", [3072], "F16"], + ["layers.38.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.38.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.38.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.38.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.38.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.38.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.38.q_norm.weight", [128], "F16"], + ["layers.39.k_norm.weight", [128], "F16"], + ["layers.39.norm1.weight", [3072], "F16"], + ["layers.39.norm2.weight", [3072], "F16"], + ["layers.39.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.39.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.39.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.39.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.39.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.39.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.39.q_norm.weight", [128], "F16"], + ["layers.4.k_norm.weight", [128], "F16"], + ["layers.4.norm1.weight", [3072], "F16"], + ["layers.4.norm2.weight", [3072], "F16"], + ["layers.4.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.4.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.4.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.4.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.4.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.4.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.4.q_norm.weight", [128], "F16"], + ["layers.40.k_norm.weight", [128], "F16"], + ["layers.40.norm1.weight", [3072], "F16"], + ["layers.40.norm2.weight", [3072], "F16"], + ["layers.40.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.40.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.40.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.40.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.40.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.40.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.40.q_norm.weight", [128], "F16"], + ["layers.41.k_norm.weight", [128], "F16"], + ["layers.41.norm1.weight", [3072], "F16"], + ["layers.41.norm2.weight", [3072], "F16"], + ["layers.41.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.41.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.41.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.41.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.41.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.41.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.41.q_norm.weight", [128], "F16"], + ["layers.5.k_norm.weight", [128], "F16"], + ["layers.5.norm1.weight", [3072], "F16"], + ["layers.5.norm2.weight", [3072], "F16"], + ["layers.5.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.5.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.5.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.5.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.5.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.5.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.5.q_norm.weight", [128], "F16"], + ["layers.6.k_norm.weight", [128], "F16"], + ["layers.6.norm1.weight", [3072], "F16"], + ["layers.6.norm2.weight", [3072], "F16"], + ["layers.6.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.6.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.6.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.6.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.6.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.6.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.6.q_norm.weight", [128], "F16"], + ["layers.7.k_norm.weight", [128], "F16"], + ["layers.7.norm1.weight", [3072], "F16"], + ["layers.7.norm2.weight", [3072], "F16"], + ["layers.7.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.7.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.7.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.7.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.7.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.7.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.7.q_norm.weight", [128], "F16"], + ["layers.8.k_norm.weight", [128], "F16"], + ["layers.8.norm1.weight", [3072], "F16"], + ["layers.8.norm2.weight", [3072], "F16"], + ["layers.8.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.8.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.8.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.8.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.8.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.8.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.8.q_norm.weight", [128], "F16"], + ["layers.9.k_norm.weight", [128], "F16"], + ["layers.9.norm1.weight", [3072], "F16"], + ["layers.9.norm2.weight", [3072], "F16"], + ["layers.9.proj.ffn_down.weight", [8192, 3072], "Q6_K"], + ["layers.9.proj.ffn_up_gate.weight", [3072, 16384], "Q4_K"], + ["layers.9.proj.k.weight", [3072, 3072], "Q4_K"], + ["layers.9.proj.o.weight", [3072, 3072], "Q4_K"], + ["layers.9.proj.q.weight", [3072, 3072], "Q4_K"], + ["layers.9.proj.v.weight", [3072, 3072], "Q6_K"], + ["layers.9.q_norm.weight", [128], "F16"], + ["rope.bands", [16], "F32"], + ["shared_block_modulation.bias", [18432], "F16"], + ["shared_block_modulation.weight", [3072, 18432], "F16"], + ["text_embedder.blocks.0.ffn_down.bias", [3072], "F16"], + ["text_embedder.blocks.0.ffn_down.weight", [8192, 3072], "F16"], + ["text_embedder.blocks.0.ffn_gate.bias", [8192], "F16"], + ["text_embedder.blocks.0.ffn_gate.weight", [3072, 8192], "F16"], + ["text_embedder.blocks.0.ffn_up.bias", [8192], "F16"], + ["text_embedder.blocks.0.ffn_up.weight", [3072, 8192], "F16"], + ["text_embedder.blocks.0.k_norm.weight", [128], "F16"], + ["text_embedder.blocks.0.norm1.weight", [3072], "F16"], + ["text_embedder.blocks.0.norm2.weight", [3072], "F16"], + ["text_embedder.blocks.0.proj.bias", [3072], "F16"], + ["text_embedder.blocks.0.proj.weight", [3072, 3072], "F16"], + ["text_embedder.blocks.0.q_norm.weight", [128], "F16"], + ["text_embedder.blocks.0.qkv.bias", [9216], "F16"], + ["text_embedder.blocks.0.qkv.weight", [3072, 9216], "F16"], + ["text_embedder.blocks.1.ffn_down.bias", [3072], "F16"], + ["text_embedder.blocks.1.ffn_down.weight", [8192, 3072], "F16"], + ["text_embedder.blocks.1.ffn_gate.bias", [8192], "F16"], + ["text_embedder.blocks.1.ffn_gate.weight", [3072, 8192], "F16"], + ["text_embedder.blocks.1.ffn_up.bias", [8192], "F16"], + ["text_embedder.blocks.1.ffn_up.weight", [3072, 8192], "F16"], + ["text_embedder.blocks.1.k_norm.weight", [128], "F16"], + ["text_embedder.blocks.1.norm1.weight", [3072], "F16"], + ["text_embedder.blocks.1.norm2.weight", [3072], "F16"], + ["text_embedder.blocks.1.proj.bias", [3072], "F16"], + ["text_embedder.blocks.1.proj.weight", [3072, 3072], "F16"], + ["text_embedder.blocks.1.q_norm.weight", [128], "F16"], + ["text_embedder.blocks.1.qkv.bias", [9216], "F16"], + ["text_embedder.blocks.1.qkv.weight", [3072, 9216], "F16"], + ["text_embedder.proj.bias", [3072], "F16"], + ["text_embedder.proj.weight", [4096, 3072], "F16"], + ["time_embedding.0.bias", [3072], "F16"], + ["time_embedding.0.weight", [256, 3072], "F16"], + ["time_embedding.2.bias", [3072], "F16"], + ["time_embedding.2.weight", [3072, 3072], "F16"], + ["video_embedder.bias", [3072], "F16"], + ["video_embedder.weight", [128, 3072], "F16"], + ["video_head_modulation.bias", [6144], "F16"], + ["video_head_modulation.weight", [3072, 6144], "F16"] + ] +} diff --git a/tests/fixtures/image/moss_image_transformer_config.json b/tests/fixtures/image/moss_image_transformer_config.json new file mode 100644 index 000000000..09ba6435c --- /dev/null +++ b/tests/fixtures/image/moss_image_transformer_config.json @@ -0,0 +1,12 @@ +{ + "_class_name": "MOSSImageTransformer", + "_diffusers_version": "0.38.0", + "dim": 3072, + "ffn_dim": 8192, + "freq_dim": 256, + "latent_channels": 128, + "num_heads": 24, + "num_layers": 42, + "rope_base": 10000.0, + "text_dim": 4096 +} diff --git a/tests/fixtures/image/moss_image_vae.header.json b/tests/fixtures/image/moss_image_vae.header.json new file mode 100644 index 000000000..e47b4a721 --- /dev/null +++ b/tests/fixtures/image/moss_image_vae.header.json @@ -0,0 +1,259 @@ +{ + "format": "safetensors", + "header_bytes": 27968, + "payload_bytes": 249493396, + "n_tensors": 251, + "tensors": [ + ["bn.num_batches_tracked", [], "I64"], + ["bn.running_mean", [128], "F32"], + ["bn.running_var", [128], "F32"], + ["decoder.conv_in.bias", [384], "F32"], + ["decoder.conv_in.weight", [384, 32, 3, 3], "F32"], + ["decoder.conv_norm_out.bias", [96], "F32"], + ["decoder.conv_norm_out.weight", [96], "F32"], + ["decoder.conv_out.bias", [3], "F32"], + ["decoder.conv_out.weight", [3, 96, 3, 3], "F32"], + ["decoder.mid_block.attentions.0.group_norm.bias", [384], "F32"], + ["decoder.mid_block.attentions.0.group_norm.weight", [384], "F32"], + ["decoder.mid_block.attentions.0.to_k.bias", [384], "F32"], + ["decoder.mid_block.attentions.0.to_k.weight", [384, 384], "F32"], + ["decoder.mid_block.attentions.0.to_out.0.bias", [384], "F32"], + ["decoder.mid_block.attentions.0.to_out.0.weight", [384, 384], "F32"], + ["decoder.mid_block.attentions.0.to_q.bias", [384], "F32"], + ["decoder.mid_block.attentions.0.to_q.weight", [384, 384], "F32"], + ["decoder.mid_block.attentions.0.to_v.bias", [384], "F32"], + ["decoder.mid_block.attentions.0.to_v.weight", [384, 384], "F32"], + ["decoder.mid_block.resnets.0.conv1.bias", [384], "F32"], + ["decoder.mid_block.resnets.0.conv1.weight", [384, 384, 3, 3], "F32"], + ["decoder.mid_block.resnets.0.conv2.bias", [384], "F32"], + ["decoder.mid_block.resnets.0.conv2.weight", [384, 384, 3, 3], "F32"], + ["decoder.mid_block.resnets.0.norm1.bias", [384], "F32"], + ["decoder.mid_block.resnets.0.norm1.weight", [384], "F32"], + ["decoder.mid_block.resnets.0.norm2.bias", [384], "F32"], + ["decoder.mid_block.resnets.0.norm2.weight", [384], "F32"], + ["decoder.mid_block.resnets.1.conv1.bias", [384], "F32"], + ["decoder.mid_block.resnets.1.conv1.weight", [384, 384, 3, 3], "F32"], + ["decoder.mid_block.resnets.1.conv2.bias", [384], "F32"], + ["decoder.mid_block.resnets.1.conv2.weight", [384, 384, 3, 3], "F32"], + ["decoder.mid_block.resnets.1.norm1.bias", [384], "F32"], + ["decoder.mid_block.resnets.1.norm1.weight", [384], "F32"], + ["decoder.mid_block.resnets.1.norm2.bias", [384], "F32"], + ["decoder.mid_block.resnets.1.norm2.weight", [384], "F32"], + ["decoder.up_blocks.0.resnets.0.conv1.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.0.conv1.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.0.resnets.0.conv2.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.0.conv2.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.0.resnets.0.norm1.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.0.norm1.weight", [384], "F32"], + ["decoder.up_blocks.0.resnets.0.norm2.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.0.norm2.weight", [384], "F32"], + ["decoder.up_blocks.0.resnets.1.conv1.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.1.conv1.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.0.resnets.1.conv2.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.1.conv2.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.0.resnets.1.norm1.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.1.norm1.weight", [384], "F32"], + ["decoder.up_blocks.0.resnets.1.norm2.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.1.norm2.weight", [384], "F32"], + ["decoder.up_blocks.0.resnets.2.conv1.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.2.conv1.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.0.resnets.2.conv2.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.2.conv2.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.0.resnets.2.norm1.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.2.norm1.weight", [384], "F32"], + ["decoder.up_blocks.0.resnets.2.norm2.bias", [384], "F32"], + ["decoder.up_blocks.0.resnets.2.norm2.weight", [384], "F32"], + ["decoder.up_blocks.0.upsamplers.0.conv.bias", [384], "F32"], + ["decoder.up_blocks.0.upsamplers.0.conv.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.1.resnets.0.conv1.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.0.conv1.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.1.resnets.0.conv2.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.0.conv2.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.1.resnets.0.norm1.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.0.norm1.weight", [384], "F32"], + ["decoder.up_blocks.1.resnets.0.norm2.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.0.norm2.weight", [384], "F32"], + ["decoder.up_blocks.1.resnets.1.conv1.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.1.conv1.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.1.resnets.1.conv2.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.1.conv2.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.1.resnets.1.norm1.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.1.norm1.weight", [384], "F32"], + ["decoder.up_blocks.1.resnets.1.norm2.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.1.norm2.weight", [384], "F32"], + ["decoder.up_blocks.1.resnets.2.conv1.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.2.conv1.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.1.resnets.2.conv2.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.2.conv2.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.1.resnets.2.norm1.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.2.norm1.weight", [384], "F32"], + ["decoder.up_blocks.1.resnets.2.norm2.bias", [384], "F32"], + ["decoder.up_blocks.1.resnets.2.norm2.weight", [384], "F32"], + ["decoder.up_blocks.1.upsamplers.0.conv.bias", [384], "F32"], + ["decoder.up_blocks.1.upsamplers.0.conv.weight", [384, 384, 3, 3], "F32"], + ["decoder.up_blocks.2.resnets.0.conv1.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.0.conv1.weight", [192, 384, 3, 3], "F32"], + ["decoder.up_blocks.2.resnets.0.conv2.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.0.conv2.weight", [192, 192, 3, 3], "F32"], + ["decoder.up_blocks.2.resnets.0.conv_shortcut.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.0.conv_shortcut.weight", [192, 384, 1, 1], "F32"], + ["decoder.up_blocks.2.resnets.0.norm1.bias", [384], "F32"], + ["decoder.up_blocks.2.resnets.0.norm1.weight", [384], "F32"], + ["decoder.up_blocks.2.resnets.0.norm2.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.0.norm2.weight", [192], "F32"], + ["decoder.up_blocks.2.resnets.1.conv1.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.1.conv1.weight", [192, 192, 3, 3], "F32"], + ["decoder.up_blocks.2.resnets.1.conv2.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.1.conv2.weight", [192, 192, 3, 3], "F32"], + ["decoder.up_blocks.2.resnets.1.norm1.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.1.norm1.weight", [192], "F32"], + ["decoder.up_blocks.2.resnets.1.norm2.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.1.norm2.weight", [192], "F32"], + ["decoder.up_blocks.2.resnets.2.conv1.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.2.conv1.weight", [192, 192, 3, 3], "F32"], + ["decoder.up_blocks.2.resnets.2.conv2.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.2.conv2.weight", [192, 192, 3, 3], "F32"], + ["decoder.up_blocks.2.resnets.2.norm1.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.2.norm1.weight", [192], "F32"], + ["decoder.up_blocks.2.resnets.2.norm2.bias", [192], "F32"], + ["decoder.up_blocks.2.resnets.2.norm2.weight", [192], "F32"], + ["decoder.up_blocks.2.upsamplers.0.conv.bias", [192], "F32"], + ["decoder.up_blocks.2.upsamplers.0.conv.weight", [192, 192, 3, 3], "F32"], + ["decoder.up_blocks.3.resnets.0.conv1.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.0.conv1.weight", [96, 192, 3, 3], "F32"], + ["decoder.up_blocks.3.resnets.0.conv2.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.0.conv2.weight", [96, 96, 3, 3], "F32"], + ["decoder.up_blocks.3.resnets.0.conv_shortcut.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.0.conv_shortcut.weight", [96, 192, 1, 1], "F32"], + ["decoder.up_blocks.3.resnets.0.norm1.bias", [192], "F32"], + ["decoder.up_blocks.3.resnets.0.norm1.weight", [192], "F32"], + ["decoder.up_blocks.3.resnets.0.norm2.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.0.norm2.weight", [96], "F32"], + ["decoder.up_blocks.3.resnets.1.conv1.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.1.conv1.weight", [96, 96, 3, 3], "F32"], + ["decoder.up_blocks.3.resnets.1.conv2.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.1.conv2.weight", [96, 96, 3, 3], "F32"], + ["decoder.up_blocks.3.resnets.1.norm1.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.1.norm1.weight", [96], "F32"], + ["decoder.up_blocks.3.resnets.1.norm2.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.1.norm2.weight", [96], "F32"], + ["decoder.up_blocks.3.resnets.2.conv1.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.2.conv1.weight", [96, 96, 3, 3], "F32"], + ["decoder.up_blocks.3.resnets.2.conv2.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.2.conv2.weight", [96, 96, 3, 3], "F32"], + ["decoder.up_blocks.3.resnets.2.norm1.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.2.norm1.weight", [96], "F32"], + ["decoder.up_blocks.3.resnets.2.norm2.bias", [96], "F32"], + ["decoder.up_blocks.3.resnets.2.norm2.weight", [96], "F32"], + ["encoder.conv_in.bias", [128], "F32"], + ["encoder.conv_in.weight", [128, 3, 3, 3], "F32"], + ["encoder.conv_norm_out.bias", [512], "F32"], + ["encoder.conv_norm_out.weight", [512], "F32"], + ["encoder.conv_out.bias", [64], "F32"], + ["encoder.conv_out.weight", [64, 512, 3, 3], "F32"], + ["encoder.down_blocks.0.downsamplers.0.conv.bias", [128], "F32"], + ["encoder.down_blocks.0.downsamplers.0.conv.weight", [128, 128, 3, 3], "F32"], + ["encoder.down_blocks.0.resnets.0.conv1.bias", [128], "F32"], + ["encoder.down_blocks.0.resnets.0.conv1.weight", [128, 128, 3, 3], "F32"], + ["encoder.down_blocks.0.resnets.0.conv2.bias", [128], "F32"], + ["encoder.down_blocks.0.resnets.0.conv2.weight", [128, 128, 3, 3], "F32"], + ["encoder.down_blocks.0.resnets.0.norm1.bias", [128], "F32"], + ["encoder.down_blocks.0.resnets.0.norm1.weight", [128], "F32"], + ["encoder.down_blocks.0.resnets.0.norm2.bias", [128], "F32"], + ["encoder.down_blocks.0.resnets.0.norm2.weight", [128], "F32"], + ["encoder.down_blocks.0.resnets.1.conv1.bias", [128], "F32"], + ["encoder.down_blocks.0.resnets.1.conv1.weight", [128, 128, 3, 3], "F32"], + ["encoder.down_blocks.0.resnets.1.conv2.bias", [128], "F32"], + ["encoder.down_blocks.0.resnets.1.conv2.weight", [128, 128, 3, 3], "F32"], + ["encoder.down_blocks.0.resnets.1.norm1.bias", [128], "F32"], + ["encoder.down_blocks.0.resnets.1.norm1.weight", [128], "F32"], + ["encoder.down_blocks.0.resnets.1.norm2.bias", [128], "F32"], + ["encoder.down_blocks.0.resnets.1.norm2.weight", [128], "F32"], + ["encoder.down_blocks.1.downsamplers.0.conv.bias", [256], "F32"], + ["encoder.down_blocks.1.downsamplers.0.conv.weight", [256, 256, 3, 3], "F32"], + ["encoder.down_blocks.1.resnets.0.conv1.bias", [256], "F32"], + ["encoder.down_blocks.1.resnets.0.conv1.weight", [256, 128, 3, 3], "F32"], + ["encoder.down_blocks.1.resnets.0.conv2.bias", [256], "F32"], + ["encoder.down_blocks.1.resnets.0.conv2.weight", [256, 256, 3, 3], "F32"], + ["encoder.down_blocks.1.resnets.0.conv_shortcut.bias", [256], "F32"], + ["encoder.down_blocks.1.resnets.0.conv_shortcut.weight", [256, 128, 1, 1], "F32"], + ["encoder.down_blocks.1.resnets.0.norm1.bias", [128], "F32"], + ["encoder.down_blocks.1.resnets.0.norm1.weight", [128], "F32"], + ["encoder.down_blocks.1.resnets.0.norm2.bias", [256], "F32"], + ["encoder.down_blocks.1.resnets.0.norm2.weight", [256], "F32"], + ["encoder.down_blocks.1.resnets.1.conv1.bias", [256], "F32"], + ["encoder.down_blocks.1.resnets.1.conv1.weight", [256, 256, 3, 3], "F32"], + ["encoder.down_blocks.1.resnets.1.conv2.bias", [256], "F32"], + ["encoder.down_blocks.1.resnets.1.conv2.weight", [256, 256, 3, 3], "F32"], + ["encoder.down_blocks.1.resnets.1.norm1.bias", [256], "F32"], + ["encoder.down_blocks.1.resnets.1.norm1.weight", [256], "F32"], + ["encoder.down_blocks.1.resnets.1.norm2.bias", [256], "F32"], + ["encoder.down_blocks.1.resnets.1.norm2.weight", [256], "F32"], + ["encoder.down_blocks.2.downsamplers.0.conv.bias", [512], "F32"], + ["encoder.down_blocks.2.downsamplers.0.conv.weight", [512, 512, 3, 3], "F32"], + ["encoder.down_blocks.2.resnets.0.conv1.bias", [512], "F32"], + ["encoder.down_blocks.2.resnets.0.conv1.weight", [512, 256, 3, 3], "F32"], + ["encoder.down_blocks.2.resnets.0.conv2.bias", [512], "F32"], + ["encoder.down_blocks.2.resnets.0.conv2.weight", [512, 512, 3, 3], "F32"], + ["encoder.down_blocks.2.resnets.0.conv_shortcut.bias", [512], "F32"], + ["encoder.down_blocks.2.resnets.0.conv_shortcut.weight", [512, 256, 1, 1], "F32"], + ["encoder.down_blocks.2.resnets.0.norm1.bias", [256], "F32"], + ["encoder.down_blocks.2.resnets.0.norm1.weight", [256], "F32"], + ["encoder.down_blocks.2.resnets.0.norm2.bias", [512], "F32"], + ["encoder.down_blocks.2.resnets.0.norm2.weight", [512], "F32"], + ["encoder.down_blocks.2.resnets.1.conv1.bias", [512], "F32"], + ["encoder.down_blocks.2.resnets.1.conv1.weight", [512, 512, 3, 3], "F32"], + ["encoder.down_blocks.2.resnets.1.conv2.bias", [512], "F32"], + ["encoder.down_blocks.2.resnets.1.conv2.weight", [512, 512, 3, 3], "F32"], + ["encoder.down_blocks.2.resnets.1.norm1.bias", [512], "F32"], + ["encoder.down_blocks.2.resnets.1.norm1.weight", [512], "F32"], + ["encoder.down_blocks.2.resnets.1.norm2.bias", [512], "F32"], + ["encoder.down_blocks.2.resnets.1.norm2.weight", [512], "F32"], + ["encoder.down_blocks.3.resnets.0.conv1.bias", [512], "F32"], + ["encoder.down_blocks.3.resnets.0.conv1.weight", [512, 512, 3, 3], "F32"], + ["encoder.down_blocks.3.resnets.0.conv2.bias", [512], "F32"], + ["encoder.down_blocks.3.resnets.0.conv2.weight", [512, 512, 3, 3], "F32"], + ["encoder.down_blocks.3.resnets.0.norm1.bias", [512], "F32"], + ["encoder.down_blocks.3.resnets.0.norm1.weight", [512], "F32"], + ["encoder.down_blocks.3.resnets.0.norm2.bias", [512], "F32"], + ["encoder.down_blocks.3.resnets.0.norm2.weight", [512], "F32"], + ["encoder.down_blocks.3.resnets.1.conv1.bias", [512], "F32"], + ["encoder.down_blocks.3.resnets.1.conv1.weight", [512, 512, 3, 3], "F32"], + ["encoder.down_blocks.3.resnets.1.conv2.bias", [512], "F32"], + ["encoder.down_blocks.3.resnets.1.conv2.weight", [512, 512, 3, 3], "F32"], + ["encoder.down_blocks.3.resnets.1.norm1.bias", [512], "F32"], + ["encoder.down_blocks.3.resnets.1.norm1.weight", [512], "F32"], + ["encoder.down_blocks.3.resnets.1.norm2.bias", [512], "F32"], + ["encoder.down_blocks.3.resnets.1.norm2.weight", [512], "F32"], + ["encoder.mid_block.attentions.0.group_norm.bias", [512], "F32"], + ["encoder.mid_block.attentions.0.group_norm.weight", [512], "F32"], + ["encoder.mid_block.attentions.0.to_k.bias", [512], "F32"], + ["encoder.mid_block.attentions.0.to_k.weight", [512, 512], "F32"], + ["encoder.mid_block.attentions.0.to_out.0.bias", [512], "F32"], + ["encoder.mid_block.attentions.0.to_out.0.weight", [512, 512], "F32"], + ["encoder.mid_block.attentions.0.to_q.bias", [512], "F32"], + ["encoder.mid_block.attentions.0.to_q.weight", [512, 512], "F32"], + ["encoder.mid_block.attentions.0.to_v.bias", [512], "F32"], + ["encoder.mid_block.attentions.0.to_v.weight", [512, 512], "F32"], + ["encoder.mid_block.resnets.0.conv1.bias", [512], "F32"], + ["encoder.mid_block.resnets.0.conv1.weight", [512, 512, 3, 3], "F32"], + ["encoder.mid_block.resnets.0.conv2.bias", [512], "F32"], + ["encoder.mid_block.resnets.0.conv2.weight", [512, 512, 3, 3], "F32"], + ["encoder.mid_block.resnets.0.norm1.bias", [512], "F32"], + ["encoder.mid_block.resnets.0.norm1.weight", [512], "F32"], + ["encoder.mid_block.resnets.0.norm2.bias", [512], "F32"], + ["encoder.mid_block.resnets.0.norm2.weight", [512], "F32"], + ["encoder.mid_block.resnets.1.conv1.bias", [512], "F32"], + ["encoder.mid_block.resnets.1.conv1.weight", [512, 512, 3, 3], "F32"], + ["encoder.mid_block.resnets.1.conv2.bias", [512], "F32"], + ["encoder.mid_block.resnets.1.conv2.weight", [512, 512, 3, 3], "F32"], + ["encoder.mid_block.resnets.1.norm1.bias", [512], "F32"], + ["encoder.mid_block.resnets.1.norm1.weight", [512], "F32"], + ["encoder.mid_block.resnets.1.norm2.bias", [512], "F32"], + ["encoder.mid_block.resnets.1.norm2.weight", [512], "F32"], + ["post_quant_conv.bias", [32], "F32"], + ["post_quant_conv.weight", [32, 32, 1, 1], "F32"], + ["quant_conv.bias", [64], "F32"], + ["quant_conv.weight", [64, 64, 1, 1], "F32"] + ] +} diff --git a/tests/fixtures/image/moss_image_vae_config.json b/tests/fixtures/image/moss_image_vae_config.json new file mode 100644 index 000000000..3c982a055 --- /dev/null +++ b/tests/fixtures/image/moss_image_vae_config.json @@ -0,0 +1,45 @@ +{ + "_class_name": "AutoencoderKLFlux2", + "_diffusers_version": "0.38.0", + "act_fn": "silu", + "batch_norm_eps": 0.0001, + "batch_norm_momentum": 0.1, + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "decoder_block_out_channels": [ + 96, + 192, + 384, + 384 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "force_upcast": true, + "in_channels": 3, + "latent_channels": 32, + "layers_per_block": 2, + "mid_block_add_attention": true, + "norm_num_groups": 32, + "out_channels": 3, + "patch_size": [ + 2, + 2 + ], + "sample_size": 1024, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "use_post_quant_conv": true, + "use_quant_conv": true +} diff --git a/tests/fixtures/image/oracle/check2_moss_image_q4_k_m.manifest.json b/tests/fixtures/image/oracle/check2_moss_image_q4_k_m.manifest.json new file mode 100644 index 000000000..45e365d42 --- /dev/null +++ b/tests/fixtures/image/oracle/check2_moss_image_q4_k_m.manifest.json @@ -0,0 +1,369 @@ +{ + "schema": 1, + "name": "check2_moss_image_q4_k_m", + "check": 2, + "family": "moss_image", + "script": "scripts/image_oracle/check2_moss_image.py", + "argv": [ + "check2_moss_image.py", + "--encoder", + "Qwen3.5-9B-Q4_K_M.gguf", + "--tokenizer", + "tok-check2", + "--name", + "check2_moss_image_q4_k_m" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": { + "encoder": { + "name": "Qwen3.5-9B-Q4_K_M.gguf", + "size": 5680522464, + "sha256": "03b74727a860a56338e042c4420bb3f04b2fec5734175f4cb9fa853daf52b7e8" + }, + "tokenizer/tokenizer.json": { + "name": "tokenizer.json", + "size": 19989325, + "sha256": "06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523" + }, + "tokenizer/tokenizer_config.json": { + "name": "tokenizer_config.json", + "size": 1123, + "sha256": "bee8eba30f0eb4af73c0fe2cd06d0f89b657d7819941c438157ec42f7c80ea87" + }, + "tokenizer/chat_template.jinja": { + "name": "chat_template.jinja", + "size": 7756, + "sha256": "a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715" + } + }, + "inputs": { + "prompts": { + "plain": "A red bicycle leaning against a weathered brick wall at sunset, a cat asleep in its basket, warm light, 35mm photograph", + "empty": "", + "non_ascii": "Un caf\u00e9 \u00e0 Montr\u00e9al en hiver, \u7a93\u8fba\u3067\u5bdd\u3066\u3044\u308b\u732b, \u4e00\u676f\u62ff\u94c1\u548c\u4e00\u672c\u4e66, \u0441\u043d\u0435\u0433 \u0437\u0430 \u043e\u043a\u043d\u043e\u043c, \ud83d\udc31\u2615", + "long": "A narrow cobblestone street in an old harbor town at dusk, lined with three story houses painted ochre, teal and faded red. Laundry hangs on lines strung between the balconies, and a bicycle with a wicker basket leans against a green door. A fishmonger closes the shutters of a small shop while gulls circle above the masts in the harbor. Warm lamplight spills from a cafe where two old men play chess at a marble table, and a cat watches them from a windowsill full of geraniums. Puddles from an afternoon rain reflect the lamps and the violet sky, and the far end of the street opens onto the sea, where a lighthouse is just switching on. In the foreground a girl in a yellow raincoat crouches to inspect a snail on the curb, her boots splashed with mud. Behind her a baker carries a tray of round loaves dusted with flour toward a delivery van whose rear doors stand open. A hand painted sign above the bakery reads Panificio in chipped gold letters. Strings of small bulbs zigzag overhead, some of them dark, and a pigeon sits on the wire between two of them. On the left, a stone staircase climbs toward a chapel with a bell tower, its steps worn into shallow bowls by centuries of feet. Moss grows in the cracks, and someone has left a bunch of wildflowers on the third step. A scooter with a dented fender is parked at an angle beside a bollard shaped like a lion. The sky shades from deep indigo at the top of the frame to a band of apricot just above the horizon, with a thin crescent moon and a single bright planet. Fishing boats with blue and white hulls rock at their moorings, nets heaped on their decks, and a man in a wool cap coils rope on the quay. Far out on the water the running lights of a ferry move slowly toward an island whose hills are already black. The air looks cool and damp, with a faint haze that softens the distant lights into halos. Shot on medium format film with a normal lens at eye level, shallow depth of field on the girl and the snail, gentle grain, muted greens and warm highlights, no text other than the bakery sign, no watermark. The composition follows the curve of the street from the lower left corner to the lighthouse at the upper right, leading the eye past each small story along the way. Every window tells something: a family eating soup under a hanging lamp, a tailor pinning the hem of a coat on a wooden dummy, an empty room with a piano and a vase of dry roses, a boy practicing the violin with his music stand turned toward the glass. Rainwater still drips from the gutters, catching the light in bright beads, and the cobbles shine like the backs of beetles. A black dog trots out of an alley with a paper bag in its mouth, looking over its shoulder as if someone might follow.", + "devanagari": "a temple in \u0928\u092e\u0938\u094d\u0924\u0947 style", + "thai": "a shop sign that reads \u0e2a\u0e27\u0e31\u0e2a\u0e14\u0e35" + }, + "max_sequence_length": 4096, + "tokenizer_class": "Qwen2Tokenizer", + "pre_tokenizer": { + "type": "Sequence", + "pretokenizers": [ + { + "type": "Split", + "pattern": { + "Regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" + }, + "behavior": "Isolated", + "invert": false + }, + { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": false + } + ] + }, + "read": "last_hidden_state, after the final norm", + "template": "the tokenizer's chat template, one user turn, add_generation_prompt and enable_thinking", + "matmul_bf16": "bf16 operands, float32 GEMM, one bf16 rounding", + "load_s": 10.9, + "layers": 32, + "gguf_codecs": { + "Q6_K": 22, + "F32": 177, + "Q4_K": 132, + "Q5_K": 48, + "Q8_0": 48 + }, + "dequant_checked": { + "Q6_K": [ + "blk.0.ffn_down.weight", + "layers.0.mlp.down_proj.weight", + "bitwise" + ], + "Q4_K": [ + "blk.4.ffn_down.weight", + "layers.4.mlp.down_proj.weight", + "bitwise" + ], + "Q5_K": [ + "blk.0.ssm_out.weight", + "layers.0.linear_attn.out_proj.weight", + "reordered" + ], + "Q8_0": [ + "blk.0.ssm_alpha.weight", + "layers.0.linear_attn.in_proj_a.weight", + "reordered" + ] + }, + "rows": { + "plain": 38, + "empty": 10, + "non_ascii": 42, + "long": 624, + "devanagari": 19, + "thai": 20 + } + }, + "results": { + "bf16_vs_fp32_plain": { + "rel_l2": 0.012346841173235992, + "token_cosine_min": 0.9995899942390147, + "worst_token": 27 + }, + "bf16_vs_fp32_empty": { + "rel_l2": 0.018317916206891317, + "token_cosine_min": 0.9995983465301802, + "worst_token": 8 + }, + "bf16_vs_fp32_non_ascii": { + "rel_l2": 0.013902518842580282, + "token_cosine_min": 0.9995989687731079, + "worst_token": 29 + }, + "bf16_vs_fp32_long": { + "rel_l2": 0.07516326836164329, + "token_cosine_min": 0.15338572738806555, + "worst_token": 288 + }, + "bf16_vs_fp32_devanagari": { + "rel_l2": 0.014019624674789555, + "token_cosine_min": 0.9996527201627962, + "worst_token": 16 + }, + "bf16_vs_fp32_thai": { + "rel_l2": 0.015967255207340315, + "token_cosine_min": 0.9995678070456726, + "worst_token": 15 + }, + "seconds": { + "plain": { + "fp32": 10.6, + "bf16": 14.2 + }, + "empty": { + "fp32": 9.8, + "bf16": 13.8 + }, + "non_ascii": { + "fp32": 10.1, + "bf16": 14.3 + }, + "long": { + "fp32": 17.2, + "bf16": 20.2 + }, + "devanagari": { + "fp32": 13.8, + "bf16": 13.9 + }, + "thai": { + "fp32": 13.9, + "bf16": 13.9 + } + } + }, + "tolerances": { + "rel_l2": { + "plain": 0.024693682346471985, + "empty": 0.03663583241378263, + "non_ascii": 0.027805037685160565, + "long": 0.15032653672328658, + "devanagari": 0.02803924934957911, + "thai": 0.03193451041468063 + }, + "token_cosine_min": { + "plain": 0.9983599769560589, + "empty": 0.9983933861207208, + "non_ascii": 0.9983958750924318, + "long": -2.386457090447738, + "devanagari": 0.9986108806511846, + "thai": 0.9982712281826904 + }, + "rule": "relative L2 = 2 x the oracle's own bf16 against fp32 relative L2 of the same prompt; per-token cosine at least 1 - 4 (1 - c), c the oracle's own smallest per-token cosine of the same prompt" + }, + "notes": [ + "The bf16 model is the float32 model rounded to bfloat16 once, as the release loads its encoder in bfloat16." + ], + "created": "2026-10-18T06:11:14Z", + "npz": "check2_moss_image_q4_k_m.npz", + "tensors": { + "features_devanagari": { + "shape": [ + 19, + 4096 + ], + "dtype": "float32", + "sha256": "5085e4297a91d3ef38b2bc58692b5b4c5fc4e227f1040b69cb0b51bde2353e6c" + }, + "features_devanagari_bf16": { + "shape": [ + 19, + 4096 + ], + "dtype": "float32", + "sha256": "b504d750a438e683aeabd34cc250b69c04ad28f4fe6a13860511ab9d8a9d1ddf" + }, + "features_empty": { + "shape": [ + 10, + 4096 + ], + "dtype": "float32", + "sha256": "b3f950d323d51789c0645b413da9174e355ea9149f751b01305b8513b6cea527" + }, + "features_empty_bf16": { + "shape": [ + 10, + 4096 + ], + "dtype": "float32", + "sha256": "b8d599ca06534d47d6c728b9b6f876018dd6723ab19ab7890a3ae29c4bef84da" + }, + "features_long": { + "shape": [ + 624, + 4096 + ], + "dtype": "float32", + "sha256": "f5458b1a5812c960fc3c62ff974e40a66d34b7bc98b5db6125e2ac6d1bf35fa8" + }, + "features_long_bf16": { + "shape": [ + 624, + 4096 + ], + "dtype": "float32", + "sha256": "cac750035ccf95c1e7cdbaa61b949f02ba0961e216ca3c28bc22f8d02ad235a3" + }, + "features_non_ascii": { + "shape": [ + 42, + 4096 + ], + "dtype": "float32", + "sha256": "d3d274436148eac1dbb85c148dc198d913c83944c043ad61d4b345277041eabc" + }, + "features_non_ascii_bf16": { + "shape": [ + 42, + 4096 + ], + "dtype": "float32", + "sha256": "561b98bbea742ccad14d43c9770b2c601e55014de2ed7fbea29b59ddcd290f82" + }, + "features_plain": { + "shape": [ + 38, + 4096 + ], + "dtype": "float32", + "sha256": "27d8b0c99bd9aa5ab204445ec70e964d6b941f45e21b14ed19e06db3dfeebead" + }, + "features_plain_bf16": { + "shape": [ + 38, + 4096 + ], + "dtype": "float32", + "sha256": "e958b65033f6337dfb48f72315da46103962425503d813f4a80a89a2965f509e" + }, + "features_thai": { + "shape": [ + 20, + 4096 + ], + "dtype": "float32", + "sha256": "1b9150b2d813eb9f436aca5ca01dd0bcb7ab1aa885fa65fb9ecd6160b85786b7" + }, + "features_thai_bf16": { + "shape": [ + 20, + 4096 + ], + "dtype": "float32", + "sha256": "43b37ec566747e83e6892b1b3c2d7f2e37b83d73c2978bc68f3e2cde966b9e1f" + }, + "ids_devanagari": { + "shape": [ + 19 + ], + "dtype": "int64", + "sha256": "1b35cfdb055d4709b7c0519682280d3c0daf5f0165459e7ca5bab60f40f78608" + }, + "ids_empty": { + "shape": [ + 10 + ], + "dtype": "int64", + "sha256": "791228b9554afb37320a489dae5451fafa841df6222f4c55607ce44d12af5e70" + }, + "ids_long": { + "shape": [ + 624 + ], + "dtype": "int64", + "sha256": "fdb2a8d4b0d78d2f68a1f8f038fde8e7ce1945c533f90e2e9aaa59d5a6a1b36a" + }, + "ids_non_ascii": { + "shape": [ + 42 + ], + "dtype": "int64", + "sha256": "17df73bdc62a885d1fdc32d051008cc7f16dad4fb06d2988de2e1b97549a529b" + }, + "ids_plain": { + "shape": [ + 38 + ], + "dtype": "int64", + "sha256": "8ea803ffe864ccf67c051ca49f80e2c1bfd70e5fec38b4e92945b2ed02ec9078" + }, + "ids_thai": { + "shape": [ + 20 + ], + "dtype": "int64", + "sha256": "dea36020e8f7a8d56545f8c6fbb877d98470d76e9ecadeb4c7885927439f722a" + } + } +} diff --git a/tests/fixtures/image/oracle/check2_moss_image_q8_0.manifest.json b/tests/fixtures/image/oracle/check2_moss_image_q8_0.manifest.json new file mode 100644 index 000000000..708a392a2 --- /dev/null +++ b/tests/fixtures/image/oracle/check2_moss_image_q8_0.manifest.json @@ -0,0 +1,351 @@ +{ + "schema": 1, + "name": "check2_moss_image_q8_0", + "check": 2, + "family": "moss_image", + "script": "scripts/image_oracle/check2_moss_image.py", + "argv": [ + "check2_moss_image.py", + "--encoder", + "Qwen3.5-9B-Q8_0.gguf", + "--tokenizer", + "tok-check2", + "--name", + "check2_moss_image_q8_0" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": { + "encoder": { + "name": "Qwen3.5-9B-Q8_0.gguf", + "size": 9527502048, + "sha256": "809626574d0cb43d4becfa56169980da2bb448f2299270f7be443cb89d0a6ae4" + }, + "tokenizer/tokenizer.json": { + "name": "tokenizer.json", + "size": 19989325, + "sha256": "06b9509352d2af50381ab2247e083b80d32d5c0aba91c272ca9ff729b6a0e523" + }, + "tokenizer/tokenizer_config.json": { + "name": "tokenizer_config.json", + "size": 1123, + "sha256": "bee8eba30f0eb4af73c0fe2cd06d0f89b657d7819941c438157ec42f7c80ea87" + }, + "tokenizer/chat_template.jinja": { + "name": "chat_template.jinja", + "size": 7756, + "sha256": "a4aee8afcf2e0711942cf848899be66016f8d14a889ff9ede07bca099c28f715" + } + }, + "inputs": { + "prompts": { + "plain": "A red bicycle leaning against a weathered brick wall at sunset, a cat asleep in its basket, warm light, 35mm photograph", + "empty": "", + "non_ascii": "Un caf\u00e9 \u00e0 Montr\u00e9al en hiver, \u7a93\u8fba\u3067\u5bdd\u3066\u3044\u308b\u732b, \u4e00\u676f\u62ff\u94c1\u548c\u4e00\u672c\u4e66, \u0441\u043d\u0435\u0433 \u0437\u0430 \u043e\u043a\u043d\u043e\u043c, \ud83d\udc31\u2615", + "long": "A narrow cobblestone street in an old harbor town at dusk, lined with three story houses painted ochre, teal and faded red. Laundry hangs on lines strung between the balconies, and a bicycle with a wicker basket leans against a green door. A fishmonger closes the shutters of a small shop while gulls circle above the masts in the harbor. Warm lamplight spills from a cafe where two old men play chess at a marble table, and a cat watches them from a windowsill full of geraniums. Puddles from an afternoon rain reflect the lamps and the violet sky, and the far end of the street opens onto the sea, where a lighthouse is just switching on. In the foreground a girl in a yellow raincoat crouches to inspect a snail on the curb, her boots splashed with mud. Behind her a baker carries a tray of round loaves dusted with flour toward a delivery van whose rear doors stand open. A hand painted sign above the bakery reads Panificio in chipped gold letters. Strings of small bulbs zigzag overhead, some of them dark, and a pigeon sits on the wire between two of them. On the left, a stone staircase climbs toward a chapel with a bell tower, its steps worn into shallow bowls by centuries of feet. Moss grows in the cracks, and someone has left a bunch of wildflowers on the third step. A scooter with a dented fender is parked at an angle beside a bollard shaped like a lion. The sky shades from deep indigo at the top of the frame to a band of apricot just above the horizon, with a thin crescent moon and a single bright planet. Fishing boats with blue and white hulls rock at their moorings, nets heaped on their decks, and a man in a wool cap coils rope on the quay. Far out on the water the running lights of a ferry move slowly toward an island whose hills are already black. The air looks cool and damp, with a faint haze that softens the distant lights into halos. Shot on medium format film with a normal lens at eye level, shallow depth of field on the girl and the snail, gentle grain, muted greens and warm highlights, no text other than the bakery sign, no watermark. The composition follows the curve of the street from the lower left corner to the lighthouse at the upper right, leading the eye past each small story along the way. Every window tells something: a family eating soup under a hanging lamp, a tailor pinning the hem of a coat on a wooden dummy, an empty room with a piano and a vase of dry roses, a boy practicing the violin with his music stand turned toward the glass. Rainwater still drips from the gutters, catching the light in bright beads, and the cobbles shine like the backs of beetles. A black dog trots out of an alley with a paper bag in its mouth, looking over its shoulder as if someone might follow.", + "devanagari": "a temple in \u0928\u092e\u0938\u094d\u0924\u0947 style", + "thai": "a shop sign that reads \u0e2a\u0e27\u0e31\u0e2a\u0e14\u0e35" + }, + "max_sequence_length": 4096, + "tokenizer_class": "Qwen2Tokenizer", + "pre_tokenizer": { + "type": "Sequence", + "pretokenizers": [ + { + "type": "Split", + "pattern": { + "Regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" + }, + "behavior": "Isolated", + "invert": false + }, + { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": false + } + ] + }, + "read": "last_hidden_state, after the final norm", + "template": "the tokenizer's chat template, one user turn, add_generation_prompt and enable_thinking", + "matmul_bf16": "bf16 operands, float32 GEMM, one bf16 rounding", + "load_s": 3.8, + "layers": 32, + "gguf_codecs": { + "Q8_0": 250, + "F32": 177 + }, + "dequant_checked": { + "Q8_0": [ + "blk.0.ffn_down.weight", + "layers.0.mlp.down_proj.weight", + "bitwise" + ] + }, + "rows": { + "plain": 38, + "empty": 10, + "non_ascii": 42, + "long": 624, + "devanagari": 19, + "thai": 20 + } + }, + "results": { + "bf16_vs_fp32_plain": { + "rel_l2": 0.018506074700880806, + "token_cosine_min": 0.9961942226532843, + "worst_token": 27 + }, + "bf16_vs_fp32_empty": { + "rel_l2": 0.018652925477714802, + "token_cosine_min": 0.9996275532206776, + "worst_token": 0 + }, + "bf16_vs_fp32_non_ascii": { + "rel_l2": 0.013433558458008636, + "token_cosine_min": 0.9996275389821867, + "worst_token": 0 + }, + "bf16_vs_fp32_long": { + "rel_l2": 0.01638260329348712, + "token_cosine_min": 0.9907086601937181, + "worst_token": 487 + }, + "bf16_vs_fp32_devanagari": { + "rel_l2": 0.014366802951649446, + "token_cosine_min": 0.9996275389821867, + "worst_token": 0 + }, + "bf16_vs_fp32_thai": { + "rel_l2": 0.021674413671501514, + "token_cosine_min": 0.9989822577265628, + "worst_token": 12 + }, + "seconds": { + "plain": { + "fp32": 10.2, + "bf16": 16.3 + }, + "empty": { + "fp32": 9.8, + "bf16": 15.9 + }, + "non_ascii": { + "fp32": 10.0, + "bf16": 16.2 + }, + "long": { + "fp32": 17.9, + "bf16": 22.2 + }, + "devanagari": { + "fp32": 15.7, + "bf16": 15.9 + }, + "thai": { + "fp32": 15.6, + "bf16": 16.3 + } + } + }, + "tolerances": { + "rel_l2": { + "plain": 0.03701214940176161, + "empty": 0.037305850955429605, + "non_ascii": 0.026867116916017272, + "long": 0.03276520658697424, + "devanagari": 0.028733605903298893, + "thai": 0.04334882734300303 + }, + "token_cosine_min": { + "plain": 0.9847768906131371, + "empty": 0.9985102128827106, + "non_ascii": 0.9985101559287468, + "long": 0.9628346407748722, + "devanagari": 0.9985101559287468, + "thai": 0.9959290309062512 + }, + "rule": "relative L2 = 2 x the oracle's own bf16 against fp32 relative L2 of the same prompt; per-token cosine at least 1 - 4 (1 - c), c the oracle's own smallest per-token cosine of the same prompt" + }, + "notes": [ + "The bf16 model is the float32 model rounded to bfloat16 once, as the release loads its encoder in bfloat16." + ], + "created": "2026-10-18T06:15:06Z", + "npz": "check2_moss_image_q8_0.npz", + "tensors": { + "features_devanagari": { + "shape": [ + 19, + 4096 + ], + "dtype": "float32", + "sha256": "71705e2f81c09f454e064b6b98202890331f47217ef9d1fd86e533e91161a714" + }, + "features_devanagari_bf16": { + "shape": [ + 19, + 4096 + ], + "dtype": "float32", + "sha256": "68d4a189aa72df8d87d8437ae447bda322a70a8cc7140dcf743e5ec90a980923" + }, + "features_empty": { + "shape": [ + 10, + 4096 + ], + "dtype": "float32", + "sha256": "62dc84e4bdef784bf5365091e6b5552df9269755fb2ab41869b6512bdeb85074" + }, + "features_empty_bf16": { + "shape": [ + 10, + 4096 + ], + "dtype": "float32", + "sha256": "d13fd0823e378e2c27d71f8f71ffc0a125dc85445ef8a27853a0bebde360a5ae" + }, + "features_long": { + "shape": [ + 624, + 4096 + ], + "dtype": "float32", + "sha256": "241a66f62651d05b60046d99f5b228ee835728fd8e05a67a18695b67f0a669b5" + }, + "features_long_bf16": { + "shape": [ + 624, + 4096 + ], + "dtype": "float32", + "sha256": "12e70e2a5e995cddf7387c295cc9f3d82cd12ad05541dc8adefe75a91e6aa66b" + }, + "features_non_ascii": { + "shape": [ + 42, + 4096 + ], + "dtype": "float32", + "sha256": "e2d2ac9e78daf5f27338cc74799756bde2bae7d8a2d4df9408ff04d6bfa2620a" + }, + "features_non_ascii_bf16": { + "shape": [ + 42, + 4096 + ], + "dtype": "float32", + "sha256": "e08142ffc531a72ee34f9107531b5a74d6c7157cb927e34f92ff491d8ab1c348" + }, + "features_plain": { + "shape": [ + 38, + 4096 + ], + "dtype": "float32", + "sha256": "012ba987671e1c79a361c5a6ce972ffbbb7c47c6a5938268d57b0f63fb345c4e" + }, + "features_plain_bf16": { + "shape": [ + 38, + 4096 + ], + "dtype": "float32", + "sha256": "f551bb467976534bd652023c051470e27006fd0426789663aae36466109e8ed5" + }, + "features_thai": { + "shape": [ + 20, + 4096 + ], + "dtype": "float32", + "sha256": "b7391180ecc5085d239d4e36926d8457b29420800295c2e64879df2eaf6a9e9c" + }, + "features_thai_bf16": { + "shape": [ + 20, + 4096 + ], + "dtype": "float32", + "sha256": "9825a8d120007fd797884586bf607167a2de84e33f15ffd28449c40780bfb893" + }, + "ids_devanagari": { + "shape": [ + 19 + ], + "dtype": "int64", + "sha256": "1b35cfdb055d4709b7c0519682280d3c0daf5f0165459e7ca5bab60f40f78608" + }, + "ids_empty": { + "shape": [ + 10 + ], + "dtype": "int64", + "sha256": "791228b9554afb37320a489dae5451fafa841df6222f4c55607ce44d12af5e70" + }, + "ids_long": { + "shape": [ + 624 + ], + "dtype": "int64", + "sha256": "fdb2a8d4b0d78d2f68a1f8f038fde8e7ce1945c533f90e2e9aaa59d5a6a1b36a" + }, + "ids_non_ascii": { + "shape": [ + 42 + ], + "dtype": "int64", + "sha256": "17df73bdc62a885d1fdc32d051008cc7f16dad4fb06d2988de2e1b97549a529b" + }, + "ids_plain": { + "shape": [ + 38 + ], + "dtype": "int64", + "sha256": "8ea803ffe864ccf67c051ca49f80e2c1bfd70e5fec38b4e92945b2ed02ec9078" + }, + "ids_thai": { + "shape": [ + 20 + ], + "dtype": "int64", + "sha256": "dea36020e8f7a8d56545f8c6fbb877d98470d76e9ecadeb4c7885927439f722a" + } + } +} diff --git a/tests/fixtures/image/oracle/check3_moss_image_1024.manifest.json b/tests/fixtures/image/oracle/check3_moss_image_1024.manifest.json new file mode 100644 index 000000000..655ec0a05 --- /dev/null +++ b/tests/fixtures/image/oracle/check3_moss_image_1024.manifest.json @@ -0,0 +1,162 @@ +{ + "schema": 1, + "name": "check3_moss_image_1024", + "check": 3, + "family": "moss_image", + "script": "scripts/image_oracle/check3_moss_image.py", + "argv": [ + "check3_moss_image.py", + "--pipeline", + "pipeline_moss_image.py", + "--dit", + "MOSS-Image-SFT-Q4_K_M.gguf", + "--config", + "config.json", + "--sizes", + "512", + "1024" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": { + "dit": { + "name": "MOSS-Image-SFT-Q4_K_M.gguf", + "size": 3702092096, + "sha256": "c6dda89c18c0c0897d09ed01a515ee31415c4baea02bb9c2dedbfca66eec50b6" + }, + "pipeline": { + "name": "pipeline_moss_image.py", + "size": 25060, + "sha256": "1dac1aca9feb1bfdfc727a401e4b67d61a505d3190706d1b7d2d784f88678b31" + }, + "config": { + "name": "config.json", + "size": 240, + "sha256": "3291c3a1c7adc22b8182a0de25ca8ec97474d72715ce84adbab09cc0be24f666" + } + }, + "inputs": { + "size": 1024, + "blocks": 42, + "sigma": 0.75, + "seed": 1024, + "features": "seeded normal, std 1.0", + "text_rows": 48, + "dequant": "exact", + "matmul_bf16": "bf16 operands, float32 GEMM, one bf16 rounding", + "timestep": "sigma x 1000 in float32 in both modes", + "gguf_codecs": { + "F16": 210, + "Q6_K": 84, + "Q4_K": 168, + "F32": 1 + } + }, + "results": { + "velocity_bf16_vs_fp32": { + "rel_l2": 0.03968968370583882, + "cosine": 0.9992164898220565, + "max_abs": 0.41244423389434814 + }, + "times": { + "fp32_s": 28.7, + "bf16_s": 217.2 + } + }, + "tolerances": { + "velocity": { + "rel_l2": 0.07937936741167764, + "cosine_min": 0.9968659592882259 + }, + "rule": "tolerance = 2 x the oracle's own bf16 against fp32 relative L2: teacher-forced for each text block, each block and the head, free running for the other rows and for the velocity" + }, + "notes": [], + "created": "2026-10-18T06:39:57Z", + "npz": "check3_moss_image_1024.npz", + "tensors": { + "features": { + "shape": [ + 48, + 4096 + ], + "dtype": "float32", + "sha256": "51f2174fad964908bd5817f252f645cdc220c2c13ca457e1ad899301a9c3ca6b" + }, + "latents": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "a4b2103b3964c8a64efd0ddead9b65e6e2e11aed1e4cb4b625caae7c425a6bb6" + }, + "sigma": { + "shape": [], + "dtype": "float32", + "sha256": "9a8208635e00348ab64aac2b759e76391fd47089e9a749bbcec770d9eb5c6421" + }, + "timestep": { + "shape": [ + 1 + ], + "dtype": "float32", + "sha256": "1d77e51c3092b797b15b02104368537eea7b222148a4d2ba2d2d578491c50ac5" + }, + "velocity": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "cbb8347e7e9ed8836f5e2c3fd337957896ba67127b5613595e1cbc86697b9a0e" + }, + "velocity_bf16": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "60afef324b6f6aaca3aae7d6b2f13bade821fa2eace5eb78772e8f7bd4db3401" + } + } +} diff --git a/tests/fixtures/image/oracle/check3_moss_image_512.manifest.json b/tests/fixtures/image/oracle/check3_moss_image_512.manifest.json new file mode 100644 index 000000000..124792e6c --- /dev/null +++ b/tests/fixtures/image/oracle/check3_moss_image_512.manifest.json @@ -0,0 +1,1189 @@ +{ + "schema": 1, + "name": "check3_moss_image_512", + "check": 3, + "family": "moss_image", + "script": "scripts/image_oracle/check3_moss_image.py", + "argv": [ + "check3_moss_image.py", + "--pipeline", + "pipeline_moss_image.py", + "--dit", + "MOSS-Image-SFT-Q4_K_M.gguf", + "--config", + "config.json", + "--sizes", + "512", + "1024" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": { + "dit": { + "name": "MOSS-Image-SFT-Q4_K_M.gguf", + "size": 3702092096, + "sha256": "c6dda89c18c0c0897d09ed01a515ee31415c4baea02bb9c2dedbfca66eec50b6" + }, + "pipeline": { + "name": "pipeline_moss_image.py", + "size": 25060, + "sha256": "1dac1aca9feb1bfdfc727a401e4b67d61a505d3190706d1b7d2d784f88678b31" + }, + "config": { + "name": "config.json", + "size": 240, + "sha256": "3291c3a1c7adc22b8182a0de25ca8ec97474d72715ce84adbab09cc0be24f666" + } + }, + "inputs": { + "size": 512, + "blocks": 42, + "sigma": 0.75, + "seed": 512, + "features": "seeded normal, std 1.0", + "text_rows": 48, + "dequant": "exact", + "matmul_bf16": "bf16 operands, float32 GEMM, one bf16 rounding", + "timestep": "sigma x 1000 in float32 in both modes", + "gguf_codecs": { + "F16": 210, + "Q6_K": 84, + "Q4_K": 168, + "F32": 1 + } + }, + "results": { + "velocity_bf16_vs_fp32": { + "rel_l2": 0.04263153404754413, + "cosine": 0.9990925493828575, + "max_abs": 0.24877294898033142 + }, + "rows_bf16_vs_fp32": { + "time_embedding": { + "rel_l2": 0.0017880038297060797, + "cosine": 0.9999986603234691, + "max_abs": 0.0019777417182922363, + "teacher_forced": false + }, + "shared_block_modulation": { + "rel_l2": 0.0016879542098240739, + "cosine": 0.9999986500267788, + "max_abs": 0.10947418212890625, + "teacher_forced": false + }, + "video_head_modulation": { + "rel_l2": 0.0015511273133948249, + "cosine": 0.99999885703462, + "max_abs": 0.017541885375976562, + "teacher_forced": false + }, + "video_embedder": { + "rel_l2": 0.0023455135392989094, + "cosine": 0.9999972492970883, + "max_abs": 0.07721519470214844, + "teacher_forced": false + }, + "text_embedder.proj": { + "rel_l2": 0.002351121441498202, + "cosine": 0.9999972361475599, + "max_abs": 0.06422615051269531, + "teacher_forced": false + }, + "text_embedder.blocks.0": { + "rel_l2": 0.005746046923318964, + "cosine": 0.9999835047040323, + "max_abs": 0.5666732788085938, + "teacher_forced": true + }, + "text_embedder.blocks.1": { + "rel_l2": 0.004199174266344991, + "cosine": 0.9999911861772524, + "max_abs": 0.7038726806640625, + "teacher_forced": true + }, + "layers.0": { + "rel_l2": 0.0033805112919488146, + "cosine": 0.999995346920703, + "max_abs": 8.157958984375, + "teacher_forced": true + }, + "layers.1": { + "rel_l2": 0.003009731834582426, + "cosine": 0.9999954920743668, + "max_abs": 5.48150634765625, + "teacher_forced": true + }, + "layers.2": { + "rel_l2": 0.004130613466698473, + "cosine": 0.9999940978660147, + "max_abs": 8.69171142578125, + "teacher_forced": true + }, + "layers.3": { + "rel_l2": 0.0031660628551223283, + "cosine": 0.9999955647403599, + "max_abs": 5.998046875, + "teacher_forced": true + }, + "layers.4": { + "rel_l2": 0.0029541544640130663, + "cosine": 0.9999957468951098, + "max_abs": 7.046875, + "teacher_forced": true + }, + "layers.5": { + "rel_l2": 0.003320484301024555, + "cosine": 0.999995223192378, + "max_abs": 8.38897705078125, + "teacher_forced": true + }, + "layers.6": { + "rel_l2": 0.0028887037271521114, + "cosine": 0.9999962639760974, + "max_abs": 11.8671875, + "teacher_forced": true + }, + "layers.7": { + "rel_l2": 0.0029597123418961917, + "cosine": 0.9999956891043652, + "max_abs": 20.08544921875, + "teacher_forced": true + }, + "layers.8": { + "rel_l2": 0.0032251165892065764, + "cosine": 0.9999958379978068, + "max_abs": 15.343994140625, + "teacher_forced": true + }, + "layers.9": { + "rel_l2": 0.002638580464974077, + "cosine": 0.9999968604093041, + "max_abs": 11.0302734375, + "teacher_forced": true + }, + "layers.10": { + "rel_l2": 0.00306328017198254, + "cosine": 0.9999955335266718, + "max_abs": 16.603515625, + "teacher_forced": true + }, + "layers.11": { + "rel_l2": 0.003321258790878133, + "cosine": 0.9999961177140309, + "max_abs": 9.247314453125, + "teacher_forced": true + }, + "layers.12": { + "rel_l2": 0.0031863883451497127, + "cosine": 0.9999961751624868, + "max_abs": 12.5003662109375, + "teacher_forced": true + }, + "layers.13": { + "rel_l2": 0.003495689074212565, + "cosine": 0.9999949698328701, + "max_abs": 10.4444580078125, + "teacher_forced": true + }, + "layers.14": { + "rel_l2": 0.005342143241925715, + "cosine": 0.9999913380821597, + "max_abs": 13.3321533203125, + "teacher_forced": true + }, + "layers.15": { + "rel_l2": 0.00574672891394024, + "cosine": 0.9999926179147978, + "max_abs": 10.822021484375, + "teacher_forced": true + }, + "layers.16": { + "rel_l2": 0.005291596892396325, + "cosine": 0.9999909897784076, + "max_abs": 9.075927734375, + "teacher_forced": true + }, + "layers.17": { + "rel_l2": 0.005607773830647339, + "cosine": 0.9999908314293308, + "max_abs": 11.46331787109375, + "teacher_forced": true + }, + "layers.18": { + "rel_l2": 0.007217745065339213, + "cosine": 0.9999826332652529, + "max_abs": 12.289306640625, + "teacher_forced": true + }, + "layers.19": { + "rel_l2": 0.0071663273319924016, + "cosine": 0.9999841525281786, + "max_abs": 14.646484375, + "teacher_forced": true + }, + "layers.20": { + "rel_l2": 0.007403384909156674, + "cosine": 0.9999822498907733, + "max_abs": 12.659027099609375, + "teacher_forced": true + }, + "layers.21": { + "rel_l2": 0.013920375934480101, + "cosine": 0.9999397632907432, + "max_abs": 17.127685546875, + "teacher_forced": true + }, + "layers.22": { + "rel_l2": 0.011032856757270797, + "cosine": 0.9999549670803568, + "max_abs": 14.04632568359375, + "teacher_forced": true + }, + "layers.23": { + "rel_l2": 0.013922649668453847, + "cosine": 0.9999276106621329, + "max_abs": 17.78741455078125, + "teacher_forced": true + }, + "layers.24": { + "rel_l2": 0.013153346148434856, + "cosine": 0.9999206767419453, + "max_abs": 13.22265625, + "teacher_forced": true + }, + "layers.25": { + "rel_l2": 0.009471184242364534, + "cosine": 0.9999588735457219, + "max_abs": 9.81866455078125, + "teacher_forced": true + }, + "layers.26": { + "rel_l2": 0.008006400623644864, + "cosine": 0.9999679638371968, + "max_abs": 5.890594482421875, + "teacher_forced": true + }, + "layers.27": { + "rel_l2": 0.006452923314511727, + "cosine": 0.9999797909453065, + "max_abs": 5.79351806640625, + "teacher_forced": true + }, + "layers.28": { + "rel_l2": 0.0076359211943152734, + "cosine": 0.999971987683924, + "max_abs": 7.8470458984375, + "teacher_forced": true + }, + "layers.29": { + "rel_l2": 0.006695458690008976, + "cosine": 0.9999794911165114, + "max_abs": 7.352752685546875, + "teacher_forced": true + }, + "layers.30": { + "rel_l2": 0.00771848633843314, + "cosine": 0.9999716311038124, + "max_abs": 8.43572998046875, + "teacher_forced": true + }, + "layers.31": { + "rel_l2": 0.008970138012250588, + "cosine": 0.9999597835853143, + "max_abs": 12.18975830078125, + "teacher_forced": true + }, + "layers.32": { + "rel_l2": 0.008930251260973673, + "cosine": 0.9999603361245712, + "max_abs": 12.3056640625, + "teacher_forced": true + }, + "layers.33": { + "rel_l2": 0.012000333981036504, + "cosine": 0.9999323362450916, + "max_abs": 16.109619140625, + "teacher_forced": true + }, + "layers.34": { + "rel_l2": 0.01350675267516355, + "cosine": 0.999909256996868, + "max_abs": 20.7421875, + "teacher_forced": true + }, + "layers.35": { + "rel_l2": 0.016121105104871667, + "cosine": 0.9998703268616204, + "max_abs": 26.9990234375, + "teacher_forced": true + }, + "layers.36": { + "rel_l2": 0.0169766962307255, + "cosine": 0.9998560027261805, + "max_abs": 25.89013671875, + "teacher_forced": true + }, + "layers.37": { + "rel_l2": 0.016693277167473215, + "cosine": 0.9998606798213903, + "max_abs": 24.6363525390625, + "teacher_forced": true + }, + "layers.38": { + "rel_l2": 0.015450623953091228, + "cosine": 0.9998812196076702, + "max_abs": 29.994384765625, + "teacher_forced": true + }, + "layers.39": { + "rel_l2": 0.013879347363074222, + "cosine": 0.9999040155383271, + "max_abs": 24.751953125, + "teacher_forced": true + }, + "layers.40": { + "rel_l2": 0.012781894670281225, + "cosine": 0.9999183273886022, + "max_abs": 26.144775390625, + "teacher_forced": true + }, + "layers.41": { + "rel_l2": 0.018279534146230302, + "cosine": 0.9998331335830456, + "max_abs": 43.5479736328125, + "teacher_forced": true + }, + "final_linear_video": { + "rel_l2": 0.004121280183839743, + "cosine": 0.9999915293975221, + "max_abs": 0.020843982696533203, + "teacher_forced": true + } + }, + "free_bf16_vs_fp32": { + "layers.0": 0.003568796111638542, + "layers.1": 0.0043838969671087065, + "layers.2": 0.005312946831784859, + "layers.3": 0.005364772768698953, + "layers.4": 0.005263227543994745, + "layers.5": 0.005855628418669753, + "layers.6": 0.005230337516982755, + "layers.7": 0.004886924745399782, + "layers.8": 0.004947188118570181, + "layers.9": 0.004642114975796911, + "layers.10": 0.005425993462861186, + "layers.11": 0.005186772440651156, + "layers.12": 0.005538870223980716, + "layers.13": 0.006101318080710389, + "layers.14": 0.007373729814392103, + "layers.15": 0.006462006530598639, + "layers.16": 0.007376730058029977, + "layers.17": 0.007617372334552163, + "layers.18": 0.010276190958139659, + "layers.19": 0.009972706445031926, + "layers.20": 0.010898405247341175, + "layers.21": 0.013974074796937014, + "layers.22": 0.014360616997252834, + "layers.23": 0.015451414285155807, + "layers.24": 0.01882663016433713, + "layers.25": 0.01765633133372524, + "layers.26": 0.021828743226979832, + "layers.27": 0.02005387110535613, + "layers.28": 0.022913333129952806, + "layers.29": 0.020662938490167863, + "layers.30": 0.02464196720165473, + "layers.31": 0.02971256731064926, + "layers.32": 0.02963242237516329, + "layers.33": 0.028980736771996383, + "layers.34": 0.03538787795881819, + "layers.35": 0.03836622252429749, + "layers.36": 0.04508697412911542, + "layers.37": 0.04382474834929958, + "layers.38": 0.045972982171921736, + "layers.39": 0.047319224173234746, + "layers.40": 0.05429789226642868, + "layers.41": 0.06379060641629042 + }, + "times": { + "fp32_s": 7.5, + "bf16_s": 20.2, + "bf16_teacher_forced_s": 21.0 + }, + "exact_free_floor": { + "rel_l2": 1.5447020847623357e-05, + "row": "free layers.41", + "measured": "gmlx's exact step on the GPU against the same step on the CPU stream, worst free row; tests/image/moss_gates.py exact_free_floor on moss-image 512c8f5a, Apple M5 Max" + }, + "exact_velocity_floor": { + "rel_l2": 1.0271884390803118e-05, + "measured": "gmlx's exact step on the GPU against the same step on the CPU stream, velocity; tests/image/moss_gates.py exact_free_floor on moss-image 512c8f5a, Apple M5 Max" + } + }, + "tolerances": { + "velocity": { + "rel_l2": 0.08526306809508825, + "cosine_min": 0.9963701975314301 + }, + "time_embedding": { + "rel_l2": 0.0035760076594121594, + "teacher_forced": false + }, + "shared_block_modulation": { + "rel_l2": 0.0033759084196481478, + "teacher_forced": false + }, + "video_head_modulation": { + "rel_l2": 0.0031022546267896497, + "teacher_forced": false + }, + "video_embedder": { + "rel_l2": 0.004691027078597819, + "teacher_forced": false + }, + "text_embedder.proj": { + "rel_l2": 0.004702242882996404, + "teacher_forced": false + }, + "text_embedder.blocks.0": { + "rel_l2": 0.011492093846637928, + "teacher_forced": true + }, + "text_embedder.blocks.1": { + "rel_l2": 0.008398348532689982, + "teacher_forced": true + }, + "layers.0": { + "rel_l2": 0.006761022583897629, + "teacher_forced": true + }, + "layers.1": { + "rel_l2": 0.006019463669164852, + "teacher_forced": true + }, + "layers.2": { + "rel_l2": 0.008261226933396947, + "teacher_forced": true + }, + "layers.3": { + "rel_l2": 0.006332125710244657, + "teacher_forced": true + }, + "layers.4": { + "rel_l2": 0.005908308928026133, + "teacher_forced": true + }, + "layers.5": { + "rel_l2": 0.00664096860204911, + "teacher_forced": true + }, + "layers.6": { + "rel_l2": 0.005777407454304223, + "teacher_forced": true + }, + "layers.7": { + "rel_l2": 0.005919424683792383, + "teacher_forced": true + }, + "layers.8": { + "rel_l2": 0.006450233178413153, + "teacher_forced": true + }, + "layers.9": { + "rel_l2": 0.005277160929948154, + "teacher_forced": true + }, + "layers.10": { + "rel_l2": 0.00612656034396508, + "teacher_forced": true + }, + "layers.11": { + "rel_l2": 0.006642517581756266, + "teacher_forced": true + }, + "layers.12": { + "rel_l2": 0.006372776690299425, + "teacher_forced": true + }, + "layers.13": { + "rel_l2": 0.00699137814842513, + "teacher_forced": true + }, + "layers.14": { + "rel_l2": 0.01068428648385143, + "teacher_forced": true + }, + "layers.15": { + "rel_l2": 0.01149345782788048, + "teacher_forced": true + }, + "layers.16": { + "rel_l2": 0.01058319378479265, + "teacher_forced": true + }, + "layers.17": { + "rel_l2": 0.011215547661294678, + "teacher_forced": true + }, + "layers.18": { + "rel_l2": 0.014435490130678425, + "teacher_forced": true + }, + "layers.19": { + "rel_l2": 0.014332654663984803, + "teacher_forced": true + }, + "layers.20": { + "rel_l2": 0.014806769818313349, + "teacher_forced": true + }, + "layers.21": { + "rel_l2": 0.027840751868960203, + "teacher_forced": true + }, + "layers.22": { + "rel_l2": 0.022065713514541594, + "teacher_forced": true + }, + "layers.23": { + "rel_l2": 0.027845299336907693, + "teacher_forced": true + }, + "layers.24": { + "rel_l2": 0.026306692296869712, + "teacher_forced": true + }, + "layers.25": { + "rel_l2": 0.018942368484729067, + "teacher_forced": true + }, + "layers.26": { + "rel_l2": 0.016012801247289727, + "teacher_forced": true + }, + "layers.27": { + "rel_l2": 0.012905846629023453, + "teacher_forced": true + }, + "layers.28": { + "rel_l2": 0.015271842388630547, + "teacher_forced": true + }, + "layers.29": { + "rel_l2": 0.013390917380017952, + "teacher_forced": true + }, + "layers.30": { + "rel_l2": 0.01543697267686628, + "teacher_forced": true + }, + "layers.31": { + "rel_l2": 0.017940276024501177, + "teacher_forced": true + }, + "layers.32": { + "rel_l2": 0.017860502521947347, + "teacher_forced": true + }, + "layers.33": { + "rel_l2": 0.024000667962073008, + "teacher_forced": true + }, + "layers.34": { + "rel_l2": 0.0270135053503271, + "teacher_forced": true + }, + "layers.35": { + "rel_l2": 0.032242210209743334, + "teacher_forced": true + }, + "layers.36": { + "rel_l2": 0.033953392461451, + "teacher_forced": true + }, + "layers.37": { + "rel_l2": 0.03338655433494643, + "teacher_forced": true + }, + "layers.38": { + "rel_l2": 0.030901247906182455, + "teacher_forced": true + }, + "layers.39": { + "rel_l2": 0.027758694726148444, + "teacher_forced": true + }, + "layers.40": { + "rel_l2": 0.02556378934056245, + "teacher_forced": true + }, + "layers.41": { + "rel_l2": 0.036559068292460604, + "teacher_forced": true + }, + "final_linear_video": { + "rel_l2": 0.008242560367679485, + "teacher_forced": true + }, + "rule": "tolerance = 2 x the oracle's own bf16 against fp32 relative L2: teacher-forced for each text block, each block and the head, free running for the other rows and for the velocity", + "exact_free": { + "rel_l2": 5e-05 + }, + "exact_velocity": { + "rel_l2": 5e-05 + } + }, + "notes": [], + "created": "2026-10-18T06:35:49Z", + "npz": "check3_moss_image_512.npz", + "tensors": { + "features": { + "shape": [ + 48, + 4096 + ], + "dtype": "float32", + "sha256": "51f2174fad964908bd5817f252f645cdc220c2c13ca457e1ad899301a9c3ca6b" + }, + "final_linear_video": { + "shape": [ + 1, + 1024, + 128 + ], + "dtype": "float32", + "sha256": "d4bc8b5ff9ae866d2c5bb2510c0bb62acd86b443075a5cc290b775dc6344536a" + }, + "in.final_norm_video": { + "shape": [ + 1, + 1024, + 3072 + ], + "dtype": "float32", + "sha256": "43116dbe95149615a633a5aefe34dd5409d33f40eaeabc7f9c16f48c0f6c57d3" + }, + "in.layers.0": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "3485428cb3984d4489ec936ef37ff6e4a27aab1935c236906f04a6f011eef0a1" + }, + "latents": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "4399e6b9bd209dc9af6f745557507a33dcfe730a81abebb0bce079d2b441f34c" + }, + "layers.0": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "3ebbb074e4f41956a62558ef1c761e864ee21819e8767ff963528ec36a0d5a26" + }, + "layers.1": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "4a4976d736543a2a338ada96c4f6a290f73dfc1c941a851923ca8ce9fc8e570c" + }, + "layers.10": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "fb545107d0a22211b62ab67efcdf17ca9ae5b37687f73b5459376f61e7b09345" + }, + "layers.11": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "934c38e5ea0744ed3b713834dcd662cba737394a23df2fe8f28fb9a55e654275" + }, + "layers.12": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "224068c7139f6370032837a2d370b8f55744f9646857187aa5b2f49ee81d2a2e" + }, + "layers.13": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "2a296aa9b8f98170ac1ff5ea37d6e83f25042ef7c712377fafab743b23ae3e3e" + }, + "layers.14": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "877561fda965a4c251405f40527f7b808bef1fdf607dd33b7d36d96359fd0b02" + }, + "layers.15": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "b240c1676eb0f543191f1da906d0194fe3d53a0883626336c9fe5a80fb3c1463" + }, + "layers.16": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "876745618e5da483c25222c34bf404cab152e30b370664cff7a3dacfbbe4a611" + }, + "layers.17": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "56e30a31015567f97e08291ae7b28eab43225178bc1656b34632ef0a7fe82fc7" + }, + "layers.18": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "fadc6b7aaca5ceee0ca119ac58de7d5db0edcca721917bb50dae23ac205d87f5" + }, + "layers.19": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "b32b0b2d915c6c44c2d2d10b37da33931cee04cf3879d9fdfd50aa6d7d678bbd" + }, + "layers.2": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "71cc5b04fb87dacb6c05b6910c2b0c4b959f80a2a27218251c1d28657221b1e9" + }, + "layers.20": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "e84b794e65acb306ffa9c9c478b4a7e48600f3a12f8f7cfa94c522254bb4ded0" + }, + "layers.21": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "068de303772c022452e455309be096cf1f6ef1d80c57957cddce50cb349aa28b" + }, + "layers.22": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "db0f22cb6652b34a9076a185a7b39685f885b66d8660623fe553d748c952f30c" + }, + "layers.23": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "bb1e06a72ffc3930dd389dff1f7ab038d2ea9c94433c56706b017508ea0863e9" + }, + "layers.24": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "6baa29c5504e9a0711eeae2029bc876b93e655ce3862a13991137e6fad629f65" + }, + "layers.25": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "190e46366a944d0f4953b6c19823ff81facafbd43cc64af339dde2283236f68b" + }, + "layers.26": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "b9973acfc2e458dcc44f07b963d21e9fc7fd273e8a569c81282e36f577f3f9a1" + }, + "layers.27": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "c45114a2d8d3bb66de05a72bb6b0b267a3b19d6d349c5881eac1baf5037e65ef" + }, + "layers.28": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "21d8b7d03e1d1481528087605c2b8ae11b833491a5e365ad0fbcb96abab1dd8a" + }, + "layers.29": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "30ce49f7a95bb4c62c8f8b7ed7002eaef5ec3e7b58d72944551a30768a0c9705" + }, + "layers.3": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "642fd9d486f5e0345765803b0f66b0de725597223004441b0b3320db07f8ebfc" + }, + "layers.30": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "89434b507bd45ce857326ccf5cc7969fe783db743dee545ec34c348bf3a7f541" + }, + "layers.31": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "2d5d13f45ebd28cfb1250d28ef791675b47930a9e6ccbe8502b8addabeb7b115" + }, + "layers.32": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "975f3333c0818e132b59c6e856e62cb618cc363d4e6d99f26e85a584f9245b2d" + }, + "layers.33": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "bf2384622da5ef1f8b930508fe4af5fac32dd69e1ba62723fd8ea3b5107ff9a4" + }, + "layers.34": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "fd5d54a5f6502faf72dfb88ddef20d2f049fb4832fc832ce4f0128896f63d172" + }, + "layers.35": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "5c9adb1b637a410870cb1a2964c3c690b76bc0c4475cd015ee8f5c192dca628a" + }, + "layers.36": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "5a29199ab141287fcd8d256802adfd92fd653410b30958c3f5487c3654fdac46" + }, + "layers.37": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "c8a565080bbeb9d250b017d54155071f232533324eabdb9530078f1c0495157d" + }, + "layers.38": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "37351b4ec934610b089a467725e7bcf197e95eec2abd531ef4d8ea46951df187" + }, + "layers.39": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "d553d547f9087778c60266d4758647e32eb3b8a8c2198468262c198c68bb1350" + }, + "layers.4": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "c0571d5931a702e48f9eb904b81d3eec87a9e2d10cf8c2497637522ed993dfb5" + }, + "layers.40": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "2c15b5648104c33ecddac7ec37308a7ddb2a97bf1a50e906a80de56adbe9de15" + }, + "layers.41": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "80f9d5f61a009c080e0a28b0db9a23dee8ae7cef4d73077d6e67318d347d6cc2" + }, + "layers.5": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "d318008e2a25f9e5a64ffc6d58cc7279518058f1ee041988de66f7beb42e2330" + }, + "layers.6": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "b17bacf906342ba00b9578d11c77bf524afd375fdfa073644564917657050095" + }, + "layers.7": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "202d3e0facd5342396545dd5a91d79fa63e91430ec79a76b1b6cc76f16837915" + }, + "layers.8": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "c10744e57f9c2a34b4e8d8943e14da1d723120d438885ce15a930d166e6cbfde" + }, + "layers.9": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "2a35cae4fa532141e55842cc578958661d65185fdd391b9504425549f1e1187b" + }, + "shared_block_modulation": { + "shape": [ + 1, + 18432 + ], + "dtype": "float32", + "sha256": "2018b1648e882852e21be9c17537a00f5c352c49d4a938e7b4fd68f60d2a98a0" + }, + "sigma": { + "shape": [], + "dtype": "float32", + "sha256": "9a8208635e00348ab64aac2b759e76391fd47089e9a749bbcec770d9eb5c6421" + }, + "text_embedder.blocks.0": { + "shape": [ + 1, + 48, + 3072 + ], + "dtype": "float32", + "sha256": "9db78b37075f6326f8a55189eb00c85c9c9924ef6532d7752cbe56f703d969ec" + }, + "text_embedder.blocks.1": { + "shape": [ + 1, + 48, + 3072 + ], + "dtype": "float32", + "sha256": "1b9a12c61a9f151b9155d610e95e706805d0c0ed25ffc8291632da83560f5764" + }, + "text_embedder.proj": { + "shape": [ + 1, + 48, + 3072 + ], + "dtype": "float32", + "sha256": "b24a5a382094fe5dc138e7282eb00fd2100f0dc809d1b4a9e7f8560f12d20450" + }, + "time_embedding": { + "shape": [ + 1, + 3072 + ], + "dtype": "float32", + "sha256": "93eac9f99ef5510c035a19eb9ff02d9623d77fdc63a5850638cd42d0f2b53586" + }, + "timestep": { + "shape": [ + 1 + ], + "dtype": "float32", + "sha256": "1d77e51c3092b797b15b02104368537eea7b222148a4d2ba2d2d578491c50ac5" + }, + "velocity": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "85299a60f4437986ff3f011d6a84e6e06f2ac5437cf0287591f07b9f3e4c8123" + }, + "velocity_bf16": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "02d83e79c6312e65e213f5d256968c1b0c222e9963112221e44aa7b093fd1412" + }, + "video_embedder": { + "shape": [ + 1, + 1024, + 3072 + ], + "dtype": "float32", + "sha256": "8e782ff4fd57af77a4633a34ccd3145bf923d4f79eba2446f8a5c01022158a12" + }, + "video_head_modulation": { + "shape": [ + 1, + 6144 + ], + "dtype": "float32", + "sha256": "72074b7e2f983d21aa6b537b0285d95ee7dfcb20e5ef9b85e842f99a84ce59ef" + } + } +} diff --git a/tests/fixtures/image/oracle/check4_moss_image.manifest.json b/tests/fixtures/image/oracle/check4_moss_image.manifest.json new file mode 100644 index 000000000..22cbed659 --- /dev/null +++ b/tests/fixtures/image/oracle/check4_moss_image.manifest.json @@ -0,0 +1,199 @@ +{ + "schema": 1, + "name": "check4_moss_image", + "check": 4, + "family": "moss_image", + "script": "scripts/image_oracle/check4_moss_image.py", + "argv": [ + "check4_moss_image.py" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": { + "pipeline": { + "name": "pipeline_moss_image.py", + "size": 25060, + "sha256": "1dac1aca9feb1bfdfc727a401e4b67d61a505d3190706d1b7d2d784f88678b31" + }, + "scheduler_config": { + "name": "scheduler_config.json", + "size": 494, + "sha256": "707322832cc507dfb8c51c0e0cd04d44ee2d30a9b036b6f3c47d6c17aa71adf0" + } + }, + "inputs": { + "steps": [ + 50, + 28, + 8, + 250, + 1000, + 822 + ], + "scheduler": { + "base_image_seq_len": 256, + "base_shift": 0.5, + "invert_sigmas": false, + "max_image_seq_len": 4096, + "max_shift": 1.15, + "num_train_timesteps": 1000, + "shift": 5.0, + "shift_terminal": null, + "stochastic_sampling": false, + "time_shift_type": "exponential", + "training_shift": 1.0, + "use_beta_sigmas": false, + "use_dynamic_shifting": false, + "use_exponential_sigmas": false, + "use_karras_sigmas": false + } + }, + "results": { + "steps_run": { + "50": 50, + "28": 28, + "8": 8, + "250": 246, + "1000": 615, + "822": 556 + }, + "grid_points_off_k_over_total": 386, + "grid_max_ulps": 1 + }, + "tolerances": { + "sigmas": "bitwise", + "count": "exact", + "timesteps": "sigma x 1000 in float32", + "rule": "each sigma equal to the release's bit for bit and the same number of steps" + }, + "notes": [], + "created": "2026-10-18T11:20:09Z", + "npz": "check4_moss_image.npz", + "tensors": { + "sigmas_1000": { + "shape": [ + 616 + ], + "dtype": "float32", + "sha256": "a2012a04b4a441b50951f4c5208462cffa733b95098aaf0d8799209640e128be" + }, + "sigmas_250": { + "shape": [ + 247 + ], + "dtype": "float32", + "sha256": "84744830a7ceff60e557e3c34090e5789608d7bbbab844229e3c9172a3489433" + }, + "sigmas_28": { + "shape": [ + 29 + ], + "dtype": "float32", + "sha256": "5f8e4e3e53d4314c4c57440c3aaa0cb7d102fecca1679bed6501c6d523d80f42" + }, + "sigmas_50": { + "shape": [ + 51 + ], + "dtype": "float32", + "sha256": "be4301af0132c0b628f0d155d50ffab5cd6714691167eee2a160e2b59fcd63a7" + }, + "sigmas_8": { + "shape": [ + 9 + ], + "dtype": "float32", + "sha256": "c229a57463f3db2e71b9595902f071390567c1a85d6ced85022c95992cdd90d1" + }, + "sigmas_822": { + "shape": [ + 557 + ], + "dtype": "float32", + "sha256": "11f5f15ede784d78b5c311166c6272570afc98003837f13bc6833b178265278c" + }, + "timesteps_1000": { + "shape": [ + 615 + ], + "dtype": "float32", + "sha256": "d354cd6b9a3be4e06ac917e1ec3fbb8d036c1d909d5a504d492bbfc53446355d" + }, + "timesteps_250": { + "shape": [ + 246 + ], + "dtype": "float32", + "sha256": "047edf44585b28115594a25b3333204c17c1a2646e83e214d83b3457b538ec14" + }, + "timesteps_28": { + "shape": [ + 28 + ], + "dtype": "float32", + "sha256": "7300eb3225f7eb822090f90a06982d0963207f07d24810b28c9b6eb49ae1fa1d" + }, + "timesteps_50": { + "shape": [ + 50 + ], + "dtype": "float32", + "sha256": "e5c8f391f1c3eb29dd3afa3fcef35fd8ab7587fa759579edfa03be23c6a23b9a" + }, + "timesteps_8": { + "shape": [ + 8 + ], + "dtype": "float32", + "sha256": "933a3da4af0be7f12e536c22fcbb88200dd6b840985fcb2e3ecdd4a95c796125" + }, + "timesteps_822": { + "shape": [ + 556 + ], + "dtype": "float32", + "sha256": "303ef9a8d6d5104f0e50e498636e6cfe1088e12b229261d8215ada4505e7387a" + }, + "training_grid": { + "shape": [ + 1000 + ], + "dtype": "float32", + "sha256": "d6f5386fddcf96144952ca78c05ec7620e4194a372c727ebd37dbe1f273d89f9" + } + } +} diff --git a/tests/fixtures/image/oracle/check4_moss_image.npz b/tests/fixtures/image/oracle/check4_moss_image.npz new file mode 100644 index 000000000..2abe8d966 Binary files /dev/null and b/tests/fixtures/image/oracle/check4_moss_image.npz differ diff --git a/tests/fixtures/image/oracle/check5_moss_image_vae.manifest.json b/tests/fixtures/image/oracle/check5_moss_image_vae.manifest.json new file mode 100644 index 000000000..e737ac4d0 --- /dev/null +++ b/tests/fixtures/image/oracle/check5_moss_image_vae.manifest.json @@ -0,0 +1,364 @@ +{ + "schema": 1, + "name": "check5_moss_image_vae", + "check": 5, + "family": "moss_image", + "script": "scripts/image_oracle/check5_flux3_klein_vae.py", + "argv": [ + "check5_flux3_klein_vae.py", + "--vae", + "diffusion_pytorch_model.safetensors", + "--family", + "moss_image", + "--sizes", + "512", + "1024" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": { + "vae": { + "name": "diffusion_pytorch_model.safetensors", + "size": 249521372, + "sha256": "c3bdf94dc67ec69b30ea9fe1fa8340085679cae64d019f7c7e4a9921f1a7f17a" + }, + "vae_config": { + "name": "config.json", + "size": 838, + "sha256": "cc77f3f0ce0357039c71018dc19edaa4fcceb0c2f7488fdc2b9c0580d6baafa1" + } + }, + "inputs": { + "sizes": [ + [ + 512, + 512 + ], + [ + 1024, + 1024 + ] + ], + "seed": 5, + "image": "synthetic_image(height, width, seed) of check5_qwen_image21_vae.py, RGB, NCHW in [-1, 1]", + "batch_norm_eps": 0.0001, + "patch_size": [ + 2, + 2 + ], + "latents": "Flux2KleinPipeline._patchify_latents of the encoder mean, then (z - running_mean) / sqrt(running_var + batch_norm_eps), in float32 from each mode's mean", + "decode_input": "latents * sqrt(running_var + eps) + running_mean, then Flux2KleinPipeline._unpatchify_latents, in float32 from the fp32 latents, cast to the VAE dtype", + "decoded": "RGB, NCHW, unclamped as the VAE returns it", + "psnr": "on x / 2 + 0.5 clamped to [0, 1], peak 1", + "source_dtypes": [ + "torch.float32", + "torch.int64" + ], + "config_revision": "fc5de0d56fc910cbca6f6625a70f5f582b26deb1", + "load": "AutoencoderKLFlux2.from_config(vae/config.json), then the file strict, in float32" + }, + "results": { + "oracle": { + "device": "cpu", + "threads": 16, + "bf16_products": "bf16 operands, float32 conv and linear, one bf16 rounding" + }, + "sizes": { + "512x512": { + "decode_bf16_vs_fp32": { + "rel_l2": 0.004622636352810195, + "cosine": 0.9999897034844557, + "max_abs": 0.06475692987442017, + "psnr_db": 57.32880164256124 + }, + "encode_bf16_vs_fp32": { + "rel_l2": 0.01739012329408501, + "cosine": 0.9998490640077118, + "max_abs": 0.19057875871658325 + }, + "encode_fp32_vs_fp64": { + "rel_l2": 1.85394467142218e-06, + "cosine": 0.999999999998282, + "max_abs": 1.1787637656662797e-05 + }, + "encode_mean_fp32_vs_fp64": { + "rel_l2": 1.8462501499014103e-06, + "cosine": 0.9999999999982965, + "max_abs": 2.1508609541726065e-05 + }, + "decode_fp32_vs_fp64": { + "rel_l2": 1.2511865637473112e-06, + "cosine": 0.9999999999992188, + "max_abs": 1.760444161930108e-05 + }, + "round_trip_fp32_psnr_db": 26.655471517058757, + "latents_std": 0.8341149687767029, + "peak_rss_gib": { + "encode": 8.34, + "decode": 12.46 + }, + "seconds": { + "fp32_encode_s": 1.1, + "bf16_encode_s": 1.6, + "fp32_decode_s": 1.4, + "bf16_decode_s": 1.7, + "float64_decode_s": 4.7 + } + }, + "1024x1024": { + "decode_bf16_vs_fp32": { + "rel_l2": 0.004803838363423286, + "cosine": 0.9999885563552089, + "max_abs": 0.08704939484596252, + "psnr_db": 56.96572023123493 + }, + "encode_bf16_vs_fp32": { + "rel_l2": 0.029780096139354124, + "cosine": 0.9995578754429438, + "max_abs": 0.4994291365146637 + }, + "encode_fp32_vs_fp64": { + "rel_l2": 3.91706261461626e-06, + "cosine": 0.999999999992329, + "max_abs": 3.675377888912035e-05 + }, + "encode_mean_fp32_vs_fp64": { + "rel_l2": 3.887321914440507e-06, + "cosine": 0.9999999999924465, + "max_abs": 6.529417830059003e-05 + }, + "decode_fp32_vs_fp64": { + "rel_l2": 4.959171167047636e-06, + "cosine": 0.999999999987892, + "max_abs": 0.00017490952150334138 + }, + "round_trip_fp32_psnr_db": 26.767989226225847, + "latents_std": 0.7904156446456909, + "peak_rss_gib": { + "encode": 20.52, + "decode": 29.74 + }, + "seconds": { + "fp32_encode_s": 5.0, + "bf16_encode_s": 15.7, + "fp32_decode_s": 6.6, + "bf16_decode_s": 11.2, + "float64_decode_s": 21.8 + } + } + } + }, + "tolerances": { + "rule": "fp32 decode max abs at twice the oracle's own fp32 against float64 max abs; fp32 encode max abs, on the model space latents, at 1e-3, or at twice the oracle's own fp32 against float64 max abs where that is larger; bf16 decode PSNR at twice the oracle's own bf16 against fp32 RMS error, its PSNR less 20 log10(2) dB; bf16 encode relative L2, on the model space latents, at twice the oracle's own", + "fp32_max_abs": 0.001, + "sizes": { + "512x512": { + "decode_fp32_max_abs": 3.520888323860216e-05, + "decode_bf16_psnr_db": 51.30820172928161, + "encode_fp32_max_abs": 0.001, + "encode_bf16_rel_l2": 0.03478024658817002 + }, + "1024x1024": { + "decode_fp32_max_abs": 0.00034981904300668276, + "decode_bf16_psnr_db": 50.94512031795531, + "encode_fp32_max_abs": 0.001, + "encode_bf16_rel_l2": 0.05956019227870825 + } + } + }, + "notes": [ + "The file holds float32 weights. The bf16 VAE is a bfloat16 copy with its BatchNorm statistics kept float32." + ], + "created": "2026-10-18T06:54:52Z", + "npz": "check5_moss_image_vae.npz", + "tensors": { + "bn_running_mean": { + "shape": [ + 128 + ], + "dtype": "float32", + "sha256": "9027fac5727854f779ebbeae3032cfce0d47a11bc85f4329eb0a317c9ad90217" + }, + "bn_running_var": { + "shape": [ + 128 + ], + "dtype": "float32", + "sha256": "8b4656821cfe8827490068f768f0dba57e09423e900b8b018da097f0c4537009" + }, + "decoded_1024x1024": { + "shape": [ + 1, + 3, + 1024, + 1024 + ], + "dtype": "float32", + "sha256": "82a36bc00e916f6b1b3e02b57c768c75cabe96e860ae4e742b4b7255048ac6ac" + }, + "decoded_512x512": { + "shape": [ + 1, + 3, + 512, + 512 + ], + "dtype": "float32", + "sha256": "5b10f3903482d938af6ca3dcb2e964f17e60b815cb570d2ea5f26fdc664bbf2b" + }, + "decoded_bf16_1024x1024": { + "shape": [ + 1, + 3, + 1024, + 1024 + ], + "dtype": "float32", + "sha256": "5b949d8b4f0bf5ba8951cd298432efd981f9193321287ed91e15079cd99a1bb2" + }, + "decoded_bf16_512x512": { + "shape": [ + 1, + 3, + 512, + 512 + ], + "dtype": "float32", + "sha256": "09a468905c8486b9a745c70fe48e7652624d204ebcdf8964ebec341c096ce40a" + }, + "encode_mean_1024x1024": { + "shape": [ + 1, + 32, + 128, + 128 + ], + "dtype": "float32", + "sha256": "e25c0df3d5be4d976cb92dd01191ec0df279372b74bf5505844aeddb41dfa8ad" + }, + "encode_mean_512x512": { + "shape": [ + 1, + 32, + 64, + 64 + ], + "dtype": "float32", + "sha256": "72c807775e530a0c4fc07a96c4ebff86251dd839d93c135f5a5643050b4a86ad" + }, + "encode_mean_bf16_1024x1024": { + "shape": [ + 1, + 32, + 128, + 128 + ], + "dtype": "float32", + "sha256": "d91533b309cd78a00fb0c5dd0050b70ed397c30d3655e64f508a6a982c6aa4cc" + }, + "encode_mean_bf16_512x512": { + "shape": [ + 1, + 32, + 64, + 64 + ], + "dtype": "float32", + "sha256": "11a390e65832643f0e9a888bd14bae293bf0158dd705b5d9fb56e47b1287a578" + }, + "image_1024x1024": { + "shape": [ + 1, + 3, + 1024, + 1024 + ], + "dtype": "float32", + "sha256": "3930a02c0be88010b140dd1b6c71a1bf518c6f75d78d0ad2fa78942f998cb1a6" + }, + "image_512x512": { + "shape": [ + 1, + 3, + 512, + 512 + ], + "dtype": "float32", + "sha256": "4020c0997410077532d908bbecca27898c4c7b7cfb2d7ba29333177c2c44ee5e" + }, + "latents_1024x1024": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "c87fce5ec93e1ab0acb0d2e56d27af75dcc2d67eb2666abe43bbda04292c5e26" + }, + "latents_512x512": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "280deb3315669ad11024f79090922a6dab13f83fd043e665629eb5225c323811" + }, + "latents_bf16_1024x1024": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "91cedc4bfe75ab2aa1e4e0fc2fb846aa4a68af95eddf503df2149e916e6f2d8f" + }, + "latents_bf16_512x512": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "ce4a39f7da605086a42796276ec3c88a8ebc3d7013acf2b46dcb6d92ad942de0" + } + } +} diff --git a/tests/fixtures/image/oracle/check6_moss_image_1024.manifest.json b/tests/fixtures/image/oracle/check6_moss_image_1024.manifest.json new file mode 100644 index 000000000..a080d8d62 --- /dev/null +++ b/tests/fixtures/image/oracle/check6_moss_image_1024.manifest.json @@ -0,0 +1,322 @@ +{ + "schema": 1, + "name": "check6_moss_image_1024", + "check": 6, + "family": "moss_image", + "script": "scripts/image_oracle/check6_moss_image.py", + "argv": [ + "check6_moss_image.py", + "--dit", + "MOSS-Image-SFT-Q4_K_M.gguf", + "--features", + "check2_moss_image_q8_0.npz", + "--sizes", + "1024" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": { + "dit": { + "name": "MOSS-Image-SFT-Q4_K_M.gguf", + "size": 3702092096, + "sha256": "c6dda89c18c0c0897d09ed01a515ee31415c4baea02bb9c2dedbfca66eec50b6" + }, + "pipeline": { + "name": "pipeline_moss_image.py", + "size": 25060, + "sha256": "1dac1aca9feb1bfdfc727a401e4b67d61a505d3190706d1b7d2d784f88678b31" + }, + "config": { + "name": "config.json", + "size": 240, + "sha256": "3291c3a1c7adc22b8182a0de25ca8ec97474d72715ce84adbab09cc0be24f666" + }, + "scheduler_config": { + "name": "scheduler_config.json", + "size": 494, + "sha256": "707322832cc507dfb8c51c0e0cd04d44ee2d30a9b036b6f3c47d6c17aa71adf0" + }, + "features": { + "name": "check2_moss_image_q8_0.npz", + "size": 24685072, + "sha256": "9f40242fe35edd0fc15bac6a2d42983e6b2d95c016d866e04441b68bbe91e3ef" + } + }, + "inputs": { + "size": 1024, + "steps": 50, + "steps_run": 50, + "blocks": 42, + "guidance_scale": 3.0, + "uncond_skip_layers": [ + 22 + ], + "negative_prompt": "", + "scheduler": { + "base_image_seq_len": 256, + "base_shift": 0.5, + "invert_sigmas": false, + "max_image_seq_len": 4096, + "max_shift": 1.15, + "num_train_timesteps": 1000, + "shift": 5.0, + "shift_terminal": null, + "stochastic_sampling": false, + "time_shift_type": "exponential", + "training_shift": 1.0, + "use_beta_sigmas": false, + "use_dynamic_shifting": false, + "use_exponential_sigmas": false, + "use_karras_sigmas": false + }, + "seed": 6, + "features": "check2_moss_image_q8_0 fp32 features", + "positive_key": "features_plain", + "negative_key": "features_empty", + "text_rows": { + "positive": 38, + "negative": 10 + }, + "kept_steps": [ + 1, + 2, + 5, + 10, + 25 + ], + "gate_size": 512, + "noise": "numpy Generator(PCG64(seed)).standard_normal(shape, dtype=float32), NCHW, shape (1, 128, size / 16, size / 16)", + "latents": "(1, 128, size / 16, size / 16) float32 in model space, after the step", + "gguf_codecs": { + "F16": 210, + "Q6_K": 84, + "Q4_K": 168, + "F32": 1 + } + }, + "results": { + "oracle": { + "device": "cpu", + "threads": 16, + "dequant": "exact", + "bf16_matmul": "bf16 operands, float32 GEMM, one bf16 rounding", + "timestep": "sigma x 1000 in float32 in both modes", + "guidance": "float32 velocities in both modes, as the release casts them" + }, + "steps_run": 50, + "fp32_step_s": [ + 56.9, + 56.8, + 56.9, + 56.8, + 56.9, + 56.8, + 56.8, + 56.8, + 56.9, + 56.8, + 56.9, + 56.8, + 56.8, + 56.8, + 56.9, + 56.9, + 56.9, + 56.8, + 56.9, + 56.9, + 56.9, + 56.8, + 56.9, + 56.8, + 56.9, + 56.8, + 56.8, + 56.8, + 56.9, + 56.9, + 56.9, + 56.8, + 56.9, + 56.8, + 56.9, + 56.8, + 56.9, + 56.8, + 56.8, + 56.8, + 56.8, + 56.9, + 56.9, + 56.8, + 56.9, + 56.9, + 56.9, + 56.9, + 56.8, + 56.9 + ], + "fp32_s": 2842.5, + "gmlx_exact_floor": { + "1": 9.279319125239364e-08, + "2": 9.760845017376464e-08, + "5": 1.1916169760375085e-07, + "10": 2.3503868223467442e-07, + "25": 2.20891604645722e-05, + "50": 0.0015016736002878234 + } + }, + "tolerances": { + "rule": "latents after a kept step and after the last. At 512: gmlx's bf16 run at relative L2 twice the oracle's own bf16 against fp32 relative L2 after that step. At any other size: gmlx's exact fp32 run against the fp32 latents at exact_rel_l2, and its bf16 run read as a share of the 512 gate of the step, not gated; the exact gate of a step is the larger of exact_rel_l2 and twice gmlx_exact_floor, the exact run from the noise moved one float32 ulp against the exact run, measured by tests/image/test_image_check6_moss_image.py on an M5 Max", + "exact_rel_l2": 0.001, + "exact_by_step": { + "1": 0.001, + "2": 0.001, + "5": 0.001, + "10": 0.001, + "25": 0.001, + "50": 0.0030033472005756468 + } + }, + "notes": [ + "gmlx_exact_floor and exact_by_step were added from the gmlx exact run on the release's schedule bit for bit; the oracle wrote the rest" + ], + "created": "2026-10-18T09:08:55Z", + "npz": "check6_moss_image_1024.npz", + "tensors": { + "features_negative": { + "shape": [ + 10, + 4096 + ], + "dtype": "float32", + "sha256": "62dc84e4bdef784bf5365091e6b5552df9269755fb2ab41869b6512bdeb85074" + }, + "features_positive": { + "shape": [ + 38, + 4096 + ], + "dtype": "float32", + "sha256": "012ba987671e1c79a361c5a6ce972ffbbb7c47c6a5938268d57b0f63fb345c4e" + }, + "latents": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "7a9f8c5e58bdbe9913ef6c19429277b43d171e6102ac51d8db86e182d0194904" + }, + "latents_step1": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "0193ffaec2e5b60137ac5ff93dc8c8dd705f13aa521f76de996337bda76872d4" + }, + "latents_step10": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "6aae656f153d6d8f081ddcca6b54b542aec6d3c420655ffe4b8cec8e2c8c098e" + }, + "latents_step2": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "8e47d4a1b37f99516577e6afec3d34304a242bee50f0f1cdac41d4c173886eae" + }, + "latents_step25": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "b99aa95a3fad21b63329fc4cf49057e1d6467ee1c6fe8b30f2e85126f3700410" + }, + "latents_step5": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "1aee320fc8ecc75a669d1dbf89ff3189c45193d3759519915ce1148acc21cb34" + }, + "noise": { + "shape": [ + 1, + 128, + 64, + 64 + ], + "dtype": "float32", + "sha256": "bcd95295bb6b6d4f2030450a235b4f3e631d78cfeabead17132ec18138429b6a" + }, + "sigmas": { + "shape": [ + 51 + ], + "dtype": "float32", + "sha256": "be4301af0132c0b628f0d155d50ffab5cd6714691167eee2a160e2b59fcd63a7" + }, + "timesteps": { + "shape": [ + 50 + ], + "dtype": "float32", + "sha256": "e5c8f391f1c3eb29dd3afa3fcef35fd8ab7587fa759579edfa03be23c6a23b9a" + } + } +} diff --git a/tests/fixtures/image/oracle/check6_moss_image_512.manifest.json b/tests/fixtures/image/oracle/check6_moss_image_512.manifest.json new file mode 100644 index 000000000..c173e384f --- /dev/null +++ b/tests/fixtures/image/oracle/check6_moss_image_512.manifest.json @@ -0,0 +1,429 @@ +{ + "schema": 1, + "name": "check6_moss_image_512", + "check": 6, + "family": "moss_image", + "script": "scripts/image_oracle/check6_moss_image.py", + "argv": [ + "check6_moss_image.py", + "--dit", + "MOSS-Image-SFT-Q4_K_M.gguf", + "--features", + "check2_moss_image_q8_0.npz", + "--sizes", + "512", + "--bf16", + "512" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": { + "dit": { + "name": "MOSS-Image-SFT-Q4_K_M.gguf", + "size": 3702092096, + "sha256": "c6dda89c18c0c0897d09ed01a515ee31415c4baea02bb9c2dedbfca66eec50b6" + }, + "pipeline": { + "name": "pipeline_moss_image.py", + "size": 25060, + "sha256": "1dac1aca9feb1bfdfc727a401e4b67d61a505d3190706d1b7d2d784f88678b31" + }, + "config": { + "name": "config.json", + "size": 240, + "sha256": "3291c3a1c7adc22b8182a0de25ca8ec97474d72715ce84adbab09cc0be24f666" + }, + "scheduler_config": { + "name": "scheduler_config.json", + "size": 494, + "sha256": "707322832cc507dfb8c51c0e0cd04d44ee2d30a9b036b6f3c47d6c17aa71adf0" + }, + "features": { + "name": "check2_moss_image_q8_0.npz", + "size": 24685072, + "sha256": "9f40242fe35edd0fc15bac6a2d42983e6b2d95c016d866e04441b68bbe91e3ef" + } + }, + "inputs": { + "size": 512, + "steps": 50, + "steps_run": 50, + "blocks": 42, + "guidance_scale": 3.0, + "uncond_skip_layers": [ + 22 + ], + "negative_prompt": "", + "scheduler": { + "base_image_seq_len": 256, + "base_shift": 0.5, + "invert_sigmas": false, + "max_image_seq_len": 4096, + "max_shift": 1.15, + "num_train_timesteps": 1000, + "shift": 5.0, + "shift_terminal": null, + "stochastic_sampling": false, + "time_shift_type": "exponential", + "training_shift": 1.0, + "use_beta_sigmas": false, + "use_dynamic_shifting": false, + "use_exponential_sigmas": false, + "use_karras_sigmas": false + }, + "seed": 6, + "features": "check2_moss_image_q8_0 fp32 features", + "positive_key": "features_plain", + "negative_key": "features_empty", + "text_rows": { + "positive": 38, + "negative": 10 + }, + "kept_steps": [ + 1, + 2, + 5, + 10, + 25 + ], + "gate_size": 512, + "noise": "numpy Generator(PCG64(seed)).standard_normal(shape, dtype=float32), NCHW, shape (1, 128, size / 16, size / 16)", + "latents": "(1, 128, size / 16, size / 16) float32 in model space, after the step", + "gguf_codecs": { + "F16": 210, + "Q6_K": 84, + "Q4_K": 168, + "F32": 1 + } + }, + "results": { + "oracle": { + "device": "cpu", + "threads": 16, + "dequant": "exact", + "bf16_matmul": "bf16 operands, float32 GEMM, one bf16 rounding", + "timestep": "sigma x 1000 in float32 in both modes", + "guidance": "float32 velocities in both modes, as the release casts them" + }, + "steps_run": 50, + "fp32_step_s": [ + 14.5, + 14.5, + 14.4, + 14.5, + 14.4, + 14.5, + 14.9, + 14.8, + 14.9, + 14.9, + 15.2, + 15.0, + 14.9, + 14.8, + 14.8, + 14.9, + 14.9, + 14.9, + 14.9, + 14.9, + 14.9, + 14.9, + 14.9, + 14.9, + 14.9, + 14.9, + 14.9, + 14.9, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8, + 14.8 + ], + "fp32_s": 739.8, + "bf16_s": 2170.0, + "latents_bf16_vs_fp32": { + "rel_l2": 0.45943544097777794, + "cosine": 0.8938588930880345, + "max_abs": 5.628334999084473 + }, + "bf16_vs_fp32_by_step": [ + 0.00017785951800384727, + 0.00027685391111303373, + 0.0003971679452805254, + 0.000495907513443672, + 0.0006704266772772455, + 0.0009635333473899233, + 0.0018884017702458875, + 0.003630256001871022, + 0.0057302323636947965, + 0.0077554798629276955, + 0.010170772795258439, + 0.012561981739524034, + 0.014651559000990749, + 0.01740407060063074, + 0.02015506988086763, + 0.02295971348274425, + 0.02558256001455884, + 0.02854916205717846, + 0.03205157774754117, + 0.03519681699508832, + 0.03905007340425514, + 0.04352978393466403, + 0.048198543535250556, + 0.05355951771710164, + 0.05906839465942481, + 0.06520537719507057, + 0.07147137523104004, + 0.07908587224377753, + 0.08700295316185488, + 0.09579588105765836, + 0.10548439071365853, + 0.11616841166438577, + 0.12779969342727338, + 0.14040115702036549, + 0.15470843296211922, + 0.17078583872646225, + 0.18852639361014067, + 0.2072039341819995, + 0.22833139066703506, + 0.25167838264063064, + 0.27666427459614973, + 0.3027972916536381, + 0.33007675785693313, + 0.3583554748313851, + 0.38475211145896104, + 0.4099225010298654, + 0.43180320608272776, + 0.4487728448752057, + 0.4589628460702799, + 0.45943544097777794 + ] + }, + "tolerances": { + "rule": "latents after a kept step and after the last. At 512: gmlx's bf16 run at relative L2 twice the oracle's own bf16 against fp32 relative L2 after that step. At any other size: gmlx's exact fp32 run against the fp32 latents at exact_rel_l2, and its bf16 run read as a share of the 512 gate of the step, not gated", + "latents_rel_l2": 0.9188708819555559, + "step1_rel_l2": 0.00035571903600769455, + "step2_rel_l2": 0.0005537078222260675, + "step5_rel_l2": 0.001340853354554491, + "step10_rel_l2": 0.015510959725855391, + "step25_rel_l2": 0.11813678931884962 + }, + "notes": [], + "created": "2026-10-18T08:21:31Z", + "npz": "check6_moss_image_512.npz", + "tensors": { + "features_negative": { + "shape": [ + 10, + 4096 + ], + "dtype": "float32", + "sha256": "62dc84e4bdef784bf5365091e6b5552df9269755fb2ab41869b6512bdeb85074" + }, + "features_positive": { + "shape": [ + 38, + 4096 + ], + "dtype": "float32", + "sha256": "012ba987671e1c79a361c5a6ce972ffbbb7c47c6a5938268d57b0f63fb345c4e" + }, + "latents": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "7e9b8ce065ca4d761e6251f3029fcd52312d088f33da62406715c405c4127ad2" + }, + "latents_bf16": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "cb7a8030e9bde907ca0e9ac2562c10d8a091a4bd3cca05c9ba1274ce7bc68adc" + }, + "latents_bf16_step1": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "b842f04987777cae9c56461ba0dd043ab78b5572f8df74183e40eec2208e6206" + }, + "latents_bf16_step10": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "38c3b0af502f32cdee9ba8067b19e2c2b0cb3458a6265332b7848b7c14735cbd" + }, + "latents_bf16_step2": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "6382f3738cb8297e9f7a8869b56183bfcc529aa5b7d37bf6465fc0f725def74a" + }, + "latents_bf16_step25": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "3bb752fb6d748ad2bf34d554fc5a1ce824407c6ed57539a981f97a94ec89abae" + }, + "latents_bf16_step5": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "67c25cf10768fc90ca16e14001119a6f125f313a47cf933cde6dbad7f6a881d2" + }, + "latents_step1": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "9637fb75aae0c9a568b726979bf2a71d2011a061a5d68ea90b8fe7761901dd31" + }, + "latents_step10": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "8ee84301e3e79f719b3b8421dbedf06519fd53d517d8136d42416dd2fbc69b21" + }, + "latents_step2": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "fac21d9385a13596034a5b63b2db9ac3b57e21cf4e10c29ebf967a8de07725c0" + }, + "latents_step25": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "34a03149a42c3932b2fdf3f686d3348195c825960138e4a0b2fee0518ff1399a" + }, + "latents_step5": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "bf76779371abb985f84ae46e9b400f25c675390cba953fd7684d6b18319cc499" + }, + "noise": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "d614e8f6414c19a6e58b7af3fb3e7e866a4c21c40c780bf93ece8c4f09d875de" + }, + "sigmas": { + "shape": [ + 51 + ], + "dtype": "float32", + "sha256": "be4301af0132c0b628f0d155d50ffab5cd6714691167eee2a160e2b59fcd63a7" + }, + "timesteps": { + "shape": [ + 50 + ], + "dtype": "float32", + "sha256": "e5c8f391f1c3eb29dd3afa3fcef35fd8ab7587fa759579edfa03be23c6a23b9a" + } + } +} diff --git a/tests/fixtures/image/oracle/check7_moss_image_synthetic_512.manifest.json b/tests/fixtures/image/oracle/check7_moss_image_synthetic_512.manifest.json new file mode 100644 index 000000000..b96a2a102 --- /dev/null +++ b/tests/fixtures/image/oracle/check7_moss_image_synthetic_512.manifest.json @@ -0,0 +1,1784 @@ +{ + "schema": 1, + "name": "check7_moss_image_synthetic_512", + "check": 7, + "family": "moss_image", + "script": "scripts/image_oracle/check7_moss_image_lora.py", + "argv": [ + "check7_moss_image_lora.py", + "--pipeline", + "pipeline_moss_image.py", + "--dit", + "MOSS-Image-SFT-Q4_K_M.gguf", + "--config", + "config.json" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": { + "dit": { + "name": "MOSS-Image-SFT-Q4_K_M.gguf", + "size": 3702092096, + "sha256": "c6dda89c18c0c0897d09ed01a515ee31415c4baea02bb9c2dedbfca66eec50b6" + }, + "pipeline": { + "name": "pipeline_moss_image.py", + "size": 25060, + "sha256": "1dac1aca9feb1bfdfc727a401e4b67d61a505d3190706d1b7d2d784f88678b31" + }, + "config": { + "name": "config.json", + "size": 240, + "sha256": "3291c3a1c7adc22b8182a0de25ca8ec97474d72715ce84adbab09cc0be24f666" + } + }, + "inputs": { + "size": 512, + "blocks": 42, + "sigma": 0.75, + "seed": 512, + "features": "seeded normal, std 1.0", + "text_rows": 48, + "dequant": "exact", + "matmul_bf16": "bf16 operands, float32 GEMM, one bf16 rounding", + "timestep": "sigma x 1000 in float32 in both modes", + "gguf_codecs": { + "F16": 210, + "Q6_K": 84, + "Q4_K": 168, + "F32": 1 + }, + "source": "check3_moss_image_512", + "strength": 1.0, + "lora_pairs": 29, + "lora_ranks": [ + 4, + 8, + 12, + 16 + ], + "lora_alphas": [ + 2.0, + 4.0, + 6.0, + 8.0 + ], + "lora_source": "synthetic, seed 7, each delta at 0.08 of its weight's Frobenius norm, ranks 4 to 16, alpha half the rank", + "targets": [ + "final_linear_video", + "layers.0.proj.ffn_down", + "layers.0.proj.ffn_up_gate", + "layers.0.proj.k", + "layers.0.proj.o", + "layers.0.proj.q", + "layers.0.proj.v", + "layers.21.proj.ffn_down", + "layers.21.proj.ffn_up_gate", + "layers.21.proj.k", + "layers.21.proj.o", + "layers.21.proj.q", + "layers.21.proj.v", + "layers.41.proj.ffn_down", + "layers.41.proj.ffn_up_gate", + "layers.41.proj.k", + "layers.41.proj.o", + "layers.41.proj.q", + "layers.41.proj.v", + "shared_block_modulation", + "text_embedder.blocks.1.ffn_down", + "text_embedder.blocks.1.ffn_gate", + "text_embedder.blocks.1.ffn_up", + "text_embedder.blocks.1.proj", + "text_embedder.blocks.1.qkv", + "text_embedder.proj", + "time_embedding.2", + "video_embedder", + "video_head_modulation" + ], + "lora_pair_alphas": { + "layers.0.proj.q": 2.0, + "layers.0.proj.k": 4.0, + "layers.0.proj.v": 6.0, + "layers.0.proj.o": 8.0, + "layers.0.proj.ffn_up_gate": 2.0, + "layers.0.proj.ffn_down": 4.0, + "layers.21.proj.q": 6.0, + "layers.21.proj.k": 8.0, + "layers.21.proj.v": 2.0, + "layers.21.proj.o": 4.0, + "layers.21.proj.ffn_up_gate": 6.0, + "layers.21.proj.ffn_down": 8.0, + "layers.41.proj.q": 2.0, + "layers.41.proj.k": 4.0, + "layers.41.proj.v": 6.0, + "layers.41.proj.o": 8.0, + "layers.41.proj.ffn_up_gate": 2.0, + "layers.41.proj.ffn_down": 4.0, + "text_embedder.blocks.1.qkv": 6.0, + "text_embedder.blocks.1.proj": 8.0, + "text_embedder.blocks.1.ffn_gate": 2.0, + "text_embedder.blocks.1.ffn_up": 4.0, + "text_embedder.blocks.1.ffn_down": 6.0, + "video_embedder": 8.0, + "text_embedder.proj": 2.0, + "time_embedding.2": 4.0, + "shared_block_modulation": 6.0, + "video_head_modulation": 8.0, + "final_linear_video": 2.0 + }, + "application": "unmerged: strength * (alpha / rank) * up(down(x)) added to each target's output in the mode's dtype" + }, + "results": { + "velocity_bf16_vs_fp32": { + "rel_l2": 0.044390082674380285, + "cosine": 0.9990144077172335, + "max_abs": 0.4622030258178711 + }, + "rows_bf16_vs_fp32": { + "time_embedding": { + "rel_l2": 0.0020187154789169685, + "cosine": 0.9999982066935266, + "max_abs": 0.001960277557373047, + "teacher_forced": false + }, + "shared_block_modulation": { + "rel_l2": 0.001988742190170364, + "cosine": 0.9999980834853716, + "max_abs": 0.10837554931640625, + "teacher_forced": false + }, + "video_head_modulation": { + "rel_l2": 0.001716889643506859, + "cosine": 0.99999856422564, + "max_abs": 0.018826007843017578, + "teacher_forced": false + }, + "video_embedder": { + "rel_l2": 0.0028831431127282433, + "cosine": 0.999995843845907, + "max_abs": 0.12289810180664062, + "teacher_forced": false + }, + "text_embedder.proj": { + "rel_l2": 0.002892886066180244, + "cosine": 0.9999958157070745, + "max_abs": 0.11258125305175781, + "teacher_forced": false + }, + "text_embedder.blocks.0": { + "rel_l2": 0.0057239684339233515, + "cosine": 0.9999836181442747, + "max_abs": 0.38885498046875, + "teacher_forced": true + }, + "text_embedder.blocks.1": { + "rel_l2": 0.004456773785192701, + "cosine": 0.9999900717114037, + "max_abs": 0.6630401611328125, + "teacher_forced": true + }, + "layers.0": { + "rel_l2": 0.0030401536787696437, + "cosine": 0.9999954769179399, + "max_abs": 6.02685546875, + "teacher_forced": true + }, + "layers.1": { + "rel_l2": 0.003103715635873704, + "cosine": 0.9999952777177121, + "max_abs": 5.62744140625, + "teacher_forced": true + }, + "layers.2": { + "rel_l2": 0.003788131612642148, + "cosine": 0.999994577721581, + "max_abs": 7.8157958984375, + "teacher_forced": true + }, + "layers.3": { + "rel_l2": 0.003119955541809922, + "cosine": 0.9999952268186282, + "max_abs": 5.55828857421875, + "teacher_forced": true + }, + "layers.4": { + "rel_l2": 0.002862047540965643, + "cosine": 0.9999959856639532, + "max_abs": 6.544189453125, + "teacher_forced": true + }, + "layers.5": { + "rel_l2": 0.0034488667694506287, + "cosine": 0.9999947177524277, + "max_abs": 8.353271484375, + "teacher_forced": true + }, + "layers.6": { + "rel_l2": 0.0028469679973086876, + "cosine": 0.9999961957712828, + "max_abs": 9.1136474609375, + "teacher_forced": true + }, + "layers.7": { + "rel_l2": 0.003225051760340546, + "cosine": 0.9999949064771451, + "max_abs": 28.864990234375, + "teacher_forced": true + }, + "layers.8": { + "rel_l2": 0.003195871196673115, + "cosine": 0.9999960215285454, + "max_abs": 11.5780029296875, + "teacher_forced": true + }, + "layers.9": { + "rel_l2": 0.0026132607115890894, + "cosine": 0.9999970479233257, + "max_abs": 11.5654296875, + "teacher_forced": true + }, + "layers.10": { + "rel_l2": 0.0031156616572957283, + "cosine": 0.999995333906672, + "max_abs": 16.512939453125, + "teacher_forced": true + }, + "layers.11": { + "rel_l2": 0.0031627321329359636, + "cosine": 0.9999962268275233, + "max_abs": 13.9443359375, + "teacher_forced": true + }, + "layers.12": { + "rel_l2": 0.002928766713158232, + "cosine": 0.9999961031017881, + "max_abs": 9.855712890625, + "teacher_forced": true + }, + "layers.13": { + "rel_l2": 0.003274534364821728, + "cosine": 0.9999950753048733, + "max_abs": 12.0478515625, + "teacher_forced": true + }, + "layers.14": { + "rel_l2": 0.004489820019516022, + "cosine": 0.9999917241605523, + "max_abs": 11.8714599609375, + "teacher_forced": true + }, + "layers.15": { + "rel_l2": 0.004613191248103328, + "cosine": 0.9999939437671664, + "max_abs": 18.96435546875, + "teacher_forced": true + }, + "layers.16": { + "rel_l2": 0.0042070955888898745, + "cosine": 0.9999922308890187, + "max_abs": 9.640380859375, + "teacher_forced": true + }, + "layers.17": { + "rel_l2": 0.004569299173512813, + "cosine": 0.9999923275489548, + "max_abs": 9.614990234375, + "teacher_forced": true + }, + "layers.18": { + "rel_l2": 0.005229238819461399, + "cosine": 0.999987550675861, + "max_abs": 14.443359375, + "teacher_forced": true + }, + "layers.19": { + "rel_l2": 0.005615343565492623, + "cosine": 0.9999874987719874, + "max_abs": 10.4100341796875, + "teacher_forced": true + }, + "layers.20": { + "rel_l2": 0.006437658395659219, + "cosine": 0.9999849892355583, + "max_abs": 11.14453125, + "teacher_forced": true + }, + "layers.21": { + "rel_l2": 0.0083457899100771, + "cosine": 0.9999722477250413, + "max_abs": 13.41864013671875, + "teacher_forced": true + }, + "layers.22": { + "rel_l2": 0.0077925201540507485, + "cosine": 0.9999736336232159, + "max_abs": 13.3956298828125, + "teacher_forced": true + }, + "layers.23": { + "rel_l2": 0.008765172724434939, + "cosine": 0.9999670786620476, + "max_abs": 13.75286865234375, + "teacher_forced": true + }, + "layers.24": { + "rel_l2": 0.009439258564747728, + "cosine": 0.9999566822583557, + "max_abs": 10.55096435546875, + "teacher_forced": true + }, + "layers.25": { + "rel_l2": 0.007385484791487392, + "cosine": 0.9999737953530788, + "max_abs": 8.49383544921875, + "teacher_forced": true + }, + "layers.26": { + "rel_l2": 0.006912184031938528, + "cosine": 0.9999761151436947, + "max_abs": 4.796142578125, + "teacher_forced": true + }, + "layers.27": { + "rel_l2": 0.005876544976859187, + "cosine": 0.9999830037856027, + "max_abs": 5.3358154296875, + "teacher_forced": true + }, + "layers.28": { + "rel_l2": 0.006600044699304974, + "cosine": 0.9999785877263496, + "max_abs": 6.535400390625, + "teacher_forced": true + }, + "layers.29": { + "rel_l2": 0.005646229044722766, + "cosine": 0.999984593661947, + "max_abs": 6.361846923828125, + "teacher_forced": true + }, + "layers.30": { + "rel_l2": 0.006837828369174909, + "cosine": 0.9999774167914105, + "max_abs": 7.477874755859375, + "teacher_forced": true + }, + "layers.31": { + "rel_l2": 0.008247466312105652, + "cosine": 0.9999660130002137, + "max_abs": 8.9798583984375, + "teacher_forced": true + }, + "layers.32": { + "rel_l2": 0.008344194965352426, + "cosine": 0.9999652473846956, + "max_abs": 9.49627685546875, + "teacher_forced": true + }, + "layers.33": { + "rel_l2": 0.00950532361336014, + "cosine": 0.9999551291958141, + "max_abs": 15.6826171875, + "teacher_forced": true + }, + "layers.34": { + "rel_l2": 0.012310363636851044, + "cosine": 0.9999242420034977, + "max_abs": 18.20947265625, + "teacher_forced": true + }, + "layers.35": { + "rel_l2": 0.015827537372845025, + "cosine": 0.9998749595508493, + "max_abs": 25.1849365234375, + "teacher_forced": true + }, + "layers.36": { + "rel_l2": 0.016685478765010673, + "cosine": 0.9998608354997787, + "max_abs": 27.2435302734375, + "teacher_forced": true + }, + "layers.37": { + "rel_l2": 0.014823509964298577, + "cosine": 0.9998901783708442, + "max_abs": 25.64599609375, + "teacher_forced": true + }, + "layers.38": { + "rel_l2": 0.014622312945486982, + "cosine": 0.9998938022034956, + "max_abs": 32.1064453125, + "teacher_forced": true + }, + "layers.39": { + "rel_l2": 0.012221662846758068, + "cosine": 0.999925511322749, + "max_abs": 31.35986328125, + "teacher_forced": true + }, + "layers.40": { + "rel_l2": 0.013119060471453468, + "cosine": 0.9999143258366655, + "max_abs": 34.70660400390625, + "teacher_forced": true + }, + "layers.41": { + "rel_l2": 0.018183323665575943, + "cosine": 0.9998346946104152, + "max_abs": 51.756072998046875, + "teacher_forced": true + }, + "final_linear_video": { + "rel_l2": 0.004440267657331908, + "cosine": 0.9999901511081131, + "max_abs": 0.025125980377197266, + "teacher_forced": true + } + }, + "free_bf16_vs_fp32": { + "layers.0": 0.0031861842375336333, + "layers.1": 0.004594587335369674, + "layers.2": 0.005362381268500324, + "layers.3": 0.005492054818950669, + "layers.4": 0.005375322967680867, + "layers.5": 0.006082104203154129, + "layers.6": 0.005669565382800687, + "layers.7": 0.005135554701206397, + "layers.8": 0.005280043116440718, + "layers.9": 0.0049243432314297585, + "layers.10": 0.00583983346217759, + "layers.11": 0.005489625022557648, + "layers.12": 0.005650423918099167, + "layers.13": 0.006338932362469926, + "layers.14": 0.007459026303118939, + "layers.15": 0.006205736862011618, + "layers.16": 0.007380737495860526, + "layers.17": 0.007358308087259622, + "layers.18": 0.010077236009820648, + "layers.19": 0.009538594987259917, + "layers.20": 0.010603625663599766, + "layers.21": 0.012958935857228618, + "layers.22": 0.013660044223886465, + "layers.23": 0.013866824224116935, + "layers.24": 0.016808228726615516, + "layers.25": 0.015926661134382892, + "layers.26": 0.020438249641896167, + "layers.27": 0.019222673949866417, + "layers.28": 0.022210265157947152, + "layers.29": 0.020340339081943277, + "layers.30": 0.024749443054652875, + "layers.31": 0.030390015872921904, + "layers.32": 0.03103268502240648, + "layers.33": 0.030879388886766206, + "layers.34": 0.038108569794760315, + "layers.35": 0.042188781228498, + "layers.36": 0.04887716233252149, + "layers.37": 0.046513691249483465, + "layers.38": 0.04864073540836466, + "layers.39": 0.05081258290168686, + "layers.40": 0.05667277614668586, + "layers.41": 0.06514794343998448 + }, + "times": { + "fp32_s": 7.6, + "bf16_s": 20.5, + "bf16_teacher_forced_s": 21.5 + }, + "adapter_effect": { + "velocity": 0.15595638414888743, + "time_embedding": 0.003370674109785744, + "shared_block_modulation": 0.002289015227805314, + "video_head_modulation": 0.0020262040484439205, + "video_embedder": 0.07837702808708451, + "text_embedder.proj": 0.07915530881381533, + "text_embedder.blocks.0": 0.08863081485281928, + "text_embedder.blocks.1": 0.11286672110132051, + "layers.0": 0.01614224007230407, + "layers.1": 0.016923150066919204, + "layers.2": 0.019610501939486488, + "layers.3": 0.018001725127289474, + "layers.4": 0.01718107582220144, + "layers.5": 0.019460064270447445, + "layers.6": 0.01832425712852528, + "layers.7": 0.019620150791070285, + "layers.8": 0.018023483904496686, + "layers.9": 0.01702209830385576, + "layers.10": 0.020012945494177366, + "layers.11": 0.019000750207336926, + "layers.12": 0.020244774096774734, + "layers.13": 0.023418389747575798, + "layers.14": 0.02587860574377009, + "layers.15": 0.02327951646405481, + "layers.16": 0.027254681907740008, + "layers.17": 0.027443791352510356, + "layers.18": 0.034712008492709, + "layers.19": 0.03432845720955072, + "layers.20": 0.037827022503225456, + "layers.21": 0.055774981722166786, + "layers.22": 0.06033407895510254, + "layers.23": 0.05954453433297213, + "layers.24": 0.06844316469700473, + "layers.25": 0.06902195326136085, + "layers.26": 0.081146491644414, + "layers.27": 0.07618318782123822, + "layers.28": 0.08520334720521558, + "layers.29": 0.0807727773159565, + "layers.30": 0.09363889188102399, + "layers.31": 0.10760542123926951, + "layers.32": 0.10032667980066186, + "layers.33": 0.09770232822775853, + "layers.34": 0.11159134446356408, + "layers.35": 0.11584416466434995, + "layers.36": 0.1236643051898582, + "layers.37": 0.12052458590864862, + "layers.38": 0.11955563175024507, + "layers.39": 0.1263884419286942, + "layers.40": 0.13891738042861596, + "layers.41": 0.1762614531213099, + "final_linear_video": 0.15595638414888743 + }, + "exact_free_floor": { + "rel_l2": 1.5160680761024458e-05, + "row": "free layers.41", + "measured": "gmlx's exact step under the adapter's patch table on the GPU against the same step on the CPU stream, worst free row; tests/image/moss_gates.py exact_free_floor on moss-image 06f68465, Apple M5 Max" + }, + "exact_velocity_floor": { + "rel_l2": 1.0271258992509532e-05, + "measured": "gmlx's exact step under the adapter's patch table on the GPU against the same step on the CPU stream, velocity; tests/image/moss_gates.py exact_free_floor on moss-image 06f68465, Apple M5 Max" + } + }, + "tolerances": { + "velocity": { + "rel_l2": 0.08878016534876057, + "cosine_min": 0.996057630868934 + }, + "time_embedding": { + "rel_l2": 0.004037430957833937, + "teacher_forced": false + }, + "shared_block_modulation": { + "rel_l2": 0.003977484380340728, + "teacher_forced": false + }, + "video_head_modulation": { + "rel_l2": 0.003433779287013718, + "teacher_forced": false + }, + "video_embedder": { + "rel_l2": 0.005766286225456487, + "teacher_forced": false + }, + "text_embedder.proj": { + "rel_l2": 0.005785772132360488, + "teacher_forced": false + }, + "text_embedder.blocks.0": { + "rel_l2": 0.011447936867846703, + "teacher_forced": true + }, + "text_embedder.blocks.1": { + "rel_l2": 0.008913547570385401, + "teacher_forced": true + }, + "layers.0": { + "rel_l2": 0.006080307357539287, + "teacher_forced": true + }, + "layers.1": { + "rel_l2": 0.006207431271747408, + "teacher_forced": true + }, + "layers.2": { + "rel_l2": 0.007576263225284296, + "teacher_forced": true + }, + "layers.3": { + "rel_l2": 0.006239911083619844, + "teacher_forced": true + }, + "layers.4": { + "rel_l2": 0.005724095081931286, + "teacher_forced": true + }, + "layers.5": { + "rel_l2": 0.0068977335389012575, + "teacher_forced": true + }, + "layers.6": { + "rel_l2": 0.005693935994617375, + "teacher_forced": true + }, + "layers.7": { + "rel_l2": 0.006450103520681092, + "teacher_forced": true + }, + "layers.8": { + "rel_l2": 0.00639174239334623, + "teacher_forced": true + }, + "layers.9": { + "rel_l2": 0.005226521423178179, + "teacher_forced": true + }, + "layers.10": { + "rel_l2": 0.006231323314591457, + "teacher_forced": true + }, + "layers.11": { + "rel_l2": 0.006325464265871927, + "teacher_forced": true + }, + "layers.12": { + "rel_l2": 0.005857533426316464, + "teacher_forced": true + }, + "layers.13": { + "rel_l2": 0.006549068729643456, + "teacher_forced": true + }, + "layers.14": { + "rel_l2": 0.008979640039032044, + "teacher_forced": true + }, + "layers.15": { + "rel_l2": 0.009226382496206656, + "teacher_forced": true + }, + "layers.16": { + "rel_l2": 0.008414191177779749, + "teacher_forced": true + }, + "layers.17": { + "rel_l2": 0.009138598347025626, + "teacher_forced": true + }, + "layers.18": { + "rel_l2": 0.010458477638922797, + "teacher_forced": true + }, + "layers.19": { + "rel_l2": 0.011230687130985247, + "teacher_forced": true + }, + "layers.20": { + "rel_l2": 0.012875316791318437, + "teacher_forced": true + }, + "layers.21": { + "rel_l2": 0.0166915798201542, + "teacher_forced": true + }, + "layers.22": { + "rel_l2": 0.015585040308101497, + "teacher_forced": true + }, + "layers.23": { + "rel_l2": 0.017530345448869878, + "teacher_forced": true + }, + "layers.24": { + "rel_l2": 0.018878517129495457, + "teacher_forced": true + }, + "layers.25": { + "rel_l2": 0.014770969582974783, + "teacher_forced": true + }, + "layers.26": { + "rel_l2": 0.013824368063877056, + "teacher_forced": true + }, + "layers.27": { + "rel_l2": 0.011753089953718375, + "teacher_forced": true + }, + "layers.28": { + "rel_l2": 0.013200089398609949, + "teacher_forced": true + }, + "layers.29": { + "rel_l2": 0.011292458089445533, + "teacher_forced": true + }, + "layers.30": { + "rel_l2": 0.013675656738349818, + "teacher_forced": true + }, + "layers.31": { + "rel_l2": 0.016494932624211305, + "teacher_forced": true + }, + "layers.32": { + "rel_l2": 0.01668838993070485, + "teacher_forced": true + }, + "layers.33": { + "rel_l2": 0.01901064722672028, + "teacher_forced": true + }, + "layers.34": { + "rel_l2": 0.024620727273702087, + "teacher_forced": true + }, + "layers.35": { + "rel_l2": 0.03165507474569005, + "teacher_forced": true + }, + "layers.36": { + "rel_l2": 0.033370957530021346, + "teacher_forced": true + }, + "layers.37": { + "rel_l2": 0.029647019928597154, + "teacher_forced": true + }, + "layers.38": { + "rel_l2": 0.029244625890973965, + "teacher_forced": true + }, + "layers.39": { + "rel_l2": 0.024443325693516135, + "teacher_forced": true + }, + "layers.40": { + "rel_l2": 0.026238120942906937, + "teacher_forced": true + }, + "layers.41": { + "rel_l2": 0.036366647331151886, + "teacher_forced": true + }, + "final_linear_video": { + "rel_l2": 0.008880535314663816, + "teacher_forced": true + }, + "rule": "tolerance = 2 x the oracle's own bf16 against fp32 relative L2: teacher-forced for each text block, each block and the head, free running for the other rows and for the velocity", + "exact_free": { + "rel_l2": 5e-05 + }, + "exact_velocity": { + "rel_l2": 5e-05 + } + }, + "notes": [ + "the adapter is applied unmerged as PEFT applies it; the gmlx side runs it under a patch table" + ], + "created": "2026-10-18T12:10:14Z", + "npz": "check7_moss_image_synthetic_512.npz", + "tensors": { + "down.final_linear_video": { + "shape": [ + 4, + 3072 + ], + "dtype": "float32", + "sha256": "229816ec33698fba4a71f8d3026d3c14266261881af17f80b3b39fc7596b5c4c" + }, + "down.layers.0.proj.ffn_down": { + "shape": [ + 8, + 8192 + ], + "dtype": "float32", + "sha256": "b5dc6cc4648b8a2a3e5e6e07dfd263b1d15a5cfd94e855b4fef0e1eba003c207" + }, + "down.layers.0.proj.ffn_up_gate": { + "shape": [ + 4, + 3072 + ], + "dtype": "float32", + "sha256": "9687d9728fd2376f4baebd89015678c4285987744b4f420e45e856efa3fba25c" + }, + "down.layers.0.proj.k": { + "shape": [ + 8, + 3072 + ], + "dtype": "float32", + "sha256": "4f1e1c4e2d12edee9e346f19fa576a335c29f57bca3989f629f9b58db365d4c4" + }, + "down.layers.0.proj.o": { + "shape": [ + 16, + 3072 + ], + "dtype": "float32", + "sha256": "c4cff7b1f4548e801e76ddb987d83b048e15b4770ea284e978189be5a84c3799" + }, + "down.layers.0.proj.q": { + "shape": [ + 4, + 3072 + ], + "dtype": "float32", + "sha256": "e6415c3573b6c76ecab2e80520335c45b92d7986a6f4fb02e18b2364ef78cfe2" + }, + "down.layers.0.proj.v": { + "shape": [ + 12, + 3072 + ], + "dtype": "float32", + "sha256": "51c86f31233e8b388da6620f0afea0415de35f9c71af21983096d16d0838a1c9" + }, + "down.layers.21.proj.ffn_down": { + "shape": [ + 16, + 8192 + ], + "dtype": "float32", + "sha256": "064025075b4230c415585396af109fb18c1122ac973ee9bab3da182408ed6e59" + }, + "down.layers.21.proj.ffn_up_gate": { + "shape": [ + 12, + 3072 + ], + "dtype": "float32", + "sha256": "c6dcdb6199c7662fcc59168c7adb374e109edf4c773d2e5fa769e86fd081a514" + }, + "down.layers.21.proj.k": { + "shape": [ + 16, + 3072 + ], + "dtype": "float32", + "sha256": "90ea71bf270710d8341dbb873fbaa006e9a95a1206469cd124d728627fe00646" + }, + "down.layers.21.proj.o": { + "shape": [ + 8, + 3072 + ], + "dtype": "float32", + "sha256": "168ea4ad5432ce95d00532261cb1f14eaa4fdd62c293752f202b1feaabb4ad5a" + }, + "down.layers.21.proj.q": { + "shape": [ + 12, + 3072 + ], + "dtype": "float32", + "sha256": "a93905d38920fcdc20ced92c9264ce139b5b235742ee9b52d3b5e5b6e336fa85" + }, + "down.layers.21.proj.v": { + "shape": [ + 4, + 3072 + ], + "dtype": "float32", + "sha256": "ca47d7c9de02c62328b5e355c784bd264a18e2bd10ebff0dc039ef268b307559" + }, + "down.layers.41.proj.ffn_down": { + "shape": [ + 8, + 8192 + ], + "dtype": "float32", + "sha256": "9a2a191c41b4da7b34b3de64c735896282f59dc115250f797efecbf0ad0bdbe5" + }, + "down.layers.41.proj.ffn_up_gate": { + "shape": [ + 4, + 3072 + ], + "dtype": "float32", + "sha256": "7febc3d410924618797042a065b1cf5c031a50680d7197ee025579f6531686ee" + }, + "down.layers.41.proj.k": { + "shape": [ + 8, + 3072 + ], + "dtype": "float32", + "sha256": "2f6f3694ff3aedc6132f98008b31ebdaf9e38617a071a2cb79b143689034bdfe" + }, + "down.layers.41.proj.o": { + "shape": [ + 16, + 3072 + ], + "dtype": "float32", + "sha256": "05be0e939d4032e7bf45bbc2fe26964ff1b85557f5f1d57058f5dea84c4027ab" + }, + "down.layers.41.proj.q": { + "shape": [ + 4, + 3072 + ], + "dtype": "float32", + "sha256": "ccf6dcf1f2d55b701b88c590335d19b9397bdb2ddcfaa4c4ea630e11276966d4" + }, + "down.layers.41.proj.v": { + "shape": [ + 12, + 3072 + ], + "dtype": "float32", + "sha256": "2bec251d9ddea951019346eb985779c9302359a18d80115f6a658f2703d96904" + }, + "down.shared_block_modulation": { + "shape": [ + 12, + 3072 + ], + "dtype": "float32", + "sha256": "7f3e974405c38bedb623475bf4bef4ad1e3e93f8567cd798ab1c2e5f0c211679" + }, + "down.text_embedder.blocks.1.ffn_down": { + "shape": [ + 12, + 8192 + ], + "dtype": "float32", + "sha256": "6a71c220ceec73b0324daa8488fca9e541a90761ddfa6aec434dc1a12a648d69" + }, + "down.text_embedder.blocks.1.ffn_gate": { + "shape": [ + 4, + 3072 + ], + "dtype": "float32", + "sha256": "bd27a8a461d4ad807b86eec85e54658a68b5ab2eae8f7fd4ec08c27cdbf0ac08" + }, + "down.text_embedder.blocks.1.ffn_up": { + "shape": [ + 8, + 3072 + ], + "dtype": "float32", + "sha256": "4f8c5f5d6b95619e3fbdd7fd9c750406e120463f5919ba1d5d069c0333202904" + }, + "down.text_embedder.blocks.1.proj": { + "shape": [ + 16, + 3072 + ], + "dtype": "float32", + "sha256": "5b1dc7d0dbfaa0d4f3ec5b63014dbbb840cbdb4ceb42ef578c6f3c93b27e7a09" + }, + "down.text_embedder.blocks.1.qkv": { + "shape": [ + 12, + 3072 + ], + "dtype": "float32", + "sha256": "fb0dc5bc211329aca5b4e1ea83753409664251af03e6b0fedc35b5c04757f10d" + }, + "down.text_embedder.proj": { + "shape": [ + 4, + 4096 + ], + "dtype": "float32", + "sha256": "a70d00dea3d02592da9db83038a0c79197247eb00a6edd2b0faee8dec5d77d78" + }, + "down.time_embedding.2": { + "shape": [ + 8, + 3072 + ], + "dtype": "float32", + "sha256": "714b80261eb07232f15a4135e3c63d7e6a847ee119c82dc13830abd3d7afb562" + }, + "down.video_embedder": { + "shape": [ + 16, + 128 + ], + "dtype": "float32", + "sha256": "dfd2ffa71b132bc50741cd7e261ebc8e9aa89b9b05944f05a412432c8bb8c79b" + }, + "down.video_head_modulation": { + "shape": [ + 16, + 3072 + ], + "dtype": "float32", + "sha256": "a88551eeeba34e487a9919140db85f59597a9db8c4a7da2a1c2a7ff296a8f0ab" + }, + "features": { + "shape": [ + 48, + 4096 + ], + "dtype": "float32", + "sha256": "51f2174fad964908bd5817f252f645cdc220c2c13ca457e1ad899301a9c3ca6b" + }, + "final_linear_video": { + "shape": [ + 1, + 1024, + 128 + ], + "dtype": "float32", + "sha256": "5c24e6b9da463bf47d602b12ad1d71949536e5ddfcc6d780b998e17a5b8b8011" + }, + "in.final_norm_video": { + "shape": [ + 1, + 1024, + 3072 + ], + "dtype": "float32", + "sha256": "39131e523f8f1d78283047250ffdafe8ee2fab7253b72b7fd30d977653d35ed1" + }, + "in.layers.0": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "bb325a22457ba5d40cfd5037ef78ee143e1e46314f66aa95e7da21c4b1e5ca55" + }, + "latents": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "4399e6b9bd209dc9af6f745557507a33dcfe730a81abebb0bce079d2b441f34c" + }, + "layers.0": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "0c8c7f95c69d7ec12b3028a0efa2e49c688bf0198e6507b85a4c5475390fed9d" + }, + "layers.1": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "06ce0afa4380da9e96711ab8095e53f101d962aaa63b4c99f69e7a8d4ae4d8cd" + }, + "layers.10": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "08f6051caefa3cdf85950259503b8420dfa156f42b5611e687d694cdbd5ac9b0" + }, + "layers.11": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "0e4f6236d98d9baf0588bb8384a43fd8ddb3c091341ceadd1486d9b07558baf4" + }, + "layers.12": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "e7408e21f1aeee72324609e4cafe1a7374d59d2f76c6aa8d12e3d8d882e99609" + }, + "layers.13": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "a9133d4386934a78cc5450d6f87012dab823d5de5761773e923387a841adcfe7" + }, + "layers.14": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "37e757afa482f2719e938574e20847322c56665dac482011bb8738cd2a38a5e1" + }, + "layers.15": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "d2bd0515669642e2f3f99803aac511133cc5cacb4d5115453bdaa3b028515afd" + }, + "layers.16": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "c1e5f181551eb476ebc6986090f992dcdd80128de7e46eedc6688cdf0a426d34" + }, + "layers.17": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "04ed2d3838d839259a8874894180829bf3215324ce20e0836e5937539cf0c52a" + }, + "layers.18": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "4df78b68567c4e115208f1fda48bb3a2e1ada3346a7199c8fb5f8e8b851c37b3" + }, + "layers.19": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "ef7450fe534bf68fa5e75123f19eba232425838deca76b025d2a1310e0d39d9f" + }, + "layers.2": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "b6d942f4c2d72fcbb0eaf0b98ae0e1d4bfd1a9ae166b44a889cbb24426c14e30" + }, + "layers.20": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "81cae6bdae12e1ff7e63301f877740b7c38fea869fd0f0126316184c0ff3efd6" + }, + "layers.21": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "57aa179c684f33d6491c1edb36f832a87e1af9627d0b50feca2ead5a6ac5995b" + }, + "layers.22": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "2a70b76e81517f711619626b59d2bb155d45d1a22a6124c5fd709bc5f04a8b6f" + }, + "layers.23": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "285ce32dc8a0915fc00dd91795a519885a467065aa9f5c0f81aaaaa241db46db" + }, + "layers.24": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "bc2dd44770d33e593d63517828d088e7d6940d980a9bf9ac987a81cb6d7e67a3" + }, + "layers.25": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "45f8c557a1278cb960a56e664eec83455f1cf34805b25ae91bc4ae07d94cfd3f" + }, + "layers.26": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "cca082266f23a08f529e52e27677e78a86ca2a1366a2cc17aab27ef04946beaa" + }, + "layers.27": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "640871f0fa2cf0a18d603e75bf71caae0c573536cd4fa06c2b0469fce52f809e" + }, + "layers.28": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "ccee1e5b9c3365e6244ccad97a7d3bb49ff9f40f5fee2862274f61bb4e54d5d6" + }, + "layers.29": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "848b9a275483d3cbf07c8613d085aebecc663ecec7d65c398650c6a60cc4d9ac" + }, + "layers.3": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "ffe3c4502acbd9407ad9004ae60445dfef193835f7be7cd202744e4bb8c8f258" + }, + "layers.30": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "0417c0e6fa561ebbb7e7dbd0ff3e4305f9e3787efe923163e36d5b4226d794b4" + }, + "layers.31": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "49dbb7fe401617f68707bdc460300be1231e73fdb31fcdaba104f06ca16dcf77" + }, + "layers.32": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "ce4cda5087ed510b5ddfc63efded9bd2a4e2828220031b9bc0f55367c992b221" + }, + "layers.33": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "2240642f089c3d26f3dcce9c254b83b0d82da1ad24b5b9dd7d2bc8da93cb7e84" + }, + "layers.34": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "9c90620b0cc23cc38d5ad504137fdf404b69619751f9ba7fb93a5b794cbfa9b6" + }, + "layers.35": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "2a8bd6cdd9e4b441c1c92f0b33798086315237a42f0e15523566115c8652d4b6" + }, + "layers.36": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "57e34a8f38ea64aef04d3c209aefec4619e61f140b8fb8e479885b04905648a2" + }, + "layers.37": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "765eec2b870c22c702dbbf20bde3941e6b8f5ed21d2ae037ac57cb6b4011d76f" + }, + "layers.38": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "34df46918ac04c9eccf63dfea7455106d2b3142274d8dded2c0f38ebc3737eea" + }, + "layers.39": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "b6a69d3a4e79a155e9e7586f9921d93abff298b5ae1a6db708d7b42321af8664" + }, + "layers.4": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "5eb34df9416e150b1bfcce2a87830aa98432beae361e3643fdc84c7159e1a1d8" + }, + "layers.40": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "a28bb73dd29d229c0f94dd9e7f684cefaeb23553c8015cd5266953cabd2cfcc8" + }, + "layers.41": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "7003bd8cbe5187ee5d86de91f6e2bdda30e0294589e74d4c8960f4db8ff788a1" + }, + "layers.5": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "ebe752bedac45e55a5915c30095df32c945520820622aa980b90cd91f8ca23c0" + }, + "layers.6": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "adda867b39d68c70892e3c49c26e0016bb83a6e2e00e50182559465dbf93d9cd" + }, + "layers.7": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "1c8c25a8fdf221792b80e28807524d5c0aaf6b393037147efc126cfb542d6c86" + }, + "layers.8": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "c4310a45bdfedb4b09def77fbb760b661908c625f84d518f4a97e3745e58a8ec" + }, + "layers.9": { + "shape": [ + 1, + 1072, + 3072 + ], + "dtype": "float32", + "sha256": "5d62b6f73c341dcddbbdd173801e2da5c70584d3e1125ce8a8570985e162a94b" + }, + "shared_block_modulation": { + "shape": [ + 1, + 18432 + ], + "dtype": "float32", + "sha256": "b8792658ee3dc24f801619914c90548b4c6d46e0a2dfb447d0256d80d3b15b02" + }, + "sigma": { + "shape": [], + "dtype": "float32", + "sha256": "9a8208635e00348ab64aac2b759e76391fd47089e9a749bbcec770d9eb5c6421" + }, + "text_embedder.blocks.0": { + "shape": [ + 1, + 48, + 3072 + ], + "dtype": "float32", + "sha256": "6002c43257ead6b015a99fa4c8c0c292506bf5186af99f0c8dcf2f0f4dfffc0a" + }, + "text_embedder.blocks.1": { + "shape": [ + 1, + 48, + 3072 + ], + "dtype": "float32", + "sha256": "7299e45b216651b8f7f92717ff7b7d7e103d124ce888133c2f1eaeff7f993aec" + }, + "text_embedder.proj": { + "shape": [ + 1, + 48, + 3072 + ], + "dtype": "float32", + "sha256": "231b30d31457f41feeb9094119d655c851236ea7a00593f00d0d0f7fb30ff053" + }, + "time_embedding": { + "shape": [ + 1, + 3072 + ], + "dtype": "float32", + "sha256": "1f374466cc191b5012c3bafb15574f220574fc85d600e1a0dd4e462dedff1395" + }, + "timestep": { + "shape": [ + 1 + ], + "dtype": "float32", + "sha256": "1d77e51c3092b797b15b02104368537eea7b222148a4d2ba2d2d578491c50ac5" + }, + "up.final_linear_video": { + "shape": [ + 128, + 4 + ], + "dtype": "float32", + "sha256": "1afabdd48d8f12b0e668a9b44229a9b1fc6b20a6cd8de1e0edcf1bdb1fedc97c" + }, + "up.layers.0.proj.ffn_down": { + "shape": [ + 3072, + 8 + ], + "dtype": "float32", + "sha256": "759139c381d1943bb75c6c7d45dbfb504293fa36a3e029df473c182a34d282fb" + }, + "up.layers.0.proj.ffn_up_gate": { + "shape": [ + 16384, + 4 + ], + "dtype": "float32", + "sha256": "5bea065e271c608837e2a0c3da156f1ec20c1e12c177c2cd299459902660a6a8" + }, + "up.layers.0.proj.k": { + "shape": [ + 3072, + 8 + ], + "dtype": "float32", + "sha256": "589f97a48ffbb1764212bb0becf58079a2d2b936c01b47794731db6b8a73e84c" + }, + "up.layers.0.proj.o": { + "shape": [ + 3072, + 16 + ], + "dtype": "float32", + "sha256": "674c2153275c6d9c95de9bf22895c8ce089d75ba21d6f8d6eaf563b918892bdd" + }, + "up.layers.0.proj.q": { + "shape": [ + 3072, + 4 + ], + "dtype": "float32", + "sha256": "74e1574bc9007bba529f407bc51707887817e54390fc6cd6b65d24dd20ee0769" + }, + "up.layers.0.proj.v": { + "shape": [ + 3072, + 12 + ], + "dtype": "float32", + "sha256": "39b5b8320d2b3a885bc9fb007a37e80a7b998fc88b1943a6477eb71c94926bbc" + }, + "up.layers.21.proj.ffn_down": { + "shape": [ + 3072, + 16 + ], + "dtype": "float32", + "sha256": "96c57498970b6adba77b17697c7dc18976037ab38b60a386e6cf53dc2ccc8fbe" + }, + "up.layers.21.proj.ffn_up_gate": { + "shape": [ + 16384, + 12 + ], + "dtype": "float32", + "sha256": "59c66eb16126cbd6965f2276af11e79f81013bafd34f4c0d541a2ab29d640861" + }, + "up.layers.21.proj.k": { + "shape": [ + 3072, + 16 + ], + "dtype": "float32", + "sha256": "cecefd2ba1280b37f45a5e2df47cb231a9d758c4c8cec3658cdce39b74b15f50" + }, + "up.layers.21.proj.o": { + "shape": [ + 3072, + 8 + ], + "dtype": "float32", + "sha256": "82ee6a782d7510207184a9c85249e7623b471c04cd72e4a3e47d629ad1dae946" + }, + "up.layers.21.proj.q": { + "shape": [ + 3072, + 12 + ], + "dtype": "float32", + "sha256": "d4bf6a0aee1df533f0e72b43c2d1b59fefd37701c4b658063334d8499f6f48d8" + }, + "up.layers.21.proj.v": { + "shape": [ + 3072, + 4 + ], + "dtype": "float32", + "sha256": "8d0f9158c1bbf12443d0fae98cbfe176fcb676a2db0f3b7c71cf2b01abc1839b" + }, + "up.layers.41.proj.ffn_down": { + "shape": [ + 3072, + 8 + ], + "dtype": "float32", + "sha256": "67e70378f8255fa3501ec674f2fc9f756452f3d2f88450fbd603d0e40c5cb9e4" + }, + "up.layers.41.proj.ffn_up_gate": { + "shape": [ + 16384, + 4 + ], + "dtype": "float32", + "sha256": "c77310a640dbab3685f234589b8dc0837624028ce41d6b3b9b7e54a61ba06250" + }, + "up.layers.41.proj.k": { + "shape": [ + 3072, + 8 + ], + "dtype": "float32", + "sha256": "f2bdf7d2c80b0417f7c9444d608c0adc5ec6351bbbdacbfc4e92aca7529914dc" + }, + "up.layers.41.proj.o": { + "shape": [ + 3072, + 16 + ], + "dtype": "float32", + "sha256": "cfe30b4eb36e5eabe6ff8a9bc5d5025e8ff8b5b17e92fc1254e95e413ff50191" + }, + "up.layers.41.proj.q": { + "shape": [ + 3072, + 4 + ], + "dtype": "float32", + "sha256": "4fe3e70f32cb6fe7ed341b26f359f610388c4d2033454e0d8db2c8dfb9ef592f" + }, + "up.layers.41.proj.v": { + "shape": [ + 3072, + 12 + ], + "dtype": "float32", + "sha256": "e07864cd63b7b3657395195b3e979253204b7666e25b1a364fe5f6dfa95880f3" + }, + "up.shared_block_modulation": { + "shape": [ + 18432, + 12 + ], + "dtype": "float32", + "sha256": "8052c8181ca91273d8a41266ed93a4492cead115ccf7e29f2231963a70d51243" + }, + "up.text_embedder.blocks.1.ffn_down": { + "shape": [ + 3072, + 12 + ], + "dtype": "float32", + "sha256": "a9516aa6414234dc22836c3f5a0adc5035b447a7500ba0178a47771d25aae722" + }, + "up.text_embedder.blocks.1.ffn_gate": { + "shape": [ + 8192, + 4 + ], + "dtype": "float32", + "sha256": "224178d0c90d39a80d79282dc7ab13958b76537b11fc197a69116afdad35f083" + }, + "up.text_embedder.blocks.1.ffn_up": { + "shape": [ + 8192, + 8 + ], + "dtype": "float32", + "sha256": "f027b7be2b82f37b21aff40dddfb35b15a1e86e413164e648657b4ee2417443b" + }, + "up.text_embedder.blocks.1.proj": { + "shape": [ + 3072, + 16 + ], + "dtype": "float32", + "sha256": "e9c42c4aae76d3cb66d1b1600bf8c32c7e7d4df1f107143db492ead01fd30f4e" + }, + "up.text_embedder.blocks.1.qkv": { + "shape": [ + 9216, + 12 + ], + "dtype": "float32", + "sha256": "cf5fbab28dd118848ac2f63ca5024485ae8074b38331e62ab598b61a55d8233c" + }, + "up.text_embedder.proj": { + "shape": [ + 3072, + 4 + ], + "dtype": "float32", + "sha256": "e6a00b564e52ac986cc65bcf76d0ddb686129123f60166cfdbc2134f78b9ddbb" + }, + "up.time_embedding.2": { + "shape": [ + 3072, + 8 + ], + "dtype": "float32", + "sha256": "48125f880ecd30b193030121fa745a3e6bb9ada485757691b264071f3bafa14b" + }, + "up.video_embedder": { + "shape": [ + 3072, + 16 + ], + "dtype": "float32", + "sha256": "3c4739649dba8e716aa8733c12b03a78e171e3fc164cc476ef8707460a64c053" + }, + "up.video_head_modulation": { + "shape": [ + 6144, + 16 + ], + "dtype": "float32", + "sha256": "6b614ed56fa23bba6944c3a5a41c080febec2c4e07ac4835cb36ddef854d04f0" + }, + "velocity": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "9362a2c80e9ad045fa4548cb9407e8b9c8033b8993195a285f81e7317bf743bb" + }, + "velocity_bf16": { + "shape": [ + 1, + 128, + 32, + 32 + ], + "dtype": "float32", + "sha256": "21c644d0aec9ae45defeca5b8fe14dbab5ca07a25a9e72a72f83f9328a106bc4" + }, + "video_embedder": { + "shape": [ + 1, + 1024, + 3072 + ], + "dtype": "float32", + "sha256": "c3830c017b383b9a444815a49883a424b85a8ec5c843e5f7515f85674e35d368" + }, + "video_head_modulation": { + "shape": [ + 1, + 6144 + ], + "dtype": "float32", + "sha256": "8da4a6455a794a4a4381c68cd0276e751d40b8e90d2617a86bab0ac1848e3c36" + } + } +} diff --git a/tests/fixtures/image/oracle/selftest_moss_image_tiny.manifest.json b/tests/fixtures/image/oracle/selftest_moss_image_tiny.manifest.json new file mode 100644 index 000000000..cb0637983 --- /dev/null +++ b/tests/fixtures/image/oracle/selftest_moss_image_tiny.manifest.json @@ -0,0 +1,364 @@ +{ + "schema": 1, + "name": "selftest_moss_image_tiny", + "check": 3, + "family": "moss_image", + "script": "scripts/image_oracle/check3_moss_image.py", + "argv": [ + "check3_moss_image.py", + "--pipeline", + "pipeline_moss_image.py", + "--self-test" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": { + "pipeline": { + "name": "pipeline_moss_image.py", + "size": 25060, + "sha256": "1dac1aca9feb1bfdfc727a401e4b67d61a505d3190706d1b7d2d784f88678b31" + } + }, + "inputs": { + "config": { + "dim": 64, + "ffn_dim": 48, + "num_heads": 2, + "num_layers": 2, + "text_dim": 24, + "freq_dim": 16, + "latent_channels": 8, + "rope_base": 10000.0 + }, + "case": { + "latent": [ + 3, + 4 + ], + "text": 5, + "sigma": 0.6875 + }, + "seed": 0 + }, + "results": { + "velocity_bf16_vs_fp32": { + "rel_l2": 0.010003348221219725, + "cosine": 0.9999547055716813, + "max_abs": 0.034577369689941406 + }, + "rows_bf16_vs_fp32": { + "time_embedding": { + "rel_l2": 0.0036691585527701353, + "cosine": 0.9999932774291955, + "max_abs": 0.004172205924987793, + "teacher_forced": false + }, + "shared_block_modulation": { + "rel_l2": 0.004006088190592911, + "cosine": 0.9999920128046587, + "max_abs": 0.006873726844787598, + "teacher_forced": false + }, + "video_head_modulation": { + "rel_l2": 0.0035652210258079416, + "cosine": 0.9999936453471964, + "max_abs": 0.005121469497680664, + "teacher_forced": false + }, + "video_embedder": { + "rel_l2": 0.0028313136262487393, + "cosine": 0.9999960481839549, + "max_abs": 0.01363682746887207, + "teacher_forced": false + }, + "text_embedder.proj": { + "rel_l2": 0.0028468714726909068, + "cosine": 0.9999959910724129, + "max_abs": 0.01065206527709961, + "teacher_forced": false + }, + "text_embedder.blocks.0": { + "rel_l2": 0.006853589186401344, + "cosine": 0.999976647456216, + "max_abs": 0.04240083694458008, + "teacher_forced": true + }, + "text_embedder.blocks.1": { + "rel_l2": 0.004318806149963958, + "cosine": 0.9999906999949248, + "max_abs": 0.025324583053588867, + "teacher_forced": true + }, + "layers.0": { + "rel_l2": 0.005086596493744148, + "cosine": 0.9999872243509609, + "max_abs": 0.03275442123413086, + "teacher_forced": true + }, + "layers.1": { + "rel_l2": 0.004460280136738237, + "cosine": 0.9999901163465567, + "max_abs": 0.032853126525878906, + "teacher_forced": true + }, + "final_linear_video": { + "rel_l2": 0.004623934932363783, + "cosine": 0.9999893499714667, + "max_abs": 0.014683723449707031, + "teacher_forced": true + } + }, + "free_bf16_vs_fp32": { + "layers.0": 0.008763644962712492, + "layers.1": 0.009593163749995323 + }, + "times": { + "fp32_s": 0.0, + "bf16_s": 0.0, + "bf16_teacher_forced_s": 0.0 + } + }, + "tolerances": { + "velocity": { + "rel_l2": 0.02000669644243945, + "cosine_min": 0.9998188222867253 + }, + "time_embedding": { + "rel_l2": 0.0073383171055402706, + "teacher_forced": false + }, + "shared_block_modulation": { + "rel_l2": 0.008012176381185822, + "teacher_forced": false + }, + "video_head_modulation": { + "rel_l2": 0.007130442051615883, + "teacher_forced": false + }, + "video_embedder": { + "rel_l2": 0.005662627252497479, + "teacher_forced": false + }, + "text_embedder.proj": { + "rel_l2": 0.0056937429453818135, + "teacher_forced": false + }, + "text_embedder.blocks.0": { + "rel_l2": 0.013707178372802688, + "teacher_forced": true + }, + "text_embedder.blocks.1": { + "rel_l2": 0.008637612299927915, + "teacher_forced": true + }, + "layers.0": { + "rel_l2": 0.010173192987488295, + "teacher_forced": true + }, + "layers.1": { + "rel_l2": 0.008920560273476473, + "teacher_forced": true + }, + "final_linear_video": { + "rel_l2": 0.009247869864727566, + "teacher_forced": true + }, + "rule": "tolerance = 2 x the oracle's own bf16 against fp32 relative L2: teacher-forced for each text block, each block and the head, free running for the other rows and for the velocity" + }, + "notes": [ + "the weights are in selftest_moss_image_tiny_weights by name in torch order, with the class's rope.bands buffer" + ], + "created": "2026-10-18T06:29:29Z", + "npz": "selftest_moss_image_tiny.npz", + "tensors": { + "features": { + "shape": [ + 5, + 24 + ], + "dtype": "float32", + "sha256": "2c535a5350cf5552afff12785b3d2903e7a75212dbf7b2dce0e5ba2a362b25e8" + }, + "final_linear_video": { + "shape": [ + 1, + 12, + 8 + ], + "dtype": "float32", + "sha256": "14b7b225469eeab47a3e882551f7d0edb409dd029f650d08fedacf235429cca7" + }, + "in.final_norm_video": { + "shape": [ + 1, + 12, + 64 + ], + "dtype": "float32", + "sha256": "0933ea67c94bb65c0aeda55dfd710d470c2deacc790be9739aa8b0046714d085" + }, + "in.layers.0": { + "shape": [ + 1, + 17, + 64 + ], + "dtype": "float32", + "sha256": "c53accdc6a1b3c24c754829306cda98613a3f4df08185b7d6d9b67a3c8c02936" + }, + "latents": { + "shape": [ + 1, + 8, + 3, + 4 + ], + "dtype": "float32", + "sha256": "7ff9d5a79d8e7811cf664445f8c7cd8d10b0f1085d8e87cccb9dd2b570b016e3" + }, + "layers.0": { + "shape": [ + 1, + 17, + 64 + ], + "dtype": "float32", + "sha256": "2c77d84aebfa763451fed9a31e21eb6f3b5256d58b2ecbdcc3a10910b1b1a7d3" + }, + "layers.1": { + "shape": [ + 1, + 17, + 64 + ], + "dtype": "float32", + "sha256": "035efc284b22f48f290eb51bf196bdb63cf99ba56dad244f2b37435540d2578d" + }, + "shared_block_modulation": { + "shape": [ + 1, + 384 + ], + "dtype": "float32", + "sha256": "ff10a725992f7b92fd2a7608f264d61bd46cbefff0a901dded05e53b3580024a" + }, + "sigma": { + "shape": [], + "dtype": "float32", + "sha256": "06ebd70f712594b2c876556943aaa43db9d26e9c3f4fa8fcfe39e589eae9940c" + }, + "text_embedder.blocks.0": { + "shape": [ + 1, + 5, + 64 + ], + "dtype": "float32", + "sha256": "a0e8825ba49f5ecf0b2cf1b6198c77148a09995517b11ce56cc1d39d8dc76829" + }, + "text_embedder.blocks.1": { + "shape": [ + 1, + 5, + 64 + ], + "dtype": "float32", + "sha256": "ae39e2d2f91929b0c860ca37b11f0d771cecf9805f736c9c801407113759ded7" + }, + "text_embedder.proj": { + "shape": [ + 1, + 5, + 64 + ], + "dtype": "float32", + "sha256": "e6ffd67fe945d607763342c4ffb234e11800be098379df32a31f81a763349de5" + }, + "time_embedding": { + "shape": [ + 1, + 64 + ], + "dtype": "float32", + "sha256": "67d6cc97267b3b5b5346a6e071956a800947439ac840a7e125c173ee775ab4a0" + }, + "timestep": { + "shape": [ + 1 + ], + "dtype": "float32", + "sha256": "723b6da32a396adc2f865ec9e3980269bc9bd75f0fc1d8a7fdb7cb3267df19dd" + }, + "velocity": { + "shape": [ + 1, + 8, + 3, + 4 + ], + "dtype": "float32", + "sha256": "5fb3bfd69a944260535d536ddef3cb0d76f5bc5ce273a9e131fd35d02684dfac" + }, + "velocity_bf16": { + "shape": [ + 1, + 8, + 3, + 4 + ], + "dtype": "float32", + "sha256": "4cbd19c2078e989802c68dcf88794f57f4883838f14d67c1e6ef1b9b79208128" + }, + "video_embedder": { + "shape": [ + 1, + 12, + 64 + ], + "dtype": "float32", + "sha256": "e718ef43813a5e225ff3913f6ad08aeed88f0c32f6ae429fe98803daaee2ea14" + }, + "video_head_modulation": { + "shape": [ + 1, + 128 + ], + "dtype": "float32", + "sha256": "4307fd53c812ae6a56f4b87ce4d39a6e05864bf3ed52e764bbaeae29dc5d8973" + } + } +} diff --git a/tests/fixtures/image/oracle/selftest_moss_image_tiny.npz b/tests/fixtures/image/oracle/selftest_moss_image_tiny.npz new file mode 100644 index 000000000..3daa99a07 Binary files /dev/null and b/tests/fixtures/image/oracle/selftest_moss_image_tiny.npz differ diff --git a/tests/fixtures/image/oracle/selftest_moss_image_tiny_weights.manifest.json b/tests/fixtures/image/oracle/selftest_moss_image_tiny_weights.manifest.json new file mode 100644 index 000000000..f98556022 --- /dev/null +++ b/tests/fixtures/image/oracle/selftest_moss_image_tiny_weights.manifest.json @@ -0,0 +1,548 @@ +{ + "schema": 1, + "name": "selftest_moss_image_tiny_weights", + "check": 3, + "family": "moss_image", + "script": "scripts/image_oracle/check3_moss_image.py", + "argv": [ + "check3_moss_image.py", + "--pipeline", + "pipeline_moss_image.py", + "--self-test" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": {}, + "inputs": { + "config": { + "dim": 64, + "ffn_dim": 48, + "num_heads": 2, + "num_layers": 2, + "text_dim": 24, + "freq_dim": 16, + "latent_channels": 8, + "rope_base": 10000.0 + }, + "case": { + "latent": [ + 3, + 4 + ], + "text": 5, + "sigma": 0.6875 + }, + "seed": 0 + }, + "results": {}, + "tolerances": {}, + "notes": [], + "created": "2026-10-18T06:29:29Z", + "npz": "selftest_moss_image_tiny_weights.npz", + "tensors": { + "final_linear_video.weight": { + "shape": [ + 8, + 64 + ], + "dtype": "float32", + "sha256": "7f0f4ab95e1e607799aa61aeac4510e73a30cf0eb62c722a499a81d998e68c53" + }, + "final_norm_video.weight": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "a99888cbee2d33093040fafb681212f24af063ba50decbc8719929a640cde2e4" + }, + "layers.0.k_norm.weight": { + "shape": [ + 32 + ], + "dtype": "float32", + "sha256": "e696e269c4bd81e50df1c75172079d3d6eaddbea759b5168f4ac27200c808004" + }, + "layers.0.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "0f55c3d8cb9e7c98b28354b5e14cd6d8eac5653fdc9206581b5f9cfc97897b80" + }, + "layers.0.norm2.weight": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "0d3dff16956d057f19c300ac90e61d22630351cbdb7123db966524650e4d9d93" + }, + "layers.0.proj.ffn_down.weight": { + "shape": [ + 64, + 48 + ], + "dtype": "float32", + "sha256": "862aa99e003ad2f04856e428084ab23fee272285e85893dc0240ae58ca6ebd62" + }, + "layers.0.proj.ffn_up_gate.weight": { + "shape": [ + 96, + 64 + ], + "dtype": "float32", + "sha256": "d74a2edacb9e3a433b81c5dcb36561c6af0866ea56680ddb52e72414a2c7e7b3" + }, + "layers.0.proj.k.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "628d18ef230d39c63b2280b4020d9054e0bbb22337c8cdaae0b6c623f3761584" + }, + "layers.0.proj.o.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "951972aa29ba01ad313887fac0f0eead607fe5355a90b9bf662cf3f9bc14673a" + }, + "layers.0.proj.q.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "dabf9bc1a95f3c5d82677f1f462d861b13387f0233215801184ea8dd853d2485" + }, + "layers.0.proj.v.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "6c93ab32cf603eeadd6daa1ef76b83655cc6cc08334b3ef4c29b7fabf1a75bb7" + }, + "layers.0.q_norm.weight": { + "shape": [ + 32 + ], + "dtype": "float32", + "sha256": "db9c20953499d798d4ea23f372d6d7cabdcfcbf835228e83d4609f75fb66285a" + }, + "layers.1.k_norm.weight": { + "shape": [ + 32 + ], + "dtype": "float32", + "sha256": "677461c337f65bf0b04c9ae94b49b67b5f9f6ad75d1df5e724cd49a24f89d912" + }, + "layers.1.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "65111d8859253797321829277b9c9b0bd18eaa342c0f47ffb2652ee67575063e" + }, + "layers.1.norm2.weight": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "b39d17d359bffe2b9a2ac16bac7d4f6784dcf6207e2eb839691ff2ed2c3f8c7c" + }, + "layers.1.proj.ffn_down.weight": { + "shape": [ + 64, + 48 + ], + "dtype": "float32", + "sha256": "657c5dd18f09867cb97b327c31a30a1d2c49e0ecac329dc4e0e5c3b77f25cb6d" + }, + "layers.1.proj.ffn_up_gate.weight": { + "shape": [ + 96, + 64 + ], + "dtype": "float32", + "sha256": "8a8cd81b4b2bfb8e7a1e1d0692279da844419c29115156373429eae105dbe1c1" + }, + "layers.1.proj.k.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "1642939e87511cc60c40f70b4f28e0d5082968602fe07c3383b3a73c8ac02a83" + }, + "layers.1.proj.o.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "1248889a53177f78a3e642f3dfa5026b84eb91f0e7d6dd03d6616fbe64b3d376" + }, + "layers.1.proj.q.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "5b05fd43168b10fa8cba2a1f03005cd67abb7a788b3d3f2beaa84ea57c41be05" + }, + "layers.1.proj.v.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "ac32f6e2aebefa4899519c40a5d06bc6a5d91247eb3fc535941e581d1f3ab861" + }, + "layers.1.q_norm.weight": { + "shape": [ + 32 + ], + "dtype": "float32", + "sha256": "119913377d9ab0125c129453ce5d0497a98a132941ce9fa1f9548cef0a3b8c38" + }, + "rope.bands": { + "shape": [ + 4 + ], + "dtype": "float32", + "sha256": "47d9d469d9542448dfc243a68cae17ace015ae732f96bd482e886c1a969cf97c" + }, + "shared_block_modulation.bias": { + "shape": [ + 384 + ], + "dtype": "float32", + "sha256": "89fa1127b9b08e143b020e5a32cca767c7bec66cb1ea839ec2a4a1418531f7bc" + }, + "shared_block_modulation.weight": { + "shape": [ + 384, + 64 + ], + "dtype": "float32", + "sha256": "23783022a044dc62a935b5b0ec096af7357f83d14d4341b9cdc8468e4a1648bb" + }, + "text_embedder.blocks.0.ffn_down.bias": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "bcdd778ab5032b9feee4e164aa172ee2515ae1808042ae07dc93d2221a10d686" + }, + "text_embedder.blocks.0.ffn_down.weight": { + "shape": [ + 64, + 48 + ], + "dtype": "float32", + "sha256": "a84bb8c1236a1cc62ff2afcbae879dfd0f82e5cdebe094ca251f37d13ae6d88a" + }, + "text_embedder.blocks.0.ffn_gate.bias": { + "shape": [ + 48 + ], + "dtype": "float32", + "sha256": "56db1e4d69f55fa292e0b9a7368dce667c73bdde4df5474eb2f591ec758fa4c2" + }, + "text_embedder.blocks.0.ffn_gate.weight": { + "shape": [ + 48, + 64 + ], + "dtype": "float32", + "sha256": "ecf8c348e5bc31771cee2aaa608f6822303bdd1db5f5ef0a644f38bedc6a132d" + }, + "text_embedder.blocks.0.ffn_up.bias": { + "shape": [ + 48 + ], + "dtype": "float32", + "sha256": "18802cb52fdbd7df157ce45faca93c082e9c3918ff6a14aade9564068af0ca20" + }, + "text_embedder.blocks.0.ffn_up.weight": { + "shape": [ + 48, + 64 + ], + "dtype": "float32", + "sha256": "f04edbca58b3e092f78133a6448c942489d2ae257348887603626994aeeaae17" + }, + "text_embedder.blocks.0.k_norm.weight": { + "shape": [ + 32 + ], + "dtype": "float32", + "sha256": "672e518c6ef599ef82ebda532407627e0bcb25ed277f2140af48f0448811ead0" + }, + "text_embedder.blocks.0.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "d41ae3581e3f64dcaaa083d4e0ba67a9fd8ca36efde1b3cfdbfdf9e4af2c5b69" + }, + "text_embedder.blocks.0.norm2.weight": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "2c3da1e35f6c4a37b8cf6849d2598c9b215a74e781cb1b0243c79234b0115945" + }, + "text_embedder.blocks.0.proj.bias": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "049bb7e0e681aa81ec430cd4e71500f74798a19fa6e57be05037b9a803020ccf" + }, + "text_embedder.blocks.0.proj.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "b346f03f6298405a227983e6e2ff6fb3a198bbfb0c5009d819c3c7d307f142ba" + }, + "text_embedder.blocks.0.q_norm.weight": { + "shape": [ + 32 + ], + "dtype": "float32", + "sha256": "110e731f01c935245a7f7d10f33461b62873327336185a182a0f4a66751362c2" + }, + "text_embedder.blocks.0.qkv.bias": { + "shape": [ + 192 + ], + "dtype": "float32", + "sha256": "db998e9acc438fbf17d4c51f4af7a17bc819cae13d8960dd6cb543d635a2c928" + }, + "text_embedder.blocks.0.qkv.weight": { + "shape": [ + 192, + 64 + ], + "dtype": "float32", + "sha256": "ee8b05d25d8e40acdc8b41d8c3590641011efcf520f84b0c65050671779338b3" + }, + "text_embedder.blocks.1.ffn_down.bias": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "dd1efaa63e48708789c4b0b358e91651accf352ec891430b3b802ac9d0de1a7b" + }, + "text_embedder.blocks.1.ffn_down.weight": { + "shape": [ + 64, + 48 + ], + "dtype": "float32", + "sha256": "08ec990296f2b0bead47f3b0b315968354b78ab81f238c163fb7863596d9b222" + }, + "text_embedder.blocks.1.ffn_gate.bias": { + "shape": [ + 48 + ], + "dtype": "float32", + "sha256": "f4bef1c3fc7005c66b98df6539c6d30b02ec1305207ddf0d2a441eecaf890823" + }, + "text_embedder.blocks.1.ffn_gate.weight": { + "shape": [ + 48, + 64 + ], + "dtype": "float32", + "sha256": "421fff86927b4656ea67e8f9d815bd2a7ba391e08b367eab98484a0b197a5a06" + }, + "text_embedder.blocks.1.ffn_up.bias": { + "shape": [ + 48 + ], + "dtype": "float32", + "sha256": "d33a19be73fe38ac857a8728133abb6fe0e01621c200e959fee9052f10b57c2b" + }, + "text_embedder.blocks.1.ffn_up.weight": { + "shape": [ + 48, + 64 + ], + "dtype": "float32", + "sha256": "1fea252b19b0d33829a199f47858ae453ce318765da908ecd00bb265d66b99bc" + }, + "text_embedder.blocks.1.k_norm.weight": { + "shape": [ + 32 + ], + "dtype": "float32", + "sha256": "3c1665ceca9e75bf691272499a51add434909ad05b26bee3d10a7f8d78afaadd" + }, + "text_embedder.blocks.1.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "4ef036673b51000a0616c113f0c871fc880b63f6e3523fc89ad3b6388631b2d3" + }, + "text_embedder.blocks.1.norm2.weight": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "bd539517831e1430b25a11aaa48b40302cb33fee4125dec327152114e0d91a73" + }, + "text_embedder.blocks.1.proj.bias": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "f7a2affc3a73f19f7dc66d546b712a6f630b04b4cb5accd2a9fa6b08fecd1e0d" + }, + "text_embedder.blocks.1.proj.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "a4b248f9c2538930d77e4d5362cb85e486c4b2b02944fe8e95200121ff4379bc" + }, + "text_embedder.blocks.1.q_norm.weight": { + "shape": [ + 32 + ], + "dtype": "float32", + "sha256": "892c80014306a1808491c7abf471d807e922542126867f4e62e9176443223667" + }, + "text_embedder.blocks.1.qkv.bias": { + "shape": [ + 192 + ], + "dtype": "float32", + "sha256": "64e383ae623c01687914ff003eb6697382a2877282ebbf750a2bb646d59462b1" + }, + "text_embedder.blocks.1.qkv.weight": { + "shape": [ + 192, + 64 + ], + "dtype": "float32", + "sha256": "d179647f34707be3dc8feb0716d90587b1c9cbed048fbb0992a69f905be7ca70" + }, + "text_embedder.proj.bias": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "e2a3ae595e066b2bb2fd2eca9ee22a5735da85bcec4f2240903da9c75fa97e2c" + }, + "text_embedder.proj.weight": { + "shape": [ + 64, + 24 + ], + "dtype": "float32", + "sha256": "c866e7fbc2c2293a712bf703be19745e95c4f318fc1e197e20ab92fe3b6c62e5" + }, + "time_embedding.0.bias": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "ff3ec5f2fe81a1d69f75b151c82d4d4ae0592e0c654179c6a50bd9fa79f329e7" + }, + "time_embedding.0.weight": { + "shape": [ + 64, + 16 + ], + "dtype": "float32", + "sha256": "7993d804edaa5bb7556c05cb57b2ac6b4c0b53ce3e9ba76c089e89c4dfca1f62" + }, + "time_embedding.2.bias": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "301e4a7d2ac1311354773b953847e3ced259444435a215a61f2b682e5d17f6ec" + }, + "time_embedding.2.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float32", + "sha256": "a6ad70d6f57ce289d5b4bc1fe3f363eee5484c6e7c0af0814a0c9167fc1a4255" + }, + "video_embedder.bias": { + "shape": [ + 64 + ], + "dtype": "float32", + "sha256": "867cf4ef934a28c8226113f4ea9e72b3442bb5db7b278c856abdf7af33e66f3d" + }, + "video_embedder.weight": { + "shape": [ + 64, + 8 + ], + "dtype": "float32", + "sha256": "1015dde7779edc916b78259ac98a5fb8cbc61788f5b18837f8565dc11aa44236" + }, + "video_head_modulation.bias": { + "shape": [ + 128 + ], + "dtype": "float32", + "sha256": "3f01c30fb6358d35233ab777516722eafba060d7fe40aa38cccd38c54688171c" + }, + "video_head_modulation.weight": { + "shape": [ + 128, + 64 + ], + "dtype": "float32", + "sha256": "b713f138df47864a3c4fd03b220c4565b01fcbd52df55313e48d4599fc6426ed" + } + } +} diff --git a/tests/fixtures/image/oracle/selftest_moss_image_tiny_weights.npz b/tests/fixtures/image/oracle/selftest_moss_image_tiny_weights.npz new file mode 100644 index 000000000..6f2b6fccd Binary files /dev/null and b/tests/fixtures/image/oracle/selftest_moss_image_tiny_weights.npz differ diff --git a/tests/fixtures/image/oracle/selftest_moss_image_vae_tiny.manifest.json b/tests/fixtures/image/oracle/selftest_moss_image_vae_tiny.manifest.json new file mode 100644 index 000000000..b72b40255 --- /dev/null +++ b/tests/fixtures/image/oracle/selftest_moss_image_vae_tiny.manifest.json @@ -0,0 +1,1451 @@ +{ + "schema": 1, + "name": "selftest_moss_image_vae_tiny", + "check": 5, + "family": "moss_image", + "script": "scripts/image_oracle/check5_flux3_klein_vae.py", + "argv": [ + "check5_flux3_klein_vae.py", + "--self-test", + "--family", + "moss_image" + ], + "oracle": { + "python": "3.12.11", + "platform": "macOS-26.5-arm64-arm-64bit", + "machine": "arm64", + "packages": { + "torch": { + "version": "2.14.1" + }, + "diffusers": { + "version": "0.41.0.dev0", + "source": "https://github.com/huggingface/diffusers/archive/e0abab83b5df05de9e7abd788643c1a7c1e42e28.tar.gz" + }, + "transformers": { + "version": "5.18.0" + }, + "tokenizers": { + "version": "0.23.2" + }, + "accelerate": { + "version": "1.15.0" + }, + "huggingface-hub": { + "version": "1.33.0" + }, + "gguf": { + "version": "0.19.0" + }, + "safetensors": { + "version": "0.8.0" + }, + "numpy": { + "version": "2.5.3" + } + }, + "hf_hub_offline": "1" + }, + "files": {}, + "inputs": { + "sizes": [ + [ + 48, + 64 + ] + ], + "seed": 5, + "image": "synthetic_image(height, width, seed) of check5_qwen_image21_vae.py, RGB, NCHW in [-1, 1]", + "batch_norm_eps": 0.0001, + "patch_size": [ + 2, + 2 + ], + "latents": "Flux2KleinPipeline._patchify_latents of the encoder mean, then (z - running_mean) / sqrt(running_var + batch_norm_eps), in float32 from each mode's mean", + "decode_input": "latents * sqrt(running_var + eps) + running_mean, then Flux2KleinPipeline._unpatchify_latents, in float32 from the fp32 latents, cast to the VAE dtype", + "decoded": "RGB, NCHW, unclamped as the VAE returns it", + "psnr": "on x / 2 + 0.5 clamped to [0, 1], peak 1", + "config": { + "block_out_channels": [ + 64, + 64, + 96 + ], + "layers_per_block": 1, + "latent_channels": 32, + "norm_num_groups": 32, + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "decoder_block_out_channels": [ + 32, + 32, + 64 + ] + }, + "weights_seed": 29, + "weights": "every float tensor by its diffusers name, float16 values the model holds in float32, under weight.", + "num_batches_tracked": 0 + }, + "results": { + "oracle": { + "device": "cpu", + "threads": 8, + "bf16_products": "bf16 operands, float32 conv and linear, one bf16 rounding" + }, + "sizes": { + "48x64": { + "decode_bf16_vs_fp32": { + "rel_l2": 0.01947363446056681, + "cosine": 0.9998104814928966, + "max_abs": 0.07079541683197021, + "psnr_db": 45.56038360250544 + }, + "encode_bf16_vs_fp32": { + "rel_l2": 0.009809563517454652, + "cosine": 0.9999519049812912, + "max_abs": 0.057882487773895264 + }, + "encode_fp32_vs_fp64": { + "rel_l2": 1.1065051661335104e-06, + "cosine": 0.9999999999993883, + "max_abs": 6.751032908480781e-06 + }, + "encode_mean_fp32_vs_fp64": { + "rel_l2": 1.4347299566782424e-06, + "cosine": 0.9999999999989712, + "max_abs": 4.632919846481176e-06 + }, + "decode_fp32_vs_fp64": { + "rel_l2": 1.8627640518728079e-06, + "cosine": 0.9999999999982652, + "max_abs": 6.114344247354353e-06 + }, + "round_trip_fp32_psnr_db": 8.16128299789587, + "latents_std": 0.8618441820144653, + "peak_rss_gib": { + "encode": 0.4, + "decode": 0.4 + }, + "seconds": { + "fp32_encode_s": 0.0, + "bf16_encode_s": 0.0, + "fp32_decode_s": 0.0, + "bf16_decode_s": 0.0, + "float64_decode_s": 0.0 + } + } + } + }, + "tolerances": { + "rule": "fp32 max abs at 5e-5 on the CPU, for the decode and the model space latents of the encode of the tiny model; bf16 decode PSNR at twice the oracle's own bf16 against fp32 RMS error, its PSNR less 20 log10(2) dB; bf16 encode relative L2, on the model space latents, at twice the oracle's own", + "fp32_max_abs": 5e-05, + "sizes": { + "48x64": { + "decode_fp32_max_abs": 5e-05, + "decode_bf16_psnr_db": 39.539783689225814, + "encode_fp32_max_abs": 5e-05, + "encode_bf16_rel_l2": 0.019619127034909305 + } + } + }, + "notes": [ + "The model is the official class at a tiny config with seeded random weights and statistics, so the numbers say nothing about image quality.", + "bn.num_batches_tracked is an int64 scalar the fixture does not hold; its value is inputs.num_batches_tracked." + ], + "created": "2026-10-18T06:52:10Z", + "npz": "selftest_moss_image_vae_tiny.npz", + "tensors": { + "bn_running_mean": { + "shape": [ + 128 + ], + "dtype": "float32", + "sha256": "f58e4b1de8eeed57cf9d601df0dd92cb513ec8bbb7cd615b95ee14d97fa5968a" + }, + "bn_running_var": { + "shape": [ + 128 + ], + "dtype": "float32", + "sha256": "7dd185cb7bdeffe4fcbd4270ff4d06ac267416c8deda633d441465fecbfc53b4" + }, + "decoded_48x64": { + "shape": [ + 1, + 3, + 48, + 64 + ], + "dtype": "float32", + "sha256": "27e76ed533b139380b7728bfd54de29e3a6b7b969504913e06d413a6ec786ff9" + }, + "decoded_bf16_48x64": { + "shape": [ + 1, + 3, + 48, + 64 + ], + "dtype": "float32", + "sha256": "23be0ad5941022848d5e6b99add6e55e734290b932b375b48810eca54c4ec431" + }, + "encode_mean_48x64": { + "shape": [ + 1, + 32, + 12, + 16 + ], + "dtype": "float32", + "sha256": "d056033301fc7f606ff894dc8c4da2b0b7417da19055177731654008399fc8c6" + }, + "encode_mean_bf16_48x64": { + "shape": [ + 1, + 32, + 12, + 16 + ], + "dtype": "float32", + "sha256": "0b9b3909cb3dc96a020f16a835de4baba5a6fb6969a0b772baa45d2753865856" + }, + "image_48x64": { + "shape": [ + 1, + 3, + 48, + 64 + ], + "dtype": "float32", + "sha256": "203924fe2b36a598eb1ea7d9f7cc5e396e02a9dfa119df8cdf152e13af34ddaf" + }, + "latents_48x64": { + "shape": [ + 1, + 128, + 6, + 8 + ], + "dtype": "float32", + "sha256": "e0bd64859cda376d56544d8494d428a7205e5821145e0ada2d0abaa94528f8c8" + }, + "latents_bf16_48x64": { + "shape": [ + 1, + 128, + 6, + 8 + ], + "dtype": "float32", + "sha256": "cc0b5cd6fe80a9dde9784e63a12eb4c0147afe039d8fa3fa790d11ab98b3e3f3" + }, + "weight.bn.running_mean": { + "shape": [ + 128 + ], + "dtype": "float16", + "sha256": "24394c92b1a4149339460fda5b7d3248f2b35263b287e5002bae0a8b8a7bada5" + }, + "weight.bn.running_var": { + "shape": [ + 128 + ], + "dtype": "float16", + "sha256": "0bd589ac45faeecb029bffef51ddd280b1548c0d8803cdc6ee3791028996feb0" + }, + "weight.decoder.conv_in.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "5226e31b8aa82f24aa7f0b6c56a8af690a84312fcf91a56258ef9be3f17c2307" + }, + "weight.decoder.conv_in.weight": { + "shape": [ + 64, + 32, + 3, + 3 + ], + "dtype": "float16", + "sha256": "b7f676fc0cda3706b91b4d3f63cf1c764a055afb9fcd349d805be23917bf2be7" + }, + "weight.decoder.conv_norm_out.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "d255771a90e1a7b688ae68c42bd56758ec4171f73e84a8863b65f67c152c6e4c" + }, + "weight.decoder.conv_norm_out.weight": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "a3394a45dc83eb4c09183a4ca2a8a59c6f70dfc00508748be82e5f3b0c5dd8ca" + }, + "weight.decoder.conv_out.bias": { + "shape": [ + 3 + ], + "dtype": "float16", + "sha256": "60c5130ee11b3d1cf70af0c5b5ef2d7541d5b0b73deadb02cbfc4dd37e7579ed" + }, + "weight.decoder.conv_out.weight": { + "shape": [ + 3, + 32, + 3, + 3 + ], + "dtype": "float16", + "sha256": "706f17506819712fbcd1b5f83bd3901c6f15aa7b74d1510530e11161d1239632" + }, + "weight.decoder.mid_block.attentions.0.group_norm.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "f3aff01fbe6bb35edad247f69f5caa8dd65757dcb8a9a55f591ae9cdfee326b1" + }, + "weight.decoder.mid_block.attentions.0.group_norm.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "b09d9ad66315727617e7b4b656173e9f2af93aea6ab488b2d2e3e80b1bf1b493" + }, + "weight.decoder.mid_block.attentions.0.to_k.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "269538923575c20e7db9d756d5dcbfc902d96b574ca2fb93b57da2d5b98e6dad" + }, + "weight.decoder.mid_block.attentions.0.to_k.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float16", + "sha256": "ba2066815732d81a0dc981e0eac079e77668a8364bcd7c8031ef395d3c879e0e" + }, + "weight.decoder.mid_block.attentions.0.to_out.0.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "cdf50a57236163a5af0b75ba1c761efb1d0e0631165b1bfd1d3b14ef1bd78084" + }, + "weight.decoder.mid_block.attentions.0.to_out.0.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float16", + "sha256": "9a8870b06cacfeaa8327cc60d68472a4e59f45c5887f6add25fc62bd5aea872b" + }, + "weight.decoder.mid_block.attentions.0.to_q.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "0b25674f9a70457842f499c15ae0ebdb503cc56caf398b3d138332fcab6e70cc" + }, + "weight.decoder.mid_block.attentions.0.to_q.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float16", + "sha256": "28b1657807d3cf030a0aec8cd808ce2a05bb4454e319e42f34b720b0e91456eb" + }, + "weight.decoder.mid_block.attentions.0.to_v.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "dbc223c747165da194dae3f3962a4b0c5ea5bd91736a8346e587442990c2141e" + }, + "weight.decoder.mid_block.attentions.0.to_v.weight": { + "shape": [ + 64, + 64 + ], + "dtype": "float16", + "sha256": "86a363c33828050f775df98cd66530d86e4bcdfbe90c90356789aa9499b0f071" + }, + "weight.decoder.mid_block.resnets.0.conv1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "b8f94aa6e3eac15d8b4a1ed4c1f22273747fec6cb40d746cc9ed151c337a4a59" + }, + "weight.decoder.mid_block.resnets.0.conv1.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "e4b3912b81409b4546b405c1d7c8ab3e38c618e9e9c5cb04f350fe9d0b95161a" + }, + "weight.decoder.mid_block.resnets.0.conv2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "eb3e89827e04a7bf584676e092dc129b01a4a5ea8d5d022026fcf7909fdfb2b1" + }, + "weight.decoder.mid_block.resnets.0.conv2.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "1756d6394346010b620034e501d7957628e73272e18a095f091f411be45ed835" + }, + "weight.decoder.mid_block.resnets.0.norm1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "20a0e751e6bbc7299fbd69a3babc59f0bc625adc0b9015a9bae4ed991c6bb9b5" + }, + "weight.decoder.mid_block.resnets.0.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "24a8b02d680ee98371af58511fab58c06e8af4cee469c8f0c0bdb28904338b43" + }, + "weight.decoder.mid_block.resnets.0.norm2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "9bbb00d05afa1465b3d6d32a2fe6fc6c8b8bddbc61a4b8bf336d17e8bc0b06b0" + }, + "weight.decoder.mid_block.resnets.0.norm2.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "17a1d5ff3b69a0cb262d8a164b8b34e0decabd50cad33fd5132d34c778bdd94a" + }, + "weight.decoder.mid_block.resnets.1.conv1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "6aaefdc6bbe9b9b7b1891032a86547f9493b50ef91d9f3a707ba5a1566c338e4" + }, + "weight.decoder.mid_block.resnets.1.conv1.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "70495e4045547789b7bff78039c91a8fa6cc688134c77cdb843e3b643df718ae" + }, + "weight.decoder.mid_block.resnets.1.conv2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "6310378a60ab6817afed2b42cd1084445eb6966187395ff87aa287e054820010" + }, + "weight.decoder.mid_block.resnets.1.conv2.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "7944efb5c67828ff8f7bb4512e91a910a5bc1343c28a8511baba078711971a69" + }, + "weight.decoder.mid_block.resnets.1.norm1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "5ade0854ce503e73316a85fb5b83d17c0a3db2441cae80215337c01a8fd3c8a0" + }, + "weight.decoder.mid_block.resnets.1.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "6274f78bda465c16b50f13dc750c784c358e3f6a41ee11c8b206c44a1d727c44" + }, + "weight.decoder.mid_block.resnets.1.norm2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "b4351388ea2cc86fa82c838e6c161d1fca592ce8da5b053bf79ef01154247d47" + }, + "weight.decoder.mid_block.resnets.1.norm2.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "d03accb609eca0031f16bebbff2519a2af36ae5befda4ef9dd67202d261933cd" + }, + "weight.decoder.up_blocks.0.resnets.0.conv1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "3d2f340cea7362aef47234698aaa4b4992b391e0aa1f6eddb0b9120d2affd81a" + }, + "weight.decoder.up_blocks.0.resnets.0.conv1.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "ef2f699346b4bcffedb65a6edb63195845c619bfbb841a4b17b0a0fcce7e396e" + }, + "weight.decoder.up_blocks.0.resnets.0.conv2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "9f1af729c47582935089ba788a15301a6b13bcf6e200ab48e18cebaeb1acda77" + }, + "weight.decoder.up_blocks.0.resnets.0.conv2.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "aac630f553271eff6bea2db4de500da6feb9ebb96382c9b503cd010c01eccf32" + }, + "weight.decoder.up_blocks.0.resnets.0.norm1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "190b52e538173ec0518f8cd6d9975325e29fff0984a0b990fc5584d2c206191e" + }, + "weight.decoder.up_blocks.0.resnets.0.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "c2e4cfa3f8062f8ae3ebff3ce1f343ecba0a6552dd0bce59dd8df752f4e22f52" + }, + "weight.decoder.up_blocks.0.resnets.0.norm2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "8c987f11fff06e20d823912752ef5da821290f139c736b6edb8cba71af6b87b3" + }, + "weight.decoder.up_blocks.0.resnets.0.norm2.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "e5d439320860e9f7a69f858fab18c715fe9b83925c5c9a46d04ad610bef8ba44" + }, + "weight.decoder.up_blocks.0.resnets.1.conv1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "e2ec92aae0f3fb20f2d01d01bee351911ff4dff162d62ea15d116fdd28c4763c" + }, + "weight.decoder.up_blocks.0.resnets.1.conv1.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "49ebc6a22d86a162f29b9b4771681e60f902d6c3266a8630e913960021578300" + }, + "weight.decoder.up_blocks.0.resnets.1.conv2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "8ce285a18f9dcf6cb53cde774d064dba93e78aa63b819ba15d0c4c76fa63e3db" + }, + "weight.decoder.up_blocks.0.resnets.1.conv2.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "fa0973407d5499cafd064869f0c199bbe838581b6ad8ebac5f079471fc833b21" + }, + "weight.decoder.up_blocks.0.resnets.1.norm1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "da2f2d6b51eb95e38a01502425eda7f5c5c7a40f8c081d5075634edeedb4431e" + }, + "weight.decoder.up_blocks.0.resnets.1.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "c8ddcea15f8f732943fccf259e8205ad5320297be219a31da8a7c2b7e57094bf" + }, + "weight.decoder.up_blocks.0.resnets.1.norm2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "cde7de3d4e2ea7d7fbdced949eb30b39c5eca49f7176afd2e3cffb46e69281cb" + }, + "weight.decoder.up_blocks.0.resnets.1.norm2.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "709f6599e3f69bc84870668ef499f1a311f086d3dcad8ccbbf98b3ef4392ebf7" + }, + "weight.decoder.up_blocks.0.upsamplers.0.conv.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "42fbed58401908644abac0c14a9c23ebb45450641417be96576bbacdc24a7a02" + }, + "weight.decoder.up_blocks.0.upsamplers.0.conv.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "f3412804e81a811701cd884d8c9665f0c0e719788dba7bb9820bc8c3aaee5e9c" + }, + "weight.decoder.up_blocks.1.resnets.0.conv1.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "6a0a6b3f609d4eab12815bf3a86a5a02d667532e2726e616f60b509bf971169c" + }, + "weight.decoder.up_blocks.1.resnets.0.conv1.weight": { + "shape": [ + 32, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "9adc00e708433c72b518115b4de703771a54593073c49d919f980139e4b6d7a6" + }, + "weight.decoder.up_blocks.1.resnets.0.conv2.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "5d899f88791b764831c9bb4bcaefd9110bf5d8b2490c93c0a412d84b6bf3571b" + }, + "weight.decoder.up_blocks.1.resnets.0.conv2.weight": { + "shape": [ + 32, + 32, + 3, + 3 + ], + "dtype": "float16", + "sha256": "fbf6c4d0fba73d1df8b5371c79c31417a57e95d9cc2bbbe5d60f56f36eb6206b" + }, + "weight.decoder.up_blocks.1.resnets.0.conv_shortcut.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "db041850fca14ade3a36926b4d5b748062e9bbc6506200cdcd97c1ef8c8d7ca1" + }, + "weight.decoder.up_blocks.1.resnets.0.conv_shortcut.weight": { + "shape": [ + 32, + 64, + 1, + 1 + ], + "dtype": "float16", + "sha256": "45dab7f958f4b719018fb9f5a9fa7e81a2c4643c9703f83720738529bb7b1311" + }, + "weight.decoder.up_blocks.1.resnets.0.norm1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "6677db810717ca375e2d0159647f56e6082bd024695dbbe8727f1b17939606f2" + }, + "weight.decoder.up_blocks.1.resnets.0.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "5c180d79bb4d3c9f88382a48426c93a6142532f955619b5ef9a9ddb5f5631dac" + }, + "weight.decoder.up_blocks.1.resnets.0.norm2.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "103992a775042e889b1d88b7787dc27e102b7b1eef81fe20363887cca283bc83" + }, + "weight.decoder.up_blocks.1.resnets.0.norm2.weight": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "4e01da7c8e8a74d8644765f888a518ab76ce6c79e9b98750d83c72b5fd5f31c5" + }, + "weight.decoder.up_blocks.1.resnets.1.conv1.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "553773f701eedeb67970022ab666f58d8b27a6d4106fbc03b305052b93198d7e" + }, + "weight.decoder.up_blocks.1.resnets.1.conv1.weight": { + "shape": [ + 32, + 32, + 3, + 3 + ], + "dtype": "float16", + "sha256": "68d768d66eb4fa0551c233618b041ee141901418bb2ed9df72b32105006bb2e4" + }, + "weight.decoder.up_blocks.1.resnets.1.conv2.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "8b99d4ae287ce38092b78d4aa745590495df249dcf05aa4c6a9026e983dc4cee" + }, + "weight.decoder.up_blocks.1.resnets.1.conv2.weight": { + "shape": [ + 32, + 32, + 3, + 3 + ], + "dtype": "float16", + "sha256": "2ba13048edfce00f1130dcdd0da59656ba1c3a088cb1be7d0a7ee00642e66f89" + }, + "weight.decoder.up_blocks.1.resnets.1.norm1.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "0bf667a286390c02baafa6e75e9107161f5a4ded68f39116d021e07a1f9daf37" + }, + "weight.decoder.up_blocks.1.resnets.1.norm1.weight": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "b2e56927bc7b871fc81bea39cdc68858caf4d13fbffda52d93f7868f2cfe888e" + }, + "weight.decoder.up_blocks.1.resnets.1.norm2.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "95ea9836cf19c293b469865d5a7987f42b58b591472f8f798436051638c37a2f" + }, + "weight.decoder.up_blocks.1.resnets.1.norm2.weight": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "8d73e62315ff80c971bccd2a34cf025418592c202138678baac7aa172bde9ecc" + }, + "weight.decoder.up_blocks.1.upsamplers.0.conv.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "6e3e19476dd124598276e16fe249fa50706cf45bcab12579085bf7f0bfca4d43" + }, + "weight.decoder.up_blocks.1.upsamplers.0.conv.weight": { + "shape": [ + 32, + 32, + 3, + 3 + ], + "dtype": "float16", + "sha256": "0b077a08b7270c9f0099daa59ab14f52dfd1bb958be108f22d8c7f92e2116586" + }, + "weight.decoder.up_blocks.2.resnets.0.conv1.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "5d147dc303b51d56984ab1d025a491c7a4725fadfaf45e8284af6fcf8121680f" + }, + "weight.decoder.up_blocks.2.resnets.0.conv1.weight": { + "shape": [ + 32, + 32, + 3, + 3 + ], + "dtype": "float16", + "sha256": "986d9c3debdfb9c416ebbbb7f5b573c7acde3d1b1f3112deb318f0d728aa1a0a" + }, + "weight.decoder.up_blocks.2.resnets.0.conv2.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "4f9b9501d4c547747606e12a69136bc0deabbc6d6f4c67a43cad616968aa3774" + }, + "weight.decoder.up_blocks.2.resnets.0.conv2.weight": { + "shape": [ + 32, + 32, + 3, + 3 + ], + "dtype": "float16", + "sha256": "75e58ba4b2e1b52154a99ecbae1892135c55c8cccb037610f74835ccf05ca581" + }, + "weight.decoder.up_blocks.2.resnets.0.norm1.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "e0b6549c0dd40a68c16e0b3910b935abcf1251e46b079371fcbb222668d5a692" + }, + "weight.decoder.up_blocks.2.resnets.0.norm1.weight": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "71e9dc6d03b44aad9892234e392e94152e79c0d380055bc7d5a54f1668c77e86" + }, + "weight.decoder.up_blocks.2.resnets.0.norm2.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "f7912e913485ca06fcde3d8daa259b649a5b16800c532a0c86505612b065e0c2" + }, + "weight.decoder.up_blocks.2.resnets.0.norm2.weight": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "e0bca95b0262ee3a39d705d723ec5284044b892822137e2896e4f253f8a2b07d" + }, + "weight.decoder.up_blocks.2.resnets.1.conv1.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "eda14c409ca19f3c2da3e1370c79e4bdb5b85aff10e914bd2f7ae26f5d5d1dd5" + }, + "weight.decoder.up_blocks.2.resnets.1.conv1.weight": { + "shape": [ + 32, + 32, + 3, + 3 + ], + "dtype": "float16", + "sha256": "594ac1ece1a90bf22a9b07afc03a96363115f439fb81fb32e2cafc02c589f4e4" + }, + "weight.decoder.up_blocks.2.resnets.1.conv2.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "8d1302b2f68c5c545c5cca6cd2aadd1c7a959aeb2210194fbb44dd12293f0895" + }, + "weight.decoder.up_blocks.2.resnets.1.conv2.weight": { + "shape": [ + 32, + 32, + 3, + 3 + ], + "dtype": "float16", + "sha256": "39d894f98c8f6d19951ab720ac08f5890b37d0c2549f343d6cb6b5245bf81ffd" + }, + "weight.decoder.up_blocks.2.resnets.1.norm1.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "ca706c5dec2d0c5e4507fdd750796c694f5d8c25ecc7a8dd0df5ce2872abdeaf" + }, + "weight.decoder.up_blocks.2.resnets.1.norm1.weight": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "6873e6efaecb622725a4ad7c068de02747b907fe7332facaed53398467b16f09" + }, + "weight.decoder.up_blocks.2.resnets.1.norm2.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "9647f8b7d945df7c3d4938efdee5572e30bfc226c3d789f6c159162af7a1aafa" + }, + "weight.decoder.up_blocks.2.resnets.1.norm2.weight": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "c36788eb61059459e5886cd05b299ca854c3560d290eef9e1bac7f2f54522f37" + }, + "weight.encoder.conv_in.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "080723073d45a4e93b94788fce66eb35b192380c17d38a5a3d87ea71ae0030d3" + }, + "weight.encoder.conv_in.weight": { + "shape": [ + 64, + 3, + 3, + 3 + ], + "dtype": "float16", + "sha256": "8ecf87dc7eaa3e1f21e14d19ca3fa749dc70b0091ab06f9deed1be04c6548897" + }, + "weight.encoder.conv_norm_out.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "8455101841f142588e4b1b096e546b4d081f9445d29187d8945b49ed4b4e7632" + }, + "weight.encoder.conv_norm_out.weight": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "d217c50e23cb78449c304a5472e7231af70909f3a0746894841e79a4267eef97" + }, + "weight.encoder.conv_out.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "7b7bf0865726da7540fb16f861ce34008c90804868bcfadb00265342df8b5bbb" + }, + "weight.encoder.conv_out.weight": { + "shape": [ + 64, + 96, + 3, + 3 + ], + "dtype": "float16", + "sha256": "ca0e14c3c8c6c6571476380aa6f22572184f3cf2976cbf6c93566bf4742d33e8" + }, + "weight.encoder.down_blocks.0.downsamplers.0.conv.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "8405977d5c12feb99bc3ffd49de618d3ab6c8b382fa5cbf6e5779be1b8feb7d2" + }, + "weight.encoder.down_blocks.0.downsamplers.0.conv.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "4d7c64468ae55fe76e513e1f2200a9e2e37dcb527a2adb2854c7bb448b825e36" + }, + "weight.encoder.down_blocks.0.resnets.0.conv1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "4f849e2e17cb50e0fd208a9a09175e25e8c3aebca0a9ac1a7551662571147e6a" + }, + "weight.encoder.down_blocks.0.resnets.0.conv1.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "2369c6988e68f277942159a0df6fdd2d7c7508363d933ada00296ce284a77c75" + }, + "weight.encoder.down_blocks.0.resnets.0.conv2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "4c5b62403a01ea9d6f3819776dd9d8e5204d3df96dd45412a6fa6a38af127dc5" + }, + "weight.encoder.down_blocks.0.resnets.0.conv2.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "557b2669346ee0d3893bd61e98dd8b3cc23da775fc6c1eecb4add1ee118d89f3" + }, + "weight.encoder.down_blocks.0.resnets.0.norm1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "1ba0618152fee497ca7ef4f295ae4767e424fa7052dd27a0f5f5bc764e5d90ec" + }, + "weight.encoder.down_blocks.0.resnets.0.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "1bba557b5c85a3776b31031d1886d75e6332254fbe2ea0374ceb19d3834f5bc3" + }, + "weight.encoder.down_blocks.0.resnets.0.norm2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "e90a64fb221bcf1cd98f526ce6594a1388481277065d80d39b2a6fb07bad0b52" + }, + "weight.encoder.down_blocks.0.resnets.0.norm2.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "85ab6adfc3c8bfd4a41e7778af61f4ba49035cfd9a6989cfdf8778dbefee5ff4" + }, + "weight.encoder.down_blocks.1.downsamplers.0.conv.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "36723a72341352d70d5122d0fd38d609d22c627b92c85ea488b9f10df94e453e" + }, + "weight.encoder.down_blocks.1.downsamplers.0.conv.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "6ebef730400a405a1ad05b4540a5a5088bc4acd24f03744d295271f869ee492c" + }, + "weight.encoder.down_blocks.1.resnets.0.conv1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "fa1648f42abb49c570ba01c11af3276c32c38a4eb55092fdfafd920f810aabf5" + }, + "weight.encoder.down_blocks.1.resnets.0.conv1.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "02921a311ea603243ba4ab3ac74e1c8f27f78dc9dff0cb502c44b8ba48a20708" + }, + "weight.encoder.down_blocks.1.resnets.0.conv2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "3380ca3d2a1d002241dc9aecbd04b18eb931b9ffa925d633b50b60110862ea35" + }, + "weight.encoder.down_blocks.1.resnets.0.conv2.weight": { + "shape": [ + 64, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "0f30a7b1bc06ab92d861a65e9c5c34b999d7cee752f0e4e6fe1f8fb1b1e73d06" + }, + "weight.encoder.down_blocks.1.resnets.0.norm1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "d275cf6667180a462096a6f209cc58c453ce495959ada058b450fa455e126a82" + }, + "weight.encoder.down_blocks.1.resnets.0.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "9ed96dcf9a2aa1a9425eb889ca20e4be88c5cb461054cc6b4bd41d181c98712e" + }, + "weight.encoder.down_blocks.1.resnets.0.norm2.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "81b158e8df409ab27948895abf30d624a67aa33b9d835b42947e7dd3babf247b" + }, + "weight.encoder.down_blocks.1.resnets.0.norm2.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "a4676833ab8fb89cd3e9a2fbf1cc32096a94fce6c067e5f2222f9b230813b100" + }, + "weight.encoder.down_blocks.2.resnets.0.conv1.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "6ae7ff486f4236b55756929755b57902904a48f4b712ac0325079bb6bb22d3ca" + }, + "weight.encoder.down_blocks.2.resnets.0.conv1.weight": { + "shape": [ + 96, + 64, + 3, + 3 + ], + "dtype": "float16", + "sha256": "47afb6c3ccc617c3c007d9dff9860f85664602e05ab5c0b69dcddb1b166349d4" + }, + "weight.encoder.down_blocks.2.resnets.0.conv2.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "35b3bd75469af2b11d61c16a314a72867e3f22a39548434a8a27da5a3f359daf" + }, + "weight.encoder.down_blocks.2.resnets.0.conv2.weight": { + "shape": [ + 96, + 96, + 3, + 3 + ], + "dtype": "float16", + "sha256": "1d07002461da962b3b564a98fe2358391f6ac8daa38c3ad4c63906a11847ba56" + }, + "weight.encoder.down_blocks.2.resnets.0.conv_shortcut.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "c013bc68bf8f790fa02d6e3e4a0dfbe81aa7dc3ccf1dc7c42cde4242d2ad2e63" + }, + "weight.encoder.down_blocks.2.resnets.0.conv_shortcut.weight": { + "shape": [ + 96, + 64, + 1, + 1 + ], + "dtype": "float16", + "sha256": "66cc2d12c2ac8d2155e746a7fa6ef3ec62a4a3c8f8449b88268e79fa49bb31a9" + }, + "weight.encoder.down_blocks.2.resnets.0.norm1.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "c013237c24303f0eef9a694e4b64f1d404094290188b8cabf0a87c8e41dc5219" + }, + "weight.encoder.down_blocks.2.resnets.0.norm1.weight": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "3f73fb9bb55f69f85ebd3e8829986b975e71c11920f473f4fba60e5d512f9f71" + }, + "weight.encoder.down_blocks.2.resnets.0.norm2.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "4ca698233f02fa565946589ed17ff027e54dcca2dc6561a69c184aeb157d86ba" + }, + "weight.encoder.down_blocks.2.resnets.0.norm2.weight": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "3c1baac396a2d8fb6572a0aabf0e4d6fe5910843a7737188b05d65d71d90e613" + }, + "weight.encoder.mid_block.attentions.0.group_norm.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "f0d6254dd4577921ecd38728edc093e57a7b6869e802be166e3b985cd6013cf0" + }, + "weight.encoder.mid_block.attentions.0.group_norm.weight": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "13d3dec4019c69a43d4e9e6c481dd463e922f70550d8afbd15c4ca475d2c9979" + }, + "weight.encoder.mid_block.attentions.0.to_k.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "397a620da5c9a0b7d43bd054f35f107384f646378b01d18fa6d8188cd6352d5d" + }, + "weight.encoder.mid_block.attentions.0.to_k.weight": { + "shape": [ + 96, + 96 + ], + "dtype": "float16", + "sha256": "c470e6a68e6da79f20e2a099c2f957c0dbb876843daebe28adb36a344d51f5ed" + }, + "weight.encoder.mid_block.attentions.0.to_out.0.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "947a66f05adc6f56c39dcedb9e9d2f4f1a334c9c139f3313d8bed34a1cb25939" + }, + "weight.encoder.mid_block.attentions.0.to_out.0.weight": { + "shape": [ + 96, + 96 + ], + "dtype": "float16", + "sha256": "ead076be670bf9526199abce1c69f98a6497aaee14826ddfbed4d9b81b3459cb" + }, + "weight.encoder.mid_block.attentions.0.to_q.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "6be4ab0e3ddd1914656a1cde202da056a86fd1e88f900632401dc2e8074f2b95" + }, + "weight.encoder.mid_block.attentions.0.to_q.weight": { + "shape": [ + 96, + 96 + ], + "dtype": "float16", + "sha256": "4ff17f98c5dbd37fcfa6a6335abaea916921faacba9757b965c85ca43698b3ee" + }, + "weight.encoder.mid_block.attentions.0.to_v.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "9dcdaa3b85adaf2f105c6feb0e3830a8de3c18eb980c7cb93e7c4844ef09bbfa" + }, + "weight.encoder.mid_block.attentions.0.to_v.weight": { + "shape": [ + 96, + 96 + ], + "dtype": "float16", + "sha256": "ee975f0729cb7e92d8e5ba641f7a5566f442da9c24144b266a43021a742c6422" + }, + "weight.encoder.mid_block.resnets.0.conv1.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "467e8718acc4b5e9cd853f77ec374b42a22c4d204b22ddf0253249693132a9c5" + }, + "weight.encoder.mid_block.resnets.0.conv1.weight": { + "shape": [ + 96, + 96, + 3, + 3 + ], + "dtype": "float16", + "sha256": "d32bc61f461694d300f0833631ce69533daaecb3fc7f6aa013fcf6466649da1b" + }, + "weight.encoder.mid_block.resnets.0.conv2.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "7f635b9878bba1dd112c60e6693cbf4565f81779ec8147953b3b8bcdffa886a6" + }, + "weight.encoder.mid_block.resnets.0.conv2.weight": { + "shape": [ + 96, + 96, + 3, + 3 + ], + "dtype": "float16", + "sha256": "ee862501a8f857b5becc8a894453ff32706ae0e72a249bcb053d1039c31ba4d9" + }, + "weight.encoder.mid_block.resnets.0.norm1.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "3982df40722fc76c4f35751a5835dd9b42f102580b2b8a9d30bdaa6b2746f09d" + }, + "weight.encoder.mid_block.resnets.0.norm1.weight": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "1ac3bdfaadec4b6bb516729ade96929418473de1ec963899e0ce283467616e02" + }, + "weight.encoder.mid_block.resnets.0.norm2.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "b3a90b7ada25b872d8286e130194bcf7bf1cf77ae1e8696359c87fd3a78101e5" + }, + "weight.encoder.mid_block.resnets.0.norm2.weight": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "ed1c0f013f4db40154822511f47d2175881edf31832e881f3b1e9f890ce50780" + }, + "weight.encoder.mid_block.resnets.1.conv1.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "3b36ee6519b86a8fee931b57dfba5f443d0a2ea0d92ca826850e735e9b92903e" + }, + "weight.encoder.mid_block.resnets.1.conv1.weight": { + "shape": [ + 96, + 96, + 3, + 3 + ], + "dtype": "float16", + "sha256": "05b10b6fc744e61e75d8bbcad086da3d151e499d0a36c07a27473de5d6a07b31" + }, + "weight.encoder.mid_block.resnets.1.conv2.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "1fb4a2d8f87704d6dbc0835f41fc6df283eb7c2f424a7ba1b30fdfd9b4d63827" + }, + "weight.encoder.mid_block.resnets.1.conv2.weight": { + "shape": [ + 96, + 96, + 3, + 3 + ], + "dtype": "float16", + "sha256": "6774bd546bc5d8545ee6b5f6926d03620ee749b62a999a11459bef25a50a6c80" + }, + "weight.encoder.mid_block.resnets.1.norm1.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "78fff97524a4182a2cc6e638d7ec4b1c9e58120959339542a5ac6fca3a4ed6cd" + }, + "weight.encoder.mid_block.resnets.1.norm1.weight": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "ca5a85fa1fb97601e70418c181dda92d714bfb0926d6e6a1efa742775d795480" + }, + "weight.encoder.mid_block.resnets.1.norm2.bias": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "f722150b43e5c49c8e1b58533b48b69db6e0ba26d3a0baa4e4825fc1b2711fb8" + }, + "weight.encoder.mid_block.resnets.1.norm2.weight": { + "shape": [ + 96 + ], + "dtype": "float16", + "sha256": "61818da55026e7151a84a32f8926c7e295a9f3c16108efa66f463b170f0c6932" + }, + "weight.post_quant_conv.bias": { + "shape": [ + 32 + ], + "dtype": "float16", + "sha256": "9b0ac5bff140e62006c2f537d38d8b2d6266c60a866760072dd181f2c6270b51" + }, + "weight.post_quant_conv.weight": { + "shape": [ + 32, + 32, + 1, + 1 + ], + "dtype": "float16", + "sha256": "2c1a2766c46eec189041f521b390274ac4178ddabd865449fa69e29cee6a773f" + }, + "weight.quant_conv.bias": { + "shape": [ + 64 + ], + "dtype": "float16", + "sha256": "58a781fc179f6823e2ffdf6eca0797da7efec62964f2f11078a2592d9c100de9" + }, + "weight.quant_conv.weight": { + "shape": [ + 64, + 64, + 1, + 1 + ], + "dtype": "float16", + "sha256": "7d6b7062b5abd7a93d950025ceac45f6600ccc11eb6fe2871d42a759b8046152" + } + } +} diff --git a/tests/fixtures/image/oracle/selftest_moss_image_vae_tiny.npz b/tests/fixtures/image/oracle/selftest_moss_image_vae_tiny.npz new file mode 100644 index 000000000..4dd4b503d Binary files /dev/null and b/tests/fixtures/image/oracle/selftest_moss_image_vae_tiny.npz differ diff --git a/tests/image/image_fixtures.py b/tests/image/image_fixtures.py index 4aa3aaae9..f685bc85e 100644 --- a/tests/image/image_fixtures.py +++ b/tests/image/image_fixtures.py @@ -176,6 +176,34 @@ def tiny_z_model(seed: int = 0): return m, {k: np.array(v, np.float64) for k, v in params} +TINY_MOSS_CONFIG = {"dim": 64, "ffn_dim": 48, "freq_dim": 16, + "latent_channels": 8, "num_heads": 2, "num_layers": 3, + "rope_base": 10000.0, "text_dim": 16, + "num_text_blocks": 2} + + +def tiny_moss_model(seed: int = 0): + """A tiny MOSS-Image DiT of three blocks with random fp32 weights and + fp32 activations.""" + import mlx.core as mx + from mlx.utils import tree_flatten + + from gmlx.image.families.moss_image.model import MossImageDiT + + m = MossImageDiT(TINY_MOSS_CONFIG) + rng = np.random.default_rng(seed) + params = [] + for k, v in tree_flatten(m.parameters()): + a = rng.standard_normal(v.shape) * 0.3 + if k.endswith(("norm1.weight", "norm2.weight")): + a += 1.0 + params.append((k, mx.array(a.astype(np.float32)))) + m.load_weights(params) + m.act_dtype = mx.float32 + m.eval() + return m + + # A small Z-Image config that both kernel routes of gmlx.image.ops take: # head dim 64 and a 256-wide residual, so a 64x64 latent gives the main # layers 1024 image rows. diff --git a/tests/image/moss_gates.py b/tests/image/moss_gates.py new file mode 100644 index 000000000..655cecba4 --- /dev/null +++ b/tests/image/moss_gates.py @@ -0,0 +1,610 @@ +"""The rows of a MOSS-Image step against a check 3 fixture of +scripts/image_oracle/check3_moss_image.py: free running from the +fixture's inputs, and teacher-forced for each text block, each block and +the head from the fixture's float32 input to it.""" +from __future__ import annotations + +import math +from contextlib import contextmanager + +import mlx.core as mx +import numpy as np + +import block_gates as bg +import qwen_gates as qg + +from gmlx.image.families.moss_image import model as mm + +TIME_ROWS = ("time_embedding", "shared_block_modulation", + "video_head_modulation") +HEAD = "final_linear_video" + + +def tiny_model(weights: dict, config: dict, act=mx.float32): + """The DiT of float32 weights by name, in torch order.""" + m = mm.MossImageDiT(config) + m.load_weights([(k, mx.array(np.asarray(v, np.float32))) + for k, v in weights.items()]) + m.act_dtype = act + m.eval() + return m + + +def _np(a) -> np.ndarray: + return np.array(a.astype(mx.float32)) + + +def rel(got, ref) -> float: + g = np.asarray(got, np.float64) + r = np.asarray(ref, np.float64) + return float(np.linalg.norm(g - r) / np.linalg.norm(r)) + + +def _inputs(arrays: dict): + lat = mx.array(arrays["latents"]) + feats = mx.array(arrays["features"])[None] + sigma = np.reshape(np.float32(arrays["sigma"]), (1,)) + return lat, feats, sigma + + +def free_rows(model, arrays: dict, mod=None) -> dict: + """Every row of one step from the fixture's inputs, and the + velocity.""" + lat, feats, sigma = _inputs(arrays) + taps: dict = {} + out = dict(zip(TIME_ROWS, model.time_vectors(sigma, 1))) + vel = model(lat, sigma, feats, taps, mod=mod) + out.update(taps) + got = {k: _np(v) for k, v in out.items()} + got["velocity"] = _np(vel) + return got + + +def fixture_modulation(model, arrays: dict) -> mm.StepModulation: + """The step's modulation from the fixture's two modulation rows, as + the reference formed them.""" + act = model.act_dtype + + def chunks(rows, n): + a = mx.array(rows).astype(act)[:, None, :] + return tuple(mx.split(a, n, axis=-1)) + + shift, scale = chunks(arrays["video_head_modulation"], 2) + return mm.StepModulation( + blocks=chunks(arrays["shared_block_modulation"], 6), shift=shift, + scale=scale) + + +def forced_rows(model, arrays: dict) -> dict: + """Each text block, each block and the head on the fixture's float32 + input to it, with the fixture's modulation.""" + act = model.act_dtype + mod = fixture_modulation(model, arrays) + n_txt = int(arrays["features"].shape[0]) + _, _, h, w = arrays["latents"].shape + out = {} + cos, sin = model.text_rope(n_txt) + for i, blk in enumerate(model.text_embedder.blocks): + src = (arrays["text_embedder.proj"] if i == 0 + else arrays[f"text_embedder.blocks.{i - 1}"]) + y = blk(mx.array(src).astype(act), cos, sin) + out[f"text_embedder.blocks.{i}"] = _np(y) + cos, sin = model.rope_tables(h, w, n_txt) + for i, blk in enumerate(model.layers): + src = arrays["in.layers.0"] if i == 0 else arrays[f"layers.{i - 1}"] + y = blk(mx.array(src).astype(act), mod.blocks, cos, sin) + out[f"layers.{i}"] = _np(y) + y = model.head(mx.array(arrays["in.final_norm_video"]).astype(act), + mod.shift, mod.scale) + out[HEAD] = _np(y) + return out + + +def measures(model, arrays: dict) -> dict: + """Relative L2 of every row against the fixture: teacher-forced for + the text blocks, the blocks and the head, free running for the rest, + and ``free `` for each teacher-forced row run free.""" + free = free_rows(model, arrays) + forced = forced_rows(model, arrays) + got = {} + for k, v in free.items(): + if k == "velocity": + got[k] = rel(v, arrays["velocity"]) + elif k in forced: + got[k] = rel(forced[k], arrays[k]) + got[f"free {k}"] = rel(v, arrays[k]) + elif k in arrays: + got[k] = rel(v, arrays[k]) + return got + + +def missed(got: dict, tolerances: dict) -> dict: + """The measures over their tolerance, by name. A ``free`` measure + has no tolerance of its own and is not gated.""" + out = {} + for k, v in got.items(): + if k.startswith("free "): + continue + tol = tolerances[k]["rel_l2"] + if not v <= tol: + out[k] = (v, tol) + return out + + +# The exact gate: the step in the oracle's fp32 arithmetic. + +def exact_floats(model, path: str) -> int: + """Every float tensor of the file into the model in float32, where the + load keeps the activation dtype. Run it before a patch table wraps + the leaves, so the names are the file's.""" + import gguf + from mlx.utils import tree_flatten + + params = dict(tree_flatten(model.parameters())) + reader = gguf.GGUFReader(path) + floats = (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16) + pairs = [] + for t in reader.tensors: + if t.tensor_type not in floats or t.name not in params: + continue + a = np.array(t.data, np.float32).reshape(params[t.name].shape) + pairs.append((t.name, mx.array(a))) + model.load_weights(pairs, strict=False) + return len(pairs) + + +@contextmanager +def _exact(model): + """The oracle's arithmetic: every K-quant product dense on its + dequantized weight, each one checked at the exit.""" + with bg.exact_products(model) as used: + yield + bg.assert_all_dense(model, used) + + +def exact_measures(model, arrays: dict) -> dict: + """``measures`` in the oracle's arithmetic.""" + with _exact(model): + return measures(model, arrays) + + +def forced_row(name: str) -> bool: + return (name.startswith(("layers.", "text_embedder.blocks.")) + or name == HEAD) + + +def free_row(name: str) -> bool: + """A free-running row of ``measures``: every measure but the + teacher-forced rows and the velocity.""" + return not forced_row(name) and name != "velocity" + + +# The exact gate's free-running rows and the velocity carry fp32 rounding +# that no other summation order reproduces. Each fixture records gmlx's +# own fp32 floor on the free rows (the worst of them) and on the +# velocity: the relative L2 between the exact step on the GPU and the +# same step on the CPU stream. The free rows and the velocity each take +# the larger of EXACT_TOL and twice their floor. The teacher-forced rows +# keep EXACT_TOL. + +def exact_free_floor(model, arrays: dict) -> dict: + """Relative L2 of each free row of ``measures`` and of the velocity + between the exact step on the GPU and on the CPU stream.""" + with _exact(model): + gpu = free_rows(model, arrays) + with mx.stream(mx.cpu): + cpu = free_rows(model, arrays) + out = {} + for k, v in gpu.items(): + name = f"free {k}" if forced_row(k) else k + if k in arrays and not forced_row(name): + out[name] = rel(v, cpu[k]) + return out + + +# The kinds of measure with a floor: manifest result and tolerance keys. +FLOORS = {"free": ("exact_free_floor", "exact_free"), + "velocity": ("exact_velocity_floor", "exact_velocity")} + + +def exact_tols(m: dict) -> dict: + """The free-row and velocity tolerances a manifest records, each of + which must follow the rule from the floor beside it.""" + out = {} + for kind, (res, key) in FLOORS.items(): + floor = m["results"][res]["rel_l2"] + tol = m["tolerances"][key]["rel_l2"] + assert tol == max(bg.EXACT_TOL, 2 * floor), (kind, tol, floor) + out[kind] = tol + return out + + +def check_floor(m: dict, floor: dict, label: str) -> dict: + """The floors measured now are the floors the manifest records. + Returns the manifest's tolerances.""" + free = {k: v for k, v in floor.items() if free_row(k)} + top = max(free, key=free.get) + now = {"free": free[top], "velocity": floor["velocity"]} + print(f"{label}: fp32 floor of the free rows {now['free']!r} at {top}, " + f"of the velocity {now['velocity']!r}") + for kind, (res, _) in FLOORS.items(): + assert res in m["results"], f"{m['name']} records no {res}" + tols = exact_tols(m) + for kind, (res, _) in FLOORS.items(): + rec = m["results"][res]["rel_l2"] + print(f"{label}: recorded {kind} floor {rec:.2e}, gate " + f"{tols[kind]:.2e}") + assert now[kind] <= 1.25 * rec, (kind, now[kind], rec) + return tols + + +def exact_gate(name: str, tols: dict) -> float: + if name == "velocity": + return tols["velocity"] + return tols["free"] if free_row(name) else bg.EXACT_TOL + + +def exact_over(got: dict, tols: dict) -> dict: + return {k: (v, exact_gate(k, tols)) for k, v in got.items() + if not v <= exact_gate(k, tols)} + + +def exact_line(label: str, got: dict, tols: dict) -> str: + def worst(keys): + k = max(keys, key=got.get) + return f"{got[k]:.2e} at {k}" + forced = [k for k in got if forced_row(k)] + free = [k for k in got if free_row(k)] + return (f"{label}: teacher-forced rows within {worst(forced)}, gate " + f"{bg.EXACT_TOL:.0e}; velocity {got['velocity']:.2e}, gate " + f"{tols['velocity']:.2e}; free rows within {worst(free)}, gate " + f"{tols['free']:.2e}") + + +def fault_line(label: str, bad: dict, tols: dict) -> tuple[str, float]: + """Where a wrong model misses the exact gate, by kind of row, and the + largest ratio of a measure to its gate.""" + parts = [] + for kind, keep in (("teacher-forced", forced_row), + ("velocity", lambda k: k == "velocity"), + ("free", free_row)): + r = {k: bad[k] / exact_gate(k, tols) for k in bad if keep(k)} + miss = [k for k, v in r.items() if v > 1] + if miss: + top = max(miss, key=r.get) + parts.append(f"{kind} {len(miss)} rows, up to {r[top]:.0f} of " + f"the gate at {top}") + top = max(bad[k] / exact_gate(k, tols) for k in bad) + return (f"{label}: misses through " + ("; ".join(parts) or "nothing"), + top) + + +# The site gate: each op of a block on the path a request takes, against +# the stock op in float32 on the same production input, within 2 n u. +# The float32 side runs on the CPU, whose float32 products are exact. A +# reading is (block, "step", op, call number) for an op of a block, with +# ``qg.HEAD`` after it for the worst head row of the q and k site, and +# (module path, call number) for a linear leaf, so the helpers of +# qwen_gates.py read it. n counts the outputs the stock ops of a site +# write in the activation dtype, and a site that a kernel route ran is +# held to the kernel's own count, ``bg.KERNEL_ROUNDINGS``. +# +# The sites of a MOSS-Image block: the two norms, which round the +# normalized value and then the weight product; the two modulations; q +# and k through the q and k seam, in the dims of ``mm.interleave``; sdpa; +# the feed-forward from its input to the rows of its down leaf, situ, +# which rounds once over the rounding of its leaf; the two gated residual +# adds, or the residual op that also takes the leaf before it; and every +# linear leaf. A text block has the norms, the q and k seam, sdpa and +# situ, and adds its residuals in line. The shift of each modulation is +# added in line, and the exact gate holds it. +OP_ROUNDINGS = { + "modulate": 2, # 1 + scale, then the product + "qk": 2, # the norm, then the cast after the fp32 rotation + "sdpa": 1, + "situ": 1, # float32, rounded once + "gated_residual": 2, # gate * y, then the sum +} +# The calls of each op in one call of a block, by kind of block. +BLOCK_OPS = {"block": {"norm": 2, "modulate": 2, "qk": 1, "sdpa": 1, + "situ": 1, "gated_residual": 2}, + "text": {"norm": 2, "qk": 1, "sdpa": 1, "situ": 1}} +# The leaves a step reads past their call: the time vectors run as +# float32 products on the dequantized weight (``ops.products.precise``). +NOT_CALLED = ("time_embedding.0", "time_embedding.2", + "shared_block_modulation", "video_head_modulation") + + +def block_names(model) -> dict: + """id of each block of ``model`` to (name, kind).""" + out = {id(b): (f"layers.{i}", "block") + for i, b in enumerate(model.layers)} + for i, b in enumerate(model.text_embedder.blocks): + out[id(b)] = (f"text_embedder.blocks.{i}", "text") + return out + + +@contextmanager +def step_sites(model, sites: dict, routed: dict | None = None, + heads: dict | None = None, built=None, + skipped: set | None = None): + """Fill ``sites`` with a reading for every op of every block and every + call of a linear leaf while steps run: key to (relative L2, gate). + ``routed`` takes the mlx-kquant op of each site a kernel route ran, + ``heads`` the place of the worst head row of each q and k site, as + (0, batch, head, token). ``built`` is the built table of the patch + table the step runs under, from ``install.build_table`` in the block + form, and gives the exact side each leaf's delta. ``skipped`` takes + the path of a leaf called on a float32 input, which runs in float32 + and has no site.""" + import mlx_kquant as kq + + from gmlx.image import ops + from gmlx.image.lora import install + from gmlx.image.ops import attention, epilogues, norms + + f = mx.float32 + act = model.act_dtype + names = block_names(model) + cur: list = [] + counts: dict = {} + calls: dict = {} + ran: list = [] + wrapped = any(isinstance(m, install.ImageLoRALinear) + for m in install.modules(model).values()) + paths = {id(m): p for p, m in install.modules(model).items() + if install.linear_dims(m) is not None} + bases = {id(m): p for p, m in install.linear_leaves(model).items()} + + def factors(path): + g = None if built is None else built.get(path) + return None if g is None else (g.a_t.astype(f), g.b_t.astype(f)) + + def exact_leaf(mod, x32, path): + """The leaf's product in float32 on its dequantized weight, with + its delta.""" + base = getattr(mod, "base", mod) + codec = getattr(base, "kquant_type", None) + w = base["weight"].astype(f) if codec is None else \ + kq.dequantize(base["weight"], base["scales"], codec, f) + y = x32 @ w.T + if "bias" in base: + y = y + base["bias"].astype(f) + fac = factors(path) + if fac is not None: + y = y + (x32 @ fac[0]) @ fac[1] + return y + + def leaf_rounds(mod, path) -> int: + base = getattr(mod, "base", mod) + n = 2 if getattr(base, "kquant_type", None) is not None \ + and "bias" in base else 1 + return n + (2 if factors(path) is not None else 0) + + def read(op, got, want, rounds, by_head=False): + key = (cur[-1], "step", op) + n = counts[key] = counts.get(key, 0) + 1 + outs = got if isinstance(got, tuple) else (got,) + same = all(o.dtype == act for o in outs) + gate = 2 * rounds * qg.SITE_UNIT * same + with mx.stream(mx.cpu): + exact = want() + rel = qg._rel_l2(got, exact) + if by_head: + e2, r2 = bg._head_sums(outs, exact if isinstance( + exact, tuple) else (exact,)) + mx.eval(e2, r2) + sites[(*key, n)] = (rel, gate) + if by_head: + worst, at = bg._worst(np.array(e2), np.array(r2)) + sites[(*key, n, qg.HEAD)] = (worst, gate) + if heads is not None: + heads[(*key, n)] = at + return (*key, n) + + def leaf_site(path, rel, gate): + n = calls[path] = calls.get(path, 0) + 1 + sites[(path, n)] = (rel, gate) + + def within(orig): + def run(blk, *a, **k): + hit = names.get(id(blk)) + if hit is None: + return orig(blk, *a, **k) + cur.append(hit[0]) + try: + return orig(blk, *a, **k) + finally: + cur.pop() + return run + + def norm(orig): + def run(nrm, x): + y = orig(nrm, x) + if cur: + read("norm", y, lambda: mx.fast.rms_norm( + x.astype(f), nrm.weight.astype(f), nrm.eps), + bg.norm_roundings(nrm.weight, x.dtype)) + return y + return run + + def modulate(orig): + def run(x, scale): + y = orig(x, scale) + if cur: + read("modulate", y, lambda: epilogues.modulate( + x.astype(f), scale.astype(f)), OP_ROUNDINGS["modulate"]) + return y + return run + + def qk(orig): + def run(q, k, q_w, k_w, cos, sin, eps): + before = len(ran) + out = orig(q, k, q_w, k_w, cos, sin, eps) + if not cur: + return out + kernel = "qk" in ran[before:] + key = read("qk", tuple(out), lambda: attention.qk_norm_rope( + q.astype(f), k.astype(f), q_w.astype(f), k_w.astype(f), cos, + sin, eps), + bg.KERNEL_ROUNDINGS["qk"] if kernel else OP_ROUNDINGS["qk"], + by_head=True) + if kernel and routed is not None: + routed[key] = bg.KERNEL_OPS["qk"] + return out + return run + + def sdpa(orig): + def run(q, k, v, mask=None): + out = orig(q, k, v, mask) + if cur: + read("sdpa", out, lambda: attention.sdpa( + q.astype(f), k.astype(f), v.astype(f), mask), + OP_ROUNDINGS["sdpa"]) + return out + return run + + def block_hidden(orig): + def run(blk, h): + y = orig(blk, h) + if not cur: + return y + lin = blk.proj["ffn_up_gate"] + path = paths[id(lin)] + + def want(): + hid, gate = mx.split(exact_leaf(lin, h.astype(f), path), 2, + axis=-1) + return mm.situ(hid, gate) + + read("situ", y, want, + OP_ROUNDINGS["situ"] + leaf_rounds(lin, path)) + return y + return run + + def text_hidden(orig): + def run(blk, h): + y = orig(blk, h) + if not cur: + return y + pu, pg = paths[id(blk.ffn_up)], paths[id(blk.ffn_gate)] + + def want(): + x32 = h.astype(f) + return mm.situ(exact_leaf(blk.ffn_up, x32, pu), + exact_leaf(blk.ffn_gate, x32, pg)) + + read("situ", y, want, OP_ROUNDINGS["situ"] + + leaf_rounds(blk.ffn_up, pu) + + leaf_rounds(blk.ffn_gate, pg)) + return y + return run + + def gated_residual(orig): + def run(x, gate, y): + out = orig(x, gate, y) + if cur: + read("gated_residual", out, lambda: epilogues.gated_residual( + x.astype(f), gate.astype(f), y.astype(f)), + OP_ROUNDINGS["gated_residual"]) + return out + return run + + def resid(orig): + # The residual route's op, ``residual + gate * (x @ W^T)`` in one + # call: the block's gated residual add and the call of its leaf. + # From 1024 rows the kernel rounds once, below that as the stock + # product and gated add do. + def run(x, leaf, residual, gate): + out = orig(x, leaf, residual, gate) + if not cur: + return out + rows = math.prod(x.shape[:-1]) + rounds = 1 if rows >= qg.RESID_ONE_ROUNDING_ROWS else 3 + + def want(): + w = kq.dequantize(leaf["weight"], leaf["scales"], + leaf.kquant_type, f) + return epilogues.gated_residual( + residual.astype(f), gate.astype(f), x.astype(f) @ w.T) + + key = read("gated_residual", out, want, rounds) + if routed is not None: + routed[key] = "quantized_matmul_residual" + leaf_site(bases[id(leaf)], *sites[key]) + return out + return run + + def leaf(orig): + def run(mod, x, *a, **k): + before = len(ran) + y = orig(mod, x, *a, **k) + path = paths.get(id(mod)) + if path is None: + return y + if x.dtype != act: + if skipped is not None: + skipped.add(path) + return y + with mx.stream(mx.cpu): + rel = qg._rel_l2(y, exact_leaf(mod, x.astype(f), path)) + rounds = leaf_rounds(mod, path) + kernel = "lora" in ran[before:] + if kernel: + rounds = bg.KERNEL_ROUNDINGS["lora"] + leaf_site(path, rel, 2 * rounds * qg.SITE_UNIT) + if kernel and routed is not None: + routed[(path, calls[path])] = bg.KERNEL_OPS["lora"] + return y + return run + + leaf_types = ({install.ImageLoRALinear} if wrapped else + {type(m) for m in install.linear_leaves(model).values()}) + hooks = [(mm.Block, "__call__", within), + (mm.TextBlock, "__call__", within), + (norms.RMSNorm, "__call__", norm), + (mm.Block, "hidden", block_hidden), + (mm.TextBlock, "hidden", text_hidden), + (ops, "modulate", modulate), (ops, "qk_norm_rope", qk), + (ops, "sdpa", sdpa), (ops, "gated_residual", gated_residual), + (ops, "resid_fused", resid)] + hooks += [(cls, "__call__", leaf) for cls in leaf_types] + saved = [] + try: + for obj, attr, hook in hooks: + saved.append((obj, attr, getattr(obj, attr))) + setattr(obj, attr, hook(saved[-1][2])) + with bg.kernel_calls(ran.append): + yield sites + finally: + for obj, attr, orig in reversed(saved): + setattr(obj, attr, orig) + + +def missing_ops(model, sites: dict, steps: int = 1, skipped=()) -> list: + """What ``steps`` steps leave unread: each block calls every op of + ``BLOCK_OPS`` for its kind, and every linear leaf but those of + ``NOT_CALLED`` and ``skipped`` is called once a step.""" + from gmlx.image.lora import install + + got: dict = {} + for block, _, op, _ in qg.op_sites(sites): + got[(block, op)] = got.get((block, op), 0) + 1 + want = {} + for name, kind in block_names(model).values(): + for op, n in BLOCK_OPS[kind].items(): + want[(name, op)] = n * steps + out = sorted(f"{k}: {got.get(k, 0)} of {n}" for k, n in want.items() + if got.get(k, 0) != n) + out += sorted(f"{k}: not a known op call" for k in got if k not in want) + leaves = {p for p in install.linear_leaves(model) + if p not in NOT_CALLED and p not in skipped} + read: dict = {} + for k in sites: + if len(k) == 2: + read[k[0]] = read.get(k[0], 0) + 1 + out += sorted(f"leaf {p}: {read.get(p, 0)} of {steps}" for p in leaves + if read.get(p, 0) != steps) + out += sorted(f"leaf {p}: not a leaf a step calls" for p in read + if p not in leaves) + return out diff --git a/tests/image/test_image_api.py b/tests/image/test_image_api.py index d9512bcfd..03ba8876d 100644 --- a/tests/image/test_image_api.py +++ b/tests/image/test_image_api.py @@ -831,3 +831,28 @@ def test_sample_takes_the_batch_from_the_noise(pipe): with pytest.raises(ValueError, match="does not match"): api.sample(pipe, _cond(0.5), height=32, width=64, noise=noise, steps=2) + + +def test_load_passes_the_borrowed_models_gguf_to_its_encoder(monkeypatch): + """A borrowed encoder model with its GGUF reads the F32 norms a + MOSS-Image encode runs on, as the CLI and serve do.""" + from gmlx.image import loader + + fam = get_family("moss_image") + ident = SimpleNamespace(family=fam, variant=None) + monkeypatch.setattr(api, "_header_bytes", lambda p, kind: 0) + monkeypatch.setattr(loader, "identify_dit", lambda *a, **k: ident) + monkeypatch.setattr(loader, "load_dit", lambda *a, **k: ( + None, SimpleNamespace(identified=ident))) + got = [] + + def borrowed(pair, spec, gguf=None): + got.append((pair, spec.fp32_norms, gguf)) + raise LookupError("stop") + + monkeypatch.setattr(api, "_borrowed", borrowed) + with pytest.raises(LookupError): + api.load("d.gguf", encoder_model=(1, 2), encoder_gguf="chat.gguf") + assert got == [((1, 2), True, "chat.gguf")] + with pytest.raises(ValueError, match="encoder_gguf"): + api.load("d.gguf", text_encoder="e.gguf", encoder_gguf="chat.gguf") diff --git a/tests/image/test_image_catalog.py b/tests/image/test_image_catalog.py index c5f8a9e05..7e4322a69 100644 --- a/tests/image/test_image_catalog.py +++ b/tests/image/test_image_catalog.py @@ -392,8 +392,12 @@ def test_sha256_file_reads_the_whole_file(tmp_path): def test_every_class_pulls_the_q8_encoder_and_the_largest_the_reference_dit(): for r in catalog.RECIPES.values(): + # No 16 GB budget holds the 9.5 GB Q8_0 file of Qwen3.5-9B. + small = {16} if r.family == "moss_image" else set() assert {catalog.DEFAULT_QUANTS[r.name][c][1] - for c in catalog.MEMORY_CLASSES} == {"Q8_0"} + for c in catalog.MEMORY_CLASSES if c not in small} == {"Q8_0"} + assert {catalog.DEFAULT_QUANTS[r.name][c][1] + for c in small} <= {"Q4_K_M"} assert catalog.DEFAULT_QUANTS[r.name][64][0] == r.dit_quants[-1] @@ -501,3 +505,17 @@ def test_a_flux2_pull_takes_the_z_image_encoder_and_the_klein_vae_on_disk( def test_a_32_gb_mac_pulls_qwen_image_at_q8_0_since_its_f16_load_does_not_fit(): assert catalog.default_quants(catalog.RECIPES["qwen-image"], 32 << 30) == ("Q8_0", "Q8_0") + + +def test_each_class_default_encoder_has_its_own_size(): + """A recipe whose class defaults take more than one encoder quant has + each file in ENCODER_BYTES, so the fit of a DiT prices the one its + class pulls.""" + for r in catalog.RECIPES.values(): + names = {r.encoder_name(catalog.DEFAULT_QUANTS[r.name][c][1]) + for c in catalog.MEMORY_CLASSES} + if len(names) > 1: + assert names <= set(catalog.ENCODER_BYTES), r.name + m = catalog.RECIPES["moss-image"] + assert catalog.companion_bytes(m, 16 << 30) == 5_680_522_464 + m.vae_bytes + assert catalog.companion_bytes(m, 64 << 30) == m.encoder_bytes + m.vae_bytes diff --git a/tests/image/test_image_check2_moss_image.py b/tests/image/test_image_check2_moss_image.py new file mode 100644 index 000000000..bbbc4250c --- /dev/null +++ b/tests/image/test_image_check2_moss_image.py @@ -0,0 +1,178 @@ +"""Check 2 of the image acceptance gate for MOSS-Image: the Qwen3.5-9B text +encoder read after its final norm. + +gmlx encodes the oracle's prompts with the text GGUF the oracle read, +for the Q8_0 and the Q4_K_M files of unsloth/Qwen3.5-9B-GGUF. The token +ids must equal the pipeline's, and the features of each prompt must stay +within the manifest's relative L2 and per-token cosine of the oracle's +fp32 features. The pipeline runs each prompt alone at its own length, +which is the encoding gmlx runs. Each wrong encoder below must fail some +reading. + +Skipped until the fixture's manifest is committed, and unless +KQUANT_TEST_GGUF_DIR holds the encoder GGUF the fixture names and +GMLX_IMAGE_FIXTURES holds the fixture. +scripts/image_oracle/check2_moss_image.py writes it.""" +from __future__ import annotations + +import os + +import mlx.core as mx +import numpy as np +import pytest +from oracle_fixture import ORACLE, fixture, fx + +pytestmark = pytest.mark.integration + +NAMES = ("check2_moss_image_q8_0", "check2_moss_image_q4_k_m") + + +def token_cosine(got: np.ndarray, ref: np.ndarray) -> np.ndarray: + g = got.astype(np.float64) + r = ref.astype(np.float64) + return (g * r).sum(-1) / (np.linalg.norm(g, axis=-1) + * np.linalg.norm(r, axis=-1)) + + +@pytest.fixture(scope="module", params=NAMES) +def check2(request): + name = request.param + if not (ORACLE / f"{name}.manifest.json").is_file(): + pytest.skip(f"{name}.manifest.json is not committed") + return fixture(name) + + +@pytest.fixture(scope="module") +def handle(gguf_index, check2): + from gmlx.image import text_encoder as te + + spec = te.encoder_spec("moss_image") + want = check2[1]["files"]["encoder"] + for p in gguf_index.get(spec.arch, []): + if (os.path.basename(p) == want["name"] + and os.path.getsize(p) == want["size"] + and fx.file_identity(p)["sha256"] == want["sha256"]): + break + else: + pytest.skip(f"{want['name']} is not under KQUANT_TEST_GGUF_DIR") + h = te.load_encoder(p, spec, log=lambda *_: None) + yield h + h.release() + + +def test_the_encoder_is_the_identified_one(handle): + from gmlx.image import text_encoder as te + + spec = handle.spec + assert len(te._inner(handle.model).layers) == spec.layers == 32 + got = te.model_identity(handle.model, spec) + assert got == {r.gguf_name: r.sha256 for r in spec.identity} + + +def test_token_ids_equal_the_pipelines(handle, check2): + arrays, m = check2 + for key, prompt in m["inputs"]["prompts"].items(): + ids = handle.token_ids(prompt) + assert ids == arrays[f"ids_{key}"].tolist(), key + assert len(ids) == m["inputs"]["rows"][key] + + +def _reading(got: np.ndarray, ref: np.ndarray, tol: dict, key: str) -> dict: + out = {"rows": int(got.shape[0]), "shape": got.shape, + "want_shape": ref.shape, + "tol_rel_l2": tol["rel_l2"][key], + "tol_cosine": tol["token_cosine_min"][key]} + if got.shape == ref.shape: + cos = token_cosine(got, ref) + out.update(rel_l2=fx.rel_l2(got, ref), + token_cosine_min=float(cos.min()), + worst_token=int(cos.argmin())) + return out + + +def _caught(r: dict, ids, want_ids) -> list: + by = [] + if list(ids) != list(want_ids): + by.append("token ids") + if r["shape"] != r["want_shape"]: + by.append("shape") + if "rel_l2" in r and not r["rel_l2"] <= r["tol_rel_l2"]: + by.append("relative L2") + if "rel_l2" in r and not r["token_cosine_min"] >= r["tol_cosine"]: + by.append("token cosine") + return by + + +def test_features_are_within_the_calibrated_tolerance(handle, check2): + arrays, m = check2 + tol = m["tolerances"] + for key, prompt in m["inputs"]["prompts"].items(): + enc = handle.encode(prompt) + got = np.array(enc.features.astype(mx.float32)) + r = _reading(got, arrays[f"features_{key}"], tol, key) + print(f"[check2] {m['name']} {key}: {r['rows']} rows, rel L2 " + f"{r.get('rel_l2', float('nan')):.3e} of gate " + f"{r['tol_rel_l2']:.3e}, min token cosine " + f"{r.get('token_cosine_min', float('nan')):.6f} floor " + f"{r['tol_cosine']:.6f}") + assert not _caught(r, [], []), (key, r) + + +def _layer_read(handle, ids, layer: int, norm: bool) -> np.ndarray: + """The output of layer ``layer`` (1-based) of the stack, with or + without the final norm.""" + from gmlx.image import text_encoder as te + + inner = te._inner(handle.model) + layers, final = inner.layers, inner.norm + inner.layers = layers[:layer] + if not norm: + inner.norm = lambda h: h + try: + h = te.final_hidden(handle.model, ids) + finally: + inner.layers, inner.norm = layers, final + return np.array(h[0].astype(mx.float32)) + + +def _thinking_off(handle, prompt: str): + import dataclasses + + from gmlx.image import text_encoder as te + + spec = handle.spec + off = dataclasses.replace(spec, template=spec.template.replace( + "\n", "\n\n\n\n")) + ids = te.token_ids(off, handle.tokenizer, prompt) + return ids, np.array(te.final_hidden(handle.model, ids)[0].astype( + mx.float32)) + + +def _ids(handle, prompt): + return handle.token_ids(prompt) + + +WRONG = { + "the read before the final norm": ( + lambda h, p: (_ids(h, p), _layer_read(h, _ids(h, p), 32, False)), + {"relative L2", "token cosine"}), + "the final norm one layer early": ( + lambda h, p: (_ids(h, p), _layer_read(h, _ids(h, p), 31, True)), + {"relative L2", "token cosine"}), + "thinking off": (_thinking_off, {"token ids", "shape"}), +} + + +@pytest.mark.parametrize("name", list(WRONG), + ids=[n.replace(" ", "_") for n in WRONG]) +def test_each_wrong_encoder_fails_a_reading(handle, check2, name): + arrays, m = check2 + make, kinds = WRONG[name] + missed = {} + for key, prompt in m["inputs"]["prompts"].items(): + ids, got = make(handle, prompt) + r = _reading(got, arrays[f"features_{key}"], m["tolerances"], key) + by = _caught(r, ids, arrays[f"ids_{key}"].tolist()) + if not set(by) & kinds: + missed[key] = (by, r) + assert not missed, (name, missed) diff --git a/tests/image/test_image_check3_moss_image.py b/tests/image/test_image_check3_moss_image.py new file mode 100644 index 000000000..0f361aafd --- /dev/null +++ b/tests/image/test_image_check3_moss_image.py @@ -0,0 +1,298 @@ +"""Check 3 of the image acceptance gate for MOSS-Image: one denoise step of +the gmlx DiT, row by row, against the release's MOSSImageTransformer on +the same GGUF, under three gates. + +scripts/image_oracle/check3_moss_image.py writes the fixtures: every row +of the step at 512 and the velocity at 1024 in fp32, with the oracle's +own bf16 against fp32 gap of each. + +The product gate runs the model as it ships, in bf16 with the kernel +routes at their default, and holds each row to twice the oracle's gap: +teacher-forced for each text block, each block and the head, free +running for the rest. The exact gate runs the model in the oracle's fp32 +arithmetic, every K-quant product dense on its dequantized weight and +every float tensor of the file in float32. It holds the teacher-forced +rows to EXACT_TOL, and the free rows and the velocity each to the larger +of EXACT_TOL and twice gmlx's own fp32 floor that the manifest records. +The site gate holds each op of every block and each linear leaf on the +path a request takes to the stock op in float32 on the same input, +within 2 n u (moss_gates.py). + +The rows need the MOSS-Image GGUF under KQUANT_TEST_GGUF_DIR and the +fixtures under GMLX_IMAGE_FIXTURES. They skip until both are there and +the manifests are committed. +""" +from __future__ import annotations + +import os +from pathlib import Path + +import mlx.core as mx +import pytest + +import moss_gates as mg +import qwen_gates as qg +from block_gates import exact_process, rerun_exact +from image_fixtures import image_dits +from oracle_fixture import ORACLE, committed, fixture, fx + +from gmlx.image import loader as image_loader +from gmlx.image import ops +from gmlx.image.families.moss_image import model as mm +from gmlx.image.ops import kq_routes + +ROWS = "check3_moss_image_512" +VELOCITY = "check3_moss_image_1024" + + +def _manifest(name: str) -> dict: + if not (ORACLE / f"{name}.manifest.json").is_file(): + pytest.skip(f"{name}.manifest.json is not committed: run " + "scripts/image_oracle/check3_moss_image.py on the " + "MOSS-Image GGUF and commit the manifest it writes") + return committed(name) + + +@pytest.fixture(scope="module") +def real(gguf_dir): + m = _manifest(ROWS) + want = m["files"]["dit"] + for p in image_dits(gguf_dir, "moss_image"): + if Path(p).name == want["name"] and os.path.getsize(p) == want["size"]: + break + else: + pytest.skip(f"{want['name']} is not under KQUANT_TEST_GGUF_DIR") + if fx.file_identity(p)["sha256"] != want["sha256"]: + pytest.fail(f"{p} is not the file the oracle read") + model, _ = image_loader.load_dit(p, log=lambda *_: None) + return model, p, list(model.layers) + + +def _depth(real, m: dict, blocks: int | None = None): + """The model with the blocks the fixture holds, or the first + ``blocks``.""" + model, _, all_blocks = real + n = m["inputs"]["blocks"] if blocks is None else blocks + assert n <= len(all_blocks) + model.layers = all_blocks[:n] + return model + + +def test_the_manifests_follow_the_rule(): + """The committed manifests alone: every row tolerance is twice the + gap beside it, the velocity's cosine floor follows its gap, and the + exact gate's free-row and velocity tolerances follow their floors.""" + for name in (ROWS, VELOCITY): + m = _manifest(name) + gap = m["results"]["velocity_bf16_vs_fp32"] + tol = m["tolerances"]["velocity"] + assert tol["rel_l2"] == 2 * gap["rel_l2"] + assert tol["cosine_min"] == pytest.approx( + 1 - 4 * (1 - gap["cosine"]), abs=1e-12) + assert m["inputs"]["dequant"] == "exact" + rows = _manifest(ROWS) + for k, g in rows["results"]["rows_bf16_vs_fp32"].items(): + assert rows["tolerances"][k]["rel_l2"] == 2 * g["rel_l2"], k + assert g["teacher_forced"] == mg.forced_row(k), k + mg.exact_tols(rows) + + +@pytest.mark.integration +def test_check3_rows_at_512(real): + arrays, m = fixture(_manifest(ROWS)["name"]) + model = _depth(real, m) + got = mg.measures(model, arrays) + out = mg.missed(got, m["tolerances"]) + ratio = {k: v / m["tolerances"][k]["rel_l2"] for k, v in got.items() + if not k.startswith("free ")} + top = max(ratio, key=ratio.get) + print(f"moss check3 512: {len(ratio)} rows, worst {got[top]:.2e} at " + f"{top}, {ratio[top]:.2f} of its gate; velocity " + f"{got['velocity']:.2e} of gate " + f"{m['tolerances']['velocity']['rel_l2']:.2e}") + assert not out, out + + +@pytest.mark.integration +def test_check3_velocity_at_1024(real): + arrays, m = fixture(_manifest(VELOCITY)["name"]) + model = _depth(real, m) + got = mg.free_rows(model, arrays)["velocity"] + rel = mg.rel(got, arrays["velocity"]) + cos = fx.cosine(got, arrays["velocity"]) + tol = m["tolerances"]["velocity"] + print(f"moss check3 1024: velocity {rel:.2e} of gate " + f"{tol['rel_l2']:.2e}, cosine {cos:.6f} floor " + f"{tol['cosine_min']:.6f}") + assert rel <= tol["rel_l2"] and cos >= tol["cosine_min"] + + +# The exact gate. + +def _situ_halves_swapped(mp, model): + real_situ = mm.situ + mp.setattr(mm, "situ", lambda hid, gate: real_situ(gate, hid)) + + +def _shift_dropped(mp, model): + real_mod = mm.MossImageDiT.modulation + + def mod(self, sigma, batch): + got = real_mod(self, sigma, batch) + b = got.blocks + return mm.StepModulation((mx.zeros_like(b[0]), *b[1:]), got.shift, + got.scale) + mp.setattr(mm.MossImageDiT, "modulation", mod) + + +def _rope_adjacent(mp, model): + mp.setattr(mm, "interleave", lambda x: x) + + +# Wrong models the exact gate must fail. +EXACT_FAULTS = {"situ halves swapped": _situ_halves_swapped, + "pre-attention shift dropped": _shift_dropped, + "rope pairs adjacent dims": _rope_adjacent} + + +@pytest.mark.integration +def test_check3_exact_gate_at_512(real, request, monkeypatch): + """Every teacher-forced row within EXACT_TOL of the oracle in its + arithmetic, the free rows and the velocity within the fixture's gates + from the floors measured now, and each wrong model of EXACT_FAULTS + over the gate.""" + arrays, m = fixture(_manifest(ROWS)["name"]) + if not exact_process(): + return rerun_exact(request, "moss check3 exact") + model = _depth(real, m) + n = mg.exact_floats(model, real[1]) + label = "moss check3 exact 512" + tol = mg.check_floor(m, mg.exact_free_floor(model, arrays), label) + got = mg.exact_measures(model, arrays) + print(mg.exact_line(f"{label}, {n} float tensors in fp32", got, tol)) + over = mg.exact_over(got, tol) + for name, apply in EXACT_FAULTS.items(): + with monkeypatch.context() as mp: + apply(mp, model) + bad = mg.exact_measures(model, arrays) + line, top = mg.fault_line(f"{label}, {name}", bad, tol) + print(line) + assert top > 10, (name, top) + assert not over, over + + +# The site gate. + +def _sites(model, arrays: dict) -> tuple: + """The sites of one free-running production step on the fixture's + inputs, the kernel sites, the head rows, the leaves called on a + float32 input and the velocity.""" + sites, routed, heads, skipped = {}, {}, {}, set() + with mg.step_sites(model, sites, routed=routed, heads=heads, + skipped=skipped): + vel = mg.free_rows(model, arrays)["velocity"] + return sites, routed, heads, skipped, vel + + +def site_line(label: str, sites: dict, routed: dict) -> str: + worst = qg.worst_by_op(sites) + kinds: dict = {} + for op in routed.values(): + kinds[op] = kinds.get(op, 0) + 1 + return (f"{label}: {len(sites)} readings, the largest over its gate " + "by op: " + + ", ".join(f"{k} {v:.2f}" for k, v in sorted(worst.items())) + + "; sites on mlx-kquant: " + + (", ".join(f"{k} {v}" for k, v in sorted(kinds.items())) + or "none")) + + +@pytest.mark.integration +def test_check3_site_gate_at_512(real): + """Every op of every block and every linear leaf a step calls on the + path a request takes is within its site gate of the stock op in + float32 on the same input, and that run's velocity passes the + product gate.""" + arrays, m = fixture(_manifest(ROWS)["name"]) + model = _depth(real, m) + sites, routed, heads, skipped, vel = _sites(model, arrays) + missing = mg.missing_ops(model, sites, skipped=skipped) + assert not missing, missing + qk = {k for k in qg.op_sites(sites) if k[2] == "qk"} + on = "qk" in kq_routes.active() + assert {k for k, v in routed.items() if v == "qk_norm_rope"} == ( + qk if on else set()) + assert set(qg.head_sites(sites)) == qk + print(site_line("moss check3 sites 512", sites, routed)) + print("moss check3 sites 512, q and k: " + + qg.qk_reading(sites, routed, heads)) + over = qg.over_gate(sites) + assert not over, over + rel = mg.rel(vel, arrays["velocity"]) + assert rel <= m["tolerances"]["velocity"]["rel_l2"], rel + + +def _situ_high(mp, model): + for cls in (mm.Block, mm.TextBlock): + real_hidden = cls.hidden + mp.setattr(cls, "hidden", + lambda self, h, f=real_hidden: f(self, h) * 1.05) + + +def _a_head_off(mp, model): + """Head 1 of the seam's output a twentieth off, on the kernel when the + qk route is in force.""" + real_qk = ops.qk_norm_rope + + def off(q, k, qw, kw, cos, sin, eps): + q2, k2 = real_qk(q, k, qw, kw, cos, sin, eps) + scale = mx.array([1.05 if h == 1 else 1.0 + for h in range(q2.shape[1])]).astype(q2.dtype) + return q2 * scale[:, None, None], k2 + mp.setattr(ops, "qk_norm_rope", off) + + +def _modulate_high(mp, model): + real_mod = ops.modulate + mp.setattr(ops, "modulate", lambda x, s: real_mod(x, s) * 1.05) + + +def _gated_residual_high(mp, model): + real_resid, real_res = ops.resid_fused, ops.gated_residual + mp.setattr(ops, "resid_fused", lambda x, leaf, r, g: + real_resid(x, leaf, r, g * 1.05)) + mp.setattr(ops, "gated_residual", + lambda x, g, y: real_res(x, g * 1.05, y)) + + +# name: (the wrong production op, the ops whose sites must miss) +SITE_FAULTS = {"situ a twentieth high": (_situ_high, {"situ"}), + "one head of q a twentieth off": (_a_head_off, {"qk"}), + "modulation a twentieth high": (_modulate_high, + {"modulate"}), + "gated residual a twentieth high": (_gated_residual_high, + {"gated_residual"})} + + +@pytest.mark.integration +@pytest.mark.parametrize("name", list(SITE_FAULTS)) +def test_check3_each_wrong_op_fails_the_site_gate(real, monkeypatch, name): + """A site reads an op on the op's own input, so a wrong op is over the + gate at its own sites. Two blocks are enough: each block reads its + own sites.""" + arrays, m = fixture(_manifest(ROWS)["name"]) + model = _depth(real, m, blocks=2) + apply, hit = SITE_FAULTS[name] + try: + with monkeypatch.context() as mp: + apply(mp, model) + sites, _, _, _, _ = _sites(model, arrays) + finally: + _depth(real, m) + over = qg.over_gate(sites) + worst = qg.worst_by_op(sites) + print(f"moss check3 sites, {name}: {len(over)} readings over the " + "gate; the largest by op " + + ", ".join(f"{k} {v:.2f}" for k, v in sorted(worst.items()))) + missed = {qg.site_op(k) for k in over if len(k) != 2} + assert hit <= missed, (name, sorted(missed)) diff --git a/tests/image/test_image_check4_moss_image.py b/tests/image/test_image_check4_moss_image.py new file mode 100644 index 000000000..b004018dd --- /dev/null +++ b/tests/image/test_image_check4_moss_image.py @@ -0,0 +1,122 @@ +"""Check 4 of the image acceptance gate for MOSS-Image: the family's +schedule against the release's MOSSImageScheduler at 50, 28 and 8 steps, +at 250 and 1000, where the snap to the training grid drops steps, and at +822, where a sigma ties between two grid points. Each sigma equals the +release's bit for bit, at the same step count. The release's timesteps +are its sigmas times 1000 in float32, as the DiT forms them. The arrays +are small and committed, so this runs on the CPU tier. +scripts/image_oracle/check4_moss_image.py writes them.""" +from __future__ import annotations + +import numpy as np +import pytest +from oracle_fixture import committed_fixture + +from gmlx.image import schedule +from gmlx.image.registry import get_family + +NAME = "check4_moss_image" + + +@pytest.fixture(scope="module") +def check4(): + return committed_fixture(NAME) + + +def _ulps(got, want) -> int: + g = np.asarray(got, np.float32).view(np.int32).astype(np.int64) + w = np.asarray(want, np.float32).view(np.int32).astype(np.int64) + return int(np.max(np.abs(g - w))) + + +def _gap(sigmas_of, arrays, n) -> tuple[int, int]: + """Steps of gmlx and of the release, and the largest ulp gap when the + counts agree.""" + got, want = sigmas_of(n), arrays[f"sigmas_{n}"] + if got.shape != want.shape: + return got.size - 1, want.size - 1 + return got.size - 1, _ulps(got, want) + + +def test_sigmas_equal_the_release_ones_bit_for_bit(check4): + arrays, m = check4 + fam = get_family("moss_image") + assert m["tolerances"]["count"] == "exact" + assert m["tolerances"]["sigmas"] == "bitwise" + assert 822 in m["inputs"]["steps"] + for n in m["inputs"]["steps"]: + got = fam.schedule(n, tokens=None, variant="sft") + want = arrays[f"sigmas_{n}"] + assert got.dtype == np.float32 and got[-1] == 0 + assert got.size - 1 == m["results"]["steps_run"][str(n)], n + assert got.shape == want.shape, n + assert _ulps(got, want) == 0, n + # The release's timestep is its sigma times 1000 in float32, as the + # DiT forms it from a sigma. + ts = want[:-1] * np.float32(1000) + assert _ulps(ts, arrays[f"timesteps_{n}"]) == 0, n + assert m["results"]["steps_run"]["1000"] < 1000 + + +def test_the_schedule_reads_the_scheduler_config(check4): + _, m = check4 + c = get_family("moss_image").constants.values + s = m["inputs"]["scheduler"] + assert s["shift"] == c["scheduler_shift"] + assert s["num_train_timesteps"] == c["num_train_timesteps"] + assert s["training_shift"] == 1.0 + + +def _unsnapped(n): + v = schedule.linspace_f32(1.0, 5.0 / 1004, n) + return np.append(schedule.shift_static(v, 5.0), np.float32(0)) + + +def _grid_k_over_1000(n): + """The snap to k / 1000 rounded once, which breaks the tie at 822 + steps the other way.""" + v = schedule.linspace_f32(1.0, 5.0 / 1004, n) + s = schedule.shift_static(v, 5.0) + grid = (np.arange(1, 1001, dtype=np.float64) / 1000).astype(np.float32) + hi = np.minimum(np.searchsorted(grid, s, side="left"), 999) + lo = np.maximum(hi - 1, 0) + pick = np.where(np.abs(s - grid[lo]) <= np.abs(grid[hi] - s), lo, hi) + snapped = grid[pick] + keep = np.concatenate([[True], snapped[1:] != snapped[:-1]]) + return np.append(snapped[keep], np.float32(0)).astype(np.float32) + + +def _no_dedup(n): + s = schedule.moss_sigmas(n) + return np.concatenate([s[:-1], np.repeat(s[-2], n + 1 - s.size), + s[-1:]]) if s.size < n + 1 else s + + +WRONG = {"sigmas not snapped": _unsnapped, "repeats kept": _no_dedup, + "the grid k / 1000": _grid_k_over_1000} + + +@pytest.mark.parametrize("name", list(WRONG)) +def test_a_wrong_schedule_misses(check4, name): + arrays, m = check4 + missed = [] + for n in m["inputs"]["steps"]: + steps, gap = _gap(WRONG[name], arrays, n) + if steps != m["results"]["steps_run"][str(n)] or gap > 0: + missed.append(n) + assert missed, name + + +def test_the_tie_at_822_steps_goes_to_the_release_point(check4): + """At 822 steps one sigma lies halfway between two points of the + release's grid, and the release takes the lower one. On k / 1000 + rounded once the same sigma goes to the upper point, a whole grid + step off, at the same count.""" + arrays, _ = check4 + want = arrays["sigmas_822"] + got = get_family("moss_image").schedule(822, tokens=None, variant="sft") + assert _ulps(got, want) == 0 + old = _grid_k_over_1000(822) + assert old.shape == want.shape + off = np.flatnonzero(old != want) + assert np.abs(old - want).max() == pytest.approx(1e-3, abs=1e-6), off diff --git a/tests/image/test_image_check5_moss_image.py b/tests/image/test_image_check5_moss_image.py new file mode 100644 index 000000000..8e55f91a2 --- /dev/null +++ b/tests/image/test_image_check5_moss_image.py @@ -0,0 +1,470 @@ +"""Check 5 of the image acceptance gate for MOSS-Image: its VAE. + +MOSS-Image decodes with the FLUX.2 autoencoder at narrower decoder widths, +[96, 192, 384, 384] where the encoder runs at [128, 256, 512, 512]. The +module, its wrong models and the helpers are those of +test_image_check5_flux3_klein.py. + +The tiny self-test runs on the MLX CPU stream: diffusers' +AutoencoderKLFlux2 at a tiny config whose decoder is narrower than its +encoder, with seeded weights and BatchNorm statistics. The tiny rows read +it as MOSS-Image's VAE at the tiny decoder widths. Its fp32 decode, encoder +mean and model space latents are held to 5e-5 max abs, each wrong model +misses that, and a GGUF packed from it loads the same weights bit for bit. +The FLUX.2 families refuse it, MOSS-Image refuses the FLUX.2 autoencoder, +and a tree at other decoder widths is refused by its widths. + +The official rows read the file under GMLX_TEST_IMAGE_DIR against the +oracle fixture at 512 and 1024 pixels: the fp32 decode and model space +latents to the manifest's max abs, the bf16 decode PSNR and latents +relative L2 to its gates, and a tiled bf16 decode at 1024 pixels to the +bf16 PSNR gate of that size. At 2048 pixels a tiled bf16 decode holds the +gate of 1024 against the tiled fp32 decode. A tiled bf16 encode at 1024 +stays within 0.05 relative L2 of the whole encode, with its round trip at +50 dB PSNR against the whole round trip. +scripts/image_oracle/check5_flux3_klein_vae.py --family moss_image writes +the fixtures. +""" +from __future__ import annotations + +import os +import re +from pathlib import Path + +import numpy as np +import pytest +from image_dir import sha256, vae_files +from image_fixtures import load_json +from oracle_fixture import ORACLE, committed, committed_fixture, fixture, fx +from test_image_check5_flux3_klein import ( + TILED_ENCODE_REL_L2, + TILED_ROUND_TRIP_DB, + WRONG, + _decode, + _gate_without_tf32, + _latents, + _tf32, + _tiny_readings, + _write_st, + to_unit, +) + +import mlx.core as mx +from mlx.utils import tree_flatten + +# The tiny rows run on a Mac of 16 GB, as most users run, and the rows on +# local files on this Mac (tests/host_box.py). + +FAMILY = "moss_image" +KLEIN = ("flux3_klein", "flux2_klein", "flux2_klein_9b", "flux2_klein_9b_kv") +NAME = "check5_moss_image_vae" +TINY = "selftest_moss_image_vae_tiny" +FLUX2_TINY = "selftest_flux3_klein_vae_tiny" +TINY_FP32_MAX_ABS = 5e-5 +HEADER = "moss_image_vae.header.json" +CONFIG = "moss_image_vae_config.json" +COUNT = "bn.num_batches_tracked" + + +def _tensors(fixture_) -> dict: + arrays, m = fixture_ + w = {k[len("weight."):]: v for k, v in arrays.items() + if k.startswith("weight.")} + w[COUNT] = np.array(m["inputs"]["num_batches_tracked"], np.int64) + return w + + +@pytest.fixture(scope="module") +def tiny(): + return committed_fixture(TINY) + + +@pytest.fixture +def tiny_file(tiny, tmp_path): + return _write_st(tmp_path / "vae.safetensors", _tensors(tiny)) + + +@pytest.fixture +def tiny_widths(tiny, monkeypatch): + """MOSS-Image's VAE read at the decoder widths of the tiny tree.""" + from gmlx.image.vae import moss_image_ae + + widths = tuple(tiny[1]["inputs"]["config"]["decoder_block_out_channels"]) + monkeypatch.setattr(moss_image_ae, "DECODER", widths) + monkeypatch.setitem(moss_image_ae.CONSTANTS, "decoder_block_out_channels", + widths) + return widths + + +@pytest.mark.host_gb(16) +def test_the_tiny_fp32_decode_and_encode_are_within_max_abs( + tiny, tiny_file, tiny_widths): + from gmlx.image.vae import load_vae, read_header + + arrays, m = tiny + cfg = m["inputs"]["config"] + assert cfg["decoder_block_out_channels"] != cfg["block_out_channels"] + assert read_header(tiny_file).layout == "diffusers" + vae = load_vae(tiny_file, family=FAMILY, dtype="fp32", conv="exact") + assert vae.config.decoder_channels == tiny_widths + assert vae.bn.running_var.dtype == mx.float32 + got = _tiny_readings(vae, arrays) + print(f"[check5] moss_image tiny fp32: {got}, gate {TINY_FP32_MAX_ABS}") + assert m["tolerances"]["fp32_max_abs"] == TINY_FP32_MAX_ABS + assert max(got.values()) <= TINY_FP32_MAX_ABS + + +@pytest.mark.host_gb(16) +@pytest.mark.parametrize("name", list(WRONG), + ids=[n.replace(" ", "_") for n in WRONG]) +def test_a_wrong_model_misses_the_tiny_gate(tiny, tiny_file, tiny_widths, + monkeypatch, name): + from gmlx.image.vae import load_vae + + arrays, _ = tiny + WRONG[name](monkeypatch) + vae = load_vae(tiny_file, family=FAMILY, dtype="fp32", conv="exact") + got = _tiny_readings(vae, arrays) + print(f"[check5] moss_image tiny, {name}: {got}") + assert max(got.values()) > 25 * TINY_FP32_MAX_ABS + + +@pytest.mark.host_gb(16) +def test_a_packed_gguf_loads_the_same_weights(tiny_file, tiny_widths, + tmp_path, capsys): + """pack-vae labels the file moss_image without --family, packs it under + --family moss-image, and the pack loads the same weights. A pack whose + config names other decoder widths is refused.""" + import json + + from gmlx.commands.image import cmd_pack_vae + from gmlx.image.vae import VAEError, load_vae, pack_vae + + bare = pack_vae(tiny_file, tmp_path / "bare.gguf", source_repo="r", + source_revision="v", config_json="{}") + assert bare.family == FAMILY and bare.lossless + config = dict(load_json(CONFIG), + decoder_block_out_channels=list(tiny_widths)) + src_config = tmp_path / "vae-config.json" + src_config.write_text(json.dumps(config)) + dst = tmp_path / "vae.gguf" + assert cmd_pack_vae([tiny_file, "--out", str(dst), "--family", + "moss-image", "--source-config", str(src_config), + "--json"]) == 0 + assert json.loads(capsys.readouterr().out)["family"] == FAMILY + a = dict(tree_flatten(load_vae(tiny_file, family=FAMILY, + dtype="fp32").parameters())) + b = dict(tree_flatten(load_vae(str(dst), family=FAMILY, + dtype="fp32").parameters())) + assert a.keys() == b.keys() + assert [k for k in a if not mx.array_equal(a[k], b[k]).item()] == [] + with pytest.raises(VAEError, match="its VAE config differs from the " + "MOSS-Image autoencoder in decoder_block_out_channels"): + pack_vae(tiny_file, tmp_path / "other.gguf", source_repo="r", + source_revision="v", config_json=json.dumps(dict( + config, decoder_block_out_channels=[32, 32, 32])), + family=FAMILY) + + +@pytest.mark.host_gb(16) +def test_the_official_config_agrees_with_the_constants(): + from gmlx.image.vae import flux2_ae, moss_image_ae + + config = load_json(CONFIG) + assert moss_image_ae.config_disagreements(config) == [] + assert moss_image_ae.CONSTANTS["decoder_block_out_channels"] \ + == moss_image_ae.DECODER == tuple(config["decoder_block_out_channels"]) + assert flux2_ae.config_disagreements(config) == [ + "decoder_block_out_channels"] + bad = dict(config, decoder_block_out_channels=config["block_out_channels"]) + assert moss_image_ae.config_disagreements(bad) == [ + "decoder_block_out_channels"] + assert flux2_ae.config_disagreements(bad) == [] + + +@pytest.mark.host_gb(16) +def test_the_official_header_fills_the_module_tree(): + from gmlx.image.vae import flux2_ae, moss_image_ae, reader + + head = load_json(HEADER) + shapes = {n: tuple(s) for n, s, _ in head["tensors"]} + dtypes = {n: d for n, _, d in head["tensors"]} + got, names = reader.diffusers_names({n: n for n in shapes}) + assert got == "diffusers" and len(names) == len(shapes) == 251 + header = reader.VaeHeader("x", "safetensors", got, names, shapes, dtypes, + {}) + cfg = reader.identify(header, FAMILY) + assert cfg.block_out_channels == (128, 256, 512, 512) + assert cfg.decoder_channels == moss_image_ae.DECODER + assert (cfg.latent_channels, cfg.layers_per_block, cfg.model_channels) \ + == (32, 2, 128) + params = dict(tree_flatten(flux2_ae.Flux2AE(cfg).parameters())) + reader.check_names(header, params, reader.VAES[FAMILY], cfg) + assert set(shapes) - set(params) == {COUNT} + # The encoder side and the BatchNorm are the FLUX.2 autoencoder's. + flux2 = {n: tuple(s) for n, s, _ in + load_json("flux3_klein_vae.header.json")["tensors"]} + assert flux2.keys() == shapes.keys() + assert [n for n in shapes if shapes[n] != flux2[n] + and not n.startswith(("decoder.", "post_quant_conv."))] == [] + + +@pytest.mark.host_gb(16) +def test_the_flux2_families_refuse_the_moss_vae(tiny_file, tiny_widths): + from gmlx.image.vae import load_vae, reader + + for family in KLEIN: + with pytest.raises(reader.VAEError, match="holds the MOSS-Image " + f"autoencoder, the moss_image VAE, and {family} " + "decodes with the FLUX.2 autoencoder"): + load_vae(tiny_file, family=family, dtype="fp32") + + +@pytest.mark.host_gb(16) +def test_moss_image_refuses_the_flux2_vae(tmp_path): + from gmlx.image.vae import load_vae, reader + + path = _write_st(tmp_path / "flux2.safetensors", + _tensors(committed_fixture(FLUX2_TINY))) + with pytest.raises(reader.VAEError, match="holds the FLUX.2 autoencoder, " + "the flux3_klein VAE, and moss_image decodes with the " + "MOSS-Image autoencoder"): + load_vae(path, family=FAMILY, dtype="fp32") + + +@pytest.mark.host_gb(16) +@pytest.mark.parametrize("family", [FAMILY, "flux3_klein", "z_image"]) +def test_a_tree_at_other_decoder_widths_is_refused_by_its_widths( + tiny, tiny_file, family): + """The tiny tree at MOSS-Image's own widths reads as a relative, and + without its quant convs and BatchNorm the FLUX.1 tree refuses it.""" + from gmlx.image.vae import flux1_ae, load_vae, reader + + widths = tiny[1]["inputs"]["config"]["decoder_block_out_channels"] + with pytest.raises(reader.VAEError, match=re.escape( + f"holds a FLUX.2 autoencoder with decoder widths {widths}, and " + f"{family} decodes with")): + load_vae(tiny_file, family=family, dtype="fp32") + flux1 = {k: v.shape for k, v in _tensors(tiny).items() + if not k.startswith(("quant_conv.", "post_quant_conv.", "bn."))} + assert flux1_ae.config_from_shapes(flux1) is None + + +@pytest.mark.host_gb(16) +def test_the_memory_row_prices_its_own_decoder(): + """The encode and the tiles are priced as the FLUX.2 autoencoder's, the + tiles at the decoder's 96 channels, and the whole decode by its own + fits, which price the fp32 decode at 512 pixels that the FLUX.1 fits + at 96 channels price under its reading of 4.08 GiB.""" + import dataclasses + + from gmlx.image import memfit + + GiB = 1 << 30 + row = memfit.vae_row(FAMILY) + assert row.measured and row.channels == {"decode": 96, "encode": 128} + assert row.whole["decode"] is memfit.MOSS_VAE_DECODE + assert row.whole["encode"] is memfit.VAE_ENCODE + assert row.tiled is memfit.VAE_TILED + for op in ("decode", "encode"): + assert memfit.tile_overlap(FAMILY, op) == 128 + for tiled in (False, True): + assert memfit.vae_peak("encode", 1024, 1024, family=FAMILY, + itemsize=2, tiled=tiled) == memfit.vae_peak( + "encode", 1024, 1024, family="flux3_klein", itemsize=2, + tiled=tiled) + assert memfit.vae_peak("decode", 2048, 2048, family=FAMILY, itemsize=2, + tiled=True) < memfit.vae_peak( + "decode", 2048, 2048, family="flux3_klein", itemsize=2, tiled=True) + for conv in ("kq", "exact"): + assert memfit.vae_peak("decode", 512, 512, family=FAMILY, itemsize=4, + conv=conv) >= (4.08 + 0.25) * GiB + flux1 = dataclasses.replace(row, whole={"decode": memfit.VAE_DECODE, + "encode": memfit.VAE_ENCODE}) + fit = flux1.whole["decode"]["exact"][4] + assert fit.bytes(512 * 512, 96, 4) < 4.08 * GiB + + +# The official file: needs it under GMLX_TEST_IMAGE_DIR and the fixture of +# the committed manifest. + +official = pytest.mark.integration + + +@pytest.fixture(scope="module") +def check5(): + if not (ORACLE / f"{NAME}.manifest.json").is_file(): + pytest.skip(f"no committed {NAME} manifest") + return fixture(NAME) + + +@pytest.fixture(scope="module") +def vaes(check5): + """The official VAE in a dtype. One is held at a time, so a pass is + priced against the memory of one VAE.""" + import gc + + from gmlx.image.vae import load_vae + + sums = {sha256(p): p for p in vae_files(FAMILY) + if p.suffix == ".safetensors"} + held = {} + + def get(dtype): + want = check5[1]["files"]["vae"]["sha256"] + if want not in sums: + pytest.skip("the fixture's VAE is not under GMLX_TEST_IMAGE_DIR") + if dtype not in held: + held.clear() + gc.collect() + mx.clear_cache() + held[dtype] = load_vae(sums[want], family=FAMILY, dtype=dtype) + return held[dtype] + + return get + + +def _full(m) -> list: + return [f"{h}x{w}" for h, w in m["inputs"]["sizes"]] + + +@pytest.mark.host_gb(16) +def test_the_manifest_is_the_moss_families(): + m = committed(NAME) + assert m["family"] == FAMILY + assert m["argv"][m["argv"].index("--family") + 1] == FAMILY + assert m["files"]["vae"]["sha256"] == load_json("manifest.json")[ + HEADER]["sha256"] + assert m["files"]["vae_config"]["sha256"] == load_json("manifest.json")[ + CONFIG]["sha256"] + + +@official +def test_the_file_is_the_inventoried_one(check5, vaes): + from gmlx.image.vae import read_header, reader + + path = vaes("bf16").source + want = load_json("manifest.json")[HEADER] + assert want["sha256"] == check5[1]["files"]["vae"]["sha256"] + assert os.path.getsize(path) == want["size"] + assert read_header(path).layout == "diffusers" + raw, _ = reader.safetensors_header(path) + raw.pop("__metadata__", None) + assert sorted([n, v["shape"], v["dtype"]] for n, v in raw.items()) \ + == load_json(HEADER)["tensors"] + config = Path(path).with_name("config.json") + if config.is_file(): + assert sha256(config) == check5[1]["files"]["vae_config"]["sha256"] + + +@official +def test_fp32_decode_and_encode_are_within_the_gates(check5, vaes, request): + arrays, m = check5 + vae = vaes("fp32") + failed = [] + for s in _full(m): + t = m["tolerances"]["sizes"][s] + dec = _decode(vae, arrays[f"latents_{s}"], tiled=False) + lat = _latents(vae, arrays[f"image_{s}"], tiled=False) + d = fx.max_abs(dec, arrays[f"decoded_{s}"]) + e = fx.max_abs(lat, arrays[f"latents_{s}"]) + print(f"[check5] moss_image {s} fp32, TF32 " + f"{'on' if _tf32() else 'off'}: decode max abs {d:.3e} of " + f"gate {t['decode_fp32_max_abs']:.3e}, latents max abs " + f"{e:.3e} of gate {t['encode_fp32_max_abs']:.3e}") + if d > t["decode_fp32_max_abs"] or e > t["encode_fp32_max_abs"]: + failed.append(s) + if _tf32(): + _gate_without_tf32(request) + else: + assert not failed + + +@official +def test_bf16_decode_psnr_and_encode(check5, vaes): + arrays, m = check5 + vae = vaes("bf16") + for s in _full(m): + t = m["tolerances"]["sizes"][s] + dec = _decode(vae, arrays[f"latents_{s}"], tiled=False) + db = fx.psnr(to_unit(dec), to_unit(arrays[f"decoded_{s}"]), 1.0) + rel = fx.rel_l2(_latents(vae, arrays[f"image_{s}"], tiled=False), + arrays[f"latents_{s}"]) + print(f"[check5] moss_image {s} bf16: decode PSNR {db:.2f} dB, gate " + f"{t['decode_bf16_psnr_db']:.2f} dB; latents rel L2 " + f"{rel:.3e}, gate {t['encode_bf16_rel_l2']:.3e}") + assert db >= t["decode_bf16_psnr_db"], s + assert rel <= t["encode_bf16_rel_l2"], s + + +@official +def test_a_tiled_bf16_decode_at_the_largest_size(check5, vaes): + arrays, m = check5 + s = _full(m)[-1] + dec = _decode(vaes("bf16"), arrays[f"latents_{s}"], tiled=True) + db = fx.psnr(to_unit(dec), to_unit(arrays[f"decoded_{s}"]), 1.0) + tol = m["tolerances"]["sizes"][s]["decode_bf16_psnr_db"] + print(f"[check5] moss_image {s} bf16 tiled decode PSNR {db:.2f} dB, " + f"gate {tol:.2f} dB") + assert s == "1024x1024" + assert db >= tol + + +@official +def test_a_tiled_bf16_decode_at_2048_holds_the_gate_of_1024(check5, vaes): + """The 1024 latents of the fixture tiled two by two to 2048 pixels, + held to the bf16 PSNR gate of 1024 against the tiled fp32 decode.""" + arrays, m = check5 + s = _full(m)[-1] + lat = np.tile(arrays[f"latents_{s}"], (1, 1, 2, 2)) + assert lat.shape[2:] == (128, 128) + ref = _decode(vaes("fp32"), lat, tiled=True) + got = _decode(vaes("bf16"), lat, tiled=True) + db = fx.psnr(to_unit(got), to_unit(ref), 1.0) + tol = m["tolerances"]["sizes"][s]["decode_bf16_psnr_db"] + print(f"[check5] moss_image 2048x2048 bf16 tiled decode PSNR {db:.2f} " + f"dB against the tiled fp32 decode, gate {tol:.2f} dB") + assert db >= tol + + +@official +def test_a_tiled_bf16_encode_at_the_largest_full_size(check5, vaes): + arrays, m = check5 + s = _full(m)[-1] + vae = vaes("bf16") + whole = _latents(vae, arrays[f"image_{s}"], tiled=False) + tiled = _latents(vae, arrays[f"image_{s}"], tiled=True) + rel = fx.rel_l2(tiled, whole) + a = _decode(vae, whole, tiled=False) + b = _decode(vae, tiled, tiled=False) + db = fx.psnr(to_unit(b), to_unit(a), 1.0) + print(f"[check5] moss_image {s} bf16 tiled encode rel L2 {rel:.3e} " + f"against the whole encode, round trip {db:.2f} dB against the " + "whole round trip") + assert s == "1024x1024" + assert rel < TILED_ENCODE_REL_L2 + assert db >= TILED_ROUND_TRIP_DB + + +@official +def test_the_gguf_pack_loads_the_same_weights_and_decodes_the_same( + check5, vaes, tmp_path): + """The F32 file packs losslessly with its config, and its pack decodes + the 512 latents bit for bit as the file does.""" + from gmlx.image.vae import load_vae, pack_vae + + arrays, m = check5 + a = vaes("bf16") + dst = tmp_path / "moss-image-vae-F32.gguf" + rep = pack_vae(a.source, dst, source_repo="OpenMOSS-Team/MOSS-Image-SFT", + source_revision=m["inputs"]["config_revision"], + config_json=Path(ORACLE.parent / CONFIG).read_text()) + assert rep.lossless and rep.family == FAMILY + b = load_vae(str(dst), family=FAMILY, dtype="bf16") + pa = dict(tree_flatten(a.parameters())) + pb = dict(tree_flatten(b.parameters())) + assert pa.keys() == pb.keys() + assert [k for k in pa if not mx.array_equal(pa[k], pb[k]).item()] == [] + s = _full(m)[0] + assert np.array_equal(_decode(a, arrays[f"latents_{s}"], tiled=False), + _decode(b, arrays[f"latents_{s}"], tiled=False)) diff --git a/tests/image/test_image_check6_moss_image.py b/tests/image/test_image_check6_moss_image.py new file mode 100644 index 000000000..cdaa81bad --- /dev/null +++ b/tests/image/test_image_check6_moss_image.py @@ -0,0 +1,379 @@ +"""Check 6 of the image acceptance gate for MOSS-Image, on latents: a +whole guided sampling run on injected noise at the release's defaults, +50 steps at guidance 3 with block 22 left out of the negative pass. gmlx +samples from the seed's noise on the oracle's features of the plain +prompt and of the empty one, which is the release's negative prompt. The +two have different lengths, and each pass runs at its own. The seed's +noise must equal the oracle's, and the schedule must be the release's +bit for bit. + +At 512 the run is the model as it ships, in bf16, and the latents after +each step the oracle kept and after the last must stay within twice the +oracle's own bf16 against fp32 gap after that step. The release's bf16 +run drifts from its fp32 run about 2000 times over 50 steps, so these +gates guard the bf16 numerics alone: a negative pass that runs every +block stays inside them. The exact gate guards the skip set. + +At 1024 the gate is the exact run: check 3's exact arithmetic, held to +the oracle's fp32 latents after each kept step and after the last at the +larger of EXACT_REL_L2 and twice gmlx's own fp32 floor after that step, +as check 3 holds its free rows. The run grows a difference about 3e4 +times over its 50 steps, so float32 rounding alone reaches about 1e-3 by +the last step. The floor is the exact run from the noise moved one +float32 ulp against the exact run, which the 1024 manifest records and +the test measures again. A negative pass that runs every block, a positive pass that +leaves out block 22 too, the prompt as the negative, and the schedule +without its snap to the training grid each fail it. The bf16 run at 1024 +must stay finite, and its readings print as a share of the 512 gate of +each step. + +Skipped unless the manifests of scripts/image_oracle/check6_moss_image.py +are committed, GMLX_IMAGE_FIXTURES holds their arrays and +KQUANT_TEST_GGUF_DIR the DiT GGUF the oracle read.""" +from __future__ import annotations + +import os +from pathlib import Path + +import mlx.core as mx +import block_gates as bg +import moss_gates as mg +import numpy as np +import pytest +from block_gates import exact_process, rerun_exact +from image_fixtures import image_dits +from oracle_fixture import ORACLE, committed, fixture, fx + +NAME = "check6_moss_image" +VARIANT = "sft" +GATE_SIZE = 512 +EXACT_SIZE = 1024 +SIZES = (GATE_SIZE, EXACT_SIZE) +EXACT_REL_L2 = 1e-3 + + +def _name(size: int) -> str: + return f"{NAME}_{size}" + + +def _have(size: int) -> bool: + return (ORACLE / f"{_name(size)}.manifest.json").is_file() + + +def _manifest(size: int) -> dict: + if not _have(size): + pytest.skip(f"{_name(size)}.manifest.json is not committed: run " + "scripts/image_oracle/check6_moss_image.py on the " + "MOSS-Image GGUF and commit the manifest") + return committed(_name(size)) + + +def _kept(m: dict) -> list: + return m["inputs"]["kept_steps"] + [m["inputs"]["steps_run"]] + + +def _gates() -> dict: + """The bf16 gate of the latents after each kept step and after the + last, at the gate size.""" + m = _manifest(GATE_SIZE) + tol = m["tolerances"] + out = {k: tol[f"step{k}_rel_l2"] for k in m["inputs"]["kept_steps"]} + out[m["inputs"]["steps_run"]] = tol["latents_rel_l2"] + return out + + +def _ulps(got, want) -> int: + g = np.asarray(got, np.float32).view(np.int32).astype(np.int64) + w = np.asarray(want, np.float32).view(np.int32).astype(np.int64) + return int(np.max(np.abs(g - w))) + + +class _Pipe: + """The pipeline of the fixtures' file, loaded once.""" + + def __init__(self, root): + self.root, self.path, self.pipe = root, None, None + + def release(self) -> None: + self.pipe = self.path = None + mx.clear_cache() + + def __call__(self, m: dict): + from gmlx.image import api + + if self.pipe is not None: + return self.pipe + want = m["files"]["dit"] + for p in image_dits(self.root, "moss_image"): + if Path(p).name == want["name"] and \ + os.path.getsize(p) == want["size"]: + break + else: + pytest.skip(f"{want['name']} is not under KQUANT_TEST_GGUF_DIR") + if fx.file_identity(p)["sha256"] != want["sha256"]: + pytest.fail(f"{p} is not the file the oracle read") + self.pipe = api.load(p, log=lambda *_: None) + assert self.pipe.variant == VARIANT + self.path = p + return self.pipe + + +@pytest.fixture(scope="module") +def pipe(gguf_dir): + return _Pipe(gguf_dir) + + +def _run(pipe, arrays: dict, m: dict, stop: int | None = None, + noise=None, kept: dict | None = None, **kw) -> tuple[dict, bool]: + """The relative L2 of the latents after each kept step and after the + last against the oracle's fp32 latents, and whether the latents + stayed finite after every step. ``stop`` ends the run after that + many steps, ``noise`` replaces the fixture's, and ``kept`` takes the + latents of the kept steps. ``kw`` goes to api.sample.""" + from gmlx.image import api + from gmlx.image.conditioning import Conditioning + + size = m["inputs"]["size"] + want = {k: arrays[f"latents_step{k}"] for k in m["inputs"]["kept_steps"]} + want[m["inputs"]["steps_run"]] = arrays["latents"] + cond = Conditioning(mx.array(arrays["features_positive"]), + f"check6 moss {size} plain") + neg = Conditioning(mx.array(arrays["features_negative"]), + f"check6 moss {size} empty") + kw.setdefault("neg", neg) + sigmas = kw.pop("sigmas", None) + if sigmas is None and stop is not None: + sigmas = pipe.sigmas(m["inputs"]["steps"]) + if stop is not None: + sigmas = sigmas[:stop + 1] + got: dict = {} + finite = [True] + + def on_step(i, x): + a = np.array(x) + finite[0] = finite[0] and bool(np.isfinite(a).all()) + if i + 1 in want: + got[i + 1] = fx.rel_l2(a, want[i + 1]) + if kept is not None: + kept[i + 1] = a + + api.sample(pipe, cond, height=size, width=size, + noise=mx.array(arrays["noise"] if noise is None else noise), + sigmas=sigmas, + on_step=on_step, **kw) + return got, finite[0] + + +def _line(label: str, got: dict, gates: dict) -> str: + return f"{label}: " + ", ".join( + f"step {k} {v:.2e} ({v / gates[k]:.2f} of its gate)" + for k, v in sorted(got.items())) + + +def test_the_manifests_follow_the_rule(): + """The committed manifests alone: each run is the release's call at + the variant's sampling row, the two prompts have different lengths, + every gate at the gate size is twice the gap beside it, and each + exact gate at the exact size follows the floor beside it.""" + from gmlx.image.registry import get_family + + fam = get_family("moss_image") + row = fam.sampling[VARIANT] + seen = 0 + for size in SIZES: + if not _have(size): + continue + seen += 1 + m = committed(_name(size)) + i, r, tol = m["inputs"], m["results"], m["tolerances"] + assert i["size"] == size + assert i["steps"] == row.steps + assert i["guidance_scale"] == row.cfg + assert tuple(i["uncond_skip_layers"]) == row.neg_skip + assert i["scheduler"]["shift"] == row.shift + assert i["negative_prompt"] == "" + assert i["blocks"] == 42 + rows = i["text_rows"] + assert rows["positive"] != rows["negative"] + assert (i["positive_key"], i["negative_key"]) == ( + "features_plain", "features_empty") + assert m["files"]["dit"]["name"] == "MOSS-Image-SFT-Q4_K_M.gguf" + if size == GATE_SIZE: + by = r["bf16_vs_fp32_by_step"] + assert len(by) == i["steps_run"] + for k in i["kept_steps"]: + assert tol[f"step{k}_rel_l2"] == 2 * by[k - 1], k + assert tol["latents_rel_l2"] == pytest.approx( + 2 * by[-1], rel=1e-12) + assert "exact_rel_l2" not in tol + else: + assert set(tol) == {"exact_rel_l2", "rule", "exact_by_step"} + assert tol["exact_rel_l2"] == EXACT_REL_L2 + floor = r["gmlx_exact_floor"] + assert set(floor) == {str(k) for k in _kept(m)} + assert tol["exact_by_step"] == { + k: max(EXACT_REL_L2, 2 * v) for k, v in floor.items()} + if not seen: + pytest.skip("no check 6 manifest of MOSS-Image is committed") + + +@pytest.mark.parametrize("size", SIZES) +def test_the_seed_gives_the_noise_of_the_oracle(size): + from gmlx.image.samplers import noise_from_seed + + m = _manifest(size) + arrays, _ = fixture(m["name"]) + noise = arrays["noise"] + assert noise.shape == (1, 128, size // 16, size // 16) + assert np.array_equal(noise_from_seed(m["inputs"]["seed"], noise.shape), + noise) + + +@pytest.mark.parametrize("size", SIZES) +def test_the_schedule_is_the_release_one(size): + """The schedule of a request is the release's bit for bit, and the + release's timesteps are its sigmas times 1000 in float32, as the DiT + forms them.""" + from gmlx.image.api import ImagePipeline + from gmlx.image.registry import get_family + + m = _manifest(size) + arrays, _ = fixture(m["name"]) + pipe = ImagePipeline(dit=None, family=get_family("moss_image"), + variant=VARIANT) + steps = pipe.defaults.steps + assert steps == m["inputs"]["steps"] + got, want = pipe.sigmas(steps, (size // 16) ** 2), arrays["sigmas"] + assert got.shape == want.shape + assert want.size - 1 == m["inputs"]["steps_run"] + assert _ulps(got, want) == 0 + ts = want[:-1] * np.float32(1000) + assert _ulps(ts, arrays["timesteps"]) == 0 + + +@pytest.mark.integration +def test_check6_latents_at_512(pipe): + m = _manifest(GATE_SIZE) + arrays, _ = fixture(m["name"]) + gates = _gates() + got, finite = _run(pipe(m), arrays, m) + print(_line(f"moss check6 {GATE_SIZE}", got, gates)) + assert finite + assert sorted(got) == sorted(gates) + over = {k: (v, gates[k]) for k, v in got.items() if v > gates[k]} + assert not over, over + + +@pytest.mark.integration +def test_check6_bf16_at_1024_stays_finite(pipe): + """The run as it ships at 1024, read against the 512 gates.""" + m = _manifest(EXACT_SIZE) + arrays, _ = fixture(m["name"]) + gates = _gates() + got, finite = _run(pipe(m), arrays, m) + print(_line(f"moss check6 bf16 {EXACT_SIZE} against the {GATE_SIZE} " + "gates", got, gates)) + assert sorted(got) == sorted(gates) + assert finite + + +def _skip_positive_too(mp): + """Block 22 left out of the positive pass as well.""" + from gmlx.image import api + + real = api.denoise_step + + def both(pipe, latents, sigma, conds, *, skip=(), **kw): + return real(pipe, latents, sigma, conds, skip=skip or (22,), **kw) + + mp.setattr(api, "denoise_step", both) + + +def _unsnapped(steps: int) -> np.ndarray: + from gmlx.image import schedule + + v = schedule.linspace_f32(1.0, 5.0 / 1004, steps) + return np.append(schedule.shift_static(v, 5.0), np.float32(0)) + + +def _floor(p, arrays: dict, m: dict, clean: dict) -> dict: + """The exact run from the noise moved one float32 ulp against the + exact run ``clean``, at each kept step.""" + up = np.nextafter(arrays["noise"], np.float32(np.inf)) + moved: dict = {} + with mg._exact(p.dit): + _run(p, arrays, m, noise=up, kept=moved) + return {k: fx.rel_l2(moved[k], clean[k]) for k in sorted(clean)} + + +def _exact_gate(m: dict, now: dict) -> dict: + """The gate of each kept step from the recorded floor, after checking + that the floor measured now gives a gate within 1.25 times of it.""" + rec = m["results"].get("gmlx_exact_floor") + shown = ", ".join(f"step {k} {v!r}" for k, v in now.items()) + print(f"moss check6 exact {EXACT_SIZE}, fp32 floor: {shown}") + assert rec is not None, f"{m['name']} records no gmlx_exact_floor" + gate = {int(k): v for k, v in m["tolerances"]["exact_by_step"].items()} + assert sorted(gate) == sorted(now) + for k, v in now.items(): + assert max(EXACT_REL_L2, 2 * v) <= 1.25 * gate[k], (k, v, gate[k]) + return gate + + +# Wrong runs the exact gate must fail: the fault, the steps it runs (all +# with None) and the arguments it gives api.sample. +EXACT_FAULTS = { + "every block in the negative pass": (None, 10, {"neg_skip": ()}), + "block 22 left out of the positive pass too": + (_skip_positive_too, 10, {}), + "the prompt as the negative": (None, 10, {"neg": "cond"}), + "the schedule not snapped": (None, None, {"sigmas": "unsnapped"}), +} + + +@pytest.mark.integration +def test_check6_exact_at_1024(pipe, request, monkeypatch): + """The run in the oracle's fp32 arithmetic within its gate of the + oracle's latents after each kept step and after the last, with every + K-quant product dense, and each wrong run of EXACT_FAULTS over the + gate at a kept step.""" + from gmlx.image import api + from gmlx.image.conditioning import ByteLRU, Conditioning + + m = _manifest(EXACT_SIZE) + arrays, _ = fixture(m["name"]) + if not exact_process(): + pipe.release() + return rerun_exact(request, "moss check6 exact") + p = pipe(m) + n = mg.exact_floats(p.dit, pipe.path) + clean: dict = {} + with mg._exact(p.dit): + got, finite = _run(p, arrays, m, kept=clean) + gate = _exact_gate(m, _floor(p, arrays, m, clean)) + print(_line(f"moss check6 exact {EXACT_SIZE}, {n} float tensors in " + "fp32", got, gate)) + assert finite and sorted(got) == sorted(gate) + over = {k: (v, gate[k]) for k, v in got.items() if v > gate[k]} + for name, (apply, stop, kw) in EXACT_FAULTS.items(): + kw = dict(kw) + if kw.get("neg") == "cond": + kw["neg"] = Conditioning(mx.array(arrays["features_positive"]), + f"check6 moss {EXACT_SIZE} plain neg") + if isinstance(kw.get("sigmas"), str): + kw["sigmas"] = _unsnapped(m["inputs"]["steps"]) + with monkeypatch.context() as mp: + # The memos of the clean run would hand the fault its + # modulations and its caption. + mp.setattr(p, "modulations", ByteLRU(api.MODULATION_BYTES)) + mp.setattr(p, "captions", ByteLRU(api.CAPTION_BYTES)) + if apply is not None: + apply(mp) + # A fault may leave a block out of both passes, so the check + # that every product ran dense is the clean run's alone. + with bg.exact_products(p.dit) as used: + bad, _ = _run(p, arrays, m, stop=stop, **kw) + assert not used.kernel, (name, used.kernel[:3]) + print(_line(f"moss check6 exact {EXACT_SIZE}, {name}", bad, gate)) + assert bad and any(v > gate[k] for k, v in bad.items()), (name, bad) + assert not over, over diff --git a/tests/image/test_image_check7_moss_image.py b/tests/image/test_image_check7_moss_image.py new file mode 100644 index 000000000..0c21f7119 --- /dev/null +++ b/tests/image/test_image_check7_moss_image.py @@ -0,0 +1,380 @@ +"""Check 7 of the image acceptance gate for MOSS-Image, step half: check +3's step at 512 with a synthetic adapter applied live under a patch +table, against the release's MOSSImageTransformer with the same adapter +applied as PEFT applies it, under the three gates of check 3. + +No community MOSS-Image adapter is known. +scripts/image_oracle/check7_moss_image_lora.py writes the adapter from a +seed: one pair on every linear layer of blocks 0, 21 and 41, the fused +proj.ffn_up_gate among them, on every linear layer of text refiner block +1, its fused qkv among them, and on six layers outside the blocks, at +ranks 4 to 16 with an alpha of half the rank, each moving its leaf by +about 8 percent. The fixture holds the pairs, and the test writes them +in a dialect. Every row is in fp32, with the oracle's own bf16 against +fp32 gap of each and check 3's rule. + +The product gate holds each row to twice the oracle's gap, and the base +model and the adapter at strength 0.8 miss it. The velocity of the +request path, api.denoise_step with the kernel routes at their default, +passes its gate too. The exact gate runs the step in the oracle's fp32 +arithmetic in a child process and holds the teacher-forced rows to +EXACT_TOL, and the free rows and the velocity to the gates the manifest +records from gmlx's floors under the table; each wrong table of +``wrong_tables`` must miss it. The site gate holds each op of every +block and each linear leaf, the adapted leaves included, to the stock op +in float32 on the same input with the table's delta. The delta half runs +on the CPU in test_image_moss_lora_map.py. + +Skipped unless the manifest is committed, GMLX_IMAGE_FIXTURES holds its +arrays and KQUANT_TEST_GGUF_DIR the DiT GGUF.""" +from __future__ import annotations + +import os +from pathlib import Path + +import mlx.core as mx +import numpy as np +import pytest + +import moss_gates as mg +import qwen_gates as qg +from block_gates import exact_process, rerun_exact +from image_fixtures import image_dits +from lora_fixtures import DIALECTS, write_adapter +from oracle_fixture import ORACLE, committed, fixture, fx + +SYNTHETIC = "check7_moss_image_synthetic_512" +SOURCE = "check3_moss_image_512" +# The dialect each gate reads the adapter in: one that stores the alpha, +# and one that does not, with the alpha in its up rows. +EXACT_DIALECT, SITE_DIALECT = "kohya", "diffusers" +FUSED = ("layers.21.proj.ffn_up_gate", "text_embedder.blocks.1.qkv") + + +def _manifest() -> dict: + if not (ORACLE / f"{SYNTHETIC}.manifest.json").is_file(): + pytest.skip(f"{SYNTHETIC}.manifest.json is not committed: run " + "scripts/image_oracle/check7_moss_image_lora.py and " + "commit the manifest it writes") + return committed(SYNTHETIC) + + +@pytest.fixture(scope="module") +def synthetic(): + _manifest() + return fixture(SYNTHETIC) + + +@pytest.fixture(scope="module") +def dit_path(gguf_dir): + want = _manifest()["files"]["dit"] + for p in image_dits(gguf_dir, "moss_image"): + if Path(p).name == want["name"] and os.path.getsize(p) == want["size"]: + if fx.file_identity(p)["sha256"] != want["sha256"]: + pytest.fail(f"{p} is not the file the oracle read") + return p + pytest.skip(f"{want['name']} is not under KQUANT_TEST_GGUF_DIR") + + +# The float tensors the exact gate's process put in float32. +_EXACT: dict = {} + + +@pytest.fixture(scope="module") +def pipe(dit_path): + from gmlx.image import api + + p = api.load(dit_path, log=lambda *_: None) + if exact_process(): + # Before a table wraps the leaves, so the names are the file's. + _EXACT["floats"] = mg.exact_floats(p.dit, dit_path) + return p + + +def synthetic_adapter(arrays: dict, m: dict, dialect: str) -> dict: + """The fixture's pairs as an adapter in the GGUF names, for + ``dialect``: a dialect that stores no alpha takes it in its up rows.""" + keeps = DIALECTS[dialect][4] + out = {} + for name, alpha in m["inputs"]["lora_pair_alphas"].items(): + down, up = arrays[f"down.{name}"], arrays[f"up.{name}"] + if not keeps: + up = (up * np.float32(alpha / down.shape[0])).astype(np.float32) + alpha = None + out[name] = (down, up, alpha) + return out + + +def _adapter(pipe, tmp, dialect: str): + from gmlx.image import lora + + arrays, m = fixture(SYNTHETIC) + path = write_adapter(Path(tmp) / f"synthetic-{dialect}.safetensors", + dialect, synthetic_adapter(arrays, m, dialect)) + ad = lora.load_adapter(pipe, path) + assert ad.text == () and len(ad.dit) == m["inputs"]["lora_pairs"] + return ad + + +def _entries(ad, strength: float) -> list: + from gmlx.image import lora + + return lora.table_entries([(ad, strength)]) + + +def _under(pipe, tid: str | None, read, arrays: dict): + """``read(dit, arrays)`` with the table ``tid`` active.""" + from gmlx.image.lora import install + + if tid is None: + return read(pipe.dit, arrays) + prev = install.activate(pipe, tid) + try: + return read(pipe.dit, arrays) + finally: + install.restore(pipe, prev) + + +def test_the_manifest_reads_check3s_step_and_follows_the_rule(): + """The committed manifest alone: the fixture takes check 3's step, + holds a pair for every target it names, the fused weights, the + refiner and the layers outside the blocks among them, every row + tolerance is twice the gap beside it, and the exact gate's free-row + and velocity tolerances follow their floors.""" + src = committed(SOURCE) + m = _manifest() + assert m["inputs"]["source"] == SOURCE + for k in ("sigma", "seed", "size", "blocks", "text_rows"): + assert m["inputs"][k] == src["inputs"][k], k + assert m["files"]["dit"] == src["files"]["dit"] + for k, g in m["results"]["rows_bf16_vs_fp32"].items(): + assert m["tolerances"][k]["rel_l2"] == 2 * g["rel_l2"], k + mg.exact_tols(m) + targets = m["inputs"]["targets"] + assert sorted(m["inputs"]["lora_pair_alphas"]) == targets + assert len(targets) == m["inputs"]["lora_pairs"] == 29 + assert set(FUSED) | {"final_linear_video", "video_embedder", + "shared_block_modulation"} <= set(targets) + assert {f"{h}.{t}" for t in targets for h in ("down", "up")} <= set( + m["tensors"]) + assert "lora" not in m["files"] + + +# The product gate. + +@pytest.mark.integration +def test_check7_adapter_step_rows_at_512(synthetic, pipe, tmp_path): + from gmlx.image import api + + arrays, m = synthetic + ad = _adapter(pipe, tmp_path, EXACT_DIALECT) + tid = api.define_patch_table(pipe, _entries(ad, m["inputs"]["strength"])) + got = _under(pipe, tid, mg.measures, arrays) + tol = m["tolerances"] + ratio = {k: v / tol[k]["rel_l2"] for k, v in got.items() + if not k.startswith("free ")} + top = max(ratio, key=ratio.get) + print(f"moss check7 512: {len(ad.dit)} pairs, {len(ratio)} rows, worst " + f"{got[top]:.2e} at {top}, {ratio[top]:.2f} of its gate; " + f"velocity {got['velocity']:.2e} of gate " + f"{tol['velocity']['rel_l2']:.2e}") + out = mg.missed(got, tol) + assert not out, out + + +@pytest.mark.integration +@pytest.mark.parametrize("wrong", ["base", "strength 0.8"]) +def test_check7_rows_see_the_adapter_and_its_strength(synthetic, pipe, + tmp_path, wrong): + from gmlx.image import api + + arrays, m = synthetic + tid = None + if wrong != "base": + ad = _adapter(pipe, tmp_path, EXACT_DIALECT) + tid = api.define_patch_table( + pipe, _entries(ad, 0.8 * m["inputs"]["strength"])) + out = mg.missed(_under(pipe, tid, mg.measures, arrays), m["tolerances"]) + worst = max(out, key=lambda k: out[k][0] / out[k][1]) if out else None + print(f"moss check7 {wrong}: {len(out)} rows miss" + + (f", worst {worst} at {out[worst][0] / out[worst][1]:.1f} of " + "its gate" if worst else "")) + assert out + + +@pytest.mark.integration +def test_check7_velocity_of_the_request_path(synthetic, pipe, tmp_path): + from gmlx.image import api + + arrays, m = synthetic + ad = _adapter(pipe, tmp_path, SITE_DIALECT) + tid = api.define_patch_table(pipe, _entries(ad, m["inputs"]["strength"])) + v = api.denoise_step(pipe, mx.array(arrays["latents"]), + float(arrays["sigma"]), + api.Conditioning(mx.array(arrays["features"])), + patch_table=tid) + got = fx.compare(np.array(v), arrays["velocity"]) + tol = m["tolerances"]["velocity"] + print(f"moss check7 request velocity: rel L2 {got['rel_l2']:.3e} of " + f"{tol['rel_l2']:.3e}, cosine {got['cosine']:.6f} of " + f"{tol['cosine_min']:.6f}") + assert got["rel_l2"] <= tol["rel_l2"] + assert got["cosine"] >= tol["cosine_min"] + + +# The exact gate. + +def _block(path: str): + parts = path.split(".") + return int(parts[1]) if parts[0] == "layers" else None + + +def wrong_tables(entries: list, strength: float) -> dict: + """Wrong tables of an adapter's ``entries`` at ``strength``: the + strength a tenth low, the pairs of the middle adapted block left out + and of the last, and the middle block's pairs on the block before it, + which the adapter leaves alone.""" + blocks = sorted({b for p, _, _ in entries + if (b := _block(p)) is not None}) + mid, last = blocks[len(blocks) // 2], blocks[-1] + assert mid - 1 not in blocks + + def scaled(k): + return [(p, band, [(key, s * k) for key, s in patches]) + for p, band, patches in entries] + + return { + "strength a tenth low": scaled(0.9), + f"block {mid} left out": [e for e in entries if _block(e[0]) != mid], + f"block {last} left out": [e for e in entries + if _block(e[0]) != last], + f"block {mid} on block {mid - 1}": [ + (p.replace(f"layers.{mid}.", f"layers.{mid - 1}.", 1) + if _block(p) == mid else p, band, patches) + for p, band, patches in entries]} + + +@pytest.mark.integration +def test_check7_exact_gate_at_512(pipe, tmp_path, request): + """Every measure within its exact gate under the right table, the + floors measured now within the manifest's, and each wrong table of + ``wrong_tables`` over the gate.""" + from gmlx.image import api + + arrays, m = fixture(_manifest()["name"]) + if not exact_process(): + return rerun_exact(request, "moss check7 exact") + assert _EXACT.get("floats", 0) > 0 + ad = _adapter(pipe, tmp_path, EXACT_DIALECT) + strength = m["inputs"]["strength"] + entries = _entries(ad, strength) + tid = api.define_patch_table(pipe, entries) + label = "moss check7 exact" + tol = mg.check_floor(m, _under(pipe, tid, mg.exact_free_floor, arrays), + label) + got = _under(pipe, tid, mg.exact_measures, arrays) + print(mg.exact_line(f"{label}, {len(ad.dit)} pairs", got, tol)) + over = mg.exact_over(got, tol) + weakest = None + for fault, table in wrong_tables(entries, strength).items(): + bad = _under(pipe, api.define_patch_table(pipe, table), + mg.exact_measures, arrays) + line, top = mg.fault_line(f"{label}, {fault}", bad, tol) + print(line) + assert top > 10, (fault, top) + if weakest is None or top < weakest[0]: + weakest = (top, fault) + print(f"{label}: the weakest wrong table is {weakest[1]} at " + f"{weakest[0]:.0f} of the gate") + assert not over, over + + +# The site gate. + +def table_sites(pipe, arrays: dict, tid: str) -> tuple: + """The sites of one free-running production step under the table + ``tid``, read against the exact ops with the table's delta.""" + from gmlx.image.lora import install + + built = install.build_table(pipe, tid, pipe.dit.act_dtype, "block") + sites, routed, heads, skipped = {}, {}, {}, set() + prev = install.activate(pipe, tid) + try: + with mg.step_sites(pipe.dit, sites, routed=routed, heads=heads, + built=built, skipped=skipped): + vel = mg.free_rows(pipe.dit, arrays)["velocity"] + finally: + install.restore(pipe, prev) + return sites, routed, skipped, vel + + +def _site_setup(pipe, tmp): + from gmlx.image import api + + arrays, m = fixture(_manifest()["name"]) + entries = _entries(_adapter(pipe, tmp, SITE_DIALECT), + m["inputs"]["strength"]) + return arrays, m, entries, api.define_patch_table(pipe, entries) + + +@pytest.mark.integration +def test_check7_site_gate_at_512(pipe, tmp_path): + """Every op of every block and every linear leaf a step calls under + the table, each adapted leaf with its delta, is within its site gate + of the stock op in float32 on the same input, and that run's velocity + passes the product gate.""" + arrays, m, entries, tid = _site_setup(pipe, tmp_path) + sites, routed, skipped, vel = table_sites(pipe, arrays, tid) + missing = mg.missing_ops(pipe.dit, sites, skipped=skipped) + assert not missing, missing + adapted = {p for p, _, _ in entries} + # An adapted leaf is held to the count of a leaf under its adapter, or + # to the lora op's own. The time vectors are read past their call, + # and the exact gate holds them. + held = {k[0] for k, (_, g) in sites.items() if len(k) == 2 + and g >= 2 * 3 * qg.SITE_UNIT} + called = adapted - skipped - set(mg.NOT_CALLED) + assert called <= held, sorted(called - held)[:4] + assert set(FUSED) <= called + lora_on = sum(v == "quantized_matmul_lora" for v in routed.values()) + print(f"moss check7 sites: {len(adapted)} adapted leaves, {lora_on} " + f"calls on the lora op, leaves on a float32 input " + f"{sorted(skipped)}") + worst = qg.worst_by_op(sites) + print(f"moss check7 sites: {len(sites)} readings, the largest over its " + "gate by op: " + + ", ".join(f"{k} {v:.2f}" for k, v in sorted(worst.items()))) + over = qg.over_gate(sites) + assert not over, over + rel = mg.rel(vel, arrays["velocity"]) + assert rel <= m["tolerances"]["velocity"]["rel_l2"], rel + + +@pytest.mark.integration +def test_check7_a_leaf_that_drops_its_delta_fails_the_site_gate( + pipe, tmp_path, monkeypatch): + """The adapter's leaves of its middle block run their base product + alone: their sites read the table's delta as error, and no other leaf + moves past its gate.""" + from gmlx.image.lora import install + + arrays, _, entries, tid = _site_setup(pipe, tmp_path) + blocks = sorted({b for p, _, _ in entries + if (b := _block(p)) is not None}) + mid = blocks[len(blocks) // 2] + dropped = {p for p, _, _ in entries if _block(p) == mid} + real = install.ImageLoRALinear.__call__ + + def call(self, x): + return self.base(x) if self._path in dropped else real(self, x) + + monkeypatch.setattr(install.ImageLoRALinear, "__call__", call) + sites, _, _, _ = table_sites(pipe, arrays, tid) + over = qg.over_gate(sites) + leaves = {k[0] for k in over if len(k) == 2} + low = min(sites[k][0] / sites[k][1] for k in sites + if len(k) == 2 and k[0] in dropped) + print(f"moss check7 sites, block {mid}'s deltas dropped: {len(over)} " + f"readings over the gate, {len(leaves)} of {len(dropped)} dropped " + f"leaves, the least at {low:.1f} of its gate") + assert leaves == dropped, sorted(dropped - leaves) diff --git a/tests/image/test_image_cli.py b/tests/image/test_image_cli.py index 27a43a64a..93e5aa27b 100644 --- a/tests/image/test_image_cli.py +++ b/tests/image/test_image_cli.py @@ -5,6 +5,7 @@ fit check are replaced. CPU.""" from __future__ import annotations +import dataclasses import json import os import re @@ -768,6 +769,53 @@ def spy(*a, **k): assert np.isfinite(latents[0]).all() +def test_bench_via_steps_on_moss_image_skips_the_negative_blocks_of_sample( + env, capsys, monkeypatch): + """A guided MOSS-Image step leaves the variant's blocks out of its + negative pass, so the bench sends the negative rows apart, with that + skip, and gives the latents of sample.""" + from image_fixtures import tiny_moss_model + + fam = get_family("moss_image") + real_defaults = api.sampler_defaults + + def defaults(family, variant): + d = real_defaults(family, variant) + return dataclasses.replace(d, neg_skip=(1,)) if d else d + + def load(dit, *, text_encoder=None, vae=None, log=print, **kw): + pipe = api.ImagePipeline( + dit=tiny_moss_model(), family=fam, variant="sft", vae=object(), + load_seconds={"dit": 0.0}, log=log) + pipe.decoded, pipe.encoded_images = [], [] + env.pipes.append(pipe) + return pipe + + rng = np.random.default_rng(7) + feats = {neg: rng.standard_normal((6 - 2 * neg, 16)).astype(np.float32) + for neg in (False, True)} + monkeypatch.setattr(api, "sampler_defaults", defaults) + monkeypatch.setattr(api, "load", load) + monkeypatch.setattr(api, "encode_text", lambda pipe, prompt, **k: + api.Conditioning(mx.array(feats[prompt == ""]), + f"moss:{prompt}")) + view = SimpleNamespace(tensors={ + "a": SimpleNamespace(shape=(4, 4), codec="Q4_K")}) + monkeypatch.setattr(image_preflight, "identify_dit", + lambda p, **k: SimpleNamespace( + family=fam, view=view, variant="sft", + config={})) + rc = umbrella_main(["image", "bench"] + env.base[2:] + [ + "--sizes", "64x48", "--steps", "3", "--runs", "1", "--warmup", "0", + "--seed", "5", "--via-steps", "--json", "--guidance", "3"]) + io = capsys.readouterr() + assert rc == 0, io.err + rec = json.loads(io.out) + assert rec["transport"]["equal_to_direct"] is True + # Two requests a step: the positive row, then the negative one. + assert [len(r["transport_ms"]) for r in rec["runs"]] == [6] + + def test_bench_via_steps_refuses_adapters(env, capsys): assert umbrella_main(["image", "bench"] + env.base[2:] + [ "--via-steps", "--adapter", "a.safetensors"]) == 2 diff --git a/tests/image/test_image_fingerprint.py b/tests/image/test_image_fingerprint.py index eab9136c7..6db74b9b2 100644 --- a/tests/image/test_image_fingerprint.py +++ b/tests/image/test_image_fingerprint.py @@ -29,11 +29,12 @@ def test_families_sort_by_order_and_name(): "qwen_image21", "flux2_klein_9b_kv", "flux2_klein_9b", "flux2_klein", - "flux3_klein"] + "flux3_klein", "moss_image"] assert get_family("z_image").title == "Z-Image" with pytest.raises(ValueError, match="known: z_image, krea2, qwen_image21, " "flux2_klein_9b_kv, flux2_klein_9b, " - "flux2_klein, flux3_klein"): + "flux2_klein, flux3_klein, " + "moss_image"): get_family("flux") assert is_image_arch("lumina2") and is_image_arch("qwen_image") assert is_image_arch("krea2") diff --git a/tests/image/test_image_flux2_family.py b/tests/image/test_image_flux2_family.py index 95834cb0e..97053d8b4 100644 --- a/tests/image/test_image_flux2_family.py +++ b/tests/image/test_image_flux2_family.py @@ -662,13 +662,15 @@ def test_the_9b_rows_leave_every_other_family_as_it_was(): """The peak rows, plans, VAE passes, recipes, default quants and sizes, VAE pairings and registry order of every family but the 9B and KV, against what they gave before those rows. A change to another - family's prices rewrites the fixture with family_prices.py.""" + family's prices rewrites the fixture with family_prices.py. MOSS-Image + came after the fixture and is left out.""" from family_prices import family_prices from gmlx.image import catalog before = load_json("memfit/family_prices.json")["prices"] - assert family_prices(skip=("flux2_klein_9b", "flux2_klein_9b_kv")) == before + assert family_prices( + skip=("flux2_klein_9b", "flux2_klein_9b_kv", "moss_image")) == before assert {"flux2_klein", "flux3_klein", "krea2", "qwen_image21", "z_image"} <= set(before) # The 9B and KV join the pairings of the FLUX.2 autoencoder and no diff --git a/tests/image/test_image_krea2_family.py b/tests/image/test_image_krea2_family.py index 5c1c76404..c8a6861e9 100644 --- a/tests/image/test_image_krea2_family.py +++ b/tests/image/test_image_krea2_family.py @@ -48,6 +48,18 @@ def test_the_raw_file_stores_the_projector_flattened(): assert fixture_view(TURBO).shape(kf.PROJECTOR) == (12, 1) +def test_a_stored_shape_that_is_not_integers_is_refused(): + """A flattened projector whose stored shape is not a list of integers + is refused with the reason, not a crash.""" + kv = {**load_json(RAW)["kv"], kf.ORIG_SHAPE: ["1", "twelve"]} + view = fixture_view(RAW, kv=kv) + reason = (f"{kf.PROJECTOR} is flattened, and its stored shape " + "['1', 'twelve'] is not a list of integers") + for kw in ({}, {"family": "krea2"}): + msg = _refused(view, **kw) + assert "Krea 2 relative" in msg and reason in msg, kw + + def test_a_file_without_the_final_mlp_is_the_plain_layout(): got = fp.identify(fixture_view(RAW, drop=kf.LAST_MLP)) assert got.layout == "plain" and not got.config["last_mlp"] diff --git a/tests/image/test_image_memfit.py b/tests/image/test_image_memfit.py index f4451d251..67050a7ac 100644 --- a/tests/image/test_image_memfit.py +++ b/tests/image/test_image_memfit.py @@ -576,8 +576,8 @@ class _TinyVae: def test_unknown_family_is_refused(z_view): with pytest.raises(ValueError, match="known: flux2_klein, " "flux2_klein_9b, flux2_klein_9b_kv, " - "flux3_klein, krea2, qwen_image21, " - "z_image"): + "flux3_klein, krea2, moss_image, " + "qwen_image21, z_image"): plan_memory(z_view, 0, 0, family="flux", height=8, width=8, budget=1e12, overcommit=False) @@ -831,7 +831,7 @@ def step(rows, conds=1): assert p.terms["step_peak"] == step(image + 640) -SHORT = ("z_image", "krea2", "qwen_image21", "flux3_klein") +SHORT = ("z_image", "krea2", "qwen_image21", "flux3_klein", "moss_image") def _short_cells(family: str) -> list[dict]: @@ -860,7 +860,8 @@ def cell(h, w, text, gib, refs=0, px=0, guided=False, prefill=None, c["guided"], c["prefill_gib"], c["state_bytes"])) return out path, short = {"krea2": (KREA2_CELLS, 36), - "flux3_klein": (FLUX3_CELLS, 80)}[family] + "flux3_klein": (FLUX3_CELLS, 80), + "moss_image": (MOSS_CELLS, 80)}[family] return [cell(c["size"], c["size"], c["text_rows"], max(c["steps_gib"]), c.get("refs", 0), c.get("ref_px", 0), c["guided"]) for c in json.loads(path.read_text())["cells"] @@ -910,7 +911,7 @@ def test_a_short_prompt_plan_covers_every_measured_cell(family): row = memfit.PEAKS[family] cells = _short_cells(family) assert len(cells) == {"z_image": 3, "krea2": 3, "qwen_image21": 79, - "flux3_klein": 13}[family] + "flux3_klein": 13, "moss_image": 11}[family] room = [_short_room(family, c) for c in cells] def formula(c, caption): @@ -964,6 +965,83 @@ def test_the_flux3_encode_term_covers_every_measured_encode(): encodes[0]["gib"] * GiB +MOSS_CELLS = QWEN_CELLS.with_name("moss_image_cells.json") + + +def _moss_room(row, c: dict) -> float: + """Bytes the row's step leaves above the cell's largest step.""" + rows = memfit.image_rows(row, c["size"], c["size"], c["text_rows"]) + plan = row.step.bytes(rows * row.width * 2, 2 if c["guided"] else 1) + return plan - max(c["steps_gib"]) * GiB + + +def test_the_moss_row_covers_every_measured_step(): + """Sizes to 2048 pixels with 80 and 4096 text rows, and one with 640, + and guided steps to 2048 pixels. A text row reads more than an image + row, so the cells of the most text sit closest to the line.""" + row = memfit.PEAKS["moss_image"] + assert row.step is memfit.MOSS_IMAGE_STEP + assert row.width == 24576 and row.caption_rows == row.text_rows == 4096 + assert memfit.prices_prompt("moss_image") + cells = json.loads(MOSS_CELLS.read_text())["cells"] + assert {c["size"] for c in cells} == {256, 512, 768, 1024, 1280, 1536, + 2048} + assert {c["text_rows"] for c in cells} == {80, 640, 4096} + guided = [(c["size"], c["text_rows"]) for c in cells if c["guided"]] + assert guided == [(512, 80), (1024, 80), (1536, 80), (1024, 4096), + (2048, 80), (2048, 4096)] + room = [_moss_room(row, c) for c in cells] + assert min(room) >= ROOM + assert min(room) <= 0.35 * GiB and max(room) <= 1.5 * GiB + near = min(cells, key=lambda c: _moss_room(row, c)) + assert (near["size"], near["text_rows"]) == (256, 4096) + + +def test_a_thinner_moss_row_misses_a_measured_cell(): + row = memfit.PEAKS["moss_image"] + cells = json.loads(MOSS_CELLS.read_text())["cells"] + + def worst(**change) -> float: + thin = dataclasses.replace(row, step=dataclasses.replace( + row.step, **change)) + return min(_moss_room(thin, c) for c in cells) + + assert worst() >= ROOM + assert worst(base=1.5 * GiB) < ROOM + assert worst(k=5.0) < ROOM + # A guided step holds both passes, more than a second row of slope 1. + assert worst(row_k=1.0) < ROOM + # FLUX.3 klein's step at MOSS-Image's rows. + assert worst(base=memfit.FLUX3_KLEIN_STEP.base, + k=memfit.FLUX3_KLEIN_STEP.k) < ROOM + # Without the text rows in the formula the most text is priced under + # its reading. + alone = dataclasses.replace(row, caption_rows=0) + assert min(_moss_room(alone, c) for c in cells) < ROOM + + +def test_the_moss_encode_term_covers_every_measured_encode(): + """Both encoder files, Q4_K_M and Q8_0, at 80 to 4096 rows.""" + row = memfit.PEAKS["moss_image"] + encodes = json.loads(MOSS_CELLS.read_text())["encode"] + assert sorted({e["file"] for e in encodes}) == [ + "Qwen3.5-9B-Q4_K_M.gguf", "Qwen3.5-9B-Q8_0.gguf"] + assert sorted({e["rows"] for e in encodes}) == [80, 640, 2048, 4096] + kw = dict(family="moss_image", height=512, width=512, + budget=BUDGET_128, overcommit=False) + room = [] + for e in encodes: + p = plan_memory(0, ENCODER, 0, caption_tokens=e["rows"], **kw) + assert p.terms["encode_peak"] == int( + row.encode_base + + e["rows"] * row.encode_row_bytes * memfit.ENCODE_LIVE) + room.append(p.terms["encode_peak"] - e["gib"] * GiB) + assert min(room) >= ROOM + # Without the base the shortest prompt is priced under its reading. + assert 80 * row.encode_row_bytes * memfit.ENCODE_LIVE < max( + e["gib"] for e in encodes if e["rows"] == 80) * GiB + + FLUX2_CELLS = QWEN_CELLS.with_name("flux2_klein_cells.json") diff --git a/tests/image/test_image_moss_encoder.py b/tests/image/test_image_moss_encoder.py new file mode 100644 index 000000000..6881c7e5d --- /dev/null +++ b/tests/image/test_image_moss_encoder.py @@ -0,0 +1,432 @@ +"""The MOSS-Image text encoder: the template against the release's chat +template, the spec, and the read after the final norm on a tiny Qwen3.5 +model, whose recurrent layers the tap cannot run. The real-file rows +need the unsloth Qwen3.5-9B GGUFs under KQUANT_TEST_GGUF_DIR.""" +from __future__ import annotations + +import hashlib +import os +from pathlib import Path +from types import SimpleNamespace + +import mlx.core as mx +import numpy as np +import pytest + +from image_fixtures import FIXTURES, load_json + +from gmlx.image import text_encoder as te + +_TEMPLATE = "moss_image_chat_template.jinja" +_PROMPTS = ["a red bicycle", "", " ", "{} {0} {{x}}", "line one\nline two", + "caf\u00e9 \u732b", " a cat\n\n", "\ta\u3000"] +_FILES = ("Qwen3.5-9B-Q8_0.gguf", "Qwen3.5-9B-Q4_K_M.gguf") + + +def _render_official(prompt: str, **kw) -> str: + from jinja2.ext import loopcontrols + from jinja2.sandbox import ImmutableSandboxedEnvironment + + env = ImmutableSandboxedEnvironment( + trim_blocks=True, lstrip_blocks=True, extensions=[loopcontrols]) + tmpl = env.from_string((FIXTURES / _TEMPLATE).read_text()) + return tmpl.render(messages=[{"role": "user", "content": prompt}], + add_generation_prompt=True, **kw) + + +@pytest.mark.parametrize("prompt", _PROMPTS) +def test_the_template_is_the_official_render_with_thinking_on(prompt): + """The release renders one user turn with the generation prompt and + enable_thinking True, and tokenizes it with no special tokens.""" + spec = te.encoder_spec("moss_image") + assert te.render(spec, prompt) == _render_official( + prompt, enable_thinking=True) + assert te.render(spec, prompt) != _render_official( + prompt, enable_thinking=False) + + +def test_the_template_fixture_is_pinned(): + row = load_json("manifest.json")[_TEMPLATE] + data = (FIXTURES / _TEMPLATE).read_bytes() + assert hashlib.sha256(data).hexdigest() == row["sha256"] + assert (row["repo"], row["file"]) == ("OpenMOSS-Team/MOSS-Image-SFT", + "tokenizer/chat_template.jinja") + from gmlx.image.registry import get_family + + assert row["revision"] == get_family("moss_image").constants.revision + + +def test_the_read_and_the_trim_are_in_the_template_id(): + import dataclasses + + spec = te.encoder_spec("moss_image") + ids = {spec.template_id, + dataclasses.replace(spec, final_norm=False, + tap_layers=(32,)).template_id, + dataclasses.replace(spec, trim=False).template_id} + assert len(ids) == 3 + + +@pytest.mark.parametrize("prompt", [" a red bicycle\n", "\ta red bicycle ", + "a red bicycle\u3000"]) +def test_surrounding_whitespace_encodes_as_the_stripped_prompt(prompt): + """The release's template trims the prompt, so gmlx's ids for a prompt + with surrounding whitespace are the ids of the stripped prompt.""" + import dataclasses + + spec = te.encoder_spec("moss_image") + tok = _CharTokenizer() + want = te.token_ids(spec, tok, "a red bicycle") + assert te.token_ids(spec, tok, prompt) == want + untrimmed = dataclasses.replace(spec, trim=False) + assert te.token_ids(untrimmed, tok, prompt) != want + + +def test_a_final_norm_spec_takes_no_tap_options(): + spec = te.encoder_spec("moss_image") + import dataclasses + + with pytest.raises(ValueError, match="after layer 32 alone"): + dataclasses.replace(spec, tap_layers=(31,)) + for kw in ({"drop": 1}, {"stack": True}, {"pad_multiple": 8}, + {"fp32_qkv": True}, {"activation": "float32"}): + with pytest.raises(ValueError, match="takes no"): + dataclasses.replace(spec, **kw) + + +# The read on a tiny Qwen3.5 model. + +_NEEDS_GPU = pytest.mark.skipif( + bool(os.environ.get("KQUANT_FORCE_CPU")), + reason="the qwen3_5 GDN forward is Metal-only") + + +class _CharTokenizer: + """One id per character, below the tiny vocabulary.""" + + def encode(self, text, add_special_tokens=True): + assert add_special_tokens is False + return [ord(c) % 120 for c in text] + + +_TINY = dict(model_type="qwen3_5", hidden_size=64, intermediate_size=128, + linear_num_value_heads=4, linear_num_key_heads=2, + linear_key_head_dim=32, linear_value_head_dim=32, + linear_conv_kernel_dim=4, num_hidden_layers=4, + num_attention_heads=4, rms_norm_eps=1e-6, vocab_size=128, + num_key_value_heads=2, max_position_embeddings=2048, + tie_word_embeddings=True, head_dim=32, + rope_parameters={"type": "default", "mrope_section": [2, 1, 1], + "rope_theta": 100000, + "partial_rotary_factor": 0.25}, + full_attention_interval=4, kv_head_layout="tiled") + + +def _built(m): + mx.eval(m.parameters()) + m.eval() + return m + + +def _tiny_text(): + """The class the text loader builds for a qwen35 GGUF.""" + from gmlx._vendor.mlx_lm.models.qwen3_5 import Model, ModelArgs + + mx.random.seed(5) + return _built(Model(ModelArgs.from_dict(dict(_TINY)))) + + +def _tiny_multimodal(): + """The language model of a Qwen3.5 load with its vision tower, which a + resident chat model may be.""" + from gmlx._vendor.mlx_vlm.models.qwen3_5.config import TextConfig + from gmlx.models.qwen35.owned import OwnedQwen3_5LanguageModel + + top = SimpleNamespace(vision_config=SimpleNamespace(spatial_merge_size=2), + image_token_id=124, video_token_id=125, + vision_start_token_id=126) + mx.random.seed(5) + return _built(OwnedQwen3_5LanguageModel(TextConfig(**_TINY), top)) + + +_BUILDS = pytest.mark.parametrize("build", [_tiny_text, _tiny_multimodal], + ids=["text", "multimodal"]) + + +def _spec(): + import dataclasses + + return dataclasses.replace(te.encoder_spec("moss_image"), hidden=64, + layers=4, tap_layers=(4,), identity=()) + + +def _logits(m, ids): + """The chat path's prefill of ``ids`` on a fresh cache.""" + out = m(mx.array([ids]), cache=m.make_cache()) + return getattr(out, "logits", out) + + +@_NEEDS_GPU +@_BUILDS +def test_the_read_is_the_chat_prefill_before_its_output_head(build): + """The features, through the tied output head, give the logits of the + chat path's prefill: the same masks, positions and recurrent state.""" + m = build() + ids = te.token_ids(_spec(), _CharTokenizer(), "a red bicycle") + got = te.final_hidden(m, ids) + assert got.shape == (1, len(ids), 64) + head = te._inner(m).embed_tokens.as_linear + assert mx.allclose(head(got), _logits(m, ids), atol=1e-4, + rtol=1e-4).item() + + +@_NEEDS_GPU +@_BUILDS +def test_an_encoding_is_the_final_norm_of_the_last_layer(build, monkeypatch): + m = build() + tok = _CharTokenizer() + spec = _spec() + enc = te.EncoderHandle(m, tok, spec, owned=False).encode("a red bicycle") + ids = list(enc.ids) + assert ids == te.token_ids(spec, tok, "a red bicycle") + assert enc.features.shape == (len(ids), 64) and enc.pads == 0 + assert enc.template_id == spec.template_id + assert mx.array_equal(enc.features, te.final_hidden(m, ids)[0]).item() + # The last layer's output before the norm is not the read. + monkeypatch.setattr(te._inner(m), "norm", lambda h: h) + before = te.final_hidden(m, ids)[0] + assert not mx.allclose(before, enc.features, atol=1e-2).item() + + +@_NEEDS_GPU +def test_a_read_leaves_a_resident_models_positions_alone(): + """A borrowed multimodal chat model keeps the positions of its last + call for its next one, and an encoding through it does not touch + them.""" + m = _tiny_multimodal() + held = mx.arange(7)[None, None] + m._position_ids, m._rope_deltas = held, mx.zeros((1, 1)) + deltas = m._rope_deltas + te.EncoderHandle(m, _CharTokenizer(), _spec(), owned=False).encode("ab") + assert m._position_ids is held and m._rope_deltas is deltas + + +# The real files. + +@pytest.mark.integration +@pytest.mark.parametrize("name", _FILES) +def test_the_unsloth_file_passes_the_header_and_identity_checks( + name, gguf_index): + spec = te.encoder_spec("moss_image") + found = [p for p in gguf_index.get(spec.arch, []) if Path(p).name == name] + if not found: + pytest.skip(f"no {name} under KQUANT_TEST_GGUF_DIR") + view = te._header_view(found[0]) + te.check_encoder_file(view, spec) + assert te.file_identity(view, spec) == { + r.gguf_name: r.sha256 for r in spec.identity} + + +@_NEEDS_GPU +@_BUILDS +def test_the_identity_reads_the_loaded_norms(build): + """The identity of a loaded encoder, which a borrowed chat model is + checked by, reads its norms under either class.""" + import dataclasses + import hashlib + + import numpy as np + + m = build() + row = te.IdentityTensor("output_norm.weight", "model.norm.weight", "") + spec = dataclasses.replace(_spec(), identity=(row,)) + bits = np.array(te._inner(m).norm.weight.astype(mx.bfloat16).view( + mx.uint16)) + assert te.model_identity(m, spec) == { + "output_norm.weight": hashlib.sha256(bits.tobytes()).hexdigest()} + + +# F32 norms. + +def test_the_f32_norms_are_moss_images_alone(): + """Only MOSS-Image encodes on the F32 norm weights; every other + family's encoder, and a chat load, runs on the norms as loaded.""" + from gmlx.image.registry import families + + got = {} + for fam in families(): + try: + got[fam.name] = te.encoder_spec(fam).fp32_norms + except te.ImageModelError: + continue + assert {k for k, v in got.items() if v} == {"moss_image"} + assert len(got) > 1 + import dataclasses + + spec = te.encoder_spec("moss_image") + assert spec.template_id != dataclasses.replace( + spec, fp32_norms=False).template_id + + +def _f32_of(m): + """F32 weights for every norm of the tiny bf16 model, each off its + bf16 value by less than half a bf16 step, so it rounds to it.""" + inner = te._inner(m) + out = {} + for i, b in enumerate(inner.layers): + for where in ("input_layernorm", "post_attention_layernorm", + "self_attn.q_norm", "self_attn.k_norm", + "linear_attn.norm"): + mod = te._module(b, where) + if mod is not None: + w = mod.weight.astype(mx.float32) + out[f"layers.{i}.{where}"] = w * (1 + 2 ** -11) + out["norm"] = inner.norm.weight.astype(mx.float32) * (1 + 2 ** -11) + for where, w in out.items(): + held = te._module(inner, where).weight + assert mx.array_equal(w.astype(mx.bfloat16), held).item(), where + return out + + +def _norm_state(m): + """Every norm module and its weight, by identity.""" + inner = te._inner(m) + mods = [mod for _, mod in inner.named_modules() + if "norm" in type(mod).__name__.lower()] + return [(id(mod), id(mod.weight)) for mod in mods] + + +@_NEEDS_GPU +@_BUILDS +def test_an_encode_on_f32_norms_restores_a_resident_models_norms(build): + """A borrowed chat model encodes on the F32 norms and gets its own back + after the encode, and after one that raises: its prefill logits are + bit for bit what they were.""" + m = build() + m.set_dtype(mx.bfloat16) + norms = _f32_of(m) + tok = _CharTokenizer() + spec = _spec() + ids = te.token_ids(spec, tok, "a red bicycle") + before, state = _logits(m, ids), _norm_state(m) + plain = te.EncoderHandle(m, tok, spec, owned=False).encode("a red bicycle") + h = te.EncoderHandle(m, tok, spec, owned=False) + h.norms = norms + seen = [] + call = te._Float32RMSNorm.__call__ + + def spy(self, x): + y = call(self, x) + seen.append((x.dtype, y.dtype)) + return y + + te._Float32RMSNorm.__call__ = spy + try: + f32 = h.encode("a red bicycle") + finally: + te._Float32RMSNorm.__call__ = call + assert seen and set(seen) == {(mx.bfloat16, mx.bfloat16)} + assert f32.features.dtype == plain.features.dtype == mx.bfloat16 + assert not mx.array_equal(f32.features, plain.features).item() + assert _norm_state(m) == state + assert mx.array_equal(_logits(m, ids), before).item() + + import gmlx.image.text_encoder as mod + + real = mod.final_hidden + + def boom(model, ids): + assert _norm_state(model) != state + raise RuntimeError("boom") + + mod.final_hidden = boom + try: + with pytest.raises(RuntimeError, match="boom"): + h.encode("a red bicycle") + finally: + mod.final_hidden = real + assert _norm_state(m) == state + assert mx.array_equal(_logits(m, ids), before).item() + + + +@_BUILDS +def test_interleaved_f32_norm_contexts_restore_the_models_own(build): + """Two contexts on one model, entered A, B and left A, B, as two + encodes on one shared model would be, leave its own norms; F32 + norms stay in place until the last one leaves.""" + m = build() + m.set_dtype(mx.bfloat16) + norms = _f32_of(m) + state = _norm_state(m) + inner = te._inner(m) + a, b = te.float32_norms(m, norms), te.float32_norms(m, norms) + a.__enter__() + b.__enter__() + a.__exit__(None, None, None) + assert isinstance(inner.layers[0].input_layernorm, te._Float32RMSNorm) + b.__exit__(None, None, None) + assert _norm_state(m) == state + assert not isinstance(inner.layers[0].input_layernorm, + te._Float32RMSNorm) + with te.float32_norms(m, norms): + with te.float32_norms(m, norms): + pass + assert _norm_state(m) != state + assert _norm_state(m) == state + other = {k: v * 1 for k, v in norms.items()} + with te.float32_norms(m, norms): + with pytest.raises(RuntimeError, match="other F32 norms"): + with te.float32_norms(m, other): + pass + assert _norm_state(m) == state + +@_NEEDS_GPU +def test_f32_norms_that_are_not_the_models_are_refused(tmp_path, monkeypatch): + m = _tiny_text() + m.set_dtype(mx.bfloat16) + norms = _f32_of(m) + rows = tuple((k, np.array(v)) for k, v in norms.items()) + monkeypatch.setattr(te, "_gguf_norms", lambda *a: rows) + p = tmp_path / "enc.gguf" + p.write_bytes(b"x") + got = te.read_norms(str(p), m, _spec()) + assert set(got) == set(norms) + bad = ((rows[0][0], rows[0][1] * 1.5),) + rows[1:] + monkeypatch.setattr(te, "_gguf_norms", lambda *a: bad) + with pytest.raises(te.EncoderMismatchError, match="is not the loaded"): + te.read_norms(str(p), m, _spec()) + + +@pytest.mark.integration +@pytest.mark.parametrize("name", _FILES) +def test_the_encoder_reads_105_f32_norms_and_a_chat_load_keeps_bf16( + name, gguf_index): + """The encoder of a real file reads the 81 offset norms and the 24 + gated norms at F32. The loaded model, and a chat load of the same + file, hold every norm in bfloat16.""" + from gmlx.load.loader import load_model + + spec = te.encoder_spec("moss_image") + found = [p for p in gguf_index.get(spec.arch, []) if Path(p).name == name] + if not found: + pytest.skip(f"no {name} under KQUANT_TEST_GGUF_DIR") + h = te.load_encoder(found[0], spec, log=lambda *_: None) + try: + roles = {} + for where, w in h.norms.items(): + assert w.dtype == mx.float32, where + roles.setdefault(where.split(".", 2)[-1], 0) + roles[where.split(".", 2)[-1]] += 1 + assert roles == {"input_layernorm": 32, "post_attention_layernorm": 32, + "self_attn.q_norm": 8, "self_attn.k_norm": 8, + "linear_attn.norm": 24, "norm": 1} + chat = load_model(found[0])[0] + for where in h.norms: + a = te._module(te._inner(h.model), where).weight + b = te._module(te._inner(chat), where).weight + assert a.dtype == b.dtype == mx.bfloat16, where + assert mx.array_equal(a, b).item(), where + del chat + finally: + h.release() diff --git a/tests/image/test_image_moss_family.py b/tests/image/test_image_moss_family.py new file mode 100644 index 000000000..baf1b92c8 --- /dev/null +++ b/tests/image/test_image_moss_family.py @@ -0,0 +1,187 @@ +"""The MOSS-Image registry row: the realrebelai files, the config against +the official transformer/config.json (check 1), the SFT and Base identity +of every published quant, relatives refused with the reason, and the line +between MOSS-Image and every other family in both directions. CPU-only. +Header fixtures hold names, shapes and codecs, no weights.""" +from __future__ import annotations + +import hashlib + +import pytest + +from image_fixtures import (FIXTURES, Q_HEADER, Z_HEADER, fixture_view, + load_json, official_config) + +from gmlx.image import fingerprint as fp +from gmlx.image.families.moss_image import family as mf +from gmlx.image.registry import get_family, is_image_arch + +SFT = "moss_image_sft_q4_k_m.header.json" +CONFIG = "moss_image_transformer_config.json" +IDENTITY = "moss_image_identity.json" +KREA2 = "krea2_turbo_q4_k_m.header.json" +OTHERS = {Z_HEADER: "Z-Image", Q_HEADER: "Qwen-Image 2.1", + KREA2: "Krea 2", + "flux3_klein_q4_k_m.header.json": "FLUX.3 klein", + "flux2_klein_4b_q4_k_m.header.json": "FLUX.2 klein", + "flux2_klein_9b_q4_k_m.header.json": "FLUX.2 klein 9B"} + + +def _refused(view, **kw) -> str: + with pytest.raises(fp.ImageModelError) as ei: + fp.identify(view, **kw) + return str(ei.value) + + +def test_the_file_is_moss_image_with_the_official_config(): + view = fixture_view(SFT) + assert view.arch == "moss_image" and is_image_arch("moss_image") + got = fp.identify(view) + assert got.family.name == "moss_image" + assert (got.prefix, got.layout) == ("", "default") + official = official_config(CONFIG) + assert {k: got.config[k] for k in official} == official + assert (got.config["head_dim"], got.config["num_bands"], + got.config["num_text_blocks"]) == (128, 16, 2) + + +def test_the_fixtures_are_pinned(): + manifest = load_json("manifest.json") + data = (FIXTURES / CONFIG).read_bytes() + assert hashlib.sha256(data).hexdigest() == manifest[CONFIG]["sha256"] + assert len(data) == manifest[CONFIG]["size"] + row = get_family("moss_image").constants + assert (row.repo, row.revision) == (manifest[CONFIG]["repo"], + manifest[CONFIG]["revision"]) + d = load_json(SFT) + assert d["n_tensors"] == len(d["tensors"]) == 463 + assert manifest[SFT]["revision"] == load_json(IDENTITY)["revision"] + + +def test_every_published_quant_is_named_by_its_identity(): + """The hash of each file of the repository, read by range, names the + variant of its file name, and the variant row lists the file.""" + fam = get_family("moss_image") + assert fam.identity_tensors == mf.IDENTITY_TENSORS + files = load_json(IDENTITY)["files"] + assert len(files) == 20 + for name, row in files.items(): + want = "sft" if name.startswith("MOSS-Image-SFT-") else "base" + assert row["identity_codecs"] == ["F16"] * 3, name + variant = fam.variants[row["sha256"]] + assert variant.name == want and name in variant.files, name + assert variant.repo == load_json(IDENTITY)["repo"] + assert fam.variant_names() == ("base", "sft") + + +@pytest.mark.parametrize("change,reason", [ + ({"drop": ("layers.41.",)}, "layers has 41 blocks, MOSS-Image has 42"), + ({"drop": ("text_embedder.blocks.1.",)}, + "text_embedder.blocks has 1 blocks, MOSS-Image has 2"), + ({"drop": ("layers.7.q_norm.",)}, "layers.7.q_norm.weight is missing"), + ({"add": {"layers.0.proj.ffn_up_gate.weight": ((3072, 12288), "Q4_K")}}, + "layers.0.proj.ffn_up_gate.weight is (3072, 12288), MOSS-Image has " + "(3072, 16384)"), + ({"add": {"text_embedder.proj.weight": ((2560, 3072), "F16")}}, + "text_embedder.proj.weight is (2560, 3072), MOSS-Image has " + "(4096, 3072)"), + ({"add": {"layers.0.proj.q.bias": ((3072,), "F16")}}, + "layers.0.proj.q.bias is not a MOSS-Image tensor"), + ({"add": {"layers.0.moe.gate.weight": ((3072, 64), "F16")}}, + "layers.0.moe.gate.weight is not a MOSS-Image tensor"), +], ids=["depth", "refiners", "missing", "ffn width", "text width", "bias", + "extra"]) +def test_a_relative_is_refused_with_the_reason(change, reason): + msg = _refused(fixture_view(SFT, **change)) + assert "MOSS-Image relative" in msg and reason in msg + + +def test_a_long_refusal_counts_what_it_does_not_show(): + msg = _refused(fixture_view(SFT, drop=("layers.40.", "layers.41."))) + assert "layers has 40 blocks" in msg + + +@pytest.mark.parametrize("prefix", ["model.diffusion_model.", + "diffusion_model."]) +def test_a_prefixed_file_loads_under_its_prefix(prefix): + view = fixture_view(SFT) + view = fp.HeaderView(kv=view.kv, tensors={ + prefix + n: t for n, t in view.tensors.items()}) + got = fp.identify(view) + assert (got.family.name, got.prefix) == ("moss_image", prefix) + + +def test_a_file_without_a_label_loads_and_another_label_is_refused(): + assert fp.identify(fixture_view(SFT, kv={})).family.name == "moss_image" + assert _refused(fixture_view(SFT, kv={"general.architecture": "flux"}) + ) == ("arch 'flux' contradicts the tensor layout, which " + "is MOSS-Image (arch 'moss_image' or none)") + + +@pytest.mark.parametrize("name", list(OTHERS)) +def test_another_family_given_as_moss_image_is_refused_by_name(name): + view = fixture_view(name) + assert fp.identify(view).family.title == OTHERS[name] + assert _refused(view, family="moss_image") == ( + f"this is a {OTHERS[name]} GGUF, not a MOSS-Image one") + + +def test_no_other_family_takes_a_moss_image_file(): + view = fixture_view(SFT) + for fam in (f for f in fp._families(None) if f.name != "moss_image"): + assert fam.fingerprint(fp.strip_prefix(view, fam.strip_prefixes)) \ + is None, fam.name + + +@pytest.mark.parametrize("name", list(OTHERS)) +def test_a_moss_image_file_given_as_another_family_is_refused(name): + view = fixture_view(SFT) + other = fp.identify(fixture_view(name)).family + assert _refused(view, family=other.name).startswith( + f"not a {other.title} file. unrecognized image GGUF: " + "arch=moss_image, 463 tensors") + + +class _RuleError(ValueError): + pass + + +def test_an_error_in_another_familys_rule_is_not_hidden(monkeypatch): + """An error that another family's rule raises on a file reaches the + caller as it is, under family="moss_image" as with no family given.""" + import dataclasses + + from gmlx.image import registry + + def broken(view): + raise _RuleError("the Krea 2 rule failed") + + fams = tuple(dataclasses.replace(f, fingerprint=broken) + if f.name == "krea2" else f for f in registry.families()) + monkeypatch.setattr(registry, "families", lambda: fams) + view = fixture_view(KREA2) + for kw in ({}, {"family": "moss_image"}): + with pytest.raises(_RuleError, match="the Krea 2 rule failed"): + fp.identify(view, **kw) + + +def test_a_file_no_family_takes_keeps_the_generic_refusal(): + view = fp.HeaderView(kv={}, tensors={ + "blocks.0.attn.qkv.weight": fp.Tensor((64, 192), "F16")}) + assert mf.fingerprint(view) is None + assert _refused(view, family="moss_image").startswith( + "not a MOSS-Image file. unrecognized image GGUF") + + +def test_a_file_of_no_known_variant_samples_as_sft(): + """A file whose identity tensors match neither release, here a header + with no data to hash, samples with the SFT row: 50 steps at guidance + 3, block 22 left out of the negative pass, as both releases do.""" + from gmlx.image import api + + got = fp.identify(fixture_view(SFT)) + assert got.variant == "unknown" + assert got.family.unknown_variant == "sft" + d = api.sampler_defaults("moss_image", got.variant) + assert d is not None + assert (d.steps, d.cfg, d.shift, d.neg_skip) == (50, 3.0, 5.0, (22,)) diff --git a/tests/image/test_image_moss_lora_map.py b/tests/image/test_image_moss_lora_map.py new file mode 100644 index 000000000..203fe0a3a --- /dev/null +++ b/tests/image/test_image_moss_lora_map.py @@ -0,0 +1,252 @@ +"""The MOSS-Image LoRA key map and the delta half of acceptance check 7: +every spelling of every target resolves to its leaf, the placed factors +equal a float64 reference per dialect, and a step of the tiny DiT under +a patch table equals the step of the DiT with the deltas merged into its +weights. CPU.""" +from __future__ import annotations + +import dataclasses + +import mlx.core as mx +import mlx.nn as nn +import numpy as np +import pytest + +from image_fixtures import fixture_view +from lora_fixtures import (DIALECTS, dialect_keys, make_pair, + reference_delta, rel_l2, write_adapter, + write_keys) + +from gmlx.image import fingerprint as fp +from gmlx.image.families.moss_image import model as mm +from gmlx.image.families.moss_image.lora_map import lora_map +from gmlx.image.lora import plan, reader +from gmlx.image.registry import get_family + +DELTA_TOL = 1e-6 +HEADER = "moss_image_sft_q4_k_m.header.json" +TINY = {"dim": 64, "ffn_dim": 48, "freq_dim": 16, "latent_channels": 8, + "num_heads": 2, "num_layers": 2, "rope_base": 10000.0, + "text_dim": 24, "num_text_blocks": 2} +H, W = 3, 4 +# The GGUF names under the prefixes ComfyUI reads, and two that it does +# not: the diffusers prefix and lycoris_. +SPELLINGS = { + "native": ("diffusion_model.", False, ".lora_A.weight", + ".lora_B.weight", False), + "native_kohya": ("lora_unet_", True, ".lora_down.weight", + ".lora_up.weight", True), + "lycoris": ("lycoris_", True, ".lora_down.weight", ".lora_up.weight", + True), +} +COMFY_READS = ("native", "native_kohya", "kohya") + + +def dims_of(cfg: dict = TINY) -> dict: + """Module path to (out, in) of every linear leaf of the built + model.""" + return {p: tuple(m.weight.shape) + for p, m in mm.MossImageDiT(cfg).named_modules() + if isinstance(m, nn.Linear)} + + +def test_the_key_map_covers_every_linear_leaf_and_only_those(): + km = lora_map(TINY) + dims = dims_of() + assert km.family == "moss_image" and km.title == "MOSS-Image" + assert set(km.targets) == set(dims) + assert all(km.targets[p] == (p, None) for p in dims) + # 2 blocks of 6 leaves, 2 refiner blocks of 5, and 7 others. + assert len(dims) == 2 * 6 + 2 * 5 + 7 + assert km.diffusers_only == frozenset() + assert km.comfy_prefixes == ("lora_unet_", "diffusion_model.") + built = dataclasses.asdict(mm.MossConfig.from_dict(TINY)) + assert lora_map(built) == km + assert get_family("moss_image").lora_map(TINY) == km + assert len({n.replace(".", "_") for n in km.targets}) == len(km.targets) + assert km.text_names == () and km.text_targets == () + + +def test_the_key_map_names_every_linear_weight_of_the_real_header(): + ident = fp.identify(fixture_view(HEADER)) + km = lora_map(ident.config) + assert ident.family.name == "moss_image" + linear = {n[:-len(".weight")] for n, t in ident.view.tensors.items() + if n.endswith(".weight") and len(t.shape) == 2} + assert linear == set(km.targets) + assert len(linear) == 42 * 6 + 2 * 5 + 7 + assert km.targets["layers.41.proj.ffn_up_gate"] == ( + "layers.41.proj.ffn_up_gate", None) + assert "layers.42.proj.q" not in km.targets + + +@pytest.mark.parametrize("dialect", [*sorted(DIALECTS), *SPELLINGS]) +def test_every_spelling_of_every_target_resolves(tmp_path, dialect): + """One pair on every linear layer, in the GGUF names, as each dialect + spells them, lands on its leaf. A spelling ComfyUI does not read for + MOSS-Image is applied with a warning.""" + row = DIALECTS.get(dialect) or SPELLINGS[dialect] + rng = np.random.default_rng(4) + want = dims_of() + by_path = {p: make_pair(rng, o, i, 2, 4.0 if row[4] else None) + for p, (o, i) in want.items()} + prefix, under, down_s, up_s, _ = row + keys = {} + for path, (down, up, alpha) in by_path.items(): + stem = prefix + (path.replace(".", "_") if under else path) + keys[stem + down_s], keys[stem + up_s] = down, up + if alpha is not None: + keys[stem + ".alpha"] = np.asarray(alpha, np.float32) + if dialect in DIALECTS: + assert keys.keys() == dialect_keys(dialect, by_path).keys() + res = reader.resolve(reader.read_lora( + write_keys(tmp_path / "all.safetensors", keys)), lora_map(TINY)) + assert res.text == () and res.unmatched == () + if dialect in COMFY_READS: + assert res.warnings == () + else: + assert len(res.warnings) == 1 + assert "ComfyUI drops them" in res.warnings[0] + assert len(res.dit) == len(want) + for t in res.dit: + down, up, alpha = by_path[t.module] + assert t.band is None + assert np.array_equal(np.array(t.pair.down), down) + assert np.array_equal(np.array(t.pair.up), up) + assert t.pair.alpha == alpha + + +def _adapter(seed: int, alpha, paths) -> dict: + dims = dims_of() + rng = np.random.default_rng(seed) + return {p: make_pair(rng, *dims[p], 2 + i % 3, alpha) + for i, p in enumerate(paths)} + + +@pytest.mark.parametrize("dialect", sorted(DIALECTS)) +def test_check7_delta_half_per_dialect(tmp_path, dialect): + """Strengths 0.5 and 1.0 and a stacked pair with a negative strength, + against float64. A dialect that stores no alpha takes the factor + 1.0.""" + km, dims = lora_map(TINY), dims_of() + has_alpha = DIALECTS[dialect][4] + paths = sorted(dims) + one = _adapter(1, 3.0 if has_alpha else None, paths[::2]) + two = _adapter(2, 1.0 if has_alpha else None, paths[::3]) + p1 = write_adapter(tmp_path / "one.safetensors", dialect, one) + p2 = write_adapter(tmp_path / "two.safetensors", dialect, two) + cases = [([(p1, one, 0.5)], "0.5"), ([(p1, one, 1.0)], "1.0"), + ([(p1, one, 0.75), (p2, two, -0.4)], "stacked")] + worst = 0.0 + for stack, label in cases: + got: dict = {} + want: dict = {} + for path, pairs, strength in stack: + res = reader.resolve(reader.read_lora(str(path)), km) + for t in res.dit: + got.setdefault(t.module, []).append(plan.Part( + t.pair.down, t.pair.up, t.pair.alpha, t.band, strength)) + for target, (down, up, alpha) in pairs.items(): + module, band = km.targets[target] + want.setdefault(module, []).append( + (down, up, alpha, band, strength)) + assert set(got) == set(want), label + for module, parts in got.items(): + lm = plan.block_diagonal(module, parts, dims[module]) + delta = np.array(lm.b, np.float64) @ np.array(lm.a, np.float64) + err = rel_l2(delta, reference_delta(want[module], dims[module])) + worst = max(worst, err) + assert err <= DELTA_TOL, (dialect, label, module, err) + print(f"moss check7 delta {dialect}: worst relative L2 {worst:.3e}") + + +def test_a_key_that_names_no_module_is_refused(tmp_path): + """A block past the last, a refiner past the last, a norm and the + RoPE bands.""" + km = lora_map(TINY) + rng = np.random.default_rng(5) + good = _adapter(6, None, ["layers.0.proj.q"]) + for i, (name, o, n) in enumerate(( + ("layers.2.proj.q", 64, 64), + ("text_embedder.blocks.2.qkv", 192, 64), + ("layers.0.norm1", 64, 64), ("rope.bands", 4, 4))): + mixed = {**good, name: make_pair(rng, o, n, 2)} + path = write_adapter(tmp_path / f"s{i}.safetensors", "diffusers", + mixed) + with pytest.raises(reader.UnmatchedKeysError, + match="1 of 2 targets name no MOSS-Image module"): + reader.resolve(reader.read_lora(path), km) + + +def _tiny(seed: int = 0): + from mlx.utils import tree_flatten + + m = mm.MossImageDiT(TINY) + rng = np.random.default_rng(seed) + params = [] + for k, v in tree_flatten(m.parameters()): + a = rng.standard_normal(v.shape) * 0.3 + if k.endswith("norm1.weight") or k.endswith("norm2.weight"): + a += 1.0 + params.append((k, mx.array(a.astype(np.float32)))) + m.load_weights(params) + m.act_dtype = mx.float32 + m.eval() + return m, params + + +# One pair on a leaf of every kind: each block leaf, the fused weights, +# the refiner and the layers that feed the modulation and caption memos. +APPLIED = ("layers.0.proj.q", "layers.1.proj.ffn_up_gate", + "layers.1.proj.o", "text_embedder.blocks.0.qkv", + "text_embedder.blocks.1.ffn_up", "text_embedder.proj", + "time_embedding.0", "shared_block_modulation", + "video_head_modulation", "final_linear_video", "video_embedder") + + +def test_a_step_under_a_patch_table_is_the_step_of_the_merged_weights( + tmp_path): + """Each pass of a guided step under the table, the negative one with + its skip set, equals the DiT whose weights hold the deltas, and the + memos of the base step do not leak into it, nor its into a later base + step.""" + from gmlx.image import api, lora + + strength = 0.7 + pairs = _adapter(7, 2.0, APPLIED) + path = write_adapter(tmp_path / "a.safetensors", "kohya", pairs) + model, params = _tiny() + pipe = api.ImagePipeline(dit=model, family=get_family("moss_image"), + variant="sft") + merged, _ = _tiny() + held = dict(params) + for p, (down, up, alpha) in pairs.items(): + w = np.array(held[f"{p}.weight"], np.float64) + w += strength * alpha / down.shape[0] * ( + up.astype(np.float64) @ down.astype(np.float64)) + held[f"{p}.weight"] = mx.array(w.astype(np.float32)) + merged.load_weights(list(held.items())) + mpipe = api.ImagePipeline(dit=merged, family=get_family("moss_image"), + variant="sft") + rng = np.random.default_rng(8) + x = mx.array(rng.standard_normal((1, 8, H, W)).astype(np.float32)) + cond = api.Conditioning(mx.array( + rng.standard_normal((5, 24)).astype(np.float32)), "moss-lora-pos") + sigma = np.float32(0.6875) + + def step(p, tid=None, skip=()): + return np.array(api.denoise_step(p, x, sigma, cond, + patch_table=tid, skip=skip), + np.float64) + + with mx.stream(mx.cpu): + base = step(pipe) + ad = lora.load_adapter(pipe, path) + assert len(ad.dit) == len(APPLIED) + tid = api.define_patch_table( + pipe, lora.table_entries([(ad, strength)])) + for skip in ((), (1,)): + got, want = step(pipe, tid, skip), step(mpipe, skip=skip) + assert rel_l2(got, want) < 1e-5, (skip, rel_l2(got, want)) + assert rel_l2(got, step(pipe, skip=skip)) > 1e-2 + assert np.array_equal(step(pipe), base) diff --git a/tests/image/test_image_moss_model.py b/tests/image/test_image_moss_model.py new file mode 100644 index 000000000..cd62bb101 --- /dev/null +++ b/tests/image/test_image_moss_model.py @@ -0,0 +1,559 @@ +"""The MOSS-Image DiT on the CPU: the module tree against the file's +tensor names, the RoPE bands against the file's stored copy, and a tiny +model in fp32 against a float64 numpy forward written from the release's +MOSSImageTransformer, with its rotate-half RoPE, its six shared vectors +and its head. The step parts give the step they stand in for. Each wrong +model of FAULTS misses the fp32 bound.""" +from __future__ import annotations + +import math + +import mlx.core as mx +import numpy as np +import pytest +from mlx.utils import tree_flatten + +import moss_gates as mg +from image_fixtures import fixture_view, load_json +from oracle_fixture import committed_fixture + +from gmlx.image import fingerprint as fp +from gmlx.image.dtypes import ImageDtypeError +from gmlx.image.families.moss_image import model as mm +from gmlx.image.loader import check_names +from gmlx.image.registry import get_family + +TINY = {"dim": 64, "ffn_dim": 48, "freq_dim": 16, "latent_channels": 8, + "num_heads": 2, "num_layers": 2, "rope_base": 10000.0, + "text_dim": 24, "num_text_blocks": 2} +H, W, L = 3, 4, 5 +SIGMA = np.float32(0.6875) +# The tiny model in fp32 on the CPU against the float64 forward: the +# right model lands within 1e-6 relative; the weakest wrong model of +# FAULTS moves the velocity by more than 25 times that. +TOL = 4e-6 + + +def _tiny(seed: int = 0, act=mx.float32): + m = mm.MossImageDiT(TINY) + rng = np.random.default_rng(seed) + params = [] + for k, v in tree_flatten(m.parameters()): + a = rng.standard_normal(v.shape) * 0.3 + if k.endswith("norm1.weight") or k.endswith("norm2.weight"): + a += 1.0 + params.append((k, mx.array(a.astype(np.float32)))) + m.load_weights(params) + m.act_dtype = act + m.eval() + return m, {k: np.array(v, np.float64) for k, v in params} + + +def _inputs(seed: int = 1, batch: int = 1): + rng = np.random.default_rng(seed) + lat = rng.standard_normal((batch, 8, H, W)).astype(np.float32) + feats = rng.standard_normal((1, L, 24)).astype(np.float32) + return lat, feats + + +# The float64 forward, written from the release. + +def _lin(p, name, x): + y = x @ p[name + ".weight"].T + return y + p[name + ".bias"] if name + ".bias" in p else y + + +def _rms(x, w): + return x / np.sqrt(np.mean(x * x, -1, keepdims=True) + 1e-6) * w + + +def _situ(hid, gate): + sig = 1 / (1 + np.exp(-gate)) + return 4 * np.tanh(gate / 4) * sig * (25 * np.tanh(hid / 25)) + + +def _rope(h, w, n): + nb = 4 + bands = (np.float32(1) / np.float32(10000) ** ( + np.arange(nb, dtype=np.float32) / np.float32(nb))).astype(np.float32) + r, c = np.meshgrid(np.arange(h), np.arange(w), indexing="ij") + img = np.stack([np.zeros(h * w), np.full(h * w, n), r.reshape(-1), + c.reshape(-1)], -1).astype(np.float32) + img[:, 2] -= (h - 1) / 2 + img[:, 3] -= (w - 1) / 2 + txt = np.zeros((n, 4), np.float32) + txt[:, 1] = np.arange(n) + phase = (np.concatenate([img, txt])[:, :, None] * bands).reshape( + h * w + n, -1).astype(np.float64) + return np.sin(phase), np.cos(phase) + + +def _rotate(x, sin, cos): + sin = np.concatenate([sin, sin], -1)[:, None] + cos = np.concatenate([cos, cos], -1)[:, None] + a, b = np.split(x, 2, -1) + return x * cos + np.concatenate([-b, a], -1) * sin + + +def _attend(q, k, v, rope, qn, kn, heads): + n = q.shape[0] + q = _rotate(_rms(q.reshape(n, heads, -1), qn), *rope) + k = _rotate(_rms(k.reshape(n, heads, -1), kn), *rope) + v = v.reshape(n, heads, -1) + s = np.einsum("qhd,khd->hqk", q, k) / np.sqrt(q.shape[-1]) + s = np.exp(s - s.max(-1, keepdims=True)) + s /= s.sum(-1, keepdims=True) + return np.einsum("hqk,khd->qhd", s, v).reshape(n, -1) + + +def reference(p, lat, sigma, feats, heads=2, skip=()): + """Velocity (B, C, H, W) float64 of the release's forward, the blocks + of ``skip`` left out.""" + b, c, h, w = lat.shape + t = np.float64(np.float32(sigma) * np.float32(1000)) + half = 8 + phase = t * 10000.0 ** (-np.arange(half, dtype=np.float64) / half) + emb = np.concatenate([np.cos(phase), np.sin(phase)]).astype( + np.float32).astype(np.float64) + hid = _lin(p, "time_embedding.0", emb) + time = _lin(p, "time_embedding.2", hid / (1 + np.exp(-hid))) + mod = _lin(p, "shared_block_modulation", time).reshape(6, -1) + head = _lin(p, "video_head_modulation", time).reshape(2, -1) + n = feats.shape[1] + sin, cos = _rope(h, w, n) + txt_rope = (sin[h * w:], cos[h * w:]) + x = _lin(p, "text_embedder.proj", feats[0].astype(np.float64)) + for i in range(2): + pre = f"text_embedder.blocks.{i}." + q, k, v = np.split(_lin(p, pre + "qkv", _rms(x, p[pre + "norm1.weight"])), + 3, -1) + x = x + _lin(p, pre + "proj", _attend( + q, k, v, txt_rope, p[pre + "q_norm.weight"], + p[pre + "k_norm.weight"], heads)) + y = _rms(x, p[pre + "norm2.weight"]) + x = x + _lin(p, pre + "ffn_down", _situ(_lin(p, pre + "ffn_up", y), + _lin(p, pre + "ffn_gate", y))) + out = [] + for j in range(b): + tok = lat[j].reshape(c, h * w).T.astype(np.float64) + z = np.concatenate([_lin(p, "video_embedder", tok), x]) + for i in range(len([k for k in p if k.endswith("proj.o.weight")])): + if i in skip: + continue + pre = f"layers.{i}." + sa, ca, ga, sf, cf, gf = mod + y = _rms(z, p[pre + "norm1.weight"]) * (1 + ca) + sa + o = _attend(_lin(p, pre + "proj.q", y), _lin(p, pre + "proj.k", y), + _lin(p, pre + "proj.v", y), (sin, cos), + p[pre + "q_norm.weight"], p[pre + "k_norm.weight"], + heads) + z = z + _lin(p, pre + "proj.o", o) * ga + y = _rms(z, p[pre + "norm2.weight"]) * (1 + cf) + sf + hg = _lin(p, pre + "proj.ffn_up_gate", y) + z = z + _lin(p, pre + "proj.ffn_down", _situ(*np.split(hg, 2, -1))) * gf + y = _rms(z[:h * w], p["final_norm_video.weight"]) * (1 + head[1]) + head[0] + out.append(_lin(p, "final_linear_video", y).T.reshape(c, h, w)) + return np.stack(out) + + +def _gap(model, p, lat, feats) -> float: + with mx.stream(mx.cpu): + got = np.array(model(mx.array(lat), np.reshape(SIGMA, (1,)), + mx.array(feats)), np.float64) + want = reference(p, lat, SIGMA, feats) + return float(np.linalg.norm(got - want) / np.linalg.norm(want)) + + +def test_the_module_tree_names_every_tensor_of_the_file(): + """Strict name equality and the torch shape of every tensor, on the + real config. Nothing is evaluated.""" + got = fp.identify(fixture_view("moss_image_sft_q4_k_m.header.json")) + model = get_family("moss_image").build(got.config) + check_names(model, got.view.tensors, "MOSS-Image") + params = dict(tree_flatten(model.parameters())) + for n, t in got.view.tensors.items(): + assert params[n].shape == tuple(reversed(t.shape)), n + + +def test_the_bands_equal_the_files_copy_bit_for_bit(): + want = np.array(load_json("moss_image_identity.json")["rope_bands"]["bits"], + np.uint32) + assert mm.rope_bands(16, 10000.0).view(np.uint32).tolist() == \ + want.tolist() + + +def test_the_coordinates_center_the_image_and_count_the_text(): + got = mm.positions(2, 3, 4) + assert got.dtype == np.float32 and got.shape == (10, 4) + assert got[:6].tolist() == [[0, 4, -0.5, -1], [0, 4, -0.5, 0], + [0, 4, -0.5, 1], [0, 4, 0.5, -1], + [0, 4, 0.5, 0], [0, 4, 0.5, 1]] + assert got[6:].tolist() == [[0, i, 0, 0] for i in range(4)] + + +def test_the_sinusoid_rounds_float64_phases_once(): + """Each entry is the double-precision cosine or sine of t times its + frequency, rounded once to float32. The same phases in float32 miss.""" + t = 300.0 + freqs = [10000.0 ** (-k / 8) for k in range(8)] + want = ([math.cos(t * f) for f in freqs] + + [math.sin(t * f) for f in freqs]) + got = np.array(mm.time_sinusoid(np.array([t], np.float32), 16, 10000.0)) + assert got[0].tolist() == np.float32(want).tolist() + single = np.float32(t) * np.float32(freqs) + assert np.cos(single).tolist() != got[0, :8].tolist() + + +def test_the_tiny_model_matches_the_release_forward_in_fp32(): + model, p = _tiny() + lat, feats = _inputs() + gap = _gap(model, p, lat, feats) + print(f"moss tiny fp32: rel L2 {gap:.2e}") + assert gap < TOL + + +@pytest.fixture(scope="module") +def release_tiny(): + """The release's class on random tiny weights, from the oracle's + --self-test.""" + arrays, m = committed_fixture("selftest_moss_image_tiny") + weights, wm = committed_fixture("selftest_moss_image_tiny_weights") + bands = weights["rope.bands"] + return arrays, m, weights, wm["inputs"]["config"], bands + + +def test_the_tiny_model_matches_the_release_class_in_fp32(release_tiny): + arrays, _, weights, cfg, bands = release_tiny + model = mg.tiny_model(weights, cfg) + assert mm.rope_bands(len(bands), cfg["rope_base"]).tobytes() == \ + bands.tobytes() + # The CPU stream keeps fp32 products exact; the GPU's TF32 would not. + with mx.stream(mx.cpu): + got = mg.measures(model, arrays) + worst = max(got, key=got.get) + print(f"moss release tiny fp32: {len(got)} measures, worst " + f"{got[worst]:.2e} at {worst}") + assert got[worst] < TOL, got + + +def test_the_tiny_model_in_bf16_is_inside_the_bf16_tolerances(release_tiny): + arrays, m, weights, cfg, _ = release_tiny + model = mg.tiny_model(weights, cfg, act=mx.bfloat16) + out = mg.missed(mg.measures(model, arrays), m["tolerances"]) + assert not out, out + + +def test_the_step_parts_give_the_step(): + model, _ = _tiny(act=mx.bfloat16) + lat, feats = (mx.array(a) for a in _inputs()) + sigma = np.reshape(SIGMA, (1,)) + whole = model(lat, sigma, feats) + mod = model.modulation(sigma, 1) + cap = model.caption(feats) + parts = model(lat, sigma, None, mod=mod, caption=cap) + assert mx.array_equal(whole, parts).item() + assert len(mod.arrays()) == 8 + assert cap.arrays()[0].shape == (1, L, 64) + with pytest.raises(ValueError, match="text features"): + model(lat, sigma, None) + + +def test_a_batch_of_two_gives_each_row_its_own_velocity(): + model, _ = _tiny() + lat, feats = _inputs(batch=2) + lat, feats = mx.array(lat), mx.array(feats) + sigma = np.reshape(SIGMA, (1,)) + with mx.stream(mx.cpu): + two = model(lat, np.array([SIGMA, SIGMA]), feats) + for i in range(2): + one = model(lat[i:i + 1], sigma, feats) + assert mx.allclose(two[i:i + 1], one, atol=1e-5).item() + + +def test_float16_is_refused(): + model, _ = _tiny() + with pytest.raises(ImageDtypeError, match="float16"): + model.act_dtype = mx.float16 + + +def test_wrong_configs_are_refused(): + with pytest.raises(ValueError, match="head dim"): + mm.MossImageDiT(dict(TINY, num_heads=3)) + with pytest.raises(ValueError, match="RoPE bands"): + mm.MossImageDiT(dict(TINY, num_bands=8)) + with pytest.raises(ValueError, match="even width"): + mm.MossImageDiT(dict(TINY, freq_dim=15)) + + +# Wrong models: each changes one convention of the forward. + +def _rope_adjacent(mp): + mp.setattr(mm, "interleave", lambda x: x) + + +def _shift_scale_swapped(mp): + real = mm.MossImageDiT.modulation + + def swapped(self, sigma, batch): + got = real(self, sigma, batch) + b = got.blocks + return mm.StepModulation((b[1], b[0], *b[2:]), got.shift, got.scale) + mp.setattr(mm.MossImageDiT, "modulation", swapped) + + +def _head_swapped(mp): + real = mm.MossImageDiT.modulation + + def swapped(self, sigma, batch): + got = real(self, sigma, batch) + return mm.StepModulation(got.blocks, got.scale, got.shift) + mp.setattr(mm.MossImageDiT, "modulation", swapped) + + +def _situ_halves_swapped(mp): + real = mm.situ + mp.setattr(mm, "situ", lambda hid, gate: real(gate, hid)) + + +def _grid_not_centered(mp): + real = mm.positions + + def pos(h, w, n): + got = real(h, w, n) + got[:h * w, 2:] += np.float32([(h - 1) / 2, (w - 1) / 2]) + return got + mp.setattr(mm, "positions", pos) + + +def _image_at_text_zero(mp): + real = mm.positions + + def pos(h, w, n): + got = real(h, w, n) + got[:h * w, 1] = 0 + return got + mp.setattr(mm, "positions", pos) + + +def _sines_first(mp): + def sinusoid(t, dim, period): + got = np.array(mm_real_sinusoid(t, dim, period)) + return mx.array(np.roll(got, dim // 2, axis=-1)) + mm_real_sinusoid = mm.time_sinusoid + mp.setattr(mm, "time_sinusoid", sinusoid) + + +FAULTS = { + "rope pairs adjacent dims": _rope_adjacent, + "block shift and scale swapped": _shift_scale_swapped, + "head shift and scale swapped": _head_swapped, + "situ halves swapped": _situ_halves_swapped, + "grid not centered": _grid_not_centered, + "image rows at text position 0": _image_at_text_zero, + "sines first": _sines_first, +} + + +@pytest.mark.parametrize("name", list(FAULTS)) +def test_each_wrong_model_fails_the_tiny_gate(monkeypatch, name): + model, p = _tiny() + lat, feats = _inputs() + with monkeypatch.context() as mp: + FAULTS[name](mp) + gap = _gap(model, p, lat, feats) + print(f"moss tiny, {name}: rel L2 {gap:.2e}") + assert gap > 25 * TOL + + +@pytest.mark.parametrize("name", list(FAULTS)) +def test_each_wrong_model_misses_the_release_class(release_tiny, monkeypatch, + name): + arrays, _, weights, cfg, _ = release_tiny + model = mg.tiny_model(weights, cfg) + with monkeypatch.context() as mp, mx.stream(mx.cpu): + FAULTS[name](mp) + gap = mg.rel(mg.free_rows(model, arrays)["velocity"], + arrays["velocity"]) + print(f"moss release tiny, {name}: velocity rel L2 {gap:.2e}") + assert gap > 25 * TOL + + +def test_every_op_of_a_tiny_bf16_step_is_within_its_rounding_bound(): + """The site gate on the tiny model in bf16: each op of each block and + each call of a linear leaf against the stock op in float32 on the same + input, within 2 n u, and no op or leaf of the step left unread.""" + import moss_gates as mg + import qwen_gates as qg + + model, _ = _tiny(act=mx.bfloat16) + model.set_dtype(mx.bfloat16) + lat, feats = (mx.array(a) for a in _inputs()) + sites: dict = {} + skipped: set = set() + with mg.step_sites(model, sites, skipped=skipped): + mx.eval(model(lat, np.reshape(SIGMA, (1,)), feats)) + assert not mg.missing_ops(model, sites, skipped=skipped) + assert not qg.over_gate(sites), qg.over_gate(sites) + assert len(qg.op_sites(sites)) == 2 * 9 + 2 * 5 + + +def test_the_sampler_runs_the_tiny_model_on_the_snapped_schedule(): + """api.sample draws noise at the DiT's latent channels and steps + through the family's schedule.""" + from gmlx.image import api + from gmlx.image.schedule import moss_sigmas + + model, _ = _tiny() + pipe = api.ImagePipeline(dit=model, family=get_family("moss_image"), + variant="sft") + seen = [] + feats = mx.array(_inputs()[1][0]) + with mx.stream(mx.cpu): + x = api.sample(pipe, api.Conditioning(feats), height=16 * H, + width=16 * W, steps=8, cfg=1.0, seed=3, + on_step=lambda i, _: seen.append(i)) + assert x.shape == (1, 8, H, W) + assert bool(mx.all(mx.isfinite(x)).item()) + assert len(seen) == len(moss_sigmas(8)) - 1 + + +def test_a_skipped_block_is_left_out_as_the_release_leaves_it(): + model, p = _tiny() + lat, feats = _inputs() + with mx.stream(mx.cpu): + got = np.array(model(mx.array(lat), np.reshape(SIGMA, (1,)), + mx.array(feats), skip=(1,)), np.float64) + want = reference(p, lat, SIGMA, feats, skip=(1,)) + whole = reference(p, lat, SIGMA, feats) + rel = lambda a, b: np.linalg.norm(a - b) / np.linalg.norm(b) # noqa: E731 + assert rel(got, want) < TOL + assert rel(whole, want) > 25 * TOL + with pytest.raises(ValueError, match="skips blocks 0 to 1"): + model(mx.array(lat), np.reshape(SIGMA, (1,)), mx.array(feats), + skip=(2,)) + + +def _cond(rows: int, seed: int): + from gmlx.image import api + + rng = np.random.default_rng(seed) + f = rng.standard_normal((rows, 24)).astype(np.float32) + return api.Conditioning(mx.array(f), f"moss-{rows}-{seed}") + + +def test_a_guided_run_skips_the_blocks_in_its_negative_pass_alone(): + """Each pass runs at its own text rows, the negative pass without the + blocks of neg_skip, and the two combine in float32 by the guide of + the release.""" + from gmlx.image import api + from gmlx.image.schedule import moss_sigmas + + model, _ = _tiny() + pipe = api.ImagePipeline(dit=model, family=get_family("moss_image"), + variant="sft") + pos, neg = _cond(5, 1), _cond(3, 2) + noise = np.random.default_rng(4).standard_normal( + (1, 8, H, W)).astype(np.float32) + s = moss_sigmas(4) + + def by_hand(skip): + x = mx.array(noise) + for i in range(len(s) - 1): + sig = np.reshape(s[i], (1,)) + vp = model(x, sig, pos.features[None]) + vn = model(x, sig, neg.features[None], skip=skip) + v = api.guide(vp, vn, 3.0) + x = x + (float(s[i + 1]) - float(s[i])) * v + return np.array(x) + + def run(**kw): + return np.array(api.sample(pipe, pos, neg=neg, cfg=3.0, + height=16 * H, width=16 * W, steps=4, + noise=noise, **kw)) + + with mx.stream(mx.cpu): + skipped, plain = run(neg_skip=(1,)), run(neg_skip=()) + want, want_plain = by_hand((1,)), by_hand(()) + assert np.allclose(skipped, want, atol=1e-5) + assert np.allclose(plain, want_plain, atol=1e-5) + assert np.abs(skipped - plain).max() > 1e-2 + with pytest.raises(ValueError, match="neg_skip needs guidance"): + api.sample(pipe, pos, cfg=1.0, height=16 * H, width=16 * W, + steps=4, noise=noise, neg_skip=(1,)) + + +def test_both_variants_guide_at_3_with_block_22_left_out_of_the_negative(): + from gmlx.image import api + + for variant in ("sft", "base"): + d = api.sampler_defaults("moss_image", variant) + assert (d.steps, d.cfg, d.neg_skip) == (50, 3.0, (22,)) + + +def test_a_dit_that_skips_no_blocks_refuses_a_skip_set(): + from gmlx.image import api + + model, _ = _tiny() + + class Plain: + cfg = model.cfg + act_dtype = mx.float32 + + def __call__(self, *a, **k): + return model(*a, **k) + + pipe = api.ImagePipeline(dit=Plain(), family=get_family("moss_image"), + variant="sft") + lat = mx.array(_inputs()[0]) + with pytest.raises(ValueError, match="DiT skips no blocks"): + api.denoise_step(pipe, lat, SIGMA, _cond(5, 1), skip=(1,)) + + +def test_a_skip_set_of_other_than_block_indices_is_refused(): + """A fractional block is refused, not cut to the block below it.""" + from gmlx.image import api + + model, _ = _tiny() + pipe = api.ImagePipeline(dit=model, family=get_family("moss_image"), + variant="sft") + noise = np.zeros((1, 8, H, W), np.float32) + for bad in ((1.5,), (True,), ("1",)): + with pytest.raises(ValueError, match="block indices"): + with mx.stream(mx.cpu): + api.sample(pipe, _cond(5, 1), neg=_cond(3, 2), cfg=3.0, + height=16 * H, width=16 * W, steps=2, + noise=noise, neg_skip=bad) + with mx.stream(mx.cpu): + api.sample(pipe, _cond(5, 1), neg=_cond(3, 2), cfg=3.0, + height=16 * H, width=16 * W, steps=2, noise=noise, + neg_skip=(np.int64(1),)) + + +def test_a_skip_set_given_as_an_array_or_a_generator_skips_its_blocks(): + """A numpy array of neg_skip and a generator of skip each leave out + the blocks they hold, as the tuple of the same blocks does.""" + from gmlx.image import api + + model, _ = _tiny() + pipe = api.ImagePipeline(dit=model, family=get_family("moss_image"), + variant="sft") + noise = np.random.default_rng(4).standard_normal( + (1, 8, H, W)).astype(np.float32) + pos, neg = _cond(5, 1), _cond(3, 2) + + def run(skip): + return np.array(api.sample(pipe, pos, neg=neg, cfg=3.0, + height=16 * H, width=16 * W, steps=2, + noise=noise, neg_skip=skip)) + + sigma = np.float32(0.75) + with mx.stream(mx.cpu): + assert np.array_equal(run(np.array([0, 1])), run((0, 1))) + want = np.array(api.denoise_step(pipe, mx.array(noise), sigma, neg, + skip=(1,))) + got = np.array(api.denoise_step(pipe, mx.array(noise), sigma, neg, + skip=(i for i in (1,)))) + whole = np.array(api.denoise_step(pipe, mx.array(noise), sigma, + neg)) + assert np.array_equal(got, want) + assert not np.array_equal(whole, want) diff --git a/tests/image/test_image_qwen_vae.py b/tests/image/test_image_qwen_vae.py index c4536a5d1..1a428b512 100644 --- a/tests/image/test_image_qwen_vae.py +++ b/tests/image/test_image_qwen_vae.py @@ -937,8 +937,8 @@ def test_other_files_and_families_are_refused(tmp_path, st_file): reader.load_vae(path, family=FAMILY, dtype=mx.float32) with pytest.raises(reader.VAEError, match="known: flux2_klein, flux2_klein_9b, " - "flux2_klein_9b_kv, flux3_klein, krea2, qwen_image21, " - "z_image"): + "flux2_klein_9b_kv, flux3_klein, krea2, moss_image, " + "qwen_image21, z_image"): reader.load_vae(st_file, family="sd15", dtype=mx.float32) # A file with another number of levels has no official time flags. for i, config in enumerate((OTHER_FLAGS["three_levels"], @@ -1024,7 +1024,7 @@ def test_the_packer_gives_a_refusal_once_when_the_families_share_it( pack_vae(video, tmp_path / "video.gguf", **kw) text = str(e.value) assert text.startswith(f"{video} matches no image family's VAE: ") - assert text.count("holds a video VAE in the Wan 2.2 layout") == 7 + assert text.count("holds a video VAE in the Wan 2.2 layout") == 8 assert "qwen_image21 decodes with" in text assert "z_image decodes with" in text assert "flux3_klein decodes with" in text diff --git a/tests/image/test_image_upsample.py b/tests/image/test_image_upsample.py index b0c6b766e..de4e65c2f 100644 --- a/tests/image/test_image_upsample.py +++ b/tests/image/test_image_upsample.py @@ -44,6 +44,10 @@ ("z_image", "z_image", "pe_prompt_template.txt", "36bdd888b953e9b7a8237dab55bee145a44e96e92e8f6e81a6e4d433a401a0b6", None, (0.0, 2048, 0.0, "off")), + # gmlx's own caption prompt, from the OpenMOSS prompt-enhancer skill. + ("moss_image", "moss_image", "caption_prompt.txt", + "389ce41e0d9f4e11df7f0c078cef7a4af772b7312b19fc7d637e5808345cbba1", + None, (0.0, 2048, 0.0, "off")), ]) def test_the_chat_turn_is_the_sources_t2i_turn(family, package, name, sha256, key, sampling): @@ -179,6 +183,20 @@ def test_z_image_takes_revised_prompt_or_the_reply_as_written(): assert upsample.upsample(fam, "x", lambda req: sign) == sign +def test_moss_image_takes_the_caption_from_its_json_or_the_reply_as_written(): + """The caption keeps its paragraphs and the text the image shows; + a reply with no object is the caption as written.""" + fam = get_family("moss_image") + caption = ('A bakery storefront under a painted sign that reads "FRESH ' + 'BREAD DAILY".\n\nLoaves fill the window in warm light.') + reply = json.dumps({"caption": caption}) + assert upsample.upsample(fam, "x", lambda req: reply) == caption + fenced = "```json\n" + reply + "\n```" + assert upsample.upsample(fam, "x", lambda req: fenced) == caption + assert upsample.upsample(fam, "x", lambda req: CAPTION) == CAPTION + with pytest.raises(upsample.UpsampleError): + upsample.upsample(fam, "x", lambda req: '{"layout": []}') + @pytest.mark.parametrize("reply", [ "<|channel>thought\nThe user wants a bike.{caption}", "<|channel|>analysis<|message|>The user wants a bike.<|end|>" diff --git a/tests/image/test_image_vae.py b/tests/image/test_image_vae.py index 4badbd538..c244bfe91 100644 --- a/tests/image/test_image_vae.py +++ b/tests/image/test_image_vae.py @@ -400,7 +400,7 @@ def test_other_files_and_families_are_refused(tmp_path): reader.load_vae(junk, family="z_image", dtype=mx.float32) with pytest.raises(reader.VAEError, match="known: flux2_klein, " "flux2_klein_9b, flux2_klein_9b_kv, flux3_klein, " - "krea2, qwen_image21, z_image"): + "krea2, moss_image, qwen_image21, z_image"): reader.load_vae(path, family="sd15", dtype=mx.float32) diff --git a/tests/image/test_image_vae_rows.py b/tests/image/test_image_vae_rows.py index d2868d0d3..4f8d9d249 100644 --- a/tests/image/test_image_vae_rows.py +++ b/tests/image/test_image_vae_rows.py @@ -102,7 +102,7 @@ def test_every_family_has_a_vae_row(): with pytest.raises(ValueError, match="no VAE memory row for image " "family 'flux'; known: flux2_klein, " "flux2_klein_9b, flux2_klein_9b_kv, flux3_klein, " - "krea2, qwen_image21, " + "krea2, moss_image, qwen_image21, " "z_image"): price("decode", 512, 512, family="flux", itemsize=2) @@ -143,6 +143,9 @@ def _official_vaes() -> dict: # reads. out["flux2_klein"] = out["flux2_klein_9b"] = out["flux3_klein"] out["flux2_klein_9b_kv"] = out["flux3_klein"] + out["moss_image"] = flux2_ae.Flux2AEConfig( + latent_channels=32, decoder_block_out_channels=(96, 192, 384, 384), + **flux) assert set(out) == set(reader.VAES) return out diff --git a/tests/serve/test_image_routes.py b/tests/serve/test_image_routes.py index ba09bb1a6..4774acd45 100644 --- a/tests/serve/test_image_routes.py +++ b/tests/serve/test_image_routes.py @@ -667,7 +667,7 @@ def test_a_shared_encoder_encodes_as_a_step_of_the_chat_engine(monkeypatch): monkeypatch.setattr(app_mod, "get_cached_model", lambda *a, **k: None) monkeypatch.setattr(runtime_mod, "runtime", SimpleNamespace(response_generator=rg)) - monkeypatch.setattr(api, "_borrowed", lambda pair, spec: SimpleNamespace( + monkeypatch.setattr(api, "_borrowed", lambda pair, spec, gguf=None: SimpleNamespace( model=pair[0], tokenizer=pair[1])) import gmlx.image.text_encoder as text_encoder monkeypatch.setattr(text_encoder, "encoder_spec", lambda fam: "SPEC") diff --git a/tests/serve/test_image_steps.py b/tests/serve/test_image_steps.py index c924903dd..a57359119 100644 --- a/tests/serve/test_image_steps.py +++ b/tests/serve/test_image_steps.py @@ -172,7 +172,8 @@ def test_capabilities_answer_without_the_version_header(host): assert caps["contract"] == {"major": 1, "minor": 0} assert caps["families"] == ["flux2_klein", "flux2_klein_9b", "flux2_klein_9b_kv", "flux3_klein", "krea2", - "qwen_image21", "z_image"] + "moss_image", "qwen_image21", "z_image"] + assert "skip_blocks" in caps["features"] assert caps["conventions"] == ["comfy", "reference"] assert caps["wire"]["container"] == "safetensors" @@ -379,9 +380,12 @@ def test_num_tokens_and_an_attention_mask_trim_the_cond(host, pipe): def test_the_convention_maps_the_timestep_and_signs_the_output( host, pipe, monkeypatch): - # ComfyUI hands every family sigma as the timestep and wants the - # velocity back, so every convention maps the identity. Its FLUX.2 - # forward alone embeds the fp32 time, at both sizes and on KV. + # ComfyUI hands every family but MOSS-Image sigma as the timestep and + # wants the velocity back, so those conventions map the identity. Its + # FLUX.2 forward alone embeds the fp32 time, at both sizes and on KV. + # MOSS-Image's model takes sigma times 1000, as the release's does. + assert steps.CONVENTIONS["comfy"]["moss_image"] == steps.Convention( + scale=0.001) for fam in ("z_image", "qwen_image21", "krea2", "flux3_klein"): assert steps.CONVENTIONS["comfy"][fam] == steps.Convention() for fam in ("flux2_klein", "flux2_klein_9b", "flux2_klein_9b_kv"): @@ -402,6 +406,104 @@ def test_the_convention_maps_the_timestep_and_signs_the_output( assert np.array_equal(out, -want) +# MOSS-Image: ComfyUI's timestep is sigma times 1000, and a step leaves +# out the blocks of its skip, as the negative pass of the release does. + +MOSS_SHAPE = (1, 8, 3, 4) # 48x64 pixels at the stride of 16 + + +@pytest.fixture(scope="module") +def mpipe(): + def build(): + from image_fixtures import tiny_moss_model + + return api.ImagePipeline(dit=tiny_moss_model(), + family=get_family("moss_image"), + variant="sft") + return _on_worker(build) + + +def _moss_velocity(pipe, x, sigma, feats, skip=()): + return _on_worker(lambda: np.array(api.denoise_step( + pipe, mx.array(x), sigma, api.Conditioning(mx.array(feats[0]), "m"), + skip=skip))) + + +@pytest.mark.parametrize("skip", [[], [1], [2, 0]]) +def test_a_moss_step_reads_comfyuis_timestep_and_leaves_out_its_skip( + host, mpipe, skip): + """A timestep of sigma times 1000 in float32 gives the velocity at + sigma bit for bit, and the rows leave out the blocks the step names.""" + _serve(host, "moss", mpipe) + c = _client() + feats = _features(4, n=6) + x = noise_from_seed(9, MOSS_SHAPE) + sigma = np.float32(0.6875) + t = np.float32(sigma * np.float32(1000)) + out = c.step("moss", None, x, [t], _rows(c, feats), skip=skip) + want = _moss_velocity(mpipe, x, sigma, feats, tuple(skip)) + assert np.array_equal(out, want) + if skip: + assert not np.array_equal(out, _moss_velocity(mpipe, x, sigma, feats)) + + +def test_a_moss_guided_loop_with_a_negative_skip_gives_the_latents_of_sample( + host, mpipe): + _serve(host, "moss", mpipe) + c = _client() + pos, neg = _features(0, n=6), _features(1, n=4) + noise = noise_from_seed(3, MOSS_SHAPE) + got = c.sample("moss", noise, SIGMAS, _rows(c, pos), cfg=3.0, + neg_rows=_rows(c, neg), neg_skip=(1,), timestep_scale=1000) + + def direct(): + x = api.sample(mpipe, api.Conditioning(mx.array(pos[0]), "p"), + neg=api.Conditioning(mx.array(neg[0]), "n"), cfg=3.0, + height=48, width=64, noise=mx.array(noise), + sigmas=SIGMAS, neg_skip=(1,)) + return np.array(x) + assert np.array_equal(got, _on_worker(direct)) + + +@pytest.mark.parametrize("model, skip, match", [ + ("moss", "1", "skip is a list of block indices"), + ("moss", [True], "skip is a list of block indices"), + ("moss", [1.0], "skip is a list of block indices"), + ("moss", [3], r"has blocks 0 to 2, and skip names \[3\]"), + ("moss", [-1, 1], r"skip names \[-1\]"), + ("zit", [1], "the Z-Image DiT skips no blocks"), +]) +def test_a_skip_that_names_no_block_of_the_dit_is_refused(host, mpipe, model, + skip, match): + _serve(host, "moss", mpipe) + c = _client() + feats = _features(4, n=6) + shape = MOSS_SHAPE if model == "moss" else SHAPE + with pytest.raises(StepsClientError, match=match): + c.container("step", {"x": np.zeros(shape, np.float32), + "timestep": np.array([500.0], np.float32)}, + {"model": model, "convention": "comfy", + "rows": _rows(c, feats), "skip": skip}) + assert host.held == 0 + + +def test_the_moss_status_names_the_packs_class_and_its_negative_skip( + host, mpipe, monkeypatch): + monkeypatch.setattr(steps, "family_of", lambda spec: ( + get_family("moss_image"), "sft")) + monkeypatch.setattr(api, "_header_bytes", lambda path, part: 7) + _serve(host, "moss", mpipe) + st = _client().status("moss") + assert st["comfy"] == {"model_class": "MOSSImage", + "image_model": "moss_image"} + assert not st["edit_capable"] + assert st["latent"] == {"channels": 128, "stride": 16, "patch": 1} + assert st["vae"] == {"channels_out": 3, "stride": 16, "space": "model"} + assert st["sampling"] == {"steps": 50, "cfg": 3.0, "shift": 5.0, + "sigmas": None, "schedule": None, + "neg_skip": [22]} + + # FLUX.3 klein: a flow map reads the time its step goes to from the # client's schedule, and the VAE routes carry model-space latents. @@ -598,7 +700,8 @@ def test_the_flux2_9b_status_names_comfyuis_flux2_as_4b_does( assert st["latent"] == {"channels": 128, "stride": 16, "patch": 1} assert st["vae"] == {"channels_out": 3, "stride": 16, "space": "model"} assert st["sampling"] == {"steps": 4, "cfg": 1.0, "shift": 1.0, - "sigmas": None, "schedule": {"rule": "flux2"}} + "sigmas": None, "schedule": {"rule": "flux2"}, + "neg_skip": []} def test_the_flux3_vae_routes_carry_model_space_latents(host, fpipes, @@ -948,7 +1051,8 @@ def test_the_flux2_status_names_comfyuis_flux2_on_model_space_latents( assert st["latent"] == {"channels": 128, "stride": 16, "patch": 1} assert st["vae"] == {"channels_out": 3, "stride": 16, "space": "model"} assert st["sampling"] == {"steps": 4, "cfg": 1.0, "shift": 1.0, - "sigmas": None, "schedule": {"rule": "flux2"}} + "sigmas": None, "schedule": {"rule": "flux2"}, + "neg_skip": []} def test_the_flux2_vae_routes_carry_model_space_latents(host, f2pipe, diff --git a/tests/test_comfyui_pack.py b/tests/test_comfyui_pack.py index 065cb5b8e..7916ddd86 100644 --- a/tests/test_comfyui_pack.py +++ b/tests/test_comfyui_pack.py @@ -67,7 +67,7 @@ def test_the_tested_comfyui_runs_every_family_the_pack_serves(): "krea2": "Krea 2", "flux2_klein": "FLUX.2 klein", "flux2_klein_9b": "FLUX.2 klein 9B", "flux2_klein_9b_kv": "FLUX.2 klein 9B KV", - "flux3_klein": "FLUX.3 klein"} + "flux3_klein": "FLUX.3 klein", "moss_image": "MOSS-Image"} def _served_families() -> set: @@ -94,6 +94,29 @@ def test_the_readme_lists_every_family_the_pack_serves_with_its_parity_rows(): (fam, name) + +def test_the_moss_image_step_rows_sit_under_headers_that_name_their_cells(): + """MOSS-Image reads its step distances from stock fp32 against the + larger of the stock bf16 and in-process gmlx ones, so its rows sit in + a table of their own whose headers name both.""" + parity = (_PACK / "README.md").read_text().split( + "### Parity with the stock loaders", 1)[1] + step = parity.split("The final image check", 1)[0] + header = None + rows = [] + for ln in step.splitlines(): + if ln.startswith("| Model |"): + header = ln + elif ln.startswith("| MOSS-Image |"): + rows.append((header, ln)) + assert rows + for header, row in rows: + assert header is not None + cols = [c.strip() for c in header.strip("|").split("|")] + assert "gmlx in process, same step" in cols, header + assert "stock bf16, same step" in cols, header + assert len(row.strip("|").split("|")) == len(cols), row + def _pack_files() -> list[str]: out = [] for p in sorted(_PACK.rglob("*")): @@ -122,7 +145,7 @@ def test_the_wheel_ships_every_pack_file(): for want in ("gmlx/comfyui_pack/LICENSE", "gmlx/comfyui_pack/__init__.py", "gmlx/comfyui_pack/version.py"): assert want in files - assert len([f for f in files if f.startswith("gmlx/comfyui_pack/example_workflows/")]) == 10 + assert len([f for f in files if f.startswith("gmlx/comfyui_pack/example_workflows/")]) == 11 def test_the_pack_keeps_its_apache_license(): @@ -203,6 +226,8 @@ def test_gmlx_pytest_collects_no_pack_suite(monkeypatch): "ReferenceLatent": [], "GetImageSize": [], "FluxKVCache": [], + "ModelSamplingSD3": ["shift"], + "SkipLayerGuidanceDiTSimple": ["double_layers", "single_layers", "start_percent", "end_percent"], } _FRONTEND_ONLY = {"control_after_generate", "upload"}