feat(fp8): enable FP8 storage for Anima

invokeai/backend/anima/anima_transformer.py

@@ -995,15 +995,27 @@ class AnimaTransformer(MiniTrainDIT):
     # `t_embedder` is the one that matters, and it is not a rounding-quality nicety: with it cast
     # to FP8 the model renders a heavily dithered image with no fine detail at all (verified
     # against a bf16 run at the same seed/steps/CFG). It feeds `adaln_lora` into every block, so
-    # its error is applied to every token everywhere. Measured, same seed each time: casting
-    # nothing = broken; `t_embedder` alone = clean, and adding either of the two below changes
-    # nothing further. They are kept as ~2MB of margin on the I/O layers, matching what diffusers
-    # skips by default for comparable DiTs. `adaln_modulation` was also tested and is deliberately
-    # NOT listed — it costs 168MB and made no difference.
+    # its error is applied to every token everywhere. It is also the most FP8-damaged group in the
+    # network by a wide margin: 38% of its weights flush to zero under unscaled e4m3fn, against
+    # 25% for the next worst group.
+    #
+    # `x_embedder` and `final_layer` are the set diffusers' `CosmosTransformer3DModel` declares
+    # (`["patch_embed", "final_layer", "norm"]`), which is apt — Anima *is* the Cosmos-Predict2
+    # DiT. They cost ~2MB and are kept as margin on the I/O layers.
+    #
+    # `adaln_modulation` (168MB) is deliberately NOT listed, but that is a cost/benefit call, not a
+    # claim that it is harmless: on a single forward the relative L2 error against bf16 drops
+    # 0.134 -> 0.091 when it is skipped, making it the largest remaining error source. The decision
+    # rests on a 35-step A/B at the same seed showing no visible difference, not on the norm. Note
+    # the errors add in quadrature and no single group dominates — skipping `t_embedder` alone only
+    # moves the total 0.152 -> 0.144; its outsized effect on the image comes from *where* its error
+    # lands, not from its share of the norm.
     _skip_layerwise_casting_patterns = [
         "t_embedder",  # timestep embedding MLP -> adaln_lora for every block
         "x_embedder",  # patch embedding (named `patch_embed` in diffusers models)
-        "final_layer",  # output projection
+        # Output head. Most of what this shields is `final_layer.adaln_modulation.*`
+        # (1.57 of its 1.70M params), not the output projection itself.
+        "final_layer",
     ]
 
     def __init__(self, *args, **kwargs):

invokeai/backend/model_manager/load/model_loaders/anima.py

@@ -76,6 +76,38 @@ def _filter_non_model_keys(sd: dict) -> dict:
     }
 
 
