import streamlit as st
import requests
from PIL import Image
from io import BytesIO
# from IPython.display import display
import base64

# helper decoder
def decode_base64_image(image_string):
    base64_image = base64.b64decode(image_string)
    buffer = BytesIO(base64_image)
    return Image.open(buffer)

# display PIL images as grid
def display_image(image=None,width=500,height=500):
    img = image.resize((width, height))
    return img

# API Gateway endpoint URL
api_url = 'https://a02q342s5b.execute-api.us-east-2.amazonaws.com/reinvent-demo-inf2-sm-20231114'


# ===========
# Define Streamlit UI elements
st.title('Stable Diffusion XL with Refiner Image Generation')



prompt = st.text_area("Enter your prompt:", 
                      "Manatee astronaut in space, sci-fi, future, cold color palette, muted colors, detailed, 8k")

negative_prompt = st.text_area("Enter your negative prompt:", 
                      "anime, cartoon, graphic, text, painting, crayon, graphite, abstract glitch, blurry")

seed = st.number_input("Random seed", value=555, placeholder="Type a number...", help="set to same value to generate same image, if other inputs are the same, change to generate a different image for same inputs")
# seed = 555

num_inference_steps = st.slider("Number of Inference Steps", 
                                min_value=1, 
                                max_value=100, 
                                value=20,
                               help="more steps might improve quality, with diminishing marginal returns. 30-50 seems best, but your mileage may vary.")

denoising_start = st.slider("Denoising Start", 
                            min_value=0.0, 
                            max_value=1.0, 
                            value=0.8,
                           help="when to stop modifying the overall image and start refining the details")



if st.button('Generate Image'):
    with st.spinner(f'Generating Image with {num_inference_steps} iterations, beginning to refine around iteration {int(num_inference_steps * denoising_start)}...'):
        # ===============
        # Example input data
        prompt_input = {
            "prompt": prompt,
            "parameters": {
                "num_inference_steps": num_inference_steps,
                "seed": seed,
                "negative_prompt": negative_prompt
                # "denoising_start": denoising_start
            }
        }

        # Make API request
        response = requests.post(api_url, json=prompt_input)

        # Process and display the response
        if response.status_code == 200:
            result = response.json()
            # st.success(f"Prediction result: {result}")
            image = display_image(decode_base64_image(result["generated_images"][0]))
            st.header("SDXL Base + Refiner")
            st.image(image, 
                 caption=f"SDXL Base + Refiner, {num_inference_steps} iterations, beginning to refine around iteration {int(num_inference_steps * denoising_start)}") 
        else:
            st.error(f"Error: {response.text}")