Implementation:AUTOMATIC1111 Stable diffusion webui DeepDanbooru Model
| Knowledge Sources | |
|---|---|
| Domains | Deep Learning, Image Tagging, Computer Vision |
| Last Updated | 2025-05-15 00:00 GMT |
Overview
Defines the PyTorch neural network architecture for the DeepDanbooru image tagging model, a deep ResNet-like CNN with 178 convolutional layers that classifies images into 9176 Danbooru tags.
Description
The DeepDanbooruModel class is a custom ResNet variant implemented as a PyTorch nn.Module. It is built from explicit sequential convolution layers (n_Conv_0 through n_Conv_178) with residual connections, ReLU activations, and padding. The forward pass processes a 512x512 image through progressively deeper feature maps (64 to 4096 channels) with strided convolutions for downsampling. The network ends with global average pooling and sigmoid activation to produce multi-label tag probabilities across 9176 Danbooru tags.
The model includes a custom load_state_dict override that extracts a tags list from the state dictionary before loading weights, allowing tag metadata to be stored alongside model parameters.
The architecture follows standard ResNet bottleneck block patterns with 1x1 projection, 3x3 spatial convolution, and 1x1 expansion layers, using explicit residual additions throughout.
Usage
Use this model when performing DeepDanbooru interrogation to automatically tag anime-style images with Danbooru tags. The model is instantiated by the DeepDanbooru module and loaded with pre-trained weights to provide the underlying inference capability for the interrogation feature in the WebUI.
Code Reference
Source Location
- Repository: AUTOMATIC1111_Stable_diffusion_webui
- File: modules/deepbooru_model.py
- Lines: 1-678
Signature
class DeepDanbooruModel(nn.Module):
def __init__(self):
def forward(self, *inputs):
def load_state_dict(self, state_dict, **kwargs):
Import
from modules.deepbooru_model import DeepDanbooruModel
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| inputs | torch.Tensor | Yes | A batch of images as a tensor with shape (N, H, W, 3) where H and W are 512. The tensor is permuted to (N, 3, H, W) internally. |
Outputs
| Name | Type | Description |
|---|---|---|
| result | torch.Tensor | A tensor of shape (N, 9176) with sigmoid-activated probabilities for each Danbooru tag. |
Usage Examples
from modules.deepbooru_model import DeepDanbooruModel
import torch
# Instantiate the model
model = DeepDanbooruModel()
# Load pre-trained weights (tags are extracted automatically)
state_dict = torch.load("deepbooru_model.pt")
model.load_state_dict(state_dict)
model.eval()
# Run inference on a 512x512 image tensor (N, H, W, C)
image_tensor = torch.randn(1, 512, 512, 3)
with torch.no_grad():
tag_probabilities = model(image_tensor)
# tag_probabilities has shape (1, 9176)
# Each value is the probability of the corresponding Danbooru tag