1 files changed, 34 insertions, 0 deletions
diff --git a/src/config.py b/src/config.py
new file mode 100644
index 0000000..4ac56c4
--- /dev/null
+++ b/src/config.py
@@ -0,0 +1,34 @@
+"""
+Configuration file for MNIST Diffusion model.
+Contains all hyperparameters and constants.
+"""
+
+import torch
+
+# Device configuration
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Image parameters (MNIST)
+IMG_SIZE = 28
+NB_CHANNEL = 1
+NB_LABEL = 10
+
+# Diffusion parameters
+DIFFU_STEPS = 1000
+
+# Noise schedule (linear beta schedule)
+BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE)
+BETA = torch.cat((torch.tensor([0.0], device=DEVICE), BETA))
+ALPHA = 1 - BETA
+ALPHA_BAR = torch.cumprod(ALPHA, dim=0)
+
+# Training parameters
+EPOCHS = 10
+BATCH_SIZE = 64
+LEARNING_RATE = 2e-4
+NUM_WORKERS = 4
+
+# Paths
+DATA_DIR = "data"
+CHECKPOINT_DIR = "checkpoints"
+PLOTS_DIR = "plots"