+# Anima's fixed transformer architecture. Kept at module level so tests can instantiate the real
+# module graph (e.g. to pin `_skip_layerwise_casting_patterns` to actual dotted module paths)
+# without duplicating these values.
+ANIMA_TRANSFORMER_CONFIG = {
+    "max_img_h": 240,
+    "max_img_w": 240,
+    "max_frames": 1,
+    "in_channels": 16,
+    "out_channels": 16,
+    "patch_spatial": 2,
+    "patch_temporal": 1,
+    "concat_padding_mask": True,
+    "model_channels": 2048,
+    "num_blocks": 28,
+    "num_heads": 16,
+    "mlp_ratio": 4.0,
+    "crossattn_emb_channels": 1024,
+    "pos_emb_cls": "rope3d",
+    # Anima reuses the Cosmos-Predict2 2B Text2Image DiT, which trains with
+    # rope_scale=(t=1.0, h=4.0, w=4.0). The NTK-scaled spatial RoPE base is mandatory; omitting it
+    # (theta=10000 on all axes) shifts every step's velocity ~7% off and compounds into degraded
+    # images. Matches diffusers CosmosTransformer3DModel rope_scale via *_extrapolation_ratio.
+    "rope_h_extrapolation_ratio": 4.0,
+    "rope_w_extrapolation_ratio": 4.0,
+    "rope_t_extrapolation_ratio": 1.0,
+    "use_adaln_lora": True,
+    "adaln_lora_dim": 256,
+    "extra_per_block_abs_pos_emb": False,
+    "image_model": "anima",
+}
+
+
 @ModelLoaderRegistry.register(base=BaseModelType.Anima, type=ModelType.Main, format=ModelFormat.Checkpoint)
 class AnimaCheckpointModel(ModelLoader):
     """Class to load Anima transformer models from single-file checkpoints.
@@ -127,34 +159,7 @@ def _load_from_singlefile(
 
         # Create an empty AnimaTransformer with Anima's default architecture parameters
         with accelerate.init_empty_weights():
-            model = AnimaTransformer(
-                max_img_h=240,
-                max_img_w=240,
-                max_frames=1,
-                in_channels=16,
-                out_channels=16,
-                patch_spatial=2,
-                patch_temporal=1,
-                concat_padding_mask=True,
-                model_channels=2048,
-                num_blocks=28,
-                num_heads=16,
-                mlp_ratio=4.0,
-                crossattn_emb_channels=1024,
-                pos_emb_cls="rope3d",
-                # Anima reuses the Cosmos-Predict2 2B Text2Image DiT, which trains with
-                # rope_scale=(t=1.0, h=4.0, w=4.0). The NTK-scaled spatial RoPE base is
-                # mandatory; omitting it (theta=10000 on all axes) shifts every step's
-                # velocity ~7% off and compounds into degraded images. Matches diffusers
-                # CosmosTransformer3DModel rope_scale via *_extrapolation_ratio.
-                rope_h_extrapolation_ratio=4.0,
-                rope_w_extrapolation_ratio=4.0,
-                rope_t_extrapolation_ratio=1.0,
-                use_adaln_lora=True,
-                adaln_lora_dim=256,
-                extra_per_block_abs_pos_emb=False,
-                image_model="anima",
-            )
+            model = AnimaTransformer(**ANIMA_TRANSFORMER_CONFIG)
 
         # Determine safe dtype
         target_device = TorchDevice.choose_torch_device()

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ControlAdapterModelDefaultSettings/ControlAdapterModelDefaultSettings.tsx

@@ -12,6 +12,7 @@ import { useTranslation } from 'react-i18next';
 import { PiCheckBold } from 'react-icons/pi';
 import { useUpdateModelMutation } from 'services/api/endpoints/models';
 import type { ControlLoRAModelConfig, ControlNetModelConfig, T2IAdapterModelConfig } from 'services/api/types';
+import { isAnimaControlNetModelConfig } from 'services/api/types';
 
 export type ControlAdapterModelDefaultSettingsFormData = {
   preprocessor: FormField<string>;
@@ -22,6 +23,13 @@ type Props = {
   modelConfig: ControlNetModelConfig | T2IAdapterModelConfig | ControlLoRAModelConfig;
 };
 
+// Only offer FP8 storage where a loader actually applies it, so the toggle never renders as a
+// no-op. ControlLoRAs are patched into the base model rather than run standalone. Anima's LLLite
+// adapters go through `AnimaControlNetLLLiteModel`, which never calls the layerwise cast - at
+// 16-63MB per adapter there is nothing worth wiring up.
+const supportsFp8Storage = (modelConfig: Props['modelConfig']): boolean =>
+  modelConfig.type !== 'control_lora' && !isAnimaControlNetModelConfig(modelConfig);
+
 export const ControlAdapterModelDefaultSettings = memo(({ modelConfig }: Props) => {
   const { t } = useTranslation();
   const canManageModels = useIsModelManagerEnabled();
@@ -42,7 +50,10 @@ export const ControlAdapterModelDefaultSettings = memo(({ modelConfig }: Props)
     (data) => {
       const body = {
         preprocessor: data.preprocessor.isEnabled ? data.preprocessor.value : null,
-        fp8_storage: data.fp8Storage.isEnabled ? data.fp8Storage.value : null,
+        // Null it out wherever the control is hidden. react-hook-form keeps unrendered fields in
+        // `defaultValues`, so without this a value persisted before the control was hidden would be
+        // re-sent verbatim on every save, with no UI left to clear it.
+        fp8_storage: supportsFp8Storage(modelConfig) && data.fp8Storage.isEnabled ? data.fp8Storage.value : null,
       };
 
       updateModel({
@@ -68,7 +79,7 @@ export const ControlAdapterModelDefaultSettings = memo(({ modelConfig }: Props)
           }
         });
     },
-    [updateModel, modelConfig.key, t, reset]
+    [updateModel, modelConfig, t, reset]
   );
 
   return (
@@ -91,7 +102,7 @@ export const ControlAdapterModelDefaultSettings = memo(({ modelConfig }: Props)
 
       <SimpleGrid columns={2} gap={8}>
         <DefaultPreprocessor control={control} name="preprocessor" />
-        {modelConfig.type !== 'control_lora' && <DefaultFp8StorageControlAdapter control={control} name="fp8Storage" />}
+        {supportsFp8Storage(modelConfig) && <DefaultFp8StorageControlAdapter control={control} name="fp8Storage" />}
       </SimpleGrid>
     </>
   );

tests/backend/model_manager/load/test_anima_fp8_wiring.py

@@ -0,0 +1,147 @@
+"""Guards for Anima's FP8-storage wiring.
+
+Two things have to hold for the `fp8_storage` toggle to do the right thing on Anima, and neither
+is observed by any other test:
+
+1. The loader has to actually call `_apply_fp8_layerwise_casting`. Without it the toggle renders
+   in the UI and silently does nothing.
+2. `AnimaTransformer._skip_layerwise_casting_patterns` has to match the modules it is meant to
+   match. The patterns are plain substrings matched against dotted module paths, so a rename in
+   the transformer disables the skip without breaking anything loudly.
+"""
+
+import re
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import accelerate
+import torch
+
+from invokeai.backend.anima.anima_transformer import AnimaTransformer
+from invokeai.backend.model_manager.configs.main import Main_Checkpoint_Anima_Config, MainModelDefaultSettings
+from invokeai.backend.model_manager.load.load_default import (
+    _FP8_DEFAULT_SKIP_PATTERNS,
+    _FP8_SUPPORTED_PYTORCH_LAYERS,
+)
+from invokeai.backend.model_manager.load.model_loaders.anima import (
+    ANIMA_TRANSFORMER_CONFIG,
+    AnimaCheckpointModel,
+)
+from invokeai.backend.model_manager.taxonomy import SubModelType
+
+# Every cast-eligible module in the real Anima graph that the declared patterns protect. Pinned
+# exhaustively rather than as a substring check, so both a rename (entries disappear) and an
+# over-broad pattern (extra entries appear) fail here.
+EXPECTED_SKIPPED_MODULES = {
+    "t_embedder.1.linear_1",
+    "t_embedder.1.linear_2",
+    "x_embedder.proj.1",
+    "final_layer.linear",
+    "final_layer.adaln_modulation.1",
+    "final_layer.adaln_modulation.2",
+}
+
+
+def _cast_eligible_modules(model: torch.nn.Module) -> list[str]:
+    """The dotted paths `_apply_fp8_to_nn_module` would consider for casting, before skipping."""
+    return [
+        name
+        for name, module in model.named_modules()
+        if isinstance(module, _FP8_SUPPORTED_PYTORCH_LAYERS) and list(module.parameters(recurse=False))
+    ]
+
+
+def _build_meta_transformer() -> AnimaTransformer:
+    with accelerate.init_empty_weights():
+        return AnimaTransformer(**ANIMA_TRANSFORMER_CONFIG)
+
+
+def test_declared_skip_patterns_pin_to_real_module_paths() -> None:
+    """The declared patterns must resolve against the real module graph, not just be strings.
+
+    `t_embedder` is the load-bearing one: casting it to FP8 renders a heavily dithered image,
+    because it feeds `adaln_lora` into every block. Renaming it in `AnimaTransformer` while the
+    pattern list still says `t_embedder` would silently reintroduce that, so match against an
+    actually-instantiated model.
+    """
+    model = _build_meta_transformer()
+    patterns = AnimaTransformer._skip_layerwise_casting_patterns
+
+    skipped = {name for name in _cast_eligible_modules(model) if any(re.search(p, name) for p in patterns)}
+
+    assert skipped == EXPECTED_SKIPPED_MODULES
+
+    # Each declared pattern earns its place — none is dead.
+    for pattern in patterns:
+        assert any(re.search(pattern, name) for name in skipped), f"pattern {pattern!r} matches nothing"
+
+
+def test_generic_skip_patterns_do_not_cover_anima() -> None:
+    """The declared list is not redundant with `_FP8_DEFAULT_SKIP_PATTERNS`.
+
+    Those defaults are written for diffusers' module naming (`norm`, `pos_embed`, `patch_embed`,
+    `proj_in/out`); this architecture names the equivalent modules differently, so they protect
+    nothing here. If a future default did start covering Anima, this test failing is the signal to
+    re-check whether the declared list is still needed.
+    """
+    model = _build_meta_transformer()
+
+    covered_by_defaults = [
+        name for name in _cast_eligible_modules(model) if any(re.search(p, name) for p in _FP8_DEFAULT_SKIP_PATTERNS)
+    ]
+
+    assert covered_by_defaults == []
+
+
+class _TinyAnimaTransformer(torch.nn.Module):
+    def __init__(self, **_kwargs) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.empty(2, 2))
+
+
+def test_single_file_loader_applies_fp8_layerwise_casting(monkeypatch, tmp_path) -> None:
+    """Regression guard for the `fp8_storage` toggle being wired up at all.
+
+    Deleting the `_apply_fp8_layerwise_casting` call from `_load_from_singlefile` leaves the whole
+    `tests/backend/model_manager` and `tests/backend/anima` suites green — nothing else observes
+    it — so the dead toggle can come straight back with CI passing.
+    """
+    import safetensors.torch
+
+    import invokeai.backend.anima.anima_transformer as anima_transformer_module
+
+    checkpoint_path = tmp_path / "anima.safetensors"
+    checkpoint_path.touch()
+    # `fp8_storage` lives under `default_settings` -- passing it as a top-level kwarg to
+    # `model_construct` would be silently discarded (the config has no such field and no extra="allow"),
+    # leaving the toggle off in a test whose whole point is that the toggle is wired up.
+    config = Main_Checkpoint_Anima_Config.model_construct(
+        path=str(checkpoint_path), default_settings=MainModelDefaultSettings(fp8_storage=True)
+    )
+
+    monkeypatch.setattr(anima_transformer_module, "AnimaTransformer", _TinyAnimaTransformer)
+    monkeypatch.setattr(safetensors.torch, "load_file", lambda _path: {"weight": torch.ones(2, 2)})
+    monkeypatch.setattr(
+        "invokeai.backend.model_manager.load.model_loaders.anima.TorchDevice.choose_torch_device",
+        lambda: torch.device("cpu"),
+    )
+    monkeypatch.setattr(
+        "invokeai.backend.model_manager.load.model_loaders.anima.TorchDevice.choose_anima_inference_dtype",
+        lambda _device: torch.float32,
+    )
+
+    cast_calls: list[tuple[object, object]] = []
+
+    def _record_cast(model, cfg, submodel):
+        cast_calls.append((cfg, submodel))
+        return model
+
+    loader = object.__new__(AnimaCheckpointModel)
+    loader._ram_cache = SimpleNamespace(make_room=MagicMock())
+    loader._apply_fp8_layerwise_casting = _record_cast
+
+    model = loader._load_from_singlefile(config)
+
+    assert isinstance(model, _TinyAnimaTransformer)
+    assert torch.equal(model.weight, torch.ones(2, 2))
+    assert cast_calls == [(config, SubModelType.Transformer)]