stable_diffusion_trainer/data_set/
mod.rs

1//! This module contains the data set types used for training.
2
3pub mod image_data_set;
4pub use image_data_set::*;
5
6use crate::prelude::*;
7
8/// A training data set.
9#[derive(Debug, Serialize, Deserialize)]
10pub struct TrainingDataSet {
11    /// The training images.
12    pub training: ImageDataSet,
13    /// The regularization images.
14    pub regularization: Option<ImageDataSet>,
15}
16
17impl TrainingDataSet {
18    /// Create a new training data set.
19    pub fn new(training: ImageDataSet) -> Self {
20        let regularization = Default::default();
21        TrainingDataSet { training, regularization }
22    }
23
24    /// Set the regularization images for the training data set.
25    pub fn with_regularization_images(mut self, regularization_images: ImageDataSet) -> Self {
26        self.regularization = Some(regularization_images);
27        self
28    }
29}