Upload 2 files

app.py

@@ -4,6 +4,9 @@ from PIL import Image
 import requests
 import traceback
 
+from huggingface_hub import login
+login(token=HF_TOKEN)
+
 class Image2Text:
     def __init__(self):
         # Load the GIT coco model
@@ -188,7 +191,7 @@ class Social_Media_Captioner:
             self.model_loaded = False
 
 
-    def inference(self, input_text: str, use_cached=True, cache_generation=True):
+    def inference(self, input_text: str, use_cached=True, cache_generation=True) -> str | None:
         if not self.model_loaded:
             raise Exception("Model not loaded")
 
@@ -230,7 +233,7 @@ class Social_Media_Captioner:
             raise Exception("Enter a valid input text to generate a valid prompt")
 
         return f"""
-            Convert the given image description to social media worthy caption
+            Convert the given image description to a appropriate metaphoric caption
             Description: {input_text}
             Caption:
             """.strip()
@@ -279,7 +282,7 @@ class Captions:
             image_description = self.image_to_text._generate_description(image, max_length=max_length_GIT)
         captions = self.LLM.inference(image_description, use_cached=use_cached_LLM, cache_generation=cache_generation_LLM)
         return captions
-    
+
 caption_generator = Captions()
 
 import gradio as gr
@@ -293,4 +296,4 @@ iface = gr.Interface(
     outputs=gr.outputs.Textbox(label="Caption")
 )
 
-iface.launch()
+iface.launch(share=True)

requirements.txt

@@ -1,7 +1,7 @@
 evaluate
 jiwer
 huggingface_hub @ git+https://github.com/huggingface/huggingface_hub.git
-gradio==3.43.1
+gradio==3.36.0
 bitsandbytes
 transformers @ git+https://github.com/huggingface/transformers.git
 peft @ git+https://github.com/huggingface/peft.git
@@ -11,3 +11,4 @@ safetensors
 torch
 xformers
 datasets
+stable-diffusion