diff --git a/.claude/skills/add-anima-model/SKILL.md b/.claude/skills/add-anima-model/SKILL.md index 7f97f2a..48b43e8 100644 --- a/.claude/skills/add-anima-model/SKILL.md +++ b/.claude/skills/add-anima-model/SKILL.md @@ -65,8 +65,8 @@ What to read out of it: the first, the community 2.9B expansion the second; the exporter strips both and refuses to find anything under a third. - **`x_embedder.proj.1.weight` is `[hidden, 68]`.** `68 // 4 - 1` is 16 latent channels, which is - the text-to-image row of ComfyUI's `model_detection.py` and **the only row the exporter has - rotary extrapolation ratios for**. Anything else it refuses by name rather than guessing, which + the Cosmos-Predict2 2B text-to-image model and **the only width the exporter has rotary + extrapolation ratios for**. Anything else it refuses by name rather than guessing, which is right: guessing those ratios is exactly how they were wrong for months (`docs/anima.md`, "Positions"). - **`final_layer.linear.weight` is `[64, hidden]`** -- 2x2 patches over those 16 channels. @@ -147,9 +147,15 @@ python3 -m venv .venv .venv/bin/pip install -r tools/requirements.txt ``` -**A ComfyUI checkout is needed only for `-test_output`.** It is the reference the test tensors are -computed with -- the arrangement `sdxl_exporter.py` has with diffusers -- and -`tools/anima_comfy.py` says what to install and where. A plain export needs none of it. +**`-test_output` needs a second venv**, from `tools/requirements-anima.txt`. The test tensors +are computed with diffusers' Anima models (`tools/anima_reference.py`) -- the arrangement +`sdxl_exporter.py` has with diffusers -- and those need diffusers 0.41 and a transformers with +Qwen3, which cannot share `.venv` with the SDXL pair. A plain export needs none of it. + +```bash +python3 -m venv .venv-anima +.venv-anima/bin/pip install -r tools/requirements-anima.txt +``` ## 3. Export @@ -381,9 +387,10 @@ produces `Aborted: out of memory`, which reads as a code failure and is not one. | symptom | cause | fix | |---|---|---| -| `this exporter only knows the rotary extrapolation ratios for the 16 channel text-to-image branch` | the denoiser's latent width is not 16, so ComfyUI's `model_detection.py` would take a different row | not a bug to route around: find the ratios for that row and add it deliberately | +| `this exporter only knows the rotary extrapolation ratios for the 16 channel text-to-image branch` | the denoiser's latent width is not 16, so it is a different base model with ratios of its own | not a bug to route around: find that model's ratios and add them deliberately | | the model will not load: a manifest key is missing | it was exported before the rotary ratios were written into manifests, and the reader does not default them | re-export. Defaulting them is how they were silently wrong for months | -| `-test_output` exits with "needs -comfyui" | the reference outputs are what ComfyUI computes; there is no second implementation here to ask | pass a ComfyUI checkout, per `tools/anima_comfy.py` | +| `-test_output` exits with "needs diffusers with Anima and transformers with Qwen3" | the exporter ran in `.venv`, whose diffusers 0.31 predates Anima | run it from `.venv-anima`, per `tools/requirements-anima.txt` | +| `decodes_a_latent_the_way_the_reference_does` is off by about 0.05 | diffusers clamps a decode to [-1, 1] and the runtime does not; the test clamps the runtime's picture first | if that clamp was removed from the test, put it back | | the export dies fetching a tokenizer | neither vocabulary is published beside Anima's weights -- they come from `Qwen/Qwen3-0.6B-Base` and `google/t5-v1_1-xxl` | pass `-qwen_tokenizer` / `-t5_tokenizer` local paths | | `OverflowError: ... which fp16 cannot hold` | a weight outside fp16's range, which bf16 has and fp16 has not | the check is right; the format would otherwise write an infinity | | a picture that is soft and unfinished, or burnt and over-contrasted | turbo's numbers on an aesthetic release, or the reverse | `suggested.steps`/`guidance` belong to the release, not to Anima | @@ -405,11 +412,10 @@ published stem (`anima-turbo-v1.1`) is deliberately a different one. Do not rena match the published spelling. ```bash -.venv/bin/python tools/anima_exporter.py \ +.venv-anima/bin/python tools/anima_exporter.py \ -dit models/... -text models/... -vae models/... \ -output models/anima-turbo-v11.safetensors \ - -test_output models/anima-turbo-v11_test.safetensors \ - -comfyui /path/to/ComfyUI + -test_output models/anima-turbo-v11_test.safetensors ``` **Regenerate it whenever `export_test_cases` changes**, and whenever anything it depends on does: @@ -418,7 +424,7 @@ hidden states, the adapter's context before and after padding, one step's veloci The corpora are 611 texts each and are what catch a package built against the wrong tokenizer revision, which no reference tensor can. -The test export runs the whole model on CPU through ComfyUI and takes a few minutes. +The test export runs the whole model on CPU through diffusers and takes under a minute. ## Disk diff --git a/docs/anima.md b/docs/anima.md index b70ac12..c1ed86a 100644 --- a/docs/anima.md +++ b/docs/anima.md @@ -2,9 +2,10 @@ Notes towards running Anima in libwaifu. The shapes were read off the published tensor headers over range requests, the conv folding was proved with a script that still runs, and the behaviour -that shapes alone do not show -- the text path, the position layout, the latent scaling -- came -from the reference implementations in ComfyUI (`comfy/ldm/anima/model.py`, -`comfy/ldm/cosmos/predict2.py`, `comfy/text_encoders/anima.py`) and diffusers. +that shapes alone do not show -- the text path, the position layout, the latent scaling -- is +diffusers' (Apache-2.0): `AnimaTextConditioner`, `CosmosTransformer3DModel`, +`AutoencoderKLQwenImage` and `AnimaModularPipeline`, which is also what the test tensors are +computed with -- see `tools/anima_reference.py`. Anima is **not an SDXL fine tune**, so none of `waifu/src/sdxl/` applies to it and the `add-sdxl-model` skill does not describe this job. It is a diffusion transformer derived from @@ -86,31 +87,22 @@ the way out, which is why the output is 64 and the input 68. ## How a block actually runs -Shapes do not give any of this and most of it is a silent error. Read off -`comfy/ldm/cosmos/predict2.py`. +Shapes do not give any of this and most of it is a silent error. diffusers' +`transformer_cosmos.py` is where to check it. -The timestep path is the one that reads backwards. `Timesteps` builds a sinusoid -- -`exponent = -log(10000) * arange(half) / half`, then `cat([cos, sin])`, cosine first -- and hands -it to `TimestepEmbedding`, whose two linears produce **only the AdaLN-LoRA term**: +The timestep path is the one that reads backwards. The timestep becomes a sinusoid -- +`exponent = -log(10000) * arange(half) / half`, then `cat([cos, sin])`, cosine first -- and the +two linears of `t_embedder` turn that into **only the AdaLN-LoRA term**, 6144 wide. They never +feed the modulation directly. What every block modulates on is `t_embedding_norm(sinusoid)` -- +the RMSNorm applied to the *input* of those linears -- and the LoRA term rides alongside. Reading +the linears' output as the timestep embedding is the obvious mistake and it is wrong. -```python -emb = self.linear_2(self.activation(self.linear_1(sample))) -if self.use_adaln_lora: - adaln_lora_B_T_3D = emb # 6144 wide - emb_B_T_D = sample # the raw sinusoid, not the projection -``` - -So `t_embedder`'s linears never feed the modulation directly. What every block modulates on is -`t_embedding_norm(sinusoid)` -- the RMSNorm applied to the *input* of those linears -- and the -6144-wide LoRA rides alongside. Reading `emb` as the projection is the obvious mistake and it is -wrong. +Each block then, three times over, for the self attention, the cross attention and the MLP: -Each block then, three times over: - -```python -shift, scale, gate = (adaln_modulation_(emb) + adaln_lora_B_T_3D).chunk(3, dim=-1) -normalized = layer_norm(x) * (1 + scale) + shift -x = x + gate * sublayer(normalized) # torch.addcmul +``` +shift, scale, gate = split3(adaln_modulation_(emb) + adaln_lora) +normalized = layer_norm(x) * (1 + scale) + shift +x = x + gate * sublayer(normalized) ``` The same `adaln_lora` is added to all three. `adaln_modulation_*` is @@ -125,17 +117,14 @@ everywhere else in this model. In the attentions, `q_norm` and `k_norm` are `RMSNorm(head_dim, eps=1e-6)` applied per head; `v` is not normalized. Nothing carries a bias. And **rotary position embedding is applied to self attention only** -- the cross attention is explicitly excluded, which makes sense once the context -is 512 padded text tokens with no geometry, but it is one line in the reference and invisible -here. The MLP is a plain `Linear -> GELU -> Linear`. +is 512 padded text tokens with no geometry, but nothing in the weights says so. The MLP is a plain `Linear -> GELU -> Linear`. ### The residual stream has to be fp32 -Straight from the reference, and the reason is not obvious: - -> The residual stream for this model has large values. To make fp16 compute_dtype work, we keep -> the residual stream in fp32, but run attention and MLP modules in fp16. An alternate method that -> clamps fp16 values "works" in the sense that it makes coherent images, but there is noticeable -> quality degradation and visual artifacts. +The residuals of this model are large enough that in fp16 they reach infinity part way through +the stack, and every step after that is NaN. So the stream between blocks is kept in fp32 while +the attentions and the MLP run in fp16. Clamping fp16 instead would keep the numbers finite but +would no longer be the model. `waifu::sdxl` runs fp16 throughout and has no reason not to. This one does. The weights are fine narrow -- the exporter checks every one against fp16's range and none of turbo v1.1 comes close -- @@ -144,11 +133,8 @@ but the accumulating `x` between blocks is not. And a wide stream is not enough on its own: **the places that step on and off it have to widen before they multiply, not after.** There are two. -```python -x = torch.addcmul(x, gate.to(residual_dtype), result.to(residual_dtype)) -``` - -Both operands are cast to float32 *before* the multiply. Gating in half and widening the product +The gated residual add, `x + gate * sublayer(...)`, casts both the gate and the sublayer's output +to float32 *before* the multiply. Gating in half and widening the product is a different computation, and the difference is an infinity. The same goes for the modulation: `norm(x) * (1 + scale) + shift` is done on the wide `x`, and only its result is narrowed for the projections that follow. @@ -314,23 +300,26 @@ what exists. | | state | agrees with the reference to | |---|---|---| | `tools/anima_exporter.py` | done | -- | -| `tools/anima_comfy.py` | done; the reference is ComfyUI itself | -- | +| `tools/anima_reference.py` | done; the reference is diffusers' Anima models | -- | | `waifu/src/anima/config.rs` | done | -- | -| `waifu/src/anima/text_encoder.rs` | done | 0.0029 | -| `waifu/src/anima/adapter.rs` | done | 0.0024 | -| `waifu/src/anima/dit.rs` | done | 0.0028 | -| `waifu/src/anima/vae.rs` | decoder done. The package carries the encoder's weights and the reference outputs to check one against, but the encoder itself is unwritten, so no img2img yet | 0.00072 | +| `waifu/src/anima/text_encoder.rs` | done | 0.0039 | +| `waifu/src/anima/adapter.rs` | done | 0.0025 | +| `waifu/src/anima/dit.rs` | done | 0.012 | +| `waifu/src/anima/vae.rs` | decoder done. The package carries the encoder's weights and the reference outputs to check one against, but the encoder itself is unwritten, so no img2img yet | 0.00070 | | `waifu/src/anima/sampler.rs` | done | -- | | the four of them together | draws a 512 by 512 picture in ten steps | -- | | tokenizers | done; both in the package | token for token, over 611 texts each | -| `waifu/src/anima/pipeline.rs` | done | 0.0031, prompt to context | +| `waifu/src/anima/pipeline.rs` | done | 0.0025, prompt to context | | package + CLI | done. `waifu draw` reads `[model] type` and takes the package to whichever model it names | | `hub.rs` `CATALOG` | done, once it was asked for: turbo v1.1 is published to both hubs as `anima:turbo`, under CircleStone Labs' licence rather than libwaifu's. See Licensing | flint: nothing, as promised. -The RMSE numbers are relative, against `tools/anima_comfy.py`, which runs ComfyUI in float32 where the -runtime runs in half. The pipeline's is the encoder and the adapter composed, from a prompt as a +The RMSE numbers are relative, against `tools/anima_reference.py`, which runs diffusers in +float32 where the runtime runs in half, measured on turbo v1.1. The decoder's is after clamping +the runtime's picture to [-1, 1], which diffusers' `decode` does and this runtime leaves to the +conversion to pixels; unclamped, a decode of a unit-normal latent differs by 0.048, all of it in +the tenth of the picture outside that range. The pipeline's is the encoder and the adapter composed, from a prompt as a string through to the padded context the denoiser reads, so it lands a little above either of them alone. The decoder is the closest of them, which is worth a second look: it is the stage that appeared to need a kernel that does not exist, and folding a causal convolution over one @@ -367,7 +356,7 @@ range, what it would otherwise write is an infinity, and what that draws is nois ### And the reference outputs `-test_output` writes a second, small package the way `sdxl_exporter.py` does, holding what -`tools/anima_comfy.py` makes ComfyUI compute for one fixed prompt: the encoder's hidden states, the adapter's +`tools/anima_reference.py` makes diffusers compute for one fixed prompt: the encoder's hidden states, the adapter's context before and after padding, the velocity of one denoising step, one decode, and then that decode's picture encoded again and decoded once more. A 16 by 16 latent, so 64 patches and a 128 by 128 image -- every block exercised, 2.4 MB to carry. @@ -413,9 +402,11 @@ gets its own base, `10000.0 * ntk_factor`, where the factor is Rotation is NeoX `rotate_half`, which is the convention flint's `rotaryEmbedding` already implements. -**The ratios are not all 1.0, and this file said for months that they were.** `model_detection.py` -branches on `in_channels`, and Anima's is 16 -- `x_embedder.proj.1.weight` is `[2048, 68]`, so -`68 // 4 - 1` -- which takes the branch asking for four times the spatial extent: +**The ratios are not all 1.0, and this file said for months that they were.** Anima's denoiser +is the 16 channel Cosmos-Predict2 2B text-to-image model -- `x_embedder.proj.1.weight` is +`[2048, 68]`, so `68 // 4 - 1` -- and that model asks for four times the spatial extent, which is +`rope_scale = [1.0, 4.0, 4.0]` (time, height, width) in the transformer config of +`circlestone-labs/Anima-Base-v1.0-Diffusers`: ``` rope_h_extrapolation_ratio = 4.0 @@ -429,37 +420,25 @@ choice, and it is the kind of error that draws a perfectly reasonable picture: m 16x16 latent, a flat 10000 against the real bases is `max|d| = 5.5e-01` on the velocity, 13% relative, at a cosine similarity of 0.9959. Close enough to look right, nowhere near right. -It was found by standing the hand transcription that used to live in `tools/anima_reference.py` -next to ComfyUI's own `comfy.ldm.anima.model.Anima`, on the same weights in float32. That -transcription is gone now, and this is why: it had been checked only by the fact that it drew a -coherent picture, the test tensors were generated from it, the runtime was built to match those, -and the suite was green throughout -- every one of those measured agreement with the -transcription rather than with Anima. **A reference that is not the upstream implementation is a -ruler nobody checks.** `tools/anima_comfy.py` replaces it by running ComfyUI itself, the same -arrangement `sdxl_exporter.py` has with diffusers, and `-test_output` now requires `-comfyui`. +It was found by putting the hand transcription that used to be the reference beside an +independent implementation, on the same weights in float32. That transcription is gone now, and +this is why: it had been checked only by the fact that it drew a coherent picture, the test +tensors were generated from it, the runtime was built to match those, and the suite was green +throughout -- every one of those measured agreement with the transcription rather than with +Anima. **A reference that is not an implementation somebody else maintains is a ruler nobody +checks.** `tools/anima_reference.py` runs diffusers' Anima models, the same arrangement +`sdxl_exporter.py` has with diffusers. Three things carried the assumption and moved together to be rid of it. `anima_exporter.py` writes the ratios into the manifest as `rope_h_extrapolation_ratio` and its two neighbours -- written rather than assumed, because assuming them is how they were wrong -- and refuses a -checkpoint whose latent width it has no row for. `DitConfig::rope_bases` turns them into one base -per axis, which `waifu/src/anima/dit.rs` now gives to `frequencies` separately rather than -sharing `rope_theta` between all three. And the test tensors were regenerated, this time by -ComfyUI. - -While the four stages were being put side by side, the other three were measured too, which had -never been done. They were right all along -- but nothing had established that, and the way the -rotary base was wrong is exactly how any of them could have been: - -| stage | ComfyUI against the old transcription | -|---|---| -| Qwen3 text encoder | `6.1e-05` | -| LLM adapter | `0.0` exactly | -| denoiser | `8.9e-06` | -| VAE decoder | `2.1e-06` | - -The last of those is worth its own line: libwaifu decodes through convolutions folded from three -dimensions to two, ComfyUI through the real `Conv3d`, and they agree to `2.1e-06` on the -published weights. The folding had only ever been proved on synthetic ones. +checkpoint whose latent width it has no ratios for. `DitConfig::rope_bases` turns them into one +base per axis, which `waifu/src/anima/dit.rs` now gives to `frequencies` separately rather than +sharing `rope_theta` between all three. And the test tensors were regenerated. + +The decoder's agreement is worth its own line: libwaifu decodes through convolutions folded from +three dimensions to two, diffusers through the real `Conv3d`, and they agree to `7.0e-04` on the +published weights in half precision. The folding had only ever been proved on synthetic ones. A model written before this has no ratios in its manifest and will not load: the reader asks for the three keys and does not default them. That is deliberate. Defaulting to 1.0 would go back diff --git a/tools/anima_comfy.py b/tools/anima_comfy.py deleted file mode 100644 index a18454a..0000000 --- a/tools/anima_comfy.py +++ /dev/null @@ -1,238 +0,0 @@ -#!/usr/bin/env python3 -# The MIT License (MIT) -# -# Copyright (c) 2026 Xiaoyang Chen -# -# Permission is hereby granted, free of charge, to any person obtaining a copy of this software -# and associated documentation files (the "Software"), to deal in the Software without -# restriction, including without limitation the rights to use, copy, modify, merge, publish, -# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the -# Software is furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all copies or -# substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING -# BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND -# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, -# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. - -"""The Anima reference, which is ComfyUI itself. - -`sdxl_exporter.py` gets its reference outputs by running the diffusers pipeline. This is the same -arrangement for Anima, against ComfyUI, which is the implementation Anima is published for and -the only one that can be called an authority on it. - -There used to be a hand transcription here instead -- `anima_reference.py`, ~450 lines of torch -copied out of `comfy/ldm/cosmos/predict2.py` by eye. It was wrong, and the way it was wrong is -the argument against ever writing another one: it assumed every rotary extrapolation ratio was -1.0 where Anima asks for 4.0 on the two spatial axes, so every position in the picture rotated on -a base of 10000 instead of 42870.9. The velocity it produced was 13% off. It still drew perfectly -reasonable pictures, the runtime was built to match it, the test tensors were generated from it, -and the suite was green the whole time -- because every one of those measured agreement with the -transcription rather than with Anima. A reference that is not the upstream implementation is a -ruler nobody checks. - -Nothing here is used at run time. It is imported by `anima_exporter.py -test_output` and by -nothing else, and it needs a ComfyUI checkout: - - git clone --depth 1 https://github.com/comfyanonymous/ComfyUI.git /path/to/ComfyUI - -Do not install ComfyUI's requirements into `.venv`: it wants `transformers>=4.50.3` and -`tools/requirements.txt` pins 4.46.3 for the SDXL exporter's tokenizer. The handful of modules -these four stages need go in a directory of their own, which nothing else has to know about: - - .venv/bin/pip install --no-deps --target /path/to/pylibs \\ - comfy-aimdo==0.5.3 comfy-kitchen==0.2.33 einops torchvision torchaudio - - PYTHONPATH=/path/to/pylibs .venv/bin/python tools/anima_exporter.py ... \\ - -comfyui /path/to/ComfyUI -test_output anima_test.waifupkg - -Everything runs in float32 on the CPU. The residual stream of this model is wide enough that -ComfyUI keeps it in float32 even when the rest is float16, so float32 throughout is what a -reference owes it. -""" - -from __future__ import annotations - -import sys - -import torch - -# What ComfyUI's `model_detection.py` asks for when `in_channels` is 16, which is what Anima's -# `x_embedder.proj.1.weight` of [2048, 68] comes to: `68 // 4 - 1`. The ratio is an NTK factor -# and not a base -- the base is `10000 * ratio ** (dim / (dim - 2))` -- and it is a property of -# the Cosmos-Predict2 2B this model was trained from, not a choice made at run time. -H_EXTRAPOLATION_RATIO = 4.0 -W_EXTRAPOLATION_RATIO = 4.0 -T_EXTRAPOLATION_RATIO = 1.0 - -# The rest of what `model_detection.py` derives for this architecture. Only two of these are read -# off the tensors -- the width and the block count -- and those are passed in. -DENOISER_CONFIG = dict( - max_img_h=240, max_img_w=240, max_frames=128, - in_channels=16, out_channels=16, - patch_spatial=2, patch_temporal=1, concat_padding_mask=True, - crossattn_emb_channels=1024, - pos_emb_cls="rope3d", pos_emb_learnable=True, pos_emb_interpolation="crop", - min_fps=1, max_fps=30, - use_adaln_lora=True, adaln_lora_dim=256, - num_heads=16, - extra_per_block_abs_pos_emb=False, - rope_h_extrapolation_ratio=H_EXTRAPOLATION_RATIO, - rope_w_extrapolation_ratio=W_EXTRAPOLATION_RATIO, - rope_t_extrapolation_ratio=T_EXTRAPOLATION_RATIO, - extra_h_extrapolation_ratio=1.0, extra_w_extrapolation_ratio=1.0, - extra_t_extrapolation_ratio=1.0, - rope_enable_fps_modulation=False, -) - -CONTEXT_LENGTH = 512 - - -def use(comfyui: str) -> None: - """Put a ComfyUI checkout on the path. Call this before anything else here.""" - # comfy.model_management reads argv at import and will take the exporter's flags for its own. - sys.argv = [sys.argv[0], "--cpu"] - if comfyui not in sys.path: - sys.path.insert(0, comfyui) - - -def _ops(): - import comfy.ops - return comfy.ops.disable_weight_init - - -def _load(model, weights: dict, what: str, ignore=()): - """Load and insist that the config and the checkpoint agree about every tensor. - - A reference that quietly runs with a randomly initialized layer is worse than no reference, - and `strict=False` is needed only for the position tables, which are buffers rebuilt at - construction rather than weights. - """ - missing, unexpected = model.load_state_dict(weights, strict=False) - missing = [k for k in missing if not any(pattern in k for pattern in ignore)] - if missing or unexpected: - raise SystemExit( - f"{what}: the config ComfyUI derives does not match this checkpoint -- " - f"{len(missing)} missing, {len(unexpected)} unexpected\n" - f" missing: {missing[:4]}\n unexpected: {unexpected[:4]}" - ) - return model.eval().float() - - -class TextEncoder: - """Qwen3-0.6B, read for its hidden states rather than for a next token.""" - - def __init__(self, weights: dict) -> None: - import comfy.text_encoders.llama as llama - model = llama.Qwen3_06B({}, torch.float32, "cpu", _ops()) - self.model = _load(model, weights, "the text encoder") - - def __call__(self, ids: torch.Tensor) -> torch.Tensor: - out = self.model(ids) - return out[0] if isinstance(out, (tuple, list)) else out - - -class Denoiser: - """The adapter and the transformer, which are one module in ComfyUI as they are in the file. - - `llm_adapter` is reached on its own as well as through `preprocess_text_embeds`, because the - package carries the context both before padding, which is the adapter's own answer, and after, - which is what the denoiser reads. A runtime can go wrong in either. - """ - - def __init__(self, weights: dict, blocks: int, width: int) -> None: - import comfy.ldm.anima.model as anima - model = anima.Anima(**DENOISER_CONFIG, num_blocks=blocks, model_channels=width, - image_model="anima", device="cpu", dtype=torch.float32, - operations=_ops()) - self.model = _load(model, weights, "the denoiser", - ignore=("pos_embedder", "_range")) - - def adapt(self, hidden: torch.Tensor, t5_ids: torch.Tensor) -> torch.Tensor: - return self.model.llm_adapter(hidden, t5_ids) - - def adapt_padded(self, hidden: torch.Tensor, t5_ids: torch.Tensor) -> torch.Tensor: - return self.model.preprocess_text_embeds(hidden, t5_ids) - - def __call__(self, latent: torch.Tensor, timestep: torch.Tensor, - context: torch.Tensor) -> torch.Tensor: - """One velocity, from a four dimensional latent. Anima draws stills, so time is one. - - `_forward` rather than `forward`, which only wraps it in ComfyUI's patcher machinery. - """ - out = self.model._forward(latent.unsqueeze(2), timestep, context, - fps=None, padding_mask=None) - return out.squeeze(2) - - -class Vae: - """The Qwen-Image VAE, in the three dimensions it is actually written in. - - libwaifu folds every `Conv3d` here onto `W[:, :, -1]` at export time, which is exact when the - temporal extent is one. Running the real thing in both directions is what makes the exported - tensors a check on that folding rather than a restatement of it. - - Both halves, because both are in the package: the decoder is what a text to image run ends - with, and the encoder is where an image to image run begins. One object rather than two, so - that a round trip is the same weights going out and coming back. - """ - - def __init__(self, weights: dict, latents_mean, latents_std) -> None: - import comfy.ldm.wan.vae as wan - # The Wan 2.1 branch of comfy/sd.py's VAE detection, read off the same keys it reads. - self.z_dim = 16 - config = dict( - dim=weights["decoder.head.0.gamma"].shape[0], - z_dim=self.z_dim, - dim_mult=[1, 2, 4, 4], - num_res_blocks=2, - attn_scales=[], - temperal_downsample=[False, True, True], - image_channels=weights["encoder.conv1.weight"].shape[1], - conv_out_channels=weights["decoder.head.2.weight"].shape[0], - dropout=0.0, - ) - self.model = _load(wan.WanVAE(**config), weights, "the VAE") - self.mean = torch.tensor(latents_mean, dtype=torch.float32).view(1, -1, 1, 1) - self.std = torch.tensor(latents_std, dtype=torch.float32).view(1, -1, 1, 1) - - def decode(self, latent: torch.Tensor) -> torch.Tensor: - """A `(N, 16, H, W)` latent to a `(N, 3, H * 8, W * 8)` picture. - - Sixteen channels normalized one at a time, where SDXL has a single factor. - """ - unscaled = latent.float() * self.std + self.mean - return self.model.decode(unscaled.unsqueeze(2)).squeeze(2) - - def encode(self, image: torch.Tensor) -> torch.Tensor: - """A `(N, 3, H, W)` picture to the `(N, 16, H / 8, W / 8)` latent that stands for it. - - The mean of the distribution rather than a draw from it: the draw is the one part of an - encoder that two implementations have no way to agree on, and image to image wants the - mean anyway. Scaled into the normalization the sampler works in, which is the decode - above run backwards, so `encode(decode(x))` and `x` are comparable numbers. - - One frame, so `WanVAE.encode` runs a single chunk and passes `feat_cache=None` down, - and with no cache `Resample`'s `downsample3d` branch never reaches `time_conv` -- which - is the branch the folded export assumes, and the reason those weights are not in the - package. `encode` returns mu itself upstream, but the shape is checked rather than - trusted: a build that handed back the undivided moments would otherwise scale - thirty-two channels by sixteen pairs and broadcast its way to a quiet, wrong answer. - """ - moments = self.model.encode(image.float().unsqueeze(2)) - if isinstance(moments, (tuple, list)): - moments = moments[0] - moments = moments.squeeze(2) - - channels = moments.shape[1] - if channels == 2 * self.z_dim: - moments = moments.chunk(2, dim=1)[0] - elif channels != self.z_dim: - raise SystemExit( - f"the VAE encoder returned {channels} channels, which is neither the " - f"{self.z_dim} of a latent nor the {2 * self.z_dim} of its moments") - - return (moments - self.mean) / self.std diff --git a/tools/anima_exporter.py b/tools/anima_exporter.py index 1e18f52..eaebe09 100644 --- a/tools/anima_exporter.py +++ b/tools/anima_exporter.py @@ -26,10 +26,10 @@ safetensors files, and the whole of what it is is written down in `docs/anima.md` -- read that first, because several of the decisions here only make sense against it. -Unlike the SDXL exporter this one does not walk a diffusers pipeline: there is no diffusers -implementation of Anima to walk. It reads the three state dictionaries directly and rewrites them, -which means the tensor names below are the contract with the runtime and nothing checks them for -us. +Unlike the SDXL exporter this one does not walk a diffusers pipeline: Anima is published as +three bare state dictionaries, not as a pipeline to walk. It reads them directly and rewrites +them, which means the tensor names below are the contract with the runtime and nothing checks them +for us. .venv/bin/python tools/anima_exporter.py \ -dit models/anima/split_files/diffusion_models/anima-turbo-v1.1.safetensors \ @@ -48,7 +48,7 @@ import torch from safetensors import safe_open -import anima_comfy as reference +import anima_reference as reference from tokenizer_exporter import encoder, read_tokenizer, tokenizer_corpus, tokenizer_json from model_exporter import Context, TensorBag, open_weights, parse_size, stem_of @@ -375,10 +375,10 @@ def generate_config(dit: dict, text: dict, blocks: int, layers: int) -> dict: hidden = dit["x_embedder.proj.1.weight"].shape[0] latent = dit["final_layer.linear.weight"].shape[0] // 4 - # ComfyUI's model_detection.py reads the rotary extrapolation ratios off the latent width: - # sixteen channels is the text-to-image branch, which asks for four times the spatial extent - # and leaves time alone. Anything else has a different row in that table, and guessing it is - # a silent error rather than a loud one. + # The rotary extrapolation ratios are those of the 16 channel Cosmos-Predict2 2B text-to-image + # model, which asks for four times the spatial extent and leaves time alone. A different + # latent width is a different base model with ratios of its own, and guessing them is a + # silent error rather than a loud one. if latent != 16: raise SystemExit( f"{latent} latent channels: this exporter only knows the rotary extrapolation " @@ -428,9 +428,10 @@ def generate_config(dit: dict, text: dict, blocks: int, layers: int) -> dict: # rectified flow, not the epsilon prediction every SDXL checkpoint carries "prediction": "flow", # How the schedule is bent and what the denoiser is handed. Neither is in any tensor: - # they are Anima's entry in ComfyUI's model table. The multiplier of one is the unusual - # half -- a flow model is normally passed a thousandfold timestep, and this one is passed - # the sigma itself. + # they are the `shift` of the scheduler config Anima's diffusers release carries, and the + # division by `num_train_timesteps` its pipeline does before calling the transformer. The + # multiplier of one is the unusual half -- a flow model is normally passed a thousandfold + # timestep, and this one is passed the sigma itself. "sampler_shift": "3.0", "sampler_multiplier": "1.0", }} @@ -489,11 +490,11 @@ def reference_encoders(qwen, t5): def export_test_cases(dit: dict, text: dict, vae: dict, qwen_ids, t5_ids) -> TensorBag: """Write what the reference makes of one prompt, so the runtime can be checked against it. - The same idea as `sdxl_exporter.export_test_cases`, against `anima_comfy` -- which is ComfyUI -- rather than - against diffusers, and covering the path in the order a picture travels it: the encoder's - hidden states, the adapter's context, one denoising step, one decode, and -- for image to - image, which starts where text to image ends -- that decode encoded again and decoded once - more. + The same idea as `sdxl_exporter.export_test_cases`, against `anima_reference` -- which is + diffusers' Anima models -- and covering the path in the order a picture travels it: the + encoder's hidden states, the adapter's context, one denoising step, one decode, and -- for + image to image, which starts where text to image ends -- that decode encoded again and decoded + once more. All of it on a 16 by 16 latent -- a 128 by 128 image, 64 patches -- which is large enough to exercise every block and small enough to carry in the repository. @@ -588,22 +589,13 @@ def main() -> int: parser.add_argument( "-test_output", type=str, default=None, help="also write reference inputs and outputs here, for the runtime to be checked " - "against. Runs the whole model on the CPU, which takes a few minutes, and needs " - "-comfyui.") - parser.add_argument( - "-comfyui", type=str, default=None, - help="a ComfyUI checkout, which is the reference the test outputs are computed with -- " - "the same arrangement sdxl_exporter.py has with diffusers. Only -test_output needs " - "it. See tools/anima_comfy.py for what to install and where.") + "against. Runs the whole model on the CPU through diffusers, which takes a few " + "minutes, and needs the venv tools/requirements-anima.txt describes.") args = parser.parse_args() - # Before anything is read, so a missing checkout costs nothing but the message. The path goes - # on sys.path here and the modules are imported stage by stage as they are needed. + # Before anything is read, so a missing diffusers costs nothing but the message. if args.test_output: - if not args.comfyui: - raise SystemExit("-test_output needs -comfyui: the reference outputs are what ComfyUI " - "computes, and there is no second implementation here to ask") - reference.use(args.comfyui) + reference.check() dit = strip_prefix(read(args.dit)) text = read(args.text) diff --git a/tools/anima_reference.py b/tools/anima_reference.py new file mode 100644 index 0000000..6072a8d --- /dev/null +++ b/tools/anima_reference.py @@ -0,0 +1,348 @@ +#!/usr/bin/env python3 +# The MIT License (MIT) +# +# Copyright (c) 2026 Xiaoyang Chen +# +# Permission is hereby granted, free of charge, to any person obtaining a copy of this software +# and associated documentation files (the "Software"), to deal in the Software without +# restriction, including without limitation the rights to use, copy, modify, merge, publish, +# distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all copies or +# substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING +# BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND +# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + +"""The Anima reference, which is diffusers. + +`sdxl_exporter.py` gets its reference outputs by running the diffusers pipeline. This is the same +arrangement for Anima: diffusers' own `AnimaTextConditioner`, `CosmosTransformer3DModel` and +`AutoencoderKLQwenImage`, the three models `AnimaModularPipeline` is built from, run on the same +weights the exporter reads. Nothing here is a transcription of the model. What is written here is +only how the published files' tensor names map onto those classes, and that is checked: every +load is strict, so a name this gets wrong is an error and not a randomly initialized layer. + +Nothing here is used at run time. It is imported by `anima_exporter.py -test_output` and by +nothing else, and it needs diffusers 0.41 and a transformers new enough to have Qwen3, which the +pair in `tools/requirements.txt` is not. They go in a venv of their own: + + python3 -m venv .venv-anima + .venv-anima/bin/pip install -r tools/requirements-anima.txt + + .venv-anima/bin/python tools/anima_exporter.py ... -test_output anima_test.waifupkg + +Everything runs in float32 on the CPU. The residual stream of this model is wide enough that it +overflows half precision, so float32 throughout is what a reference owes it. +""" + +from __future__ import annotations + +import torch + +# The rotary position extrapolation ratios, time first, as `rope_scale` in the transformer config +# of `circlestone-labs/Anima-Base-v1.0-Diffusers` and of diffusers' Cosmos-2.0 2B text-to-image +# entry, which Anima's denoiser is trained from. Each is an NTK factor and not a base -- the base +# is `10000 * ratio ** (dim / (dim - 2))` -- and they are a property of the model, not a choice +# made at run time. +T_EXTRAPOLATION_RATIO = 1.0 +H_EXTRAPOLATION_RATIO = 4.0 +W_EXTRAPOLATION_RATIO = 4.0 + +# The rest of that config. Only the width and the block count are read off the tensors. +DENOISER_CONFIG = dict( + in_channels=16, + out_channels=16, + attention_head_dim=128, + mlp_ratio=4.0, + text_embed_dim=1024, + adaln_lora_dim=256, + max_size=(128, 240, 240), + patch_size=(1, 2, 2), + rope_scale=(T_EXTRAPOLATION_RATIO, H_EXTRAPOLATION_RATIO, W_EXTRAPOLATION_RATIO), + concat_padding_mask=True, + extra_pos_embed_type=None, +) + +CONTEXT_LENGTH = 512 + +# How the denoiser's names in the published file become diffusers' names, applied in this order +# as substring replacements. The order matters: `blocks` has to become `transformer_blocks` before +# anything inside a block is renamed, and `final_layer.adaln_modulation` must not be caught by the +# per-block `adaln_modulation_*` patterns, which it is not because those carry an underscore. +DENOISER_RENAMES = ( + ("t_embedder.1", "time_embed.t_embedder"), + ("t_embedding_norm", "time_embed.norm"), + ("blocks", "transformer_blocks"), + ("adaln_modulation_self_attn.1", "norm1.linear_1"), + ("adaln_modulation_self_attn.2", "norm1.linear_2"), + ("adaln_modulation_cross_attn.1", "norm2.linear_1"), + ("adaln_modulation_cross_attn.2", "norm2.linear_2"), + ("adaln_modulation_mlp.1", "norm3.linear_1"), + ("adaln_modulation_mlp.2", "norm3.linear_2"), + ("self_attn", "attn1"), + ("cross_attn", "attn2"), + ("q_proj", "to_q"), + ("k_proj", "to_k"), + ("v_proj", "to_v"), + ("output_proj", "to_out.0"), + ("q_norm", "norm_q"), + ("k_norm", "norm_k"), + ("mlp.layer1", "ff.net.0.proj"), + ("mlp.layer2", "ff.net.2"), + ("x_embedder.proj.1", "patch_embed.proj"), + ("final_layer.adaln_modulation.1", "norm_out.linear_1"), + ("final_layer.adaln_modulation.2", "norm_out.linear_2"), + ("final_layer.linear", "proj_out"), +) + +# Buffers and bookkeeping some releases carry beside the weights. diffusers rebuilds the position +# tables at construction, and the rest is training state. +DENOISER_SKIPPED = ("pos_embedder.", "accum_", "_extra_state") + +ADAPTER_PREFIX = "llm_adapter." + + +def check() -> None: + """Fail early, before anything is read, if this environment cannot run the reference.""" + try: + import diffusers + from diffusers import AnimaTextConditioner # noqa: F401 + from transformers import Qwen3Model # noqa: F401 + except ImportError as error: + raise SystemExit( + f"-test_output needs diffusers with Anima and transformers with Qwen3 ({error}). " + f"They do not fit beside the pair tools/requirements.txt pins; install " + f"tools/requirements-anima.txt into a venv of its own and run the exporter from it.") + del diffusers + + +def _load(model, weights: dict, what: str): + """Load and insist that the model and the checkpoint agree about every tensor. + + A reference that quietly runs with a randomly initialized layer is worse than no reference. + """ + missing, unexpected = model.load_state_dict(weights, strict=False) + if missing or unexpected: + raise SystemExit( + f"{what}: the diffusers model does not match this checkpoint -- " + f"{len(missing)} missing, {len(unexpected)} unexpected\n" + f" missing: {missing[:4]}\n unexpected: {unexpected[:4]}" + ) + return model.eval().float() + + +class TextEncoder: + """Qwen3-0.6B, read for its last hidden state rather than for a next token. + + The last hidden state, after the final norm, is what `AnimaTextEncoderStep` hands the adapter. + """ + + def __init__(self, weights: dict) -> None: + from transformers import Qwen3Config, Qwen3Model + + weights = {name.removeprefix("model."): tensor for name, tensor in weights.items() + if not name.startswith("lm_head.")} + vocab, hidden = weights["embed_tokens.weight"].shape + head_dim = weights["layers.0.self_attn.q_norm.weight"].shape[0] + config = Qwen3Config( + vocab_size=vocab, + hidden_size=hidden, + intermediate_size=weights["layers.0.mlp.gate_proj.weight"].shape[0], + num_hidden_layers=1 + max(int(name.split(".")[1]) for name in weights + if name.startswith("layers.")), + num_attention_heads=weights["layers.0.self_attn.q_proj.weight"].shape[0] // head_dim, + num_key_value_heads=weights["layers.0.self_attn.k_proj.weight"].shape[0] // head_dim, + head_dim=head_dim, + max_position_embeddings=32768, + rms_norm_eps=1e-6, + rope_theta=1000000.0, + attention_bias=False, + tie_word_embeddings=False, + ) + self.model = _load(Qwen3Model(config), weights, "the text encoder") + + def __call__(self, ids: torch.Tensor) -> torch.Tensor: + return self.model(input_ids=ids).last_hidden_state + + +class Denoiser: + """The adapter and the transformer, which diffusers keeps as two models and the file as one. + + The adapter is reached both before padding, which is its own answer, and after, which is what + the denoiser reads: the package carries both, and a runtime can go wrong in either. + """ + + def __init__(self, weights: dict, blocks: int, width: int) -> None: + from diffusers import AnimaTextConditioner, CosmosTransformer3DModel + + adapter = {name.removeprefix(ADAPTER_PREFIX): tensor for name, tensor in weights.items() + if name.startswith(ADAPTER_PREFIX)} + model_dim = adapter["blocks.0.self_attn.q_proj.weight"].shape[0] + vocab, target_dim = adapter["embed.weight"].shape + self.adapter = _load(AnimaTextConditioner( + source_dim=adapter["blocks.0.cross_attn.k_proj.weight"].shape[1], + target_dim=target_dim, + model_dim=model_dim, + num_layers=1 + max(int(name.split(".")[1]) for name in adapter + if name.startswith("blocks.")), + num_attention_heads=model_dim // adapter["blocks.0.self_attn.q_norm.weight"].shape[0], + target_vocab_size=vocab, + min_sequence_length=CONTEXT_LENGTH, + ), adapter, "the adapter") + + dit = {} + for name, tensor in weights.items(): + if name.startswith(ADAPTER_PREFIX) or any(s in name for s in DENOISER_SKIPPED): + continue + for old, new in DENOISER_RENAMES: + name = name.replace(old, new) + if name in dit: + raise SystemExit(f"the denoiser: two tensors were renamed to {name}") + dit[name] = tensor + config = dict(DENOISER_CONFIG, num_layers=blocks, + num_attention_heads=width // DENOISER_CONFIG["attention_head_dim"]) + self.model = _load(CosmosTransformer3DModel(**config), dit, "the denoiser") + + def adapt(self, hidden: torch.Tensor, t5_ids: torch.Tensor) -> torch.Tensor: + # The conditioner pads its answer with zeros out to the context length at the very end, + # so the rows before the padding are the adapter's own output. + return self.adapt_padded(hidden, t5_ids)[:, : t5_ids.shape[1]] + + def adapt_padded(self, hidden: torch.Tensor, t5_ids: torch.Tensor) -> torch.Tensor: + padded = self.adapter(source_hidden_states=hidden, target_input_ids=t5_ids) + if padded.shape[1] != max(CONTEXT_LENGTH, t5_ids.shape[1]): + raise SystemExit(f"the adapter returned {padded.shape[1]} rows of context") + return padded + + def __call__(self, latent: torch.Tensor, timestep: torch.Tensor, + context: torch.Tensor) -> torch.Tensor: + """One velocity, from a four dimensional latent. Anima draws stills, so time is one. + + `timestep` is the sigma itself, which is what `AnimaLoopBeforeDenoiser` passes: the + scheduler's timestep divided by its thousand training steps. The padding mask is the + all-zero one the pipeline builds, since nothing this runtime draws is padded. + """ + n, _, h, w = latent.shape + padding = latent.new_zeros(1, 1, h * 8, w * 8) + out = self.model(hidden_states=latent.unsqueeze(2), timestep=timestep, + encoder_hidden_states=context, padding_mask=padding, + return_dict=False)[0] + return out.squeeze(2) + + +def _vae_name(name: str) -> str: + """A tensor name in the published Qwen-Image VAE to the one `AutoencoderKLQwenImage` uses.""" + + def residual(name: str) -> str: + for old, new in ((".residual.0.", ".norm1."), (".residual.2.", ".conv1."), + (".residual.3.", ".norm2."), (".residual.6.", ".conv2."), + (".shortcut.", ".conv_shortcut.")): + name = name.replace(old, new) + return name + + def middle(name: str) -> str: + for old, new in ((".middle.0.", ".mid_block.resnets.0."), + (".middle.1.", ".mid_block.attentions.0."), + (".middle.2.", ".mid_block.resnets.1.")): + name = name.replace(old, new) + return residual(name) + + prefixes = ( + ("conv1.", "quant_conv."), + ("conv2.", "post_quant_conv."), + ("encoder.conv1.", "encoder.conv_in."), + ("decoder.conv1.", "decoder.conv_in."), + ("encoder.head.0.", "encoder.norm_out."), + ("encoder.head.2.", "encoder.conv_out."), + ("decoder.head.0.", "decoder.norm_out."), + ("decoder.head.2.", "decoder.conv_out."), + ) + for old, new in prefixes: + if name.startswith(old): + return new + name.removeprefix(old) + if name.startswith("encoder.downsamples."): + return residual("encoder.down_blocks." + name.removeprefix("encoder.downsamples.")) + if name.startswith("decoder.upsamples."): + # Four to a resolution, flattened: three residual blocks and then the upsampler, except + # the last resolution, which has no upsampler after it. + index, rest = name.removeprefix("decoder.upsamples.").split(".", 1) + index = int(index) + if index in (3, 7, 11): + return residual(f"decoder.up_blocks.{(index - 3) // 4}.upsamplers.0.{rest}") + return residual(f"decoder.up_blocks.{index // 4}.resnets.{index % 4}.{rest}") + if name.startswith(("encoder.middle.", "decoder.middle.")): + return middle(name) + return residual(name) + + +class Vae: + """The Qwen-Image VAE, in the three dimensions it is actually written in. + + libwaifu folds every `Conv3d` here onto `W[:, :, -1]` at export time, which is exact when the + temporal extent is one. Running the real thing in both directions is what makes the exported + tensors a check on that folding rather than a restatement of it. + + Both halves, because both are in the package: the decoder is what a text to image run ends + with, and the encoder is where an image to image run begins. One object rather than two, so + that a round trip is the same weights going out and coming back. + """ + + def __init__(self, weights: dict, latents_mean, latents_std) -> None: + from diffusers import AutoencoderKLQwenImage + + self.z_dim = 16 + renamed = {} + for name, tensor in weights.items(): + new = _vae_name(name) + if new in renamed: + raise SystemExit(f"the VAE: two tensors were renamed to {new}") + renamed[new] = tensor + model = AutoencoderKLQwenImage( + base_dim=weights["decoder.head.0.gamma"].shape[0], + z_dim=self.z_dim, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0, + input_channels=weights["encoder.conv1.weight"].shape[1], + latents_mean=list(latents_mean), + latents_std=list(latents_std), + ) + self.model = _load(model, renamed, "the VAE") + self.mean = torch.tensor(latents_mean, dtype=torch.float32).view(1, -1, 1, 1) + self.std = torch.tensor(latents_std, dtype=torch.float32).view(1, -1, 1, 1) + + def decode(self, latent: torch.Tensor) -> torch.Tensor: + """A `(N, 16, H, W)` latent to a `(N, 3, H * 8, W * 8)` picture. + + Sixteen channels normalized one at a time, where SDXL has a single factor. diffusers + clamps the picture to [-1, 1] on the way out, and so the reference is a clamped picture: + a runtime compared with it has to clamp its own the same way first. + """ + unscaled = latent.float() * self.std + self.mean + return self.model.decode(unscaled.unsqueeze(2), return_dict=False)[0].squeeze(2) + + def encode(self, image: torch.Tensor) -> torch.Tensor: + """A `(N, 3, H, W)` picture to the `(N, 16, H / 8, W / 8)` latent that stands for it. + + The mean of the distribution rather than a draw from it: the draw is the one part of an + encoder that two implementations have no way to agree on, and image to image wants the + mean anyway. Scaled into the normalization the sampler works in, which is the decode + above run backwards, so `encode(decode(x))` and `x` are comparable numbers. + + One frame, so the encoder runs a single chunk, and on the first chunk the cache is empty + and the downsamplers never reach `time_conv` -- which is the branch the folded export + assumes, and the reason those weights are not in the package. + """ + posterior = self.model.encode(image.float().unsqueeze(2), return_dict=False)[0] + mean = posterior.mode().squeeze(2) + if mean.shape[1] != self.z_dim: + raise SystemExit(f"the VAE encoder returned {mean.shape[1]} channels, not the " + f"{self.z_dim} of a latent") + return (mean - self.mean) / self.std diff --git a/tools/requirements-anima.txt b/tools/requirements-anima.txt new file mode 100644 index 0000000..9bdb796 --- /dev/null +++ b/tools/requirements-anima.txt @@ -0,0 +1,24 @@ +# What `anima_exporter.py -test_output` needs, in a venv of its own: +# +# python3 -m venv .venv-anima +# .venv-anima/bin/pip install -r tools/requirements-anima.txt +# +# The reference outputs are what diffusers' Anima models compute -- see tools/anima_reference.py +# -- and those arrived in diffusers 0.41, with a Qwen3 text encoder that needs a transformers far +# newer than the 4.46.3 tools/requirements.txt pins for the SDXL reference. The two pairs cannot +# share a venv, so this one is for Anima and nothing else. A plain export, without -test_output, +# runs in either. +# +# The pins are the set the Anima reference tensors are generated with; a different diffusers is a +# different reference. torch is unpinned for the same reason it is in tools/requirements.txt, and +# torchvision is here because diffusers' Cosmos transformer resizes the padding mask with it. +torch +torchvision +diffusers==0.41.0 +transformers==5.19.0 +huggingface-hub==1.33.0 +tokenizers==0.23.3 +accelerate +safetensors +sentencepiece +protobuf diff --git a/tools/requirements.txt b/tools/requirements.txt index dcb2ad9..9b7ae10 100644 --- a/tools/requirements.txt +++ b/tools/requirements.txt @@ -17,7 +17,8 @@ # transformers at all -- `tokenizer_exporter.py` asks it for the `tokenizer.json` and stores it. # So the pins are free to move, and moving them is its own job: the reference tensors in every # `_test.waifupkg` were generated by this pair, and a newer diffusers is a different reference. -# Worth doing, because it is also what keeps ComfyUI out of this venv -- see the note at the end. +# Worth doing, because it is also what would let the Anima reference share this venv -- see the +# note at the end. # # torch is unpinned: any build that matches the machine's CUDA will do, and the exporter runs on # the CPU anyway. @@ -59,25 +60,10 @@ numpy # scope for a rearrange the layers this uses never reach. Neither is pinned and neither has a # dependency of its own, so neither can move the pair above -- install them with `--no-deps` and # check `transformers.__version__` afterwards if you doubt it. -# -# einops also appears in the ComfyUI note below, where it goes into a directory of its own. That -# is about ComfyUI, not about einops: the package is harmless here, and what must stay out is the -# transformers ComfyUI drags with it. munch einops -# ComfyUI is not in here, and must not be installed into this venv. -# -# It is the reference implementation for Anima the way diffusers is for SDXL, so -# `anima_exporter.py -test_output` needs it -- but it asks for transformers >= 4.50.3, and the pin -# above is 4.46.3 because bpe_exporter.py imports something 5.x removed. Installing ComfyUI here -# breaks the SDXL exporter. -# -# It goes in a checkout of its own, with the few modules it needs in a directory of their own: -# -# git clone --depth 1 https://github.com/comfyanonymous/ComfyUI.git /path/to/ComfyUI -# .venv/bin/pip install --no-deps --target /path/to/pylibs \ -# comfy-aimdo==0.5.3 comfy-kitchen==0.2.33 einops torchvision torchaudio -# -# and is reached with PYTHONPATH=/path/to/pylibs and -comfyui /path/to/ComfyUI. See -# tools/anima_comfy.py. +# The Anima reference is not in here. `anima_exporter.py -test_output` runs diffusers' Anima +# models, which need diffusers 0.41 and a transformers with Qwen3 -- neither fits the pair above, +# and installing them here breaks the SDXL exporter. They have a venv of their own, described in +# tools/requirements-anima.txt. diff --git a/waifu/src/anima/config.rs b/waifu/src/anima/config.rs index e6c99ab..2780652 100644 --- a/waifu/src/anima/config.rs +++ b/waifu/src/anima/config.rs @@ -79,7 +79,7 @@ pub struct DitConfig { /// They are read from the package rather than assumed here because assuming them is how they /// came to be wrong. Anima asks for four times the spatial extent and leaves time alone, and /// a flat base on all three axes is 13% off on the velocity while still drawing a picture - /// that looks perfectly reasonable. It was caught by running ComfyUI beside it. + /// that looks perfectly reasonable. It was caught by running a reference beside it. pub rope_h_ratio: f32, pub rope_w_ratio: f32, pub rope_t_ratio: f32, diff --git a/waifu/src/anima/dit.rs b/waifu/src/anima/dit.rs index d7aefab..9655a65 100644 --- a/waifu/src/anima/dit.rs +++ b/waifu/src/anima/dit.rs @@ -340,8 +340,8 @@ fn write(config: &DitConfig, float_type: DType, g: &Graph) { let d = config.hidden_size; // The residual stream runs in float32 whatever the sublayers run in. This model's residuals // are large enough that in half they reach infinity part way through the first block and every - // step after it is NaN -- which is what ComfyUI means by keeping the stream wide and running - // the attention and the feed forward narrow. + // step after it is NaN -- so the stream is kept wide while the attention and the feed forward + // run narrow. let mut x = g.cast( Linear::graph( &g.subgraph("x_embedder"), diff --git a/waifu/src/sdxl/sampler.rs b/waifu/src/sdxl/sampler.rs index 6d6641f..b1a20a8 100644 --- a/waifu/src/sdxl/sampler.rs +++ b/waifu/src/sdxl/sampler.rs @@ -125,9 +125,9 @@ impl EulerSampler { /// Which is the other way round from the way diffusers takes the same two numbers: there, /// `num_inference_steps` is the whole schedule and `strength` cuts it down, so asking for /// thirty steps at 0.8 runs twenty-four. The arithmetic is the same and the knob is the same; - /// what differs is which of the two numbers is the one you can count on. ComfyUI counts on - /// this one, and so does this, because a step count that quietly means something else as soon - /// as another box is touched is a step count nobody can read. + /// what differs is which of the two numbers is the one you can count on. This counts on the + /// step count, because a step count that quietly means something else as soon as another box + /// is touched is a step count nobody can read. /// /// The latent to start from is the encoded picture plus unit noise times `sigmas()[start()]`. pub fn from_image( @@ -359,7 +359,7 @@ mod tests { #[test] fn the_strength_is_how_far_down_the_schedule_the_tail_begins() { // Twenty of forty at a half, twenty of twenty-five at 0.8: the last num_steps of a - // schedule num_steps/strength long, which is the arithmetic ComfyUI does. + // schedule num_steps/strength long. let half = from_image(20, 0.5); assert_eq!(half.len(), 40); assert_eq!(half.start(), 20); @@ -368,8 +368,8 @@ mod tests { assert_eq!(most.len(), 25); assert_eq!(most.start(), 5); - // Where the arithmetic is not exact the tail is what is left of the floor, which is what - // ComfyUI's int(steps / denoise) gives: 20 / 0.75 is 26 and not 27. + // Where the arithmetic is not exact the tail is what is left of the floor: 20 / 0.75 is + // 26 and not 27. let rough = from_image(20, 0.75); assert_eq!(rough.len(), 26); assert_eq!(rough.start(), 6); @@ -382,14 +382,10 @@ mod tests { } #[test] - fn the_schedule_is_the_one_comfyui_would_have_built() { - // Its KSampler::set_steps, which is where this shape comes from: - // - // new_steps = int(steps / denoise) - // sigmas = calculate_sigmas(new_steps)[-(steps + 1):] - // - // The pairs below are that arithmetic run in python, so a change here that quietly moved - // the tail would have to disagree with the tool people are coming from. + fn the_schedule_is_the_floor_of_the_steps_over_the_strength() { + // The whole schedule is floor(num_steps / strength) long and the run is its last + // num_steps. The pairs below are that arithmetic worked by hand, so a change here that + // quietly moved the tail has to disagree with something written down. for (steps, strength, whole) in [ (20, 1.0, 20), (20, 0.9, 22), diff --git a/waifu/tests/anima.rs b/waifu/tests/anima.rs index 7674293..8a6da32 100644 --- a/waifu/tests/anima.rs +++ b/waifu/tests/anima.rs @@ -17,13 +17,12 @@ // DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -//! Anima's parts against what ComfyUI produces for the same inputs, by way of the test -//! tensors `tools/anima_comfy.py` writes into the package. +//! Anima's parts against what diffusers produces for the same inputs, by way of the test +//! tensors `tools/anima_reference.py` writes into the package. //! -//! The reference is ComfyUI run directly on the same weights, so agreeing with it here is -//! agreeing with the model rather than with a second opinion of our own. None of what passes -//! between these stages can be judged by -//! looking -- a hidden state, a context, a velocity -- and an error in any of them shows up only +//! The reference is diffusers' Anima models run directly on the same weights, so agreeing with +//! it here is agreeing with the model rather than with a second opinion of our own. None of what +//! passes between these stages can be judged by looking -- a hidden state, a context, a velocity -- and an error in any of them shows up only //! as a picture that is subtly not the prompt, so they are compared as numbers. //! //! Each stage is measured against the reference's own inputs where it can be, so that what a @@ -254,7 +253,8 @@ fn one_step_matches_the_reference() { println!("dit velocity rmse = {rmse}"); // Twenty-eight blocks of fp16 against a float32 reference, and this model's residual stream - // is large enough that ComfyUI keeps its own in float32 for exactly this reason. + // is large enough that it overflows half precision, which is why the runtime keeps it in + // float32. assert!(rmse < 5e-2, "the velocity drifted by {rmse}"); } @@ -307,6 +307,21 @@ fn decodes_a_latent_the_way_the_reference_does() { // An eighth of the side, three channels out of sixteen: a 16 by 16 latent is 128 by 128. assert_eq!(image.shape(), vec![1, 3, 128, 128]); + // The reference decode clamps its picture to [-1, 1] and this runtime leaves that to the + // conversion to pixels, so the runtime's is clamped here before the two are compared. On a + // latent off a unit normal a tenth of the picture lies outside that range. + let pixels: Vec = image + .to_device(Device::Cpu) + .unwrap() + .cast(DType::Float) + .unwrap() + .to_vec_f32() + .unwrap() + .into_iter() + .map(|value| value.clamp(-1.0, 1.0)) + .collect(); + let image = Tensor::from_f32(&[1, 3, 128, 128], &pixels).unwrap(); + let reference = cases["test_case.decoded"].clone(); let rmse = relative_rmse(&image, &reference); println!("vae decoded rmse = {rmse}");