Metadata-Version: 2.4
Name: latent-diffusion-tools
Version: 0.1.1
Summary: Latent diffusion model training and inference tools.
License-Expression: MIT
Project-URL: Homepage, https://github.com/CompVis/latent-diffusion
Project-URL: Source, https://github.com/CompVis/latent-diffusion
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0.0
Requires-Dist: numpy>=1.23.0
Requires-Dist: Pillow>=9.0.0
Requires-Dist: tqdm>=4.66.0
Requires-Dist: einops>=0.6.0
Requires-Dist: omegaconf>=2.3.0
Provides-Extra: train
Requires-Dist: torchvision>=0.15.0; extra == "train"
Requires-Dist: pytorch-lightning>=2.0.0; extra == "train"
Provides-Extra: test
Requires-Dist: build>=1.2.2; extra == "test"
Requires-Dist: twine>=5.1.0; extra == "test"
Dynamic: license-file

# latent-diffusion-tools

`latent-diffusion-tools` is a compact latent diffusion training and inference package for Python experiments.

It provides:

- a small latent diffusion model builder with checkpoint-aware setup
- a command-line fine-tuning loop for image folders and `.npz` arrays
- a command-line sample generation workflow
- Python APIs for model construction, training, and prediction

## Install

For the remote notebook workflow where dependencies are prepared separately:

```python
!python -m pip install --no-deps latent-diffusion-tools
```

For a local environment where dependency resolution is desired:

```powershell
pip install latent-diffusion-tools
```

## Build a Model

```python
from latent_diffusion_tools import model

net = model(pretrained=True, root="./ckpt/")
```

`pretrained=True` builds the latent diffusion module, prepares the bundled checkpoint, and records the prepared checkpoint path on `net.checkpoint_path`.

Use `pretrained=False` when training from scratch.

```python
from latent_diffusion_tools import model

net = model(pretrained=False)
```

## Training

Training input can be an image directory or a directory of `.npz` files.

Image files may be `.png`, `.jpg`, `.jpeg`, `.bmp`, `.tif`, or `.tiff`.

Each `.npz` file should include one of these array keys:

- `image`
- `img`
- `data`
- `x`

Command line:

```powershell
latent-diffusion-train ^
  --data-dir .\train_images ^
  --output-dir .\runs\latent_diffusion ^
  --epochs 5 ^
  --batch-size 4 ^
  --lr 1e-4 ^
  --image-size 64 ^
  --checkpoint-root .\ckpt
```

The command uses the bundled checkpoint setup by default and writes `latest.pt` to the output directory.

Python usage:

```python
from latent_diffusion_tools import fit

net, history = fit(
    data_dir="./train_images",
    output_dir="./runs/latent_diffusion",
    epochs=5,
    batch_size=4,
    lr=1e-4,
    image_size=64,
    checkpoint_root="./ckpt",
)
```

## Inference

Generate samples with the default model setup:

```powershell
latent-diffusion-infer ^
  --output-dir .\runs\latent_diffusion_samples ^
  --count 8 ^
  --image-size 64 ^
  --steps 16 ^
  --root .\ckpt
```

Use a fine-tuned checkpoint produced by `latent-diffusion-train`:

```powershell
latent-diffusion-infer ^
  --output-dir .\runs\latent_diffusion_samples ^
  --checkpoint .\runs\latent_diffusion\latest.pt ^
  --count 8 ^
  --image-size 64 ^
  --steps 16 ^
  --root .\ckpt
```

Python usage:

```python
from latent_diffusion_tools import predict

paths = predict(
    output_dir="./runs/latent_diffusion_samples",
    count=8,
    image_size=64,
    steps=16,
    root="./ckpt",
)
print(paths)
```

## Checkpoint Utilities

The model API prepares the bundled checkpoint automatically when `pretrained=True`.

Manual preparation is also available:

```python
from latent_diffusion_tools import prepare_checkpoint

checkpoint_path = prepare_checkpoint(root="./ckpt")
print(checkpoint_path)
```

## Attribution

This package is inspired by the latent diffusion training and sampling workflow from CompVis latent-diffusion.

- Upstream repository: <https://github.com/CompVis/latent-diffusion>
- Upstream license: MIT License

The implementation here is a compact training and inference package for practical experiments.
