kohya-ss / kohya-ss/sd-scripts

New format for caching latent and Text Encoder outputs

Open
#1,750 22 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
7.2k
Forks
1.2k
Avg merge
11m
Merged PRs (30d)
2

Description

# Overview
- Policy: File management per image, supporting multiple architectures and data types
- Format: Using `.safetensors` format

---

# Latent Cache

## Filename Format
```
{original_image_base_name}_{width}x{height}_{architecture}.safetensors
```
Example: `image1_1024x0768_sd.safetensors`

## Internal File Structure

### Tensor Data
The following tensors are stored as needed:
```
latents_{latents_height}x{latents_width}_{data_type}
latents_flipped_{latents_height}x{latents_width}_{data_type}
alpha_mask_{latents_height}x{latents_width}
```

The data type can be one of the following: `fp32`, `bf16`, `fp16`, `float8_e4m3fn`, `float8_e5m2`.

The priority of types is `fp32>bf16=fp16>float8_e4m3fn=fp8_e5m2`, and if a higher-level type has already been stored, a new one will not be stored. fp16 and bf16, and multiple fp8 may coexist.

Also, when saving a higher-level type, the lower-level type is deleted.

`{latents_height}x{latents_width}` is used for multi-resolution training.

### Metadata (JSON format)
```json
{
"architecture": "sd", // or "sdxl" etc.
"width": "1024",
"height": "768"
"crop_ltrb_{latents_height}x{latents_width}": "0,10,0,10", // comma separated values
"format_version": "1.0.0"
}
```

## Key Features
1. File Management
- Individual file management per image
- Multiple data types consolidated into a single file
- Quick image size retrieval through filename parsing (`{width}_{height}` part) without opening the file

2. Type Support
- Support for multiple data types (fp32, fp16, bf16, fp8)
- Support for multiple fp8 implementations (E4M3, E5M2)

3. Functionality
- Includes flipped data for flip augmentation
- Alpha mask retention

---

# Text Encoder Cache

## Filename Format
```
{original_image_base_name}_{architecture}_{max_token_length}{_m}_te.safetensors
```
`_m` indicates that an attention mask has been applied. `max_token_length` is the length for the main Text Encoder (SD3, FLUX for T5XXL).

Example: `image1_sd3_512_te.safetensors`, `image2_sdxl_77_te_safetensors`, `image1_flux_512_m_te.safetensors`

## Internal File Structure

### Tensor Data
Depending on the architecture, the following is an example:
```
lg_out_1_fp16
t5_out_1_fp16
t5_attn_mask_1
lg_out_2_fp16
t5_out_2_fp16
...
```

The key has a suffix at the end that indicates the data type. The type and priority are the same as for latent.

The attention mask is optional. The attention mask is a long tensor, so it doesn't have the data type suffix.

There is always a suffix indicating the index of the caption before the masked or the data type suffix. The index is 1-based.

The tensor data doesn't have a batch dimension.

### Metadata (JSON format)
```json
{
"architecture": "sd", // or "sdxl" etc.
"caption1": "a caption for the image",
"caption2": "another caption for the image", // "capiton2", "caption3" ... optional
"format_version": "1.0.0"
}
```

## Key Features
1. File Management
- Individual file management per image
- Multiple data types consolidated into a single file

2. Type Support
- Support for multiple data types (fp32, fp16, bf16, fp8)
- Support for multiple fp8 implementations (E4M3, E5M2)

3. Functionality
- Includes optional attention mask.
- Supports multiple captions.

---

`.safetensors` may be zipped in the future, in which case the extension will be `.zip`.

Please feel free to share your thoughts and suggestions on this specification.

EDIT: Fix fp8 dtype to `float8_e4m3fn` and `float8_e5m2`, `captions` as list. Use `x` instead of `_` between the height and the width to match the existing name. Add `crop_ltrb` to the metadata. Fix the metadata because safetensors can have `Dict[str, str]` only. Change format version to `major.minor.patch`. Add the max token length to the file name for the text encoder cache. The tensor in the text encoder cache doesn't have the batch dimension, because it is cumbersome to read the entire tensor for reading one of the captions, or update part of the tensor. If an attention mask is applied, it is saved to a separate file. Define what model `max_token_length` refers to.

Contributor guide

No contributing guide indexed for this repository

Research direction

No files, tests, or entry points are identified. Start by reviewing the proposed latent-cache and text-encoder-cache schemas, then clarify the implementation scope and acceptance criteria; done should mean an agreed format specification and a defined path for adopting it.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.