Ostris AI Toolkit

Gradient Checkpointing

When train.gradient_checkpointing: true is set, the toolkit calls model.enable_gradient_checkpointing(). Your network should implement this using torch.utils.checkpoint.checkpoint(..., use_reentrant=False). Crucially, do not gate this check on self.training; it should run whenever torch.is_grad_enabled() is true.

Quantization

When quantize: true is used, nn::Linear layers are replaced with optimum.quanto quantized layers.

• Dimension Constraint: The quanto matmul kernel only accepts 2D or 3D activations. If your network applies a Linear layer over a 4D tensor (e.g., (B, L, D, N)), you must reshape to 3D before the call and back to 4D after.
• Performance Tip: If a frozen sub-model (like a vision tower) contains Conv3d, it may be slow due to lack of fast bf16 kernels. If it is not needed for your specific task, consider dropping it (e.g., text_encoder.model.visual = None).

When implementing a new BaseModel subclass, you must follow this execution order:

1. Load: load_model() is called to build the transformer, text encoder(s), tokenizer(s), VAE, and scheduler. These are stored on self for other methods to access.
2. Caching (Optional): The trainer may call encode_images() (for latent caching) and get_prompt_embeds() (for text-embed caching) before training starts.
3. Train Step:
   • Latents are retrieved from cache or encode_images().
   • add_noise() (from BaseModel) mixes noise and timesteps.
   • condition_noisy_latents(noisy_latents, batch) is your hook to inject control/reference conditioning.
   • get_noise_prediction(latent_model_input, timestep, text_embeddings) performs the forward pass under autograd.
   • Loss is calculated using get_loss_target(noise=..., batch=...).
4. Sampling Previews: generate_images() calls get_prompt_embeds(), then get_generation_pipeline() once, and generate_single_image(...) per prompt. Note: The pipeline receives embeds only, never raw text.
5. Saving: Full fine-tunes use save_model(). For LoRA, use convert_lora_weights_before_save/load() to map keys to the public convention (typically using the diffusion_model. prefix).

The manager is designed to be self-contained and robust across different hardware and environments:

• Self-Contained Logic: The installation logic resides within the repository itself. External launchers (like install.sh or desktop launchers) simply call this CLI.
• Hardware-Specific Specs: Hardware-to-spec mapping (e.g., CUDA 13, ROCm, Mac, or CPU) is handled via spec.py. It manages specific torch pins and accelerator wheels (like flash-attn, NATTEN, or triton) based on the detected platform.
• Environment Isolation: To prevent system conflicts, the manager installs dependencies into a local virtual environment (.venv/). It also manages local, git-ignored versions of tools like FFmpeg (.ffmpeg/), Node.js (.node/), uv (.uv/), and MinGit (.mingit/) inside the repository.
• Update Safety: The update command performs a fast-forward git pull. If the repository has untracked files (a "dirty tree"), the update will abort by default to prevent overwriting local work. After pulling, the manager re-executes itself to ensure the new code performs the dependency sync and migrations.
• State Management: Environment state (such as requirements hashes and applied migrations) is stored within the virtual environment in aitk_manager_state.json. Deleting the .venv/ directory will reset the manager state.

Training is performed by passing a YAML configuration file to run.py.

Workflow:

1. Prepare Config: Copy an example config from config/examples/ to the config folder and rename it (e.g., config/examples/train_lora_flux_24gb.yaml $\rightarrow$ config/my_training.yml).
2. Edit Config: Modify the .yml file according to the comments provided within the file.
3. Execute Training: Run the following command:

   python run.py config/your_config.yml

Important Notes:

• Outputs: A folder containing checkpoints and images will be created based on the names specified in your config file.
• Resuming: You can stop training with Ctrl+C. When you restart using the same config, it will resume from the last checkpoint.
• Warning: Do not press Ctrl+C while the system is actively saving a checkpoint, as this will likely corrupt the file.

python run.py config/my_training.yml

To add a new model architecture so it can be used for training, follow these three steps:

1. Register the class: The toolkit scans extensions/ and extensions_built_in/ for a module-level AI_TOOLKIT_MODELS list. Import your model class into extensions_built_in/diffusion_models/__init__.py and append it to that list. Alternatively, create a new folder under extensions/ with its own AI_TOOLKIT_MODELS list.
2. Match the architecture: Ensure your class has an arch attribute (e.g., arch = "example"). This string is matched against the model.arch key in your training configuration YAML.
3. Expose in Web UI: To make the model selectable in the UI, add an entry to ui/src/app/jobs/new/options.ts. You can copy the structure from an existing architecture like ideogram4.

