Update app.py

app.py

@@ -1,39 +1,18 @@
-# import os
-
-# os.system('pip install --upgrade pip')
-
 import streamlit as st
 from transformers import BlipForConditionalGeneration, AutoTokenizer
-import torchvision.transforms as transforms
 import torch
 from PIL import Image
-from io import BytesIO
-
-# # Load the fine-tuned model
-# model_path = '/content/model_after_5_epochs.pth'
-# model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
-# model.load_state_dict(torch.load(model_path))
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# model.to(device)
-# tokenizer = AutoTokenizer.from_pretrained("Salesforce/blip-image-captioning-base")
-
+import torchvision.transforms as transforms
 
+# Load the fine-tuned model and tokenizer
 model = BlipForConditionalGeneration.from_pretrained("MLInAi/CartoonCaptionGen")
 tokenizer = AutoTokenizer.from_pretrained("MLInAi/CartoonCaptionGen")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Function to generate caption for the uploaded image
-# def generate_caption(image):
-#     # Preprocess the image
-#     image = Image.open(image).convert("RGB")
-#     image = image.resize((224, 224))  # Resize the image to match model input size
-#     image_tensor = torch.tensor([torch.Tensor(image)]).permute(0, 3, 1, 2).to(device)
-
-#     # Generate caption
-#     output = model.generate(pixel_values=image_tensor)
-#     caption = tokenizer.decode(output[0], skip_special_tokens=True)
-#     return caption
-
+# Fixed prompt
+FIXED_PROMPT = "Generate a funny caption from image"
 
+# Function to generate caption for the uploaded image with the fixed prompt
 def generate_caption(image):
     # Preprocess the image
     image = Image.open(image).convert("RGB")
@@ -44,13 +23,14 @@ def generate_caption(image):
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     ])
-    image_tensor = transform(image).unsqueeze(0)
-
-    # Generate caption
-    output = model.generate(pixel_values=image_tensor)
+    image_tensor = transform(image).unsqueeze(0).to(device)
+    
+    # Generate caption with the fixed prompt
+    input_text = FIXED_PROMPT + " "
+    output = model.generate(pixel_values=image_tensor, input_text=input_text)
     caption = tokenizer.decode(output[0], skip_special_tokens=True)
     return caption
-   
+
 # Streamlit app
 st.title("Image Caption Generator")
 
@@ -61,6 +41,6 @@ if uploaded_image is not None:
     st.write("")
     st.write("Generating caption...")
 
-    # Generate caption for the uploaded image
+    # Generate caption for the uploaded image with the fixed prompt
     caption = generate_caption(uploaded_image)
     st.write("Caption:", caption)