latent-diffusion-tools 0.1.1

Latent diffusion model training and inference tools.

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 -m pip install --no-deps latent-diffusion-tools

For a local environment where dependency resolution is desired:

pip install latent-diffusion-tools

Build a Model

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.

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:

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:

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:

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:

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:

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:

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.