update app.py

app.py

@@ -38,7 +38,36 @@ if "vision_encoder.logit_scale"in model_state_dict:
     del model_state_dict["vision_encoder.visual.ln_post.weight"]
     del model_state_dict["vision_encoder.visual.ln_post.bias"]
 flamingo.load_state_dict(model_state_dict, strict=True)
+def get_outputs(
+    model,
+    batch_images,
+    attention_mask,
+    max_generation_length,
+    min_generation_length,
+    num_beams,
+    length_penalty,
+    input_ids,
+    image_start_index_list=None,
+    image_nums=None,
+    bad_words_ids=None,
+):
+    #  and torch.cuda.amp.autocast(dtype=torch.float16)
+    with torch.inference_mode():
+        outputs = model.generate(
+            batch_images,
+            input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=max_generation_length,
+            min_length=min_generation_length,
+            num_beams=num_beams,
+            length_penalty=length_penalty,
+            image_start_index_list=image_start_index_list,
+            image_nums=image_nums,
+            bad_words_ids=bad_words_ids,
+        )
 
+    outputs = outputs[:, len(input_ids[0]) :]
+    return outputs
 
 def generate(
     idx,