diff --git a/invokeai/backend/anima/anima_transformer.py b/invokeai/backend/anima/anima_transformer.py index 670c3f70cca..8f42d30b8fe 100644 --- a/invokeai/backend/anima/anima_transformer.py +++ b/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): diff --git a/invokeai/backend/model_manager/load/model_loaders/anima.py b/invokeai/backend/model_manager/load/model_loaders/anima.py index a2c334ee6d6..97782b07ffd 100644 --- a/invokeai/backend/model_manager/load/model_loaders/anima.py +++ b/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() diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ControlAdapterModelDefaultSettings/ControlAdapterModelDefaultSettings.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ControlAdapterModelDefaultSettings/ControlAdapterModelDefaultSettings.tsx index 0bb177ddbe1..864c45c1024 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ControlAdapterModelDefaultSettings/ControlAdapterModelDefaultSettings.tsx +++ b/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; @@ -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) - {modelConfig.type !== 'control_lora' && } + {supportsFp8Storage(modelConfig) && } ); diff --git a/tests/backend/model_manager/load/test_anima_fp8_wiring.py b/tests/backend/model_manager/load/test_anima_fp8_wiring.py new file mode 100644 index 00000000000..7bcdda873cb --- /dev/null +++ b/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)]