Add Model Architecture

Adding a new model architecture

Forge/A1111's backend only understands SD1.x / SD2 / SDXL / Flux / SD3-style UNet checkpoints. Brand-new base models (DiT transformers with their own text encoder + VAE) won't load. This skill is the method for bringing one in — first as a standalone diffusers/torch pipeline (fast to iterate + easy to verify), then wired natively into the Forge backend so the checkpoint loads and generates from the WebUI/API like any other model. See anima_infer.py (standalone) and backend/diffusion_engine/anima.py + backend/nn/anima.py + backend/nn/anima_hf_register.py (native) + ANIMA.md for a complete worked example: the Anima model = NVIDIA Cosmos-Predict2-2B DiT + Qwen3 text encoder + LLM adapter + Wan2.1 VAE.

Do the standalone port first: it isolates the model math from Forge's loader/patcher/sampler plumbing, so when the native path misbehaves you already have a trusted reference to diff against.

The one rule that matters

Verify every stage bit-exact against a reference implementation before trusting it. Most of the work is not writing the forward pass — it's finding the one subtle place where your port silently diverges. Bugs hide in the stages you assume are correct. (In the Anima port the only bug was the LLM adapter doing out_proj(norm(x)) instead of norm(out_proj(x)) — everything else was already bit-exact. It cost a day because it was never directly diffed.)

Step 1 — Identify the architecture

• HuggingFace model card: the base_model: / base_model:finetune: tags name the real architecture (e.g. nvidia/Cosmos-Predict2-2B-Text2Image). The library_name (diffusion-single-file) and a comfyui tag mean "ComfyUI-native, not diffusers".
• Civitai: the baseModel field (e.g. Anima, Illustrious, NoobAI, Pony) is the architecture family. Read it via the Civitai API with your own API token (Civitai may need a proxy in some regions).
• Note the components: a DiT usually ships a diffusion model + a separate text encoder + a separate VAE. Read the model card "installing" section — it lists which file goes in diffusion_models/, text_encoders/, vae/.

Step 2 — Inspect the checkpoint structure

Read the safetensors header (8-byte length prefix + JSON) without loading tensors: list keys, shapes, dtype; find the top-level prefix (net.), count transformer blocks, infer hidden size / heads / head_dim (from q_norm length), patch size (from the input/output projection: in_ch * patch^2 * t_patch), latent channels, context dim (from cross-attn k_proj in-features). This tells you the exact config.

Step 3 — Reuse what already exists

• VAE: try diffusers.AutoencoderKL*.from_single_file(path, ...). Qwen-Image VAE = Wan 2.1 VAE → AutoencoderKLWan.from_single_file loads it (auto key conversion).
• Text encoder: if it's a standard LLM (Qwen3, T5, Llama), load it with transformers by stripping the key prefix (model.) into the matching *Model + config. Match rope_theta, rms_norm_eps, head/kv counts exactly.
• Tokenizer: download the matching tokenizer; verify token IDs equal the reference's.
• DiT transformer: usually NOT in diffusers as-is (custom). Port it (Step 4).

Step 4 — Port the custom transformer

• Get the reference implementation. For ComfyUI-native models it's comfy/ldm/<family>/... + comfy/text_encoders/<family>.py + comfy/supported_models.py (config) + comfy/model_base.py (forward wiring) + comfy/latent_formats.py (latent mean/std). Read them; don't trust a summary for the forward math.
• Build nn.Modules whose submodule names exactly match the checkpoint keys, then load_state_dict(strict=True) → must be 0 missing / 0 unexpected. That validates the structure.
• Replicate the forward precisely: patchify order, AdaLN modulation (and any low-rank adaln_lora added before chunking), block order (self-attn → cross-attn → MLP), QK-norm placement (before RoPE), RoPE axes/theta/NTK, final layer + unpatchify.

Step 5 — Sampler, schedule, normalization, timestep

• Latent normalization (latent_format): the DiT works in (latent - mean)/std space; decode applies z*std + mean. Use the reference's exact mean/std.
• Timestep convention: flow models often feed timestep = sigma ∈ (0,1] (multiplier 1.0), NOT sigma*1000. Check model_sampling.timestep() + process_timestep().
• Sampler: rectified flow (CONST): denoised = x - v*sigma, Euler step x += (σ_next−σ)·v. Respect the model's shift.

