"""One controlled trainer: diffusion, flow, and two equally budgeted flow follow-ups.""" import argparse import copy import hashlib import json from pathlib import Path import time import numpy as np import torch from torch import nn from torch.nn import functional as F from digit_model import DigitUNet from quality import dataset from .processes import cosine_schedule,generate,training_pair def file_hash(path):return hashlib.sha256(Path(path).read_bytes()).hexdigest() def state_hash(model): digest=hashlib.sha256() for key,value in sorted(model.state_dict().items()): digest.update(key.encode());digest.update(value.detach().cpu().numpy().tobytes()) return digest.hexdigest() def write_config(path,config):path.write_text(json.dumps(config,indent=2)) def main(): parser=argparse.ArgumentParser() parser.add_argument('--stage',choices=['diffusion','flow','pairs','reflow','continued'],required=True) parser.add_argument('--root',type=Path,default=Path('results/unified')) parser.add_argument('--seed',type=int,default=42) parser.add_argument('--steps',type=int,default=20000) parser.add_argument('--batch-size',type=int,default=128) parser.add_argument('--pairs',type=int,default=20000) parser.add_argument('--teacher-nfe',type=int,default=128) parser.add_argument('--resume',action='store_true') args=parser.parse_args() torch.set_num_threads(4) if not torch.cuda.is_available():raise RuntimeError('Use a GPU allocation; training on login nodes is disabled.') device=torch.device('cuda') torch.manual_seed(args.seed) np.random.seed(args.seed) model=DigitUNet(num_classes=10).to(device) out=args.root/args.stage;out.mkdir(parents=True,exist_ok=True) teacher_path=args.root/'flow/checkpoint.pt' if args.stage=='pairs': checkpoint=torch.load(teacher_path,weights_only=True,map_location=device) if checkpoint['config']['steps'] != checkpoint['config']['planned_steps']: raise RuntimeError('Finish the base-flow training plan before generating teacher pairs.') model.load_state_dict(checkpoint['ema']);model.eval().requires_grad_(False) rng=torch.Generator(device=device).manual_seed(1234) starts=[];ends=[];ys=[] started=time.perf_counter() for offset in range(0,args.pairs,args.batch_size): size=min(args.batch_size,args.pairs-offset) z0=torch.randn(size,1,28,28,device=device,generator=rng) labels=torch.arange(offset,offset+size,device=device)%10 z1=generate(model,z0,labels,'fm',args.teacher_nfe) starts.append(z0.cpu());ends.append(z1.cpu());ys.append(labels.cpu()) if offset%(args.batch_size*10)==0:print('teacher pairs',offset,args.pairs,flush=True) torch.cuda.synchronize() elapsed=time.perf_counter()-started torch.save({'start':torch.cat(starts),'end':torch.cat(ends),'labels':torch.cat(ys)},out/'pairs.pt') write_config(out/'metadata.json',{'pairs':args.pairs,'teacher_nfe':args.teacher_nfe,'seconds':elapsed, 'teacher_checkpoint':str(teacher_path),'teacher_sha256':file_hash(teacher_path),'seed':1234, 'label_counts':torch.cat(ys).bincount(minlength=10).tolist()}) print('Pair bank completed',elapsed,flush=True);return extra=args.stage in ['reflow','continued'] teacher_sha=None if extra: checkpoint=torch.load(teacher_path,weights_only=True,map_location=device) model.load_state_dict(checkpoint['ema']) teacher_sha=file_hash(teacher_path) initial_state_hash=state_hash(model) ema=copy.deepcopy(model).eval().requires_grad_(False) optimizer=torch.optim.AdamW(model.parameters(),lr=1e-4 if extra else 2e-4,weight_decay=1e-4) data_rng=torch.Generator(device=device).manual_seed(144 if extra else 100) noise_rng=torch.Generator(device=device).manual_seed(244 if extra else 200) time_rng=torch.Generator(device=device).manual_seed(344 if extra else 300) if args.stage=='reflow': bank=torch.load(args.root/'pairs/pairs.pt',weights_only=True) starts,ends,bank_y=[bank[key].to(device) for key in ['start','end','labels']] pair_meta=json.loads((args.root/'pairs/metadata.json').read_text()) if pair_meta['teacher_sha256']!=teacher_sha:raise RuntimeError('Pair bank was generated by a different teacher.') classes=[torch.where(bank_y==label)[0] for label in range(10)] lengths=torch.tensor([len(c) for c in classes],device=device) lookup=torch.zeros(10,int(lengths.max()),device=device,dtype=torch.long) for label,indices in enumerate(classes):lookup[label,:len(indices)]=indices else: images,labels=dataset('train');images,labels=images.to(device),labels.to(device) classes=[torch.where(labels==label)[0] for label in range(10)] lengths=torch.tensor([len(c) for c in classes],device=device) lookup=torch.zeros(10,int(lengths.max()),device=device,dtype=torch.long) for label,indices in enumerate(classes):lookup[label,:len(indices)]=indices abar=cosine_schedule(1000,device) losses=[];start_step=0;previous_seconds=0 class_counts=torch.zeros(10,device=device,dtype=torch.long) if args.resume: checkpoint=torch.load(out/'checkpoint.pt',weights_only=True,map_location=device) if checkpoint['config']['planned_steps']!=args.steps:raise ValueError('Keep planned --steps unchanged when resuming.') model.load_state_dict(checkpoint['model']);ema.load_state_dict(checkpoint['ema']) optimizer.load_state_dict(checkpoint['optimizer']) start_step=checkpoint['config']['steps'];previous_seconds=checkpoint['config']['training_seconds'] for rng,key in [(data_rng,'data_rng'),(noise_rng,'noise_rng'),(time_rng,'time_rng')]:rng.set_state(checkpoint[key].cpu()) torch.set_rng_state(checkpoint['rng_state'].cpu());torch.cuda.set_rng_state(checkpoint['cuda_rng_state'].cpu()) losses=np.load(out/'loss.npy').tolist() class_counts=torch.tensor(checkpoint['config']['class_presentations'],device=device) started=time.perf_counter() def save(step): torch.cuda.synchronize() config={'stage':args.stage,'seed':args.seed,'steps':step,'planned_steps':args.steps,'batch_size':args.batch_size, 'parameters':sum(p.numel() for p in model.parameters()),'architecture':'DigitUNet + 10-class embedding', 'image_shape':[1,28,28],'normalization':'[-1,1]','conditioning':'class labels, no CFG or label dropout', 'class_sampling':'uniform over 10 labels','optimizer':'AdamW','weight_decay':1e-4,'gradient_clip':1, 'learning_rate':1e-4 if extra else 2e-4,'lr_schedule':'constant' if extra else '1000 warmup + cosine to 2e-5', 'ema_max_decay':.999,'training_seconds':previous_seconds+time.perf_counter()-started, 'last_200_loss':float(np.mean(losses[-200:])),'hardware':torch.cuda.get_device_name(), 'torch_version':str(torch.__version__),'precision':'FP32','teacher_sha256':teacher_sha, 'initialization':'same teacher EMA' if extra else 'same seeded initialization', 'initial_weights_sha256':initial_state_hash,'class_presentations':class_counts.tolist(), 'dataset_presentations':step*args.batch_size,'source_protocol':'unified/protocol.json'} checkpoint={'model':model.state_dict(),'ema':ema.state_dict(),'optimizer':optimizer.state_dict(),'config':config, 'data_rng':data_rng.get_state(),'noise_rng':noise_rng.get_state(),'time_rng':time_rng.get_state(), 'rng_state':torch.get_rng_state(),'cuda_rng_state':torch.cuda.get_rng_state()} temp=out/'checkpoint.new.pt';torch.save(checkpoint,temp);temp.replace(out/'checkpoint.pt') write_config(out/'metadata.json',config);np.save(out/'loss.npy',np.asarray(losses)) return config for step in range(start_step+1,args.steps+1): if not extra: rate=min(step/1000,1.)*(.1+.9*.5*(1+np.cos(np.pi*max(step-1000,0)/max(args.steps-1000,1)))) for group in optimizer.param_groups:group['lr']=float(2e-4*rate) y=torch.randint(10,(args.batch_size,),device=device,generator=data_rng) u=torch.rand(args.batch_size,device=device,generator=data_rng) indices=lookup[y,(u*lengths[y]).long()] if args.stage=='reflow': data,noise=ends[indices],starts[indices] assert torch.equal(bank_y[indices],y) else: data=images[indices] noise=torch.randn(data.shape,device=device,generator=noise_rng) class_counts+=y.bincount(minlength=10) t=torch.rand(args.batch_size,device=device,generator=time_rng) xt,target,time_input=training_pair(data,noise,t,'diffusion' if args.stage=='diffusion' else 'flow',abar) loss=F.mse_loss(model(xt,time_input,y),target) optimizer.zero_grad(set_to_none=True);loss.backward();nn.utils.clip_grad_norm_(model.parameters(),1.);optimizer.step() decay=.999 if extra else min(.999,(1+step)/(10+step)) with torch.no_grad(): for average,current in zip(ema.parameters(),model.parameters()):average.lerp_(current,1-decay) losses.append(float(loss.detach())) if step%500==0:print(args.stage,step,'loss',np.mean(losses[-200:]),'seconds',time.perf_counter()-started,flush=True) if step%2000==0:save(step) print(json.dumps(save(args.steps),indent=2),flush=True) if __name__=='__main__':main()