24fromtypingimportList2526importtorch27importtorch.nnasnn2829fromlabml\_nn.diffusion.stable\_diffusion.model.autoencoderimportAutoencoder30fromlabml\_nn.diffusion.stable\_diffusion.model.clip\_embedderimportCLIPTextEmbedder31fromlabml\_nn.diffusion.stable\_diffusion.model.unetimportUNetModel

34classDiffusionWrapper(nn.Module):

42def\_\_init\_\_(self,diffusion\_model:UNetModel):43super().\_\_init\_\_()44self.diffusion\_model=diffusion\_model

46defforward(self,x:torch.Tensor,time\_steps:torch.Tensor,context:torch.Tensor):47returnself.diffusion\_model(x,time\_steps,context)

50classLatentDiffusion(nn.Module):

60model:DiffusionWrapper61first\_stage\_model:Autoencoder62cond\_stage\_model:CLIPTextEmbedder

64def\_\_init\_\_(self,65unet\_model:UNetModel,66autoencoder:Autoencoder,67clip\_embedder:CLIPTextEmbedder,68latent\_scaling\_factor:float,69n\_steps:int,70linear\_start:float,71linear\_end:float,72):

84super().\_\_init\_\_()

87self.model=DiffusionWrapper(unet\_model)

89self.first\_stage\_model=autoencoder90self.latent\_scaling\_factor=latent\_scaling\_factor

92self.cond\_stage\_model=clip\_embedder

95self.n\_steps=n\_steps

98beta=torch.linspace(linear\_start\*\*0.5,linear\_end\*\*0.5,n\_steps,dtype=torch.float64)\*\*299self.beta=nn.Parameter(beta.to(torch.float32),requires\_grad=False)

101alpha=1.-beta

103alpha\_bar=torch.cumprod(alpha,dim=0)104self.alpha\_bar=nn.Parameter(alpha\_bar.to(torch.float32),requires\_grad=False)

106@property107defdevice(self):

111returnnext(iter(self.model.parameters())).device

113defget\_text\_conditioning(self,prompts:List[str]):

117returnself.cond\_stage\_model(prompts)

119defautoencoder\_encode(self,image:torch.Tensor):

126returnself.latent\_scaling\_factor\*self.first\_stage\_model.encode(image).sample()

128defautoencoder\_decode(self,z:torch.Tensor):

134returnself.first\_stage\_model.decode(z/self.latent\_scaling\_factor)

136defforward(self,x:torch.Tensor,t:torch.Tensor,context:torch.Tensor):

145returnself.model(x,t,context)