bug fixes

app.py

@@ -3,9 +3,17 @@ import gradio as gr
 from prefix_clip import download_pretrained_model, generate_caption
 from gpt2_story_gen import generate_story
 
+coco_weights = 'coco_weights.pt'
+conceptual_weights = 'conceptual_weights.pt'
+download_pretrained_model('coco', file_to_save=coco_weights)
+download_pretrained_model('conceptual', file_to_save=conceptual_weights)
+
 
 def main(pil_image, genre, model, use_beam_search=False):
-    model_file = "pretrained_weights.pt"
+    if model.lower()=='coco':
+        model_file = coco_weights
+    elif model.lower()=='conceptual':
+        model_file = conceptual_weights
 
     download_pretrained_model(model.lower(), file_to_save=model_file)
 
@@ -14,7 +22,7 @@ def main(pil_image, genre, model, use_beam_search=False):
         pil_image=pil_image,
         use_beam_search=use_beam_search,
     )
-    story = generate_story(image_caption, image, genre.lower())
+    story = generate_story(image_caption, pil_image, genre.lower())
     return story