Update app.py

app.py

@@ -1,25 +1,31 @@
 import streamlit as st
-from transformers import pipeline,AutoFeatureExtractor
+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
 model_id = "meta-llama/Llama-3.2-11B-Vision"
-@st.cache_resource
+
+# Initialize feature extractor and pipeline
+@st.experimental_singleton
 def initialize_pipeline():
-    return pipeline(
-        "image-text-to-text",
+    feature_extractor = AutoFeatureExtractor.from_pretrained(model_id, use_auth_token=HF_TOKEN)
+    model_pipeline = pipeline(
+        "image-to-text",
         model=model_id,
-        model_kwargs={"torch_dtype": torch.bfloat16, "use_auth_token": HF_TOKEN}
+        feature_extractor=feature_extractor,
+        model_kwargs={"torch_dtype": torch.float32, "use_auth_token": HF_TOKEN},
     )
+    return feature_extractor, model_pipeline
 
-# 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)
+# Preprocess function using feature extractor
+def preprocess_image(image, feature_extractor):
+    """Preprocess the image for the model."""
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    return inputs.pixel_values
 
 # Streamlit UI
 st.title("Image and Text to Text Generation")
@@ -31,29 +37,21 @@ input_text = st.text_input("Enter your text input (optional):")
 if st.button("Generate"):
     if uploaded_file:
         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}
+            # Load and preprocess image
+            image = Image.open(uploaded_file).convert("RGB")
+            feature_extractor, model_pipeline = initialize_pipeline()
+            preprocessed_image = preprocess_image(image, feature_extractor)
 
-            # Run the model and get the response
-            response = model_pipeline(**inputs)
+            # Run the pipeline
+            response = model_pipeline(images=preprocessed_image, text=input_text)
             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("Ensure 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 to proceed")
+        st.error("Please upload an image to proceed.")
+
+