|
|
| import torch, torch.nn as nn, numpy as np, json, math, os, argparse |
| torch.set_default_device('cuda') |
| LAT=64 |
| class Dec(nn.Module): |
| def __init__(s): |
| super().__init__(); s.net=nn.Sequential(nn.Conv2d(LAT,192,3,1,1),nn.GELU(),nn.ConvTranspose2d(192,96,4,2,1),nn.GELU(),nn.ConvTranspose2d(96,48,4,2,1),nn.GELU(),nn.ConvTranspose2d(48,3,4,2,1)) |
| def forward(s,z): return s.net(z) |
| class TagEnc(nn.Module): |
| def __init__(s, vocab, emb=48, d=256): |
| super().__init__(); s.emb=nn.Embedding(vocab,emb); s.gru=nn.GRU(emb,d,batch_first=True) |
| def forward(s,seq,lens): |
| packed=torch.nn.utils.rnn.pack_padded_sequence(s.emb(seq).to('cuda'),lens.cpu(),batch_first=True,enforce_sorted=False) |
| return s.gru(packed)[1].squeeze(0) |
| def t_emb(t,dim): |
| half=dim//2; f=torch.exp(-math.log(10000)*torch.arange(half,device='cuda')/half); a=t[:,None].float()*f[None,:] |
| return torch.cat([torch.cos(a),torch.sin(a)],1) |
| class CDiff(nn.Module): |
| def __init__(s,cin=LAT,tagdim=256,base=256): |
| super().__init__(); s.temb=nn.Sequential(nn.Linear(64,base),nn.GELU(),nn.Linear(base,cin)); s.tagproj=nn.Linear(tagdim,cin) |
| s.d1=nn.Sequential(nn.Conv2d(cin,base,3,1,1),nn.GroupNorm(8,base),nn.GELU(),nn.Conv2d(base,base,3,1,1),nn.GroupNorm(8,base),nn.GELU()) |
| s.d2=nn.Sequential(nn.Conv2d(base,base,3,2,1),nn.GroupNorm(8,base),nn.GELU()); s.mid=nn.Sequential(nn.Conv2d(base,base,3,1,1),nn.GroupNorm(8,base),nn.GELU()) |
| s.u2=nn.Sequential(nn.ConvTranspose2d(base,base,3,2,1,1),nn.GroupNorm(8,base),nn.GELU()); s.outc=nn.Conv2d(base,cin,3,1,1) |
| def forward(s,x,t,tag): |
| te=s.temb(t_emb(t,64)).unsqueeze(-1).unsqueeze(-1); tg=s.tagproj(tag).unsqueeze(-1).unsqueeze(-1) |
| x=x+te+tg; h1=s.d1(x); h2=s.mid(s.d2(h1)); h=s.u2(h2)+h1; return s.outc(h) |
| def load_model(ckpt): |
| ck=torch.load(ckpt,weights_only=False,map_location='cuda'); tok2id=ck['tok2id'] |
| dec=Dec(); te=TagEnc(vocab=len(tok2id)+1); cd=CDiff(LAT) |
| dec.load_state_dict(ck['dec']); te.load_state_dict(ck['te']); cd.load_state_dict(ck['cd']) |
| dec.eval(); te.eval(); cd.eval() |
| pc_mean=ck.get('pc_mean',torch.zeros(1,64,1,1)).to('cuda'); pc_std=ck.get('pc_std',torch.ones(1,64,1,1)).to('cuda') |
| return dec,te,cd,tok2id,pc_mean,pc_std |
| def generate(prompt, dec,te,cd,tok2id,pc_mean,pc_std, n=3, steps=30, cfg=4.0, seed=0): |
| MAXL=40 |
| def encp(t): |
| toks=[tok2id[x.strip()] for x in t.split(',') if x.strip() in tok2id][:MAXL]; return toks+[0]*(MAXL-len(toks)) |
| torch.manual_seed(seed); dev='cuda' |
| TMAX=200; beta=torch.linspace(1e-4,0.02,TMAX,device=dev); alpha=1-beta; abar=torch.cumprod(alpha,0) |
| toks=torch.tensor([encp(prompt)]*n,device=dev); lens=torch.tensor([max(1,min(len([x for x in encp(prompt) if x!=0]),MAXL))]*n) |
| utok=torch.zeros_like(toks); ulen=torch.ones_like(lens) |
| st=torch.linspace(TMAX-1,0,steps).long(); zn=torch.randn(n,LAT,12,12,device=dev) |
| with torch.no_grad(): |
| c=te(toks,lens); ucond=te(utok,ulen) |
| for idx in range(len(st)): |
| t=st[idx]; tn=st[idx+1] if idx+1<len(st) else -1; tt=torch.tensor([t.item()]*n,device=dev) |
| ep_c=cd(zn,tt,c); ep_u=cd(zn,tt,ucond); ep=ep_u+cfg*(ep_c-ep_u) |
| ab_t=abar[t].view(-1,1,1,1) |
| if tn>=0: |
| ab_tn=abar[tn].view(-1,1,1,1); z0=(zn-torch.sqrt(1-ab_t)*ep)/torch.sqrt(ab_t) |
| zn=torch.sqrt(ab_tn)*z0+torch.sqrt(torch.clamp(1-ab_tn,min=0))*ep |
| else: zn=(zn-torch.sqrt(1-ab_t)*ep)/torch.sqrt(ab_t) |
| z=zn*pc_std+pc_mean; imgs=dec(z) |
| return imgs.clamp(0,1) |
| if __name__=='__main__': |
| ap=argparse.ArgumentParser(); ap.add_argument('--ckpt',required=True); ap.add_argument('--prompt',required=True); ap.add_argument('--out',default='out'); ap.add_argument('-n',type=int,default=3); ap.add_argument('--steps',type=int,default=30); ap.add_argument('--cfg',type=float,default=4.0) |
| a=ap.parse_args() |
| dec,te,cd,tok,pm,ps=load_model(a.ckpt) |
| imgs=generate(a.prompt,dec,te,cd,tok,pm,ps,n=a.n,steps=a.steps,cfg=a.cfg) |
| from PIL import Image |
| os.makedirs(a.out,exist_ok=True) |
| for i in range(a.n): |
| im=imgs[i].permute(1,2,0).cpu().numpy() |
| Image.fromarray((np.clip(im,0,1)*255).astype(np.uint8)).save(f"{a.out}/gen_{i}.png") |
| print(f"saved {a.n} images to {a.out}") |
|
|
