Add model options

app.py

@@ -26,13 +26,12 @@ Caption:
 """)
 
 hf_api_key = os.environ["HF_API_KEY"]
-prompt_node = PromptNode(model_name_or_path="tiiuae/falcon-7b-instruct", api_key=hf_api_key, default_prompt_template=prompt_template, model_kwargs={"trust_remote_code":True})
-
 captioning_pipeline = Pipeline()
-captioning_pipeline.add_node(component=image_to_text, name="image_to_text", inputs=["File"])
-captioning_pipeline.add_node(component=prompt_node, name="prompt_node", inputs=["image_to_text"])
 
-def generate_caption(image_file_paths):
+def generate_caption(image_file_paths, model_name):
+    prompt_node = PromptNode(model_name_or_path=model_name, api_key=hf_api_key, default_prompt_template=prompt_template, model_kwargs={"trust_remote_code":True})
+    captioning_pipeline.add_node(component=image_to_text, name="image_to_text", inputs=["File"])
+    captioning_pipeline.add_node(component=prompt_node, name="prompt_node", inputs=["image_to_text"])
     caption = captioning_pipeline.run(file_paths=[image_file_paths])
     print(caption)
     return caption["results"][0]
@@ -40,10 +39,10 @@ def generate_caption(image_file_paths):
 with gr.Blocks(theme="soft") as demo:
     gr.Markdown(value=description)
     image = gr.Image(type="filepath")
-    # model_name = gr.Dropdown(["tiiuae/falcon-7b-instruct", "gpt-3.5-turbo", "gpt-4", "gpt-4-32k", "command", "command-light", "base", "base-light"], value="tiiuae/falcon-7b-instruct", label="Choose your model!")
+    model_name = gr.Dropdown(["tiiuae/falcon-7b-instruct", "tiiuae/falcon-7b", "EleutherAI/gpt-neox-20b", "HuggingFaceH4/starchat-beta", "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5", "timdettmers/guanaco-33b-merged", "bigscience/bloom"], value="tiiuae/falcon-7b-instruct", label="Choose your model!")
     submit_btn = gr.Button("✨ Captionate ✨")
     caption = gr.Textbox(label="Caption")
-    submit_btn.click(fn=generate_caption, inputs=[image], outputs=[caption])
+    submit_btn.click(fn=generate_caption, inputs=[image, model_name], outputs=[caption])
 
 if __name__ == "__main__":
     demo.launch()