From 00d9363172e8c6affca051d517ed12a03ee3a92c Mon Sep 17 00:00:00 2001 From: Brett Kuprel Date: Fri, 15 Jul 2022 21:21:16 -0400 Subject: [PATCH] fixed bug on replicate --- min_dalle/models/dalle_bart_decoder.py | 8 ++++---- replicate_predictor.py | 3 ++- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/min_dalle/models/dalle_bart_decoder.py b/min_dalle/models/dalle_bart_decoder.py index f2ee357..9ab707c 100644 --- a/min_dalle/models/dalle_bart_decoder.py +++ b/min_dalle/models/dalle_bart_decoder.py @@ -161,16 +161,16 @@ class DalleBartDecoder(nn.Module): ) decoder_state = self.final_ln(decoder_state) logits = self.lm_head(decoder_state) - temperature = settings[0] - top_k = settings[1].to(torch.long) - supercondition_factor = settings[2] + temperature = settings[[0]] + top_k = settings[[1]].to(torch.long) + supercondition_factor = settings[[2]] logits = logits[:, -1, : 2 ** 14] logits: FloatTensor = ( logits[:image_count] * (1 - supercondition_factor) + logits[image_count:] * supercondition_factor ) logits_sorted, _ = logits.sort(descending=True) - is_kept = logits >= logits_sorted[:, top_k: top_k + 1] + is_kept = logits >= logits_sorted[:, top_k - 1] logits -= logits_sorted[:, [0]] logits /= temperature logits.exp_() diff --git a/replicate_predictor.py b/replicate_predictor.py index 8e702e4..3bac013 100644 --- a/replicate_predictor.py +++ b/replicate_predictor.py @@ -11,7 +11,8 @@ class ReplicatePredictor(BasePredictor): self.model = MinDalle( is_mega=True, is_reusable=True, - dtype=torch.float32 + dtype=torch.float32, + device='cuda' ) def predict(