Chore: Parameterize Depth Pro's DINOv2/decoder tree, then add a tiny test

[ianscrivener]

Another tweak arising from tiny test work

---

PRD: Parameterize Depth Pro's DINOv2/decoder tree, then add a tiny test

Background

Depth Pro (src/mflux/models/depth_pro/) is a ViT encoder/decoder depth-estimation
model — a DINOv2 backbone plus a multires convolutional decoder and FOV head. It is
architecturally unrelated to the diffusion-transformer models (FLUX.1, FIBO, Z-Image,
etc.), but shares the same underlying problem as FLUX.1: it predates the
"constructor takes dimension kwargs" convention the newer ports follow, and every
class in its tree hardcodes dimensions inline with zero constructor arguments:

# dino_v2/dino_vision_transformer.py (illustrative — see file for exact contents)
class DinoVisionTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.cls_token = ...      # hardcoded shape (1, 1, 1024)
        self.pos_embed = ...      # hardcoded shape (1, 577, 1024)
        self.blocks = [TransformerBlock() for _ in range(24)]  # ViT-L, 1024-dim

# dino_v2/patch_embed.py
class PatchEmbed(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Conv2d(3, 1024, kernel_size=16, stride=16)

# dino_v2/mlp.py
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1024, 4096)
        self.fc2 = nn.Linear(4096, 1024)

DepthProModel, DepthProEncoder, TransformerBlock are all zero-constructor-arg
too. DinoVisionTransformer (ViT-L, 24 layers, 1024-dim) is instantiated twice
(patch_encoder and image_encoder).

DepthProEncoder/MultiresConvDecoder hardcode conv channel counts (256/512/1024)
and call UpSampleBlock/FeatureFusionBlock2d/ResidualBlock with fixed dims baked
into the call site — even though those three leaf classes do accept dim kwargs
themselves. That's the key difference from FIBO's situation: in FIBO, composing real
leaves at toy dims worked around a blocked top-level class. Here, there is no level
of the tree where that composition bottoms out cleanly — the parameterizable leaves
are always invoked from hardcoded parents, so recomposing them still requires
touching (or bypassing) the hardcoded parent classes.

This is why the Depth Pro tiny model-saving test is currently blocked, and more
severely than FLUX.1: FLUX.1 at least has isolated components (T5Encoder,
CLIPEncoder) that are self-contained dead ends; Depth Pro's hardcoding runs through
its only component (get_components() is a single mega-component named
"depth_pro", mapping onto the whole DepthProModel), so there's nothing to route
around.

Why this is its own workstream, not a quick fix

Same reasoning as the FLUX.1/FIBO PRD (prd-parameterize-flux-timestep-norm.md):
mflux is a line-by-line numerical port. Any change to DinoVisionTransformer,
TransformerBlock, MLP, PatchEmbed, or the hardcoded call sites in
DepthProEncoder/MultiresConvDecoder touches Depth Pro's actual forward pass and
needs to be proven bit-identical at the default (production) size before it can be
trusted — not just "looks like a mechanical refactor." Given ~7 files are involved
(more than the 2-file FLUX.1/FIBO shared-class fix), this is a larger and more
error-prone refactor, and the verification pass is proportionally bigger.

Goal

1. Thread dimension parameters through the Depth Pro tree so a toy-sized instance
   can be constructed. Minimum viable surface, in dependency order:
   • PatchEmbed(in_chans: int = 3, embed_dim: int = 1024, patch_size: int = 16)
   • MLP(dim: int = 1024, hidden_dim: int = 4096)
   • TransformerBlock(dim: int = 1024, num_heads: int, mlp_hidden_dim: int, ...)
     — check Attention's existing signature first; it may already take dims and
     just need TransformerBlock to pass them through, similar to the FIBO
     single-block situation.
   • DinoVisionTransformer(embed_dim: int = 1024, num_blocks: int = 24, ...) —
     needs cls_token/pos_embed shapes to derive from embed_dim and a
     configurable img_size/patch grid (pos_embed length depends on patch count).
   • DepthProEncoder(...) / MultiresConvDecoder(...) — thread channel-count
     params through to the UpSampleBlock/FeatureFusionBlock2d/ResidualBlock
     calls that already accept them.
   • DepthProModel(...) — top-level wiring, likely just forwarding kwargs down.
     All new parameters default to current production values (no behavior change at
     defaults).
