arrow_back Projects
PERSONAL•2026•Duccio Meconcelli

Stable Diffusion 3.5 Feature Extractor

PythonPyTorchStable Diffusion 3.5MMDiTSafetensorsFeature Extraction
codeCode ↗launchView Resource ↗
§ 00 // ABSTRACT

A Python library for extracting internal hidden states and attention maps from Stable Diffusion 3.5 MM-DiT blocks for synthetic image detection, model attribution, and diffusion interpretability.

Stable Diffusion 3.5 Feature Extractor (sd35-feature-extractor) is a research toolkit that hooks into the Multimodal Diffusion Transformer (MM-DiT) backbone of Stable Diffusion 3.5 to extract internal latent representations—hidden states and joint attention maps—for synthetic image detection, generative model attribution, and mechanistic interpretability.

Stable Diffusion 3.5 Feature Extractor overview
Feature extraction pipeline across the MM-DiT joint transformer blocks of Stable Diffusion 3.5

Motivation & Inversion Pipeline

Standard pixel-space classifiers operate exclusively on the final rendered image, where generator-specific traces are often subtle or destroyed by post-processing. By encoding an input image into the Stable Diffusion 3.5 VAE latent space and running a single forward pass through the MM-DiT transformer at a controlled diffusion timestep $t \in [0, 1000]$, intermediate activations expose structural and semantic regularities across both the image and text conditioning streams.

Decoupling feature extraction from downstream classification allows high-dimensional representations to be computed once and reused across linear probes, MLP classifiers, layer-wise attribution studies, and cross-dataset generalization benchmarks.

MM-DiT Hook Architecture & Controls

The library instruments the 38 joint transformer blocks of Stable Diffusion 3.5 Large (as well as the Medium and Large Turbo variants) without modifying model weights:

  • Granular Layer & Stream Selection: Extracts any subset of image-branch hidden states (hidden_x, blocks 0–37), text-branch context states (hidden_context, blocks 0–36), and joint attention weights (attention), along with cross-layer averages (middle_avg).
  • Precomputed Triple-Encoder Conditioning: Supports offline extraction of text conditioning tensors (c_crossattn of shape [1, 154, 4096] and pooled y of shape [1, 2048]) from CLIP-L, CLIP-G, and T5-XXL, eliminating the need to keep all three text encoders in VRAM during image-batch processing.
  • Spatial Mean Pooling & Safetensors Serialization: Optional spatial mean pooling compresses per-layer image tokens from [1, 1024, 1536] to [1, 1536] (float16), serializing all selected streams into a single namespaced .safetensors archive per image.
output_features/
├── experiment_config.json
├── execution.log
├── real/
│   └── image001.safetensors   # hidden_x__*, hidden_context__*, attention__*
└── fake/
    └── image001.safetensors

Bias Mitigation & Benchmark Suite

A recurring failure mode in deepfake detection is shortcut learning on compression artifacts or resizing pipelines rather than generative priors. To support controlled ablation studies, the extractor integrates:

  • Preprocessing Modes: Configurable input pipelines including imagenet_style, brutal_resize, crop_100_then_resize, and native-resolution pass-through (none).
  • On-the-Fly JPEG Augmentation: Stochastic JPEG recompression during extraction to neutralize format asymmetries between real and synthetic datasets.
  • Automated Benchmark Acquisition: Interactive TUI and CLI downloaders (download_tui.py) for model checkpoints and standard detection benchmarks, including GenImage, Tiny GenImage, Unbiased GenImage, OpenFake (1.9M images across 20+ generators), Echo-4o-Image, and NanoBanana.

Resources