model:
  arch: "example"
  name_or_path: "/path/to/weights"
  quantize: true        # optional: qfloat8 the transformer
  quantize_te: true    # optional: qfloat8 the text encoder
train:
  gradient_checkpointing: true

When implementing attention, do not import flash_attn or xformers at the top level of your module, as this will cause ImportError on machines without them. Instead, make them optional and selected at runtime.

Recommended Pattern:

1. Guarded Import: Use a try-except block to set a flag.
2. Default to Native: Use torch.nn.functional.scaled_dot_product_attention (SDPA) as the default backend.
3. Runtime Selection: Provide a set_attention_backend("native"|"flash") method on the parent model that propagates the flag to individual attention modules.
4. Branching: Inside the forward pass of the attention module, branch on the flag rather than swapping the module instance (to avoid losing trained weights).

try:
    from flash_attn import flash_attn_varlen_func
    _FLASH_ATTN_AVAILABLE = True
except ImportError:
    flash_attn_varlen_func = None
    _FLASH_ATTN_AVAILABLE = False

To implement a model that uses an image for conditioning (like an instruct or edit model), follow these patterns:

1. Latent Conditioning: In condition_noisy_latents, encode batch.control_tensor (expected shape (B, 3, H, W) in [0, 1]) using the VAE. Attach these to the noisy latents by either adding extra channels via torch.cat(..., dim=1) or adding extra sequence tokens. You must slice the prediction back down in get_noise_prediction before returning it.
2. Visual Language (VL) Conditioning: If the text encoder requires the control image, set self.encode_control_in_text_embeddings = True. This ensures get_prompt_embeds receives control_images.
3. Multiple Reference Images: Set self.has_multiple_control_images = True to use batch.control_tensor_list.
4. Preview Generation: In generate_single_image, load the file path from gen_config.ctrl_img and apply the same conditioning logic used during training to ensure previews are accurate.

To run AI Toolkit on DGX OS, you must use Python 3.11. It is recommended to use miniconda to create a virtual environment to avoid conflicts with the system Python installation.

1. Set up Python 3.11 via Miniconda

First, download and install the latest Miniconda for Linux aarch64:

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh
chmod u+x Miniconda3-latest-Linux-aarch64.sh
./Miniconda3-latest-Linux-aarch64.sh

After installation, restart your bash or SSH session. To prevent the 'base' environment from activating automatically, you can run:

conda config --set auto_activate_base false

Create and activate a dedicated environment for AI Toolkit:

conda create --name ai-toolkit python=3.11
conda activate ai-toolkit

2. Install PyTorch

Install the specific PyTorch versions required for the environment:

pip3 install torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130

3. Install remaining requirements

Install the dependencies listed in the DGX requirements file:

pip3 install -r dgx_requirements.txt

# Summary of installation steps
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh
chmod u+x Miniconda3-latest-Linux-aarch64.sh
./Miniconda3-latest-Linux-aarch64.sh
conda create --name ai-toolkit python=3.11
conda activate ai-toolkit
pip3 install torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130
pip3 install -r dgx_requirements.txt

To train a PRXPixel model, use the prx_pixel architecture in your configuration. Ensure your name_or_path points to a directory containing the required subfolders: transformer/, text_encoder/, tokenizer/, and scheduler/.

Important Dataset Requirement: Datasets must be bucketed to multiples of 16px (calculated as vae_scale_factor * patch_size).

Optimization Options:

• quantize: true: Enables qfloat8 quantization for the transformer.
• quantize_te: true: Enables qfloat8 quantization for the Qwen3-VL text encoder.
• gradient_checkpointing: true: Recommended for training to save memory.

model:
  arch: "prx_pixel"
  name_or_path: "/path/to/prxpixel-t2i"   # diffusers folder: transformer/, 
                                          # text_encoder/, tokenizer/, scheduler/
  quantize: true        # optional: qfloat8 the transformer
  quantize_te: true     # optional: qfloat8 the Qwen3-VL text encoder
train:
  gradient_checkpointing: true
sample:
  guidance_scale: 5.0
  sample_steps: 28