Spaces:
Sleeping
Sleeping
File size: 2,994 Bytes
cad10d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """
train.py — Training entry point.
Usage:
# Phase 1: Train VAE
python scripts/train.py --phase vae
# Phase 2: Train diffusion U-Net
python scripts/train.py --phase diffusion --vae-checkpoint checkpoints/vae_final.pt
# Resume training
python scripts/train.py --phase diffusion --resume checkpoints/diffusion_step_50000.pt
"""
import argparse
import os
import sys
# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
from omegaconf import OmegaConf
from data.dataset import create_dataloader
from training.trainer import Trainer
def main():
parser = argparse.ArgumentParser(description="Train the image editing model")
parser.add_argument("--config", type=str, default="config/default.yaml", help="Config file path")
parser.add_argument("--phase", type=str, required=True, choices=["vae", "diffusion"], help="Training phase")
parser.add_argument("--vae-checkpoint", type=str, default=None, help="VAE checkpoint for diffusion phase")
parser.add_argument("--resume", type=str, default=None, help="Resume from checkpoint")
parser.add_argument("--max-samples", type=int, default=None, help="Limit dataset size (for testing)")
parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
args = parser.parse_args()
# Load config (default + optional override)
base_config = OmegaConf.load("config/default.yaml")
override_config = OmegaConf.load(args.config)
config = OmegaConf.to_container(OmegaConf.merge(base_config, override_config), resolve=True)
if args.max_samples:
config["data"]["max_samples"] = args.max_samples
print(f"Device: {args.device}")
print(f"Phase: {args.phase}")
print(f"Config: {args.config}")
if args.device == "cuda" and torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB")
# Create dataloader
dataloader = create_dataloader(config, is_train=True, phase=args.phase, split="train")
print(f"Dataset size: {len(dataloader.dataset)}")
# Create trainer
trainer = Trainer(config, device=args.device)
if args.phase == "vae":
trainer.train_vae(dataloader, resume_from=args.resume)
elif args.phase == "diffusion":
if not args.vae_checkpoint:
# Look for default VAE checkpoint
default_vae = os.path.join(config["paths"]["checkpoint_dir"], "vae_final.pt")
if os.path.exists(default_vae):
args.vae_checkpoint = default_vae
else:
print("ERROR: --vae-checkpoint required for diffusion phase")
print("Train the VAE first: python scripts/train.py --phase vae")
sys.exit(1)
trainer.train_diffusion(dataloader, args.vae_checkpoint, resume_from=args.resume)
if __name__ == "__main__":
main()
|