image2story / app.py
bipin's picture
Update app.py
837545d
import gradio as gr
from huggingface_hub import hf_hub_download
from prefix_clip import generate_caption
from gpt2_story_gen import generate_story
conceptual_weights = hf_hub_download(repo_id="akhaliq/CLIP-prefix-captioning-conceptual-weights", filename="conceptual_weights.pt")
coco_weights = hf_hub_download(repo_id="akhaliq/CLIP-prefix-captioning-COCO-weights", filename="coco_weights.pt")
def main(pil_image, genre, model, n_stories, use_beam_search=False):
if model.lower()=='coco':
model_file = coco_weights
elif model.lower()=='conceptual':
model_file = conceptual_weights
image_caption = generate_caption(
model_path=model_file,
pil_image=pil_image,
use_beam_search=use_beam_search,
)
story = generate_story(image_caption, pil_image, genre.lower(), n_stories)
return story
if __name__ == "__main__":
title = "Image to Story"
article = "Combines the power of [clip prefix captioning](https://github.com/rmokady/CLIP_prefix_caption) with [gpt2 story generator](https://huggingface.co/pranavpsv/genre-story-generator-v2) to create stories of different genres from image"
description = "Drop an image and generate stories of different genre based on that image"
interface = gr.Interface(
main,
title=title,
description=description,
article=article,
inputs=[
gr.inputs.Image(type="pil", source="upload", label="Input"),
gr.inputs.Dropdown(
type="value",
label="Story genre",
choices=[
"superhero",
"action",
"drama",
"horror",
"thriller",
"sci_fi",
],
),
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"]],
enable_queue=True,
)
interface.launch()