import io
import time
import torch
import numpy as np
import streamlit as st

from PIL import Image

from diffusers import DiffusionPipeline, StableDiffusionPipeline
from transformers import BlipProcessor, BlipForConditionalGeneration

processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")

model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id)

def image2text(raw_image):
    inputs = processor(raw_image, return_tensors="pt")
    out = model.generate(**inputs)
    return processor.decode(out[0], skip_special_tokens=True)

@st.cache_resource(ttl="10h", show_spinner=False)
def text2image(prompt, negative_prompt,):
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
    ).images[0]
    return image

def main(pil_image, prompt, n_prompt,):

    # Generate using own prompt
    # runwayml/stable-diffusion-v1-5 model
    if st.button("Generate using own prompt", type="primary", key=5) and prompt and n_prompt:
        with st.spinner('Wait for it...'):
            start = time.time()
            image = text2image(prompt, n_prompt,)
            if image is None:
                st.write("Something go wrong try again")
            else:
                st.image(image)
        st.success('Done!')
        st.write(f"It's taken {time.time() - start}")

    # The prompt from photo
    # Salesforce/blip-image-captioning-large model
    if st.button("The prompt from photo", type="primary", key=4) and pil_image:
        with st.spinner('Wait for it...'):
            prompt = image2text(pil_image,)
            if prompt is None:
                st.write("Something go wrong try again")
            else:
                st.write(prompt)
        st.success('Done!')

    # Clears all st.cache_resource caches
    if st.button("Clear All"):
        st.cache_resource.clear()
        st.success("Done!")
            
    st.divider()

if __name__ == "__main__":
    st.header("This application will help you generate your own images.")
    st.divider()
    st.write("If you're unsure which prompt to use for generating your masterpiece, upload a photo you like, and the generative model will automatically determine the necessary prompt.")
    
    pil_image = None
    uploaded_file = st.file_uploader("Choose an image", type=["png", "jpg", "jpeg"])
    if uploaded_file is not None:
        pil_image = Image.open(uploaded_file)
        st.image(pil_image)
        
    st.divider()
    prompt = st.text_input("Write your prompt",)
    n_prompt = st.text_input("Write your negative prompt",)
    st.divider()
    
    main(pil_image, prompt, n_prompt)