ComfyUI/docs/agents/models.md

# Models, Device, and Memory

Detailed rules referenced from [AGENTS.md](../../AGENTS.md).

- Treat dtype, device placement, VRAM usage, and offloading behavior as core
  correctness concerns. Check CPU, CUDA, ROCm, MPS, DirectML, XPU, NPU, and low
  VRAM implications when touching shared execution or loading code.

## Autograd and Model Freezing

- Do not add `torch.no_grad`, `torch.inference_mode`, or inference-mode helper
  wrappers in ComfyUI code. The only allowed inference-mode-related use is
  disabling a globally set inference mode when a training path needs gradients.
- Do not add freeze, unfreeze, or trainability toggles to model classes. ComfyUI
  models are always treated as frozen for inference, so explicit freeze
  functionality is redundant and should not be added.
- Remove training-only behavior such as dropout from inference model code, but
  preserve checkpoint and state-dict compatibility when doing so. If deleting a
  module would change state-dict keys, module ordering, or checkpoint loading
  behavior, replace it with a no-op such as `nn.Identity` instead of removing the
  slot outright.

## Use the Existing Optimized Operations

- Prefer native ComfyUI formats and existing quantization/offload helpers over
  adding parallel code paths. Use `comfy.quant_ops`, `comfy.model_management`,
  `comfy.memory_management`, `comfy.pinned_memory`, `comfy_aimdo`, and
  `comfy-kitchen` helpers where they already solve the problem.
- Model implementations must use an existing optimized Comfy Kitchen or
  ComfyUI operation whenever one supports the required math and tensor layout
  without changing expected dtype, device, memory, or interface behavior. This
  is the default implementation requirement, not an optional follow-up
  optimization.
- Before implementing model math, inspect the operations already exposed by
  Comfy Kitchen, `comfy.quant_ops`, and existing ComfyUI model helpers. Check
  for optimized single, paired, fused, layout-specific, and quantized variants
  before writing a local implementation or composing lower-level torch ops.
- Use the compatible optimized operation first and adapt the model's inputs to
  its documented layout while preserving the model's exact math. If several
  optimized variants apply, benchmark representative model shapes and select
  the fastest valid path.
- Add or retain a local implementation only when no existing optimized
  operation supports the required math, layout, dtype, device, autograd, or
  patch contract. Keep differentiable or patch-compatible fallbacks when the
  optimized inference operation does not provide those contracts.
- Use the existing ComfyUI cast, offload, and cleanup helpers for parameters
  passed to optimized operations. Preserve model-specific epsilon, scaling,
  layout, dtype, device, and output-shape behavior.
- Prefer ComfyUI's shared optimized kernels and backend dispatchers over
  handwritten implementations of the same operation. Remove duplicate local
  kernels and adapt inputs to the shared operation's documented layout while
  preserving the model's original math and output contract.
- All models should use the optimized attention function selected by ComfyUI.
  Treat optimized backend functions, dispatch helpers, and capability-selected
  callables as opaque. Higher-level code must not inspect function identity,
  names, modules, or implementation details to decide behavior.
- Apply the same opacity rule to similar patterns beyond attention: callers
  should depend on the documented interface and result contract, not on which
  backend implementation was selected underneath.
- Do not use custom inference ops that only duplicate an existing op while
  upcasting to float32, such as custom RMSNorm variants. Use the generic ComfyUI
  ops and/or native torch ops instead.

## Model Classes and Constructors

- If a model class `__init__` has an `operations` parameter, assume
  `operations` is never `None`. Do not add fallback branches or default torch
  ops for a missing `operations` object.
- Do not add unnecessary parameters to model, model block, or model ops related
  classes. Constructor and forward signatures should carry only values that are
  actually needed by that object for inference.
- Reuse existing model classes, blocks, ops, and helper modules when appropriate.
  Before implementing a new version of a model component, search the existing
  model code for a class or helper that already provides the behavior.

## Model Detection

- Model detection code that inspects linear weight shapes should only use the
  first dimension. The second dimension may be half the original size for
  NVFP4 or other 4-bit quantized models.
- A model-detection signature must guard every state-dict key it dereferences.
  Do not partially match a format and then raise an incidental `KeyError` while
  extracting its configuration.
- Order model-detection checks from established or more-specific signatures to
  newer or broader signatures. Put a broad new detector near the generic
  fallback when giving it higher precedence could steal another model family.

## Tensors and Python Values

- Avoid adding `einops` usage in core inference code. Use native torch tensor
  ops such as `reshape`, `view`, `permute`, `transpose`, `flatten`, `unflatten`,
  `unsqueeze`, and `squeeze` instead.
