Update app.py

app.py

@@ -7,25 +7,8 @@ from PIL import Image
 import torch.nn as nn
 from transformers import Blip2Processor, Blip2ForConditionalGeneration, InstructBlipProcessor, InstructBlipForConditionalGeneration
 from my_model.object_detection import detect_and_draw_objects
+from my_model.captioner.image_captioning import get_caption
 
-def load_caption_model(blip2=False, instructblip=True):
-
-    if blip2:
-        processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b", load_in_8bit=True,torch_dtype=torch.float16)
-        model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", load_in_8bit=True,torch_dtype=torch.float16)
-        if torch.cuda.device_count() > 1:
-            model = nn.DataParallel(model)
-            model.to('cuda')
-        #model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16, device_map="auto")
-        
-    if instructblip:
-        model = InstructBlipForConditionalGeneration.from_pretrained("Salesforce/instructblip-vicuna-7b", load_in_8bit=True,torch_dtype=torch.float16)
-        if torch.cuda.device_count() > 1:
-            model = nn.DataParallel(model)
-            model.to('cuda')
-        processor = InstructBlipProcessor.from_pretrained("Salesforce/instructblip-vicuna-7b", load_in_8bit=True,torch_dtype=torch.float16)
-
-    return model, processor
 
     
 
@@ -54,7 +37,16 @@ image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
 # Text input for the question
 question = st.text_input("Enter your question about the image:")
 
+if st.button('Generate Caption'):
+    if image is not None:
+        # Display the image
+        st.image(image, use_column_width=True)
+        caption = get_caption(image)
+        st.write(caption)
 
+    else:
+        st.write("Please upload an image and enter a question.")
+        
 if st.button("Get Answer"):
     if image is not None and question:
         # Display the image