diff options
Diffstat (limited to 'main.py')
| -rw-r--r-- | main.py | 36 |
1 files changed, 18 insertions, 18 deletions
@@ -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) |