Create a model

Build a small 2D regression model from Python:

import torch
from widitapp import WiDiTApp

app = WiDiTApp()
model = app.model(
    dim=2,
    use_diffusion=False,
    hidden_size=64,
    depth=1,
    num_heads=4,
    patch_size=2,
    window_size=4,
    use_flash_attention=False,
)
model.eval()
with torch.no_grad():
    prediction = model(torch.zeros(1, 1, 32, 32))

Use dim=3 for volumes, or unet=True to build a U-Net. Inputs use channel-first shapes: (N, C, H, W) for batched 2D images and (N, C, D, H, W) for batched 3D volumes. Choose spatial sizes compatible with the model’s patch and window settings.

Create your own training app

As in Supercat and Dosefusion, implement datasets() and decorate it with Cluey’s @method. Return a pair of PyTorch datasets: training first, validation second. The inherited train tool collects the options from model(), dataloaders(), and your datasets() method.

Save this example as myapp.py. It reads tensor pairs from train.pt and validation.pt in a directory you provide:

from pathlib import Path

import torch
from cluey import method
from torch.utils.data import TensorDataset
from widitapp import WiDiTApp


class MyApp(WiDiTApp):
    @method
    def datasets(self, data_dir: Path = None, **kwargs) -> tuple:
        """Load paired input and target tensors."""
        if data_dir is None:
            raise ValueError("Provide data_dir (CLI: --data-dir).")
        data_dir = Path(data_dir)

        def load_split(name):
            inputs, targets = torch.load(
                data_dir / name, map_location="cpu", weights_only=True
            )
            return TensorDataset(inputs.float(), targets.float())

        return load_split("train.pt"), load_split("validation.pt")


if __name__ == "__main__":
    MyApp.tools()

Each file should contain a tuple (inputs, targets) saved with torch.save((inputs, targets), path). For single-channel 2D data, both tensors have shape (N, 1, H, W). Use separate samples for training and validation. Preprocess and normalize your data consistently before saving it; WiDiTApp does not read image formats or normalize intensities for you.

Inspect the generated options and train a small regression model:

python myapp.py train --help
python myapp.py train \
    --data-dir /path/to/tensor-pairs \
    --no-use-diffusion \
    --dim 2 \
    --hidden-size 64 \
    --depth 1 \
    --num-heads 4 \
    --no-use-flash-attention \
    --batch-size 4 \
    --epochs 40 \
    --results-dir results \
    --run-name example

You can also call MyApp().train(data_dir=Path("/path/to/tensor-pairs"), use_diffusion=False, ...) from Python with the corresponding keyword arguments.

Dataset and training modes

• Each dataset item must return (input, target) or (input, target, timestep). Individual images omit the batch dimension; the data loader adds it. Samples in a batch must have compatible shapes.

• With use_diffusion=True (the default), the target is noised and the input is passed to the model as conditioned=input. The training loop samples diffusion timesteps; any dataset-provided timestep is replaced. WiDiT expects the conditioning input and target to have the same shape, so resize paired data to match before training.

• With use_diffusion=False, the model predicts the target directly from the input using mean squared error. An optional dataset timestep is decremented by one before being passed to the model.

• Return (training_dataset, None) to omit validation. The current training loop saves checkpoints during validation only, so provide a validation dataset when you need saved models.

Training data is shuffled; validation data is not. Both loaders retain the last incomplete batch. Override dataloaders() if you need custom sampling or collation.

Common options

Python parameter names use underscores; their CLI equivalents use hyphens, for example batch_size becomes --batch-size. Boolean options also have a negative form, such as --no-use-diffusion.

The default output-channel counts assume single-channel targets. For multi-channel regression set out_channels to the target channel count; diffusion with learned variance needs twice that count. A preset defines its own spatial dimensionality and core architecture. Loading a checkpoint uses the saved model configuration rather than the model-building options.

For finer control, call widitapp.training.train directly with a model and data loaders. It additionally accepts precision (fp16 by default, bf16, fp32, or no), loss_fn (mse, smoothl1/huber, or a PyTorch loss module), wandb_config, and wandb_log_artifacts. These arguments are not forwarded by WiDiTApp.train. The custom loss applies to regression; diffusion uses its own objective.

Checkpoints and prediction

Training maintains an exponential moving average (EMA) of model weights and evaluates that model after each epoch. With validation, outputs are written to:

results/<run_name>/
    log.txt
    checkpoints/
        best.pt

best.pt stores the EMA model with the best validation loss. If the initial validation loss is non-finite, a fallback checkpoint is saved until a finite improvement is available. There is currently no separate final checkpoint or optimizer-state resume: checkpoint=... initializes a new training run with saved model weights.

Load a checkpoint with WiDiTApp().model(checkpoint=Path(".../best.pt")) or widit.load_model(path). For regression, move the loaded model and inputs to the same device, call model.eval(), and predict under torch.no_grad(). Diffusion inference requires an iterative sampling loop, rather than a single forward pass; widitapp.diffusion.create_diffusion supplies the diffusion process and sampling methods.

To expose prediction in your app, override predict() with @main from Cluey. Implement input loading, preprocessing, checkpoint loading, regression or diffusion inference, and output writing for your data format.

Package your command-line app

If your class lives in myapp/apps.py, add these entry points to your app’s Poetry configuration and install that app with python -m pip install -e .:

[tool.poetry.scripts]
myapp = "myapp.apps:MyApp.main"
myapp-tools = "myapp.apps:MyApp.tools"

myapp invokes the method decorated with @main (your prediction method), while myapp-tools train invokes inherited training. The training-only example above can use just the myapp-tools entry point until prediction is implemented. Decorate additional commands with Cluey’s @tool.

Example applications

The following projects demonstrate how to extend WiDiTApp:

• Supercat (supercat/apps.py): provides 2D and 3D super-resolution datasets, tiled prediction, and additional porosity tools. Its supercat/pretrain.py provides further subclasses for image and video pretraining.

• Dosefusion (dosefusion/apps.py): builds datasets from a CSV with partition selection and implements volume prediction with diffusion and fusion sampling options. Its data preparation and prediction code live in dosefusion/data.py and dosefusion/models.py.

In both projects, pyproject.toml registers the app’s .main and .tools entry points. Use their dataset and prediction implementations as examples when adapting the base class to your own files and task.