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

also consider accepting tokenizer, so dalle2 forward pass can just be invoked as DALLE2(<prompt string>)

Browse Source
This commit is contained in:
Phil Wang
2022-04-12 10:29:29 -07:00
parent de75a8af76
commit 771fe0d0d2
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
dalle2_pytorch/dalle2_pytorch.py
View File

@@ -180,12 +180,14 @@ class DALLE2(nn.Module):
*, *,
clip, clip,
prior, prior,
decoder decoder,
tokenizer = None
): ):
super().__init__() super().__init__()
assert isinstance(clip), CLIP assert isinstance(clip), CLIP
assert isinstance(prior), DiffusionPrior assert isinstance(prior), DiffusionPrior
assert isinstance(decoder), Decoder assert isinstance(decoder), Decoder
self.tokenizer = tokenizer
@torch.no_grad() @torch.no_grad()
def forward( def forward(