import torch


# =========================
# Device
# =========================

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# =========================
# Dataset
# =========================

IMAGE_SIZE = 28
CHANNELS = 1


# =========================
# DDPM
# =========================

TIMESTEPS = 1000

BETA_START = 1e-4
BETA_END = 0.02


# =========================
# Model
# =========================

TIME_EMBED_DIM = 128

BASE_CHANNELS = 64
MID_CHANNELS = 128


# =========================
# Training
# =========================

BATCH_SIZE = 128
EPOCHS = 20

LEARNING_RATE = 2e-4
WEIGHT_DECAY = 1e-4


# =========================
# Checkpoints
# =========================

CHECKPOINT_DIR = "checkpoints"

MODEL_NAME = "ddpm_fashion_mnist.safetensors"


# =========================
# Sampling
# =========================

NUM_SAMPLES = 16


# =========================
# Random Seed
# =========================

SEED = 42