2. Prove numerical parity at default (production) dimensions against the PyTorch
   reference, per mflux-debugging (export-then-compare).
3. Once (1)-(2) land, write tests/model_saving/test_tiny_model_saving_depth_pro.py
   following the mflux-model-tiny-test skill — single mega-component, so
   _tiny_components() returns one key ("depth_pro") holding one fully-toy
   DepthProModel.

Non-goals

• Not touching FLUX.1's T5Encoder/CLIPEncoder or the shared
  AdaLayerNormZeroSingle/TimestepEmbedder fix — that's a separate PRD
  (prd-parameterize-flux-timestep-norm.md), unrelated model family.
• Not changing Depth Pro's torch_checkpoint loading mode or its single-component
  weight definition shape — only making the underlying classes constructible small.
• Not attempting a partial/leaf-only workaround (à la FIBO) — the investigation found
  no clean composition point, so don't force one; do the real parameterization.

Verification plan (per mflux-debugging skill)

1. Pick one deterministic repro: a fixed input image, fixed resolution.
2. Export reference tensors from the PyTorch/diffusers-adjacent Depth Pro reference
   (Apple's ml-depth-pro) at these checkpoints:
   • Patch embedding output (both patch_encoder and image_encoder paths)
   • DinoVisionTransformer output (first block, last block, final norm)
   • DepthProEncoder/MultiresConvDecoder fused feature maps
   • Final depth map + FOV head output
3. Run the same repro in mflux after the parameterization change (at default/
   production dims) and export matching tensors.
4. Compare per mflux-debugging's standard approach: shape+dtype match,
   max_abs_diff/mean_abs_diff/max_rel_diff against atol=1e-5, rtol=1e-5
   (fp32) or relaxed bf16/fp16 tolerances as appropriate; inspect actual values,
   not just summary stats.
5. Run mflux-save-depth end-to-end on a fixed input, before/after the change,
   confirm the output depth map is unchanged (or within known fp precision noise).
6. This model has no diffusion/RNG loop, so the usual seed/latent-injection RNG
   warning from mflux-debugging mostly doesn't apply — but confirm there's no
   dropout or other stochastic path active at inference before assuming determinism.

Acceptance criteria

• PatchEmbed, MLP, TransformerBlock, DinoVisionTransformer,
  DepthProEncoder, MultiresConvDecoder, DepthProModel all accept
  dimension parameters, defaulting to current production values.
• Tensor-level parity confirmed and documented against the PyTorch/Apple
  reference at default dimensions.
• End-to-end depth-map parity confirmed on a fixed input image.
• just test-fast / just test green, no regressions.
• tests/model_saving/test_tiny_model_saving_depth_pro.py written and green,
  single "depth_pro" component, fully toy-sized (no unshrinkable floor).
• .claude/notes/tiny-tests.md coverage matrix updated: depth_pro from
  "❌ blocked" to "✅".

References

• .claude/notes/tiny-tests.md — "Depth Pro — blocked, worse than FLUX.1" section
  (investigation findings, files read, exact param-count estimate).
• .cursor/skills/mflux-model-tiny-test/SKILL.md — tiny test pattern.
• .cursor/skills/mflux-debugging/SKILL.md — export-then-compare verification
  workflow.
• prd-parameterize-flux-timestep-norm.md — sibling PRD for the FLUX.1/FIBO shared
  timestep/norm classes; same category of problem, different model family, kept
  separate since the affected code and verification surface don't overlap.
• Files read during investigation (starting points for the fix): src/mflux/models/ depth_pro/weights/depth_pro_weight_definition.py, depth_pro_initializer.py,
  model/depth_pro.py, model/depth_pro_model.py, model/encoder/ depth_pro_encoder.py, model/encoder/upsample_block.py, model/decoder/ multires_conv_decoder.py, model/decoder/feature_fusion_block_2d.py, model/ decoder/residual_block.py, model/head/fov_head.py, model/dino_v2/ dino_vision_transformer.py, model/dino_v2/transformer_block.py, model/dino_v2/ patch_embed.py, model/dino_v2/mlp.py, model/dino_v2/attention.py.

[fxd0h]

Unparking note: the freeze this was filed under (#588) was about landing 0.19.0, and with 0.19.1 out too it's over. The marker can come off whenever someone picks this up.