Update app.py

app.py

@@ -1,51 +1,69 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline,AutoFeatureExtractor
 from PIL import Image
 import torch
 
 # Load Hugging Face token
 HF_TOKEN = st.secrets["Hf_token"]
 
-
 # Load the model and pipeline
 model_id = "meta-llama/Llama-3.2-11B-Vision"
 
-# Initialize pipeline for image-text-to-text
-pipeline = pipeline(
-    "image-text-to-text",
-    model=model_id,
-    model_kwargs={"torch_dtype": torch.bfloat16, "use_auth_token": HF_TOKEN}
-)
-def preprocess_image(image):
-    return image.convert("RGB") 
-    
-st.title("Image-Text-to-Text Inference")
-st.write(f"**Using model:** {model_id}")
+# Secret Token
+HF_TOKEN = st.secrets["hf_token"]
+
+# Model and Task Configuration
+MODEL_ID = "meta-llama/Llama-3.2-11B-Vision"  # Replace with the correct model ID
+
+# Initialize the pipeline
+@st.cache_resource
+def initialize_pipeline():
+    return pipeline(
+        "image-text-to-text",
+        model=MODEL_ID,
+        model_kwargs={"torch_dtype": torch.bfloat16, "use_auth_token": HF_TOKEN}
+    )
 
-# Text Input
-input_text = st.text_input("Enter your prompt (optional):")
+# Preprocess image function
+def preprocess_image(image, size=(224, 224)):
+    """Converts the image to RGB and resizes it to the required dimensions."""
+    return image.convert("RGB").resize(size)
 
-# Image Input
-uploaded_file = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
+# Streamlit UI
+st.title("Image and Text to Text Generation")
+st.write(f"**Using model:** {MODEL_ID}")
+
+uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
+input_text = st.text_input("Enter your text input (optional):")
 
 if st.button("Generate"):
     if uploaded_file:
-        # Preprocess the uploaded image
-        image = preprocess_image(Image.open(uploaded_file))
-
-        # Prepare inputs
-        inputs = {"images": [image]}  # Wrap the image in a list
-        if input_text:
-            inputs["text"] = input_text  # Include text if provided
-
-        # Generate response
-        response = pipeline(**inputs)
-        
-        # Display results
-        st.write("Generated Response:")
-        st.write(response)  # Assuming the pipeline returns text directly
+        try:
+            # Preprocess image
+            image = Image.open(uploaded_file)
+            preprocessed_image = preprocess_image(image)
+
+            # Initialize pipeline
+            model_pipeline = initialize_pipeline()
+
+            # Create inputs for the pipeline
+            inputs = {"images": [preprocessed_image], "text": input_text}
+
+            # Run the model and get the response
+            response = model_pipeline(**inputs)
+            st.write("Generated Response:")
+            st.write(response)
+        except ValueError as ve:
+            if str(ve) == "The number of image token (0) should be the same as in the number of provided images (1)":
+                st.error("Make sure your image is correctly preprocessed and passed to the model.")
+            else:
+                st.error(f"Error: {ve}")
+        except Exception as e:
+            st.error(f"Error: {e}")
     else:
-        st.error("Please upload an image.")
+        st.error("Please upload an image to proceed.")
+
+