add mask

app.py

@@ -63,6 +63,10 @@ def caption_next_word(latent, model, tokenizer, prefix='a photo of a'):
         x_l = cross_attn2(x_l, latent)
 
     pred = model.to_logits(x_l)
+    pred[:, :, 103] = -100
+    pred[:, :, 101] = -100
+    pred[:, :, 100] = -100
+    pred[:, :, 0] = -100
     next_word = pred.argmax(dim=-1)[0, -1]
     next_word = tokenizer.decode(next_word)
     
@@ -73,6 +77,8 @@ def caption(max_len, latent, model, tokenizer, prefix='a photo of a'):
     while len(words) < max_len:
         next_word = caption_next_word(latent, model, tokenizer, prefix=' '.join(words))
         words.append(next_word)
+        if next_word == '[SEP]':
+            break
     return ' '.join(words)