import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from PIL import Image
import requests
import os

# Load the image classification pipeline
@st.cache_resource
def load_image_classification_pipeline():
    return pipeline("image-classification", model="Shresthadev403/food-image-classification")

pipe_classification = load_image_classification_pipeline()

# Load the Meta-Llama model and tokenizer for text generation
@st.cache_resource
def load_llama_pipeline():
    # Retrieve Hugging Face token from environment variables
    token = os.getenv("HF_AUTH_TOKEN")
    tokenizer = AutoTokenizer.from_pretrained(
        "meta-llama/Llama-3.2-3B-Instruct",
        use_auth_token=token
    )
    model = AutoModelForCausalLM.from_pretrained(
        "meta-llama/Llama-3.2-3B-Instruct",
        use_auth_token=token
    )
    return pipeline("text-generation", model=model, tokenizer=tokenizer)

pipe_llama = load_llama_pipeline()

# Function to generate ingredients using Meta-Llama
def get_ingredients(food_name):
    prompt = f"List the main ingredients typically used to prepare {food_name}:"
    response = pipe_llama(prompt, max_length=50, num_return_sequences=1)
    return response[0]['generated_text']

# Streamlit app
st.title("Food Image Classification with Ingredients Generation")
st.write("Upload an image to classify the type of food and get its ingredients!")

# Upload image
uploaded_file = st.file_uploader("Choose a food image...", type=["jpg", "png", "jpeg"])

if uploaded_file is not None:
    # Display the uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)
    st.write("Classifying...")
    
    # Make predictions
    predictions = pipe_classification(image)
    
    # Display only the top prediction
    top_food = predictions[0]['label']
    confidence = predictions[0]['score']
    st.subheader("Top Prediction")
    st.write(f"**{top_food}** with confidence {confidence:.2f}")
    
    # Generate and display ingredients for the top prediction
    st.subheader("Ingredients")
    try:
        ingredients = get_ingredients(top_food)
        st.write(ingredients)
    except Exception as e:
        st.write("Could not generate ingredients. Please try again later.")