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

def load_image(input_type, uploaded_file=None, image_url=""):
    """
    Loads an image from an uploaded file or URL.
    """
    if input_type == "Upload Image" and uploaded_file is not None:
        return Image.open(io.BytesIO(uploaded_file.getvalue()))
    elif input_type == "Image URL" and image_url:
        try:
            response = requests.get(image_url)
            return Image.open(io.BytesIO(response.content))
        except Exception as e:
            st.error(f"Error loading image from URL: {e}")
            return None

def image_to_caption(image, input_type, uploaded_file, image_url):
    """
    Generates a caption for the given image.
    """
    image_to_text_pipeline = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
    if input_type == "Upload Image" and uploaded_file:
        return image_to_text_pipeline(uploaded_file.getvalue())[0]['generated_text']
    elif input_type == "Image URL" and image_url:
        return image_to_text_pipeline(image_url)[0]['generated_text']
    return ""

def generate_text_from_caption(caption):
    """
    Generates text based on the provided caption.
    """
    model = AutoModelForCausalLM.from_pretrained('pranavpsv/genre-story-generator-v2')
    tokenizer = AutoTokenizer.from_pretrained('pranavpsv/genre-story-generator-v2')
    input_ids = tokenizer.encode(caption, return_tensors='pt')
    generated_outputs = model.generate(input_ids, max_length=100, num_return_sequences=1)
    return tokenizer.decode(generated_outputs[0], skip_special_tokens=True)

def main():
    st.title('Image to Story Converter')

    # User interface for input selection
    input_type = st.radio("Select input type:", ("Upload Image", "Image URL"))
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) if input_type == "Upload Image" else None
    image_url = st.text_input("Enter the image URL here:", "") if input_type == "Image URL" else ""

    # Load image based on input selection
    image = load_image(input_type, uploaded_file, image_url)
    if image:
        st.image(image, caption='Selected Image', use_column_width=True)

    # Process image and generate text
    if st.button('Generate Caption and Continue'):
        if image:
            with st.spinner("Processing..."):
                # Convert image to text
                caption = image_to_caption(image, input_type, uploaded_file, image_url)
                st.success(f'Caption: {caption}')

                # Generate additional text based on the caption
                generated_text = generate_text_from_caption(caption)
                st.text_area("Generated Story:", generated_text, height=200)
        else:
            st.error("Please upload an image or enter an image URL first.")

if __name__ == "__main__":
    main()