20fromtypingimportList2122importtorch23importtorch.utils.data24importtorchvision25fromPILimportImage2627fromlabmlimportlab,tracker,experiment,monit28fromlabml.configsimportBaseConfigs,option29fromlabml\_helpers.deviceimportDeviceConfigs30fromlabml\_nn.diffusion.ddpmimportDenoiseDiffusion31fromlabml\_nn.diffusion.ddpm.unetimportUNet

34classConfigs(BaseConfigs):

41device:torch.device=DeviceConfigs()

44eps\_model:UNet

46diffusion:DenoiseDiffusion

49image\_channels:int=3

51image\_size:int=32

53n\_channels:int=64

56channel\_multipliers:List[int]=[1,2,2,4]

58is\_attention:List[int]=[False,False,False,True]

61n\_steps:int=1\_000

63batch\_size:int=64

65n\_samples:int=16

67learning\_rate:float=2e-5

70epochs:int=1\_000

73dataset:torch.utils.data.Dataset

75data\_loader:torch.utils.data.DataLoader

78optimizer:torch.optim.Adam

80definit(self):

82self.eps\_model=UNet(83image\_channels=self.image\_channels,84n\_channels=self.n\_channels,85ch\_mults=self.channel\_multipliers,86is\_attn=self.is\_attention,87).to(self.device)

90self.diffusion=DenoiseDiffusion(91eps\_model=self.eps\_model,92n\_steps=self.n\_steps,93device=self.device,94)

97self.data\_loader=torch.utils.data.DataLoader(self.dataset,self.batch\_size,shuffle=True,pin\_memory=True)

99self.optimizer=torch.optim.Adam(self.eps\_model.parameters(),lr=self.learning\_rate)

102tracker.set\_image("sample",True)

104defsample(self):

108withtorch.no\_grad():

110x=torch.randn([self.n\_samples,self.image\_channels,self.image\_size,self.image\_size],111device=self.device)

114fort\_inmonit.iterate('Sample',self.n\_steps):

116t=self.n\_steps-t\_-1

118x=self.diffusion.p\_sample(x,x.new\_full((self.n\_samples,),t,dtype=torch.long))

121tracker.save('sample',x)

123deftrain(self):

129fordatainmonit.iterate('Train',self.data\_loader):

131tracker.add\_global\_step()

133data=data.to(self.device)

136self.optimizer.zero\_grad()

138loss=self.diffusion.loss(data)

140loss.backward()

142self.optimizer.step()

144tracker.save('loss',loss)

146defrun(self):

150for\_inmonit.loop(self.epochs):

152self.train()

154self.sample()

156tracker.new\_line()

158experiment.save\_checkpoint()

161classCelebADataset(torch.utils.data.Dataset):

166def\_\_init\_\_(self,image\_size:int):167super().\_\_init\_\_()

170folder=lab.get\_data\_path()/'celebA'

172self.\_files=[pforpinfolder.glob(f'\*\*/\*.jpg')]

175self.\_transform=torchvision.transforms.Compose([176torchvision.transforms.Resize(image\_size),177torchvision.transforms.ToTensor(),178])

180def\_\_len\_\_(self):

184returnlen(self.\_files)

186def\_\_getitem\_\_(self,index:int):

190img=Image.open(self.\_files[index])191returnself.\_transform(img)

194@option(Configs.dataset,'CelebA')195defceleb\_dataset(c:Configs):

199returnCelebADataset(c.image\_size)

202classMNISTDataset(torchvision.datasets.MNIST):

207def\_\_init\_\_(self,image\_size):208transform=torchvision.transforms.Compose([209torchvision.transforms.Resize(image\_size),210torchvision.transforms.ToTensor(),211])212213super().\_\_init\_\_(str(lab.get\_data\_path()),train=True,download=True,transform=transform)

215def\_\_getitem\_\_(self,item):216returnsuper().\_\_getitem\_\_(item)[0]

219@option(Configs.dataset,'MNIST')220defmnist\_dataset(c:Configs):

224returnMNISTDataset(c.image\_size)

227defmain():

229experiment.create(name='diffuse',writers={'screen','labml'})

232configs=Configs()

235experiment.configs(configs,{236'dataset':'CelebA',# 'MNIST'237'image\_channels':3,# 1,238'epochs':100,# 5,239})

242configs.init()

245experiment.add\_pytorch\_models({'eps\_model':configs.eps\_model})

248withexperiment.start():249configs.run()

253if\_\_name\_\_=='\_\_main\_\_':254main()