| """Main training entry point for LuminaRS.""" |
| import os, torch |
| from torch.utils.data import DataLoader |
| from torchvision import transforms |
| from transformers import CLIPTextModel, CLIPTokenizer |
| from diffusers import AutoencoderKL |
| from datasets import load_dataset |
| from luminars.config import LuminaRSConfig |
| from luminars.model import LuminaRS |
| from luminars.flow_loss import flow_loss |
| from luminars.train import setup_stage1, setup_stage2, setup_stage3, freeze, count_trainable |
|
|
| STYLE_LABELS = ["Abstract","Action Painting","Analytical Cubism","Art Nouveau","Baroque","Color Field Painting","Contemporary Realism","Cubism","Early Renaissance","Expressionism","Fauvism","High Renaissance","Impressionism","Mannerism","Minimalism","Naive Art","Neoclassicism","Northern Renaissance","Pointillism","Pop Art","Post-Impressionism","Realism","Rococo","Romanticism","Symbolism","Synthetic Cubism","Ukiyo-e"] |
|
|
| def get_dataloader(batch_size=4, num_workers=2): |
| ds = load_dataset("huggan/wikiart", split="train") |
| tfm = transforms.Compose([transforms.Resize(512), transforms.CenterCrop(512), transforms.ToTensor(), transforms.Normalize([.5]*3, [.5]*3)]) |
| class ArtDS(torch.utils.data.Dataset): |
| def __init__(self, ds, tfm): self.ds=ds; self.tfm=tfm |
| def __len__(self): return len(self.ds) |
| def __getitem__(self, i): |
| it = self.ds[i]; img = self.tfm(it["image"].convert("RGB")); cap = f"a painting in {STYLE_LABELS[it['style']]} style" |
| return img, cap |
| return DataLoader(ArtDS(ds,tfm), batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True, drop_last=True) |
|
|
| def train(stage=1, epochs=10, lr=1e-4, batch_size=4, save_dir="./checkpoints"): |
| os.makedirs(save_dir, exist_ok=True); cfg = LuminaRSConfig() |
| dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| vae = AutoencoderKL.from_pretrained(cfg.vae_pretrained).to(dev); freeze(vae) |
| tok = CLIPTokenizer.from_pretrained(cfg.clip_pretrained) |
| clip = CLIPTextModel.from_pretrained(cfg.clip_pretrained).to(dev); freeze(clip) |
| model = LuminaRS(cfg).to(dev) |
| if stage==1: setup_stage1(model) |
| elif stage==2: setup_stage2(model); lr*=0.1 |
| elif stage==3: setup_stage3(model); lr*=0.01 |
| print(f"Stage {stage}: {count_trainable(model)/1e6:.1f}M train, lr={lr:.1e}") |
| opt = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=lr, weight_decay=0.01) |
| dl = get_dataloader(batch_size); scaler = torch.amp.GradScaler("cuda"); model.train() |
| for ep in range(epochs): |
| for step, (imgs, caps) in enumerate(dl): |
| imgs=imgs.to(dev); tks=tok(caps,padding="max_length",max_length=77,truncation=True,return_tensors="pt").input_ids.to(dev) |
| with torch.no_grad(): z=vae.encode(imgs).latent_dist.sample()*vae.config.scaling_factor; te=clip(tks).last_hidden_state |
| with torch.amp.autocast("cuda"): loss=flow_loss(model,z,te) |
| scaler.scale(loss).backward(); scaler.unscale_(opt) |
| torch.nn.utils.clip_grad_norm_(model.parameters(),1.0) |
| scaler.step(opt); scaler.update(); opt.zero_grad() |
| if step%50==0: print(f"[S{stage} E{ep} {step}] loss={loss.item():.4f}") |
| torch.save(model.state_dict(), f"{save_dir}/s{stage}_e{ep}.pt") |
| return model |
|
|
| if __name__ == "__main__": train(stage=1, epochs=5, lr=1e-4) |
|
|
