sd_vae_2_dino_v2_bridge

SD/E2E-SD VAE to DINOv2 Bridge

GET Dino v2 features from SD VAE latents

Architecture

Our design cleanly separates the modality-specific layers from the spatial processing body:

Latent Adapter: A lightweight convolutional stem (VAE-specific). Maps 4-channel VAE latents up to the bridge's working width. Bridge Backbone: A standard transformer body (width 768, depth 8) that remains VAE-agnostic. Token Head: A linear projection that maps transformer outputs to the expected DINO patch targets (e.g., 64 tokens of 768 dim).

Note: This decoupled design means for future models like FLUX, you can swap out just the Latent Adapter (to handle 16-channel latents) while freezing/reusing the learned bridge backbone.

Loss Function

The training utilizes a composite, geometry-aware loss function (bridge_loss) designed to prioritize structural and directional alignment over raw magnitude matching:

• Cosine Loss (Weight: 1.0): 1.0 - cosine_similarity(pred, target). The primary driver, focusing heavily on matching the semantic direction of the DINOv2 features.
• MSE Norm (Weight: 0.25): Standard MSE applied after L2-normalizing the predictions and targets.
• MSE Raw (Weight: 0.05): Standard MSE applied to the raw values. Keeps the scale grounded without letting magnitude differences dominate the gradients.

Training History

Current Quality Read

The bridge is learning correctly.

Strong signal" Retrieval@1 reached ~98.8% on the held-out validation subset.

This means the predicted features preserve enough image identity / semantic structure to retrieve the matching true DINO target among validation candidates.

However, raw patch cosine is still: ~0.759

So the bridge is not yet a perfect DINO replacement. It is already useful for ranking / retrieval-like proxy supervision, but should be improved before being treated as a high-fidelity DINO teacher.

Suggested target before production use as a serious DINO proxy:

• val cosine: 0.85+
• val NMSE: <0.25
• retrieval@1: remain >0.95 on larger external eval

How to Use

Prepare Data: Pre-pack your SD latents ([N, 4, 32, 32]) and DINOv2 features ([N, 64, 768]) into memory-mappable .npy files. Ensure you are targeting the 8x8 patch grid, excluding the DINO CLS token. Configure: Update paths and training knobs in the @dataclass class CFG (Cell 3). This serves as the single source of truth for the run. Run All: The notebook will handle package installation, wandb logging, dataset splitting, and mixed-precision (AMP) training automatically.

Basic

import torch

ckpt = torch.load("checkpoints/best.pt", map_location="cpu")


state_dict = ckpt["model"] if "model" in ckpt else ckpt

model = DinoBridgeV3(
    in_ch=4,
    target_tokens=64,
    target_dim=768,
    adapter_mid_channels=256,
    adapter_out_channels=512,
    adapter_depth=2,
    width=768,
    depth=8,
    heads=12,
    mlp_ratio=4.0,
    dropout=0.02,
)

model.load_state_dict(state_dict, strict=True)
model.eval().cuda()

Advanced

import torch
import torch.nn.functional as F

best_ckpt_path = OUT_DIR / "best.pt"
print(f"Loading checkpoint from: {best_ckpt_path}")

inference_model = DinoBridgeV3(
    in_ch=in_ch,
    target_tokens=TARGET_TOKENS,
    target_dim=TARGET_DIM,
    adapter_mid_channels=cfg.adapter_mid_channels,
    adapter_out_channels=cfg.adapter_out_channels,
    adapter_depth=cfg.adapter_depth,
    width=cfg.model_width,
    depth=cfg.model_depth,
    heads=cfg.model_heads,
    mlp_ratio=cfg.mlp_ratio,
    dropout=cfg.dropout,
).to(device)


if best_ckpt_path.exists():
    ckpt = torch.load(best_ckpt_path, map_location=device)
    inference_model.load_state_dict(ckpt["model"])
    print("Weights loaded successfully.")
else:
    print("Checkpoint not found. Make sure you have completed at least one training epoch.")

inference_model.eval()

# In practice, this would be the output from your VAE encoder: latent = vae.encode(image)
sample_latent, sample_target = val_ds[0]
sample_latent = sample_latent.unsqueeze(0).to(device).float() # Add batch dim: [1, C, H, W]

with torch.no_grad():
    with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=(device=="cuda" and cfg.amp)):
        pred_dino_features = inference_model(sample_latent)

print("\n--- Inference Results ---")
print("Input latent shape:", sample_latent.shape)
print("Predicted DINO features shape:", pred_dino_features.shape)
print("Ground truth DINO features shape:", sample_target.shape)

# 5. Quick comparison to ground truth (Cosine Similarity)
pred_norm = F.normalize(pred_dino_features[0].float(), dim=-1)
target_norm = F.normalize(sample_target.to(device).float(), dim=-1)
sim = F.cosine_similarity(pred_norm, target_norm, dim=-1).mean().item()
print(f"Average Cosine Similarity for this sample: {sim:.4f}")

Short HF model-card table

The training log shows validation improving consistently from epoch 1 to epoch 6: cosine rose from 0.624136 to 0.759174, NMSE fell from 0.610722 to 0.420825, and Retrieval@1 rose from 0.751818 to 0.988182. The run was interrupted during the final latest.pt save after epoch_006.pt had already been saved, so best.pt / epoch_006.pt are the safer checkpoints. :contentReference[oaicite:0]{index=0}