examples.infer

Running inference.

Below is a code example of an indference script. Simply plug in your own checkpoint and start running inference. To extend this further, you might want to look into compiling your model to ONNX, TensorRT, OpenVino or other formats.

from torchvision import datasets
from torchvision.transforms import v2

from diffusion_models.diffusion_trainer import DiffusionTrainer
from diffusion_models.gaussian_diffusion.beta_schedulers import (
  LinearBetaScheduler,
)
from diffusion_models.gaussian_diffusion.gaussian_diffuser import (
  GaussianDiffuser,
)
from diffusion_models.models.SimpleUnet import SimpleUnet
from diffusion_models.utils.schemas import LogConfiguration
from diffusion_models.utils.schemas import TrainingConfiguration


if __name__ == "__main__":
  image_size = 64
  image_channels = 3

  training_configuration = TrainingConfiguration(
    batch_size=256,
    learning_rate=2 * 10e-4,
    number_of_epochs=500,
    training_name="ReworkedFrameworkBase",
    checkpoint_rate=100,
    mixed_precision_training=False,
    # gradient_clip=0.1,
  )
  log_configuration = LogConfiguration(
    log_rate=10,
    image_rate=635,
    number_of_images=5,
  )
  model = SimpleUnet(
    image_channels=image_channels,