“update”

app.py

@@ -92,6 +92,9 @@ def generate(
     all_ids = set(range(flamingo.lang_encoder.lm_head.out_features))
     bad_words_ids = list(all_ids - set(loc_token_ids))
     bad_words_ids = [[b] for b in bad_words_ids]
+    loc_word_ids = list(set(loc_token_ids))
+    loc_word_ids = [[b] for b in loc_word_ids]
+
     min_loc_token_id = min(loc_token_ids)
     max_loc_token_id = max(loc_token_ids)
     image_ori = image
@@ -103,9 +106,11 @@ def generate(
     if idx == 1:
         prompt = [f"<|#image#|>{tokenizer.pad_token*vis_embed_size}<|#endofimage#|><|#obj#|>{text.rstrip('.')}<|#loc#|>"]
         bad_words_ids = None
+        max_generation_length = 5
     else:
         prompt = [f"<|#image#|>{tokenizer.pad_token * vis_embed_size}<|#endofimage#|>{text.rstrip('.')}"]
-        bad_words_ids = None
+        bad_words_ids = loc_word_ids
+        max_generation_length = 100
     encodings = tokenizer(
         prompt,
         padding="longest",
@@ -122,7 +127,7 @@ def generate(
         model=flamingo,
         batch_images=batch_images,
         attention_mask=attention_mask,
-        max_generation_length=5,
+        max_generation_length=max_generation_length,
         min_generation_length=4,
         num_beams=1,
         length_penalty=1.0,