img_dir = r"datasets/FinalDataset_128"
prompt_path = r"datasets/FinalDatasetPrompts.txt"
vae_path = "stabilityai/sd-vae-ft-ema"
clip_path = "openai/clip-vit-base-patch32"
lpips_path = "vgg_weights/vgg.pth"



unet_params = {
    "in_channels": 4,  # 4 would be input/output channels due to VAE, 3 if operating in pixel space
    "channels": (128, 256, 512, 1024),
    "n_groups": 16,  # Number of normalization groups for groupnorm
    "dropout": 0.1,  # Dropout probability
    "T": 1000,  # Number of timesteps
    "t_embd": 128,  # Embedding dimension to timesteps
    "n_embd": 512,  # Fixed at 512 for CLIP Embeddings
    "n_heads": 16,  # Number of attention heads
    "n_layers": 3,  # Number of layers per Encoder/Mid/Decoder blocks
}


diffusion_params = {
    "beta1": 0.00085,  # Lowest noise level
    "beta2": 0.01200,  # Highest noise level
    "guidance_scale": 4.5,  # CFG Scale
}


training_params = {
    "batch_size": 64,
    "grad_accum_steps": 4,  # Number of batches to accumulate
    "epochs": 500,
    "t_max": 250,  # Dictates over how many epochs to adjust optimization lr. If t_max < epochs, usually the case, then the final (epochs - t_max) epochs would be at minimum lr
    "lr": 1e-4,  # Starting lr
    "min_lr": 1e-6,  # Ending (minimum) lr
    "pl_coeff": 0.25,  # Perceptual loss weight
    "epoch_save_interval": 5,  # Model checkpointing interval
    "model_dir": "saved_models"  # Checkpoint folder name
}