aljaz/DALLE2-pytorch
1
1
Fork 0
mirror of https://github.com/lucidrains/DALLE2-pytorch.git synced 2025-12-19 17:54:20 +01:00
Code Issues Packages Projects Releases Wiki Activity

per-fect

Browse Source
This commit is contained in:
Phil Wang
2022-04-14 08:30:07 -07:00
parent 7fb3f695d5
commit 82464d7bd3
1 changed files with 4 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
dalle2_pytorch/dalle2_pytorch.py
View File

@@ -285,17 +285,16 @@ class DiffusionPriorNetwork(nn.Module):
def forward_with_cond_scale( def forward_with_cond_scale(
self, self,
x,
*args, *args,
cond_scale = 1., cond_scale = 1.,
**kwargs **kwargs
): ):
logits = self.forward(x, *args, **kwargs) logits = self.forward(*args, **kwargs)
if cond_scale == 1: if cond_scale == 1:
return logits return logits
null_logits = self.forward(x, *args, cond_drop_prob = 1., **kwargs) null_logits = self.forward(*args, cond_drop_prob = 1., **kwargs)
return null_logits + (logits - null_logits) * cond_scale return null_logits + (logits - null_logits) * cond_scale
def forward( def forward(
@@ -675,17 +674,16 @@ class Unet(nn.Module):
def forward_with_cond_scale( def forward_with_cond_scale(
self, self,
x,
*args, *args,
cond_scale = 1., cond_scale = 1.,
**kwargs **kwargs
): ):
logits = self.forward(x, *args, **kwargs) logits = self.forward(*args, **kwargs)
if cond_scale == 1: if cond_scale == 1:
return logits return logits
null_logits = self.forward(x, *args, cond_drop_prob = 1., **kwargs) null_logits = self.forward(*args, cond_drop_prob = 1., **kwargs)
return null_logits + (logits - null_logits) * cond_scale return null_logits + (logits - null_logits) * cond_scale
def forward( def forward(