- Do not use tensors as general-purpose Python data structures. Keep metadata,
  bookkeeping, counters, flags, shape math, padding math, index planning, memory
  estimates, and control-flow decisions in plain Python values unless the data
  must participate directly in tensor computation. Do not create tensors for
  structural metadata that is only used for Python-side control flow. Sequence
  lengths, cumulative offsets, split indices, window counts, slice boundaries,
  and repeat counts should be kept as Python ints/lists from the point they are
  computed. Do not build them as CPU/GPU tensors and then cast, move, validate,
  or convert them back to Python for `split`, `tensor_split`, indexing plans,
  loops, or cache keys. Avoid creating temporary tensors just to use tensor
  methods for scalar or structural calculations.

## Dtype and Device

- Avoid unnecessary casts and transfers. Preserve the intended compute dtype,
  storage dtype, bias dtype, and original tensor shape metadata.
- Do not cast the result of an optimized backend operation back to its input
  dtype unless that backend's documented result contract requires normalization.
  In particular, trust the selected optimized-attention implementation to honor
  its dtype contract.
- Avoid defensive shape and configuration checks that merely replace the clear
  failure from the tensor operation immediately below them. Add explicit
  validation only when it provides materially better context at a real boundary
  or prevents silent incorrect output.
- Assume inputs to the main model forward are already in the compute dtype by
  default, except integer inputs such as some model timestep tensors. Do not add
  defensive or convenience casts in model code; it is better for invalid dtype
  plumbing to error clearly than to hide it with unnecessary casts.
- Raw model parameters that are not owned by an op and may be initialized in a
  dtype different from the compute dtype should be cast at use in forward or
  inference code with `comfy.ops.cast_to_input` or
  `comfy.model_management.cast_to` to avoid dtype mismatches.
- Model code should not care what dtype it is initialized in, and model
  `__init__` methods should not contain workarounds for specific dtypes. Dtype
  workaround code, such as making a model work with fp16 compute, belongs in the
  execution or model-management layer that owns compute policy.
- Model code should not perform unnecessary device-to-CPU or CPU-to-device
  transfers. New allocations must be created on the correct device and dtype;
  never allocate on CPU and then move to GPU, or allocate in one dtype and then
  convert to another.

## Latent Layout

- Keep model-native latent layout handling inside the model or latent-format
  owner, not in helper nodes. Do not collapse, expand, pack, or unpack latent
  dimensions in nodes or other caller-side adapters just to satisfy a model
  forward; the model path should consume and return the native latent shape for
  that model family.
- DiT models should accept latent dimensions that are not exact patch-size
  multiples. Use `comfy.ldm.common_dit.pad_to_patch_size` on every patchified
  target or reference input, then crop only the target output back to its
  original dimensions.

## Memory and Caches

- Model code itself should not perform memory management. Loading, unloading,
  offloading, device movement, VRAM policy, cache lifetime, and cleanup belong
  in the relevant model-management and execution layers, not inside model
  implementations.
- Do not add global, module-level, class-level, singleton, or model-owned stores
  for tensors or other large memory that persist across executions. Temporary
  caches must be scoped to a single execution or forward/encode/decode call:
  allocate them in the owning top-level call, pass them explicitly through the
  call stack, and let them be discarded when that call returns.
- Follow the Wan VAE temporal cache pattern for temporary caches: create a local
  cache such as `feat_map` for the encode/decode operation, pass it into the
  blocks that need it, and do not retain it on the model or in global state.
- When slicing large tensors, copy the slice if the sliced tensor's lifetime
  exceeds the current function scope. Do not keep a long-lived view into a large
  backing tensor when a smaller copy would release memory sooner.
- Avoid caches that persist across different executions as much as possible.
  Persistent caches are acceptable only when they use a very minimal amount of
  memory and have a clear ownership and invalidation story.

## Initialization

- In model init code, prefer `torch.empty` for parameter/buffer placeholders
  that are populated from the model state dict instead of zero-initializing with
  `torch.zeros` or similar. If an allocation is not loaded from the state dict
  and is useless for inference, do not include it.
- `nn.Parameter` tensors that are stored in and populated from the model state
  dict should be initialized with `torch.empty`, not with zero, random, or
  otherwise meaningful initialization.
- Model initialization should describe module structure, not fabricate
  checkpoint-owned tensor contents. Parameters and buffers that are loaded from
  the state dict must not be manually initialized, reassigned, or filled with
  fallback values unless that value is actually used when no checkpoint key
  exists.

## Optimization

- Use fused or compound torch operations such as `addcmul` when they naturally
  match the math. Reducing Python and torch dispatch overhead is a valid
  optimization when it does not obscure the code or change dtype/device
  behavior.
- When optimizing, favor small measurable changes: fewer allocations, fewer
  device transfers, less peak memory, better batching, or use of a faster
  existing backend op.