updated streamlit app

.ipynb_checkpoints/Dockerfile-checkpoint

@@ -30,7 +30,7 @@ RUN mkdir -p /app/cache && chmod -R 777 /app/cache
 ENV HF_HOME=/app/cache
 
 # Expose the necessary ports
-EXPOSE 8501 8502
+EXPOSE 7860 8502
 
 # Create a supervisor configuration file
 RUN mkdir -p /etc/supervisor/conf.d/

.ipynb_checkpoints/start-checkpoint.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
 # Start FastAPI
-uvicorn app:app --host 0.0.0.0 --port 8502 &
+# uvicorn app:app --host 0.0.0.0 --port 8502 &
 
 # Start Streamlit
 streamlit run streamlit_app.py --server.port=7860 --server.address=0.0.0.0

.ipynb_checkpoints/streamlit_app-checkpoint.py

@@ -1,6 +1,21 @@
 import streamlit as st
 import requests
 from PIL import Image
+from transformers import AutoProcessor, Blip2ForConditionalGeneration
+import torch
+import io
+
+
+@st.cache_resource
+def load_model():
+    model = Blip2ForConditionalGeneration.from_pretrained("ybelkada/blip2-opt-2.7b-fp16-sharded")
+    model.load_adapter('blip-cpu-model')
+    processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    return model, processor
+
+model, processor = load_model()
 
 st.title("Image Captioning with Fine-Tuned BLiPv2 Model")
 
@@ -11,8 +26,15 @@ if uploaded_file is not None:
     st.image(image, caption="Uploaded Image", use_column_width=True)
     
     files = {"file": uploaded_file.getvalue()}
-    response = requests.post("http://0.0.0.0:8502/generate-caption/", files=files)
-    caption = response.json().get("caption")
+    print("Sending API request")
+    # response = requests.post("http://0.0.0.0:8502/generate-caption/", files=files)
+    # caption = response.json().get("caption")
+    
+    inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
+    
+    with torch.no_grad():
+        caption_ids = model.generate(**inputs, max_length=128)
+        caption = processor.decode(caption_ids[0], skip_special_tokens=True)
     
     st.write("Generated Caption:")
     st.write(f"**{caption}**")

Dockerfile

@@ -30,7 +30,7 @@ RUN mkdir -p /app/cache && chmod -R 777 /app/cache
 ENV HF_HOME=/app/cache
 
 # Expose the necessary ports
-EXPOSE 8501 8502
+EXPOSE 7860 8502
 
 # Create a supervisor configuration file
 RUN mkdir -p /etc/supervisor/conf.d/

start.sh

@@ -1,7 +1,7 @@
 #!/bin/bash
 
 # Start FastAPI
-uvicorn app:app --host 0.0.0.0 --port 8502 &
+# uvicorn app:app --host 0.0.0.0 --port 8502 &
 
 # Start Streamlit
 streamlit run streamlit_app.py --server.port=7860 --server.address=0.0.0.0

streamlit_app.py

@@ -1,6 +1,21 @@
 import streamlit as st
 import requests
 from PIL import Image
+from transformers import AutoProcessor, Blip2ForConditionalGeneration
+import torch
+import io
+
+
+@st.cache_resource
+def load_model():
+    model = Blip2ForConditionalGeneration.from_pretrained("ybelkada/blip2-opt-2.7b-fp16-sharded")
+    model.load_adapter('blip-cpu-model')
+    processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    return model, processor
+
+model, processor = load_model()
 
 st.title("Image Captioning with Fine-Tuned BLiPv2 Model")
 
@@ -12,8 +27,14 @@ if uploaded_file is not None:
     
     files = {"file": uploaded_file.getvalue()}
     print("Sending API request")
-    response = requests.post("http://0.0.0.0:8502/generate-caption/", files=files)
-    caption = response.json().get("caption")
+    # response = requests.post("http://0.0.0.0:8502/generate-caption/", files=files)
+    # caption = response.json().get("caption")
+    
+    inputs = processor(images=image, return_tensors="pt").to(device, torch.float16)
+    
+    with torch.no_grad():
+        caption_ids = model.generate(**inputs, max_length=128)
+        caption = processor.decode(caption_ids[0], skip_special_tokens=True)
     
     st.write("Generated Caption:")
     st.write(f"**{caption}**")