added clip ranking

app.py

@@ -14,7 +14,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, genre.lower())
+    story = generate_story(image_caption, image, genre.lower())
     return story
 
 

gpt2_story_gen.py

@@ -1,11 +1,27 @@
-from transformers import pipeline
+from transformers import pipeline, CLIPProcessor, CLIPModel
 
 
-def generate_story(image_caption, genre):
-    story_gen = pipeline("text-generation", "pranavpsv/genre-story-generator-v2")
+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)
 
+    story_gen = pipeline(
+        "text-generation", 
+        "pranavpsv/genre-story-generator-v2"
+        )
+    
     input = f"<BOS> <{genre}> {image_caption}"
-    story = story_gen(input)[0]["generated_text"]
-    story = f"{story.strip(input)}"
+    stories = [story_gen(input)[0]['generated_text'].strip(input) for i in range(3)]
+    clip_ranker_inputs = clip_ranker_processor(
+        text=stories, 
+        images=image, 
+        return_tensors='pt', 
+        padding=True
+        )
+    clip_ranker_outputs = clip_ranker_model(**clip_ranker_inputs)
+    logits_per_image = outputs.logits_per_image
+    probs = logits_per_image.softmax(dim=1)
+    story = stories[torch.argmax(probs).item()]
 
     return story