added more than single story option

app.py

@@ -9,7 +9,7 @@ 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):
+def main(pil_image, genre, model, n_stories, use_beam_search=False):
     if model.lower()=='coco':
         model_file = coco_weights
     elif model.lower()=='conceptual':
@@ -20,7 +20,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, pil_image, genre.lower())
+    story = generate_story(image_caption, pil_image, genre.lower(), n_stories)
     return story
 
 
@@ -48,7 +48,8 @@ if __name__ == "__main__":
                     "sci_fi",
                 ],
             ),
-            gr.inputs.Radio(choices=["coco", "conceptual"], label="Model")
+            gr.inputs.Radio(choices=["coco", "conceptual"], label="Model"),
+            gr.inputs.Dropdown(choices=[1, 2, 3], label="No. of stories", type="value"),
         ],
         outputs=gr.outputs.Textbox(label="Generated story"),
         examples=[["car.jpg", "drama", "conceptual"], ["gangster.jpg", "action", "coco"]],

gpt2_story_gen.py

@@ -1,11 +1,7 @@
-from transformers import pipeline, CLIPProcessor, CLIPModel
-import torch
+from transformers import pipeline
 
 
-def generate_story(image_caption, image, genre):
-    clip_ranker_checkpoint = "openai/clip-vit-base-patch32"
-    clip_ranker_processor = CLIPProcessor.from_pretrained(clip_ranker_checkpoint)
-    clip_ranker_model = CLIPModel.from_pretrained(clip_ranker_checkpoint)
+def generate_story(image_caption, image, genre, n_stories):
 
     story_gen = pipeline(
         "text-generation", 
@@ -13,17 +9,6 @@ def generate_story(image_caption, image, genre):
         )
     
     input = f"<BOS> <{genre}> {image_caption}"
-    stories = [story_gen(input)[0]['generated_text'].strip(input) for i in range(3)]
-    clip_ranker_inputs = clip_ranker_processor(
-        text=stories, 
-        images=image, 
-        truncation=True,
-        return_tensors='pt', 
-        padding=True
-        )
-    clip_ranker_outputs = clip_ranker_model(**clip_ranker_inputs)
-    logits_per_image = clip_ranker_outputs.logits_per_image
-    probs = logits_per_image.softmax(dim=1)
-    story = stories[torch.argmax(probs).item()]
+    stories = '\n'.join([f"Story {i+1}\n{story_gen(input)[0]['generated_text'].strip(input)}" for i in range(n_stories)])
 
-    return story
+    return stories