From f1be02af8c33d4136ad7c90dd43158a8ff6e5af1 Mon Sep 17 00:00:00 2001
From: gdamms <damguillotin@gmail.com>
Date: Thu, 5 Feb 2026 19:55:58 +0100
Subject: better plots, semi working autoencoder

---
 main.py | 36 ++++++++++++++++++------------------
 1 file changed, 18 insertions(+), 18 deletions(-)

(limited to 'main.py')

diff --git a/main.py b/main.py
index 46c8a50..1c05580 100644
--- a/main.py
+++ b/main.py
@@ -27,9 +27,9 @@ def main():
         formatter_class=argparse.RawDescriptionHelpFormatter,
         epilog=__doc__
     )
-    
+
     subparsers = parser.add_subparsers(dest="command", help="Available commands")
-    
+
     # Train diffusion model
     train_parser = subparsers.add_parser("train", help="Train diffusion model")
     train_parser.add_argument("--epochs", type=int, default=10, help="Number of epochs")
@@ -38,7 +38,7 @@ def main():
     train_parser.add_argument("--attention", action="store_true", help="Use self-attention")
     train_parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint")
     train_parser.add_argument("--name", type=str, default=None, help="Run name")
-    
+
     # Train autoencoder
     ae_parser = subparsers.add_parser("train-ae", help="Train autoencoder")
     ae_parser.add_argument("--epochs", type=int, default=10, help="Number of epochs")
@@ -46,14 +46,14 @@ def main():
     ae_parser.add_argument("--batch-size", type=int, default=64, help="Batch size")
     ae_parser.add_argument("--latent-channels", type=int, default=1, help="Latent channels")
     ae_parser.add_argument("--checkpoint", type=str, default=None, help="Resume from checkpoint")
-    
+
     # Sample from model
     sample_parser = subparsers.add_parser("sample", help="Generate samples")
     sample_parser.add_argument("--checkpoint", type=str, default="checkpoints/diffusion_latest.pt",
                                help="Path to model checkpoint")
     sample_parser.add_argument("--n-samples", type=int, default=10, help="Samples per class")
     sample_parser.add_argument("--attention", action="store_true", help="Use attention in model")
-    
+
     # Visualize diffusion
     viz_parser = subparsers.add_parser("visualize", help="Visualize diffusion process")
     viz_parser.add_argument("--checkpoint", type=str, default="checkpoints/diffusion_latest.pt",
@@ -62,16 +62,16 @@ def main():
     viz_parser.add_argument("--forward", action="store_true", help="Visualize forward diffusion")
     viz_parser.add_argument("--backward", action="store_true", help="Visualize backward diffusion")
     viz_parser.add_argument("--all", action="store_true", help="Run all visualizations")
-    
+
     args = parser.parse_args()
-    
+
     if args.command is None:
         parser.print_help()
         return
-    
+
     # Set multiprocessing start method
     torch.multiprocessing.set_start_method("spawn", force=True)
-    
+
     if args.command == "train":
         from src.train_diffusion import train_diffusion
         train_diffusion(
@@ -82,7 +82,7 @@ def main():
             checkpoint_path=args.checkpoint,
             run_name=args.name,
         )
-    
+
     elif args.command == "train-ae":
         from src.train_autoencoder import train_autoencoder
         train_autoencoder(
@@ -92,22 +92,22 @@ def main():
             latent_channels=args.latent_channels,
             checkpoint_path=args.checkpoint,
         )
-    
+
     elif args.command == "sample":
         import os
         from src.config import DEVICE
         from src.sample import generate_grid
         from src.utils import load_checkpoint
         from models import UNetMNIST
-        
+
         model = UNetMNIST(use_attention=args.attention).to(DEVICE)
         if os.path.exists(args.checkpoint):
             model = load_checkpoint(model, os.path.basename(args.checkpoint))
         else:
             print(f"Warning: Checkpoint {args.checkpoint} not found.")
-        
+
         generate_grid(model, n_per_class=args.n_samples)
-    
+
     elif args.command == "visualize":
         import os
         from src.config import DEVICE
@@ -118,17 +118,17 @@ def main():
         )
         from src.utils import load_checkpoint
         from models import UNetMNIST
-        
+
         model = UNetMNIST(use_attention=args.attention).to(DEVICE)
         if os.path.exists(args.checkpoint):
             model = load_checkpoint(model, os.path.basename(args.checkpoint))
-        
+
         if args.forward or args.all:
             visualize_forward_diffusion()
-        
+
         if args.backward or args.all:
             visualize_backward_diffusion(model)
-        
+
         if args.all:
             generate_grid(model)
 
-- 
cgit v1.3.1