Step 6 — Verify bit-exact, stage by stage (do NOT skip)

Clone the reference runtime into a temp dir (e.g. ComfyUI), make a throwaway venv, symlink the model files into its model dirs, and diff each stage on identical inputs:

1. tokenizer IDs (both encoders)
2. text-encoder hidden states
3. any adapter between encoder and DiT ← easy to get wrong, easy to forget to test
4. RoPE tables / the rope-apply kernel
5. the full DiT forward on a fixed (x, sigma, context)
6. the final latent + the decoded image

cos should be ~1.0 and max-abs-diff ~0 at every stage. The first stage that isn't is the bug. Then delete the temp runtime (don't leave multi-GB junk).

Step 7 — Wire it natively into Forge

Once the standalone pipeline is bit-exact, integrate so the checkpoint loads in the UI/API. Five touch points (Anima example in parentheses):

1. Arch detection + config — teach huggingface_guess to recognise the checkpoint and emit its config. repositories/huggingface_guess is a git-ignored vendor copy, so do NOT rely on editing it in place (a fresh clone loses the edit). Instead write a version-controlled, idempotent runtime-registration module that injects:
   • a detection.detect_unet_config wrapper that returns the DiT config when it sees a signature key (net.llm_adapter.blocks.0.self_attn.q_proj.weight), else defers;
   • a latent.<Format> class (the mean/std latent format);
   • a model_list.<Arch>(BASE) config (unet_config, sampling_settings shift/multiplier, latent_format, supported_inference_dtypes, model_type() → ModelType.FLOW, clip_target()), appended to model_list.models. (backend/nn/anima_hf_register.py; import + call register() from the engine module before it references model_list.<Arch>.) Make every step guard on existence so it is a no-op when the vendor copy is already patched.
2. DiT loader routing — in backend/loader.py add the diffusers cls_name (AnimaTransformer2DModel) to the transformer branch and map it to your Integrated<Arch> nn.Module. Add a backend/huggingface/<Arch>/model_index.json stub naming that cls_name so the guess resolves to it.
3. Diffusion engine — backend/diffusion_engine/<arch>.py, a ForgeDiffusionEngine subclass with matched_guesses = [model_list.<Arch>]. It loads the extra components the single-file DiT doesn't contain (text encoder, tokenizers, adapter, VAE) from an assets dir; implements get_learned_conditioning (return the crossattn tensor if there is no pooled vector — compile_conditions then takes the crossattn-only path, no 'vector' KeyError), encode_first_stage / decode_first_stage (VAE + latent_format), and sets the predictor on the UnetPatcher.
4. Predictor — a Prediction* in backend/modules/k_prediction.py for the sampler.
5. Register the engine — import it in backend/loader.py and add to possible_models.

Step 7.5 — Running a checkpoint too big for VRAM (quantization)

If the bf16 weights exceed the GPU (e.g. a 20B model on 24 GB), do NOT reimplement the DiT with Forge ops — reuse Forge's built-in quantization path:

• Feed Forge a quantized single-file checkpoint (fp8_e4m3fn / nf4 / gguf). The loader (backend/loader.py::load_huggingface_component) auto-detects the dtype and builds the model inside using_forge_operations(..., bnb_dtype=...), which monkeypatches torch.nn.Linear/ Conv for the duration of construction. Because your IntegratedX subclasses the diffusers model, its internal nn.Linear layers become Forge's quantization-aware ops automatically — no model-specific quant code. The bf16 checkpoint still loads unchanged on a bigger card (same path, dtype auto-detected), so "get the quantized one working" also gets bf16 working.
• Non-Linear fp8 params bite. A checkpoint that quantizes everything leaves RMSNorm / LayerNorm / Embedding weights in fp8, and those do raw multiplies/lookups Forge doesn't patch → RuntimeError: Promotion for Float8 Types is not supported. In the engine, after the model is loaded, cast every fp8 param that is NOT a Linear/Conv weight back up to the compute dtype (iterate .modules(), skip nn.Linear/nn.Conv*, recast the rest). The Linear weights stay quantized so the memory win survives.
• Memory-heavy VAE. A 3D/causal VAE decode can OOM next to a multi-B DiT even when sampling fit. Enable vae.enable_tiling() + enable_slicing() and free the DiT from VRAM before decode (memory_management.free_memory(1e30, device, free_all=True)).
• The Comfy-Org repackaged single-files use diffusers key naming (diffusers' single-file mapping is often the identity), so no key conversion is needed — verify with a 0-missing/0-unexpected load_state_dict against your IntegratedX.

Step 8 — Debug native generation with stage prints, then remove them

The standalone port is your oracle. If the native image is black/NaN/mush, gate temporary prints behind an env var and compare the native (x_in std, timestep, context std, latent std, decoded min/max) against the standalone at the same step. Black image = NaN somewhere; uniform-nonzero = a scaling/normalization mismatch. Remove the prints once it works.

The NaN-black-image trap (rectified flow + k-diffusion)

ForgeScheduleLinker.get_sigmas(n) builds the sampling schedule as append_zero(predictor.sigma(linspace(len(sigmas)-1, 0, n))) — i.e. it calls your predictor's sigma(t) on buffer indices t ∈ [0, len-1], then appends a terminal 0. The k-diffusion Euler loop runs for i in range(len(sigmas)-1), so it evaluates the model at every entry except the appended terminal. If your sigma(index) returns exactly 0 at index == 0 (e.g. time_snr_shift(shift, t/1000) with t=0), the last sampled sigma is 0, and to_d = (x - denoised)/sigma = 0/0 = NaN → black image. Fix: make sigma(index) map index 0 to sigma_min (>0), e.g. time_snr_shift(shift, (t+1)/timesteps), so only the never-evaluated terminal is 0. Symptoms to recognise: DiT forward is clean at every step (nan=False) but the final latent going into decode is nan=True; and the printed timestep of the last step is exactly 0.0. (Do NOT "fix" this by making timestep()/sigma() mutual identities on the raw sigma — that makes the scheduler read indices as sigmas, giving sigma_max≈999 and blowing up x_in.)

Pitfalls seen in practice

• out_proj(norm(x)) vs norm(out_proj(x)) — order matters (RMSNorm ⊄ Linear).
• RoPE "split-half" ([:d/2],[d/2:]) vs interleaved; HF rotate_half with duplicated cos/sin is equivalent to split-half — verify against the actual kernel.
• Context zero-padded to a fixed length (e.g. 512) with no attention mask: this is often how the model was trained — keep it; removing it makes things worse.
• 3D/causal VAE decode is memory-heavy; run it on CPU or tile it to avoid OOM next to a multi-B DiT on a 24 GB card.
• A symptom where img2img refines fine but text-to-image is mush points at the conditioning/context path (the DiT and VAE are probably fine) — diff the text→context stages.
• Editing repositories/huggingface_guess (or anything under repositories/) directly: it's git-ignored, so the change vanishes on a fresh clone. Register the arch at runtime from a tracked module instead (Step 7.1).
• fp16 on a DiT with a large residual stream (Cosmos-Predict2) overflows to NaN — declare supported_inference_dtypes = [torch.bfloat16, torch.float32] so Forge never picks fp16.
• Velocity SIGN. When wrapping a diffusers transformer, read the pipeline's denoising loop, not just the transformer: Z-Image's pipeline does noise_pred = -noise_pred right before scheduler.step. Forge's CONST predictor uses the model output directly as the flow derivative, so any such negation must be folded into the wrapper. Symptom: no NaN, but the latent std grows every step (diverges) and decodes to pure colour noise; the fix flips it so std tracks the diffusers per-step trajectory. Diagnose by printing per-step latent std for both your engine and the reference pipeline (callback_on_step_end) — they must match.
• Timestep convention can be inverted between arch families even at the same shift: Anima's Cosmos DiT wants t = sigma (t=1 is noise); Z-Image's NextDiT wants t = 1 - sigma (t=0 is noise, matching diffusers' (1000 - sigma*1000)/1000). Encode this in the predictor's timestep(sigma), and check it against the reference by printing the timestep the reference feeds its transformer at the first and last step.
• List/variable-length conditioning vs Forge's fixed-size cond tensor: models like Z-Image (NextDiT) / Qwen-Image take a list of variable-length, masked text embeddings. Forge can only carry a single padded tensor per cond, so zero-pad to a fixed length in get_learned_conditioning and recover each item's true length in the DiT wrapper by trimming trailing all-zero rows. Do NOT just pad and attend over the padding without a mask unless the model was trained that way (Anima was; Z-Image was not — padding-without-mask diverges ~30%).