import streamlit as st
import glib as glib
from io import BytesIO

def check_password():
    """Returns `True` if the user had the correct password."""

    def password_entered():
        """Checks whether a password entered by the user is correct."""
        if st.session_state["password"] == st.secrets["password"]:
            st.session_state["password_correct"] = True
            # del st.session_state["password"]  # don't store password
        else:
            st.session_state["password_correct"] = False

    if "password_correct" not in st.session_state:
        # First run, show input for password.
        st.text_input(
            "Password", type="password", on_change=password_entered, key="password"
        )
        return False
    elif not st.session_state["password_correct"]:
        # Password not correct, show input + error.
        st.text_input(
            "Password", type="password", on_change=password_entered, key="password"
        )
        st.error("😕 Password incorrect")
        return False
    else:
        # Password correct.
        return True

def s3_uploader(checkbox, filename, filebytes):
    if checkbox:
        glib.uploadFileToS3(filename, filebytes)
    st.session_state.filename = ''

some_prompts = {
    "default" : "Type or Select prompt. Something like - pencil sketch of pizza and some cake",
    "prompt-1": "a person in 8k resolution, photorealistic masterpiece by Aaron Horkey and Jeremy Mann, intricately detailed fluid gouache painting by Jean Baptiste, professional photography, natural lighting, volumetric lighting, maximalist, concept art, intricately detailed, complex, elegant, expansive cover",
    "prompt-2": "anime style",
    "prompt-3": "Portrait photo side profile, looking away, serious eyes, 50mm portrait photography, hard rim lighting photography–beta –ar 2:3 –beta –upbeta",
    "prompt-4": "High Quality, Intricately Detailed, Hyper-Realistic Lawyer Portrait Photography, Volumetric Lighting, Full Character, 4k, In Workwear, boca",
    "prompt-5": "Street style photo portrait, clear edge definition, unique and one-of-a-kind pieces, light brown and light amber, Fujifilm X-T4, Sony FE 85mm f/1.4 GM, —quality 2 —s 750 —v 5.2",
}

# Perform some crude state management
if "input_img" not in st.session_state:
    st.session_state.input_img = ''
if "gen_img" not in st.session_state:
    st.session_state.gen_img = ''


# start rendering the App
if not check_password():
    st.stop()

st.subheader("GenAI Application (beta)")

input_img = ''
input_slider = 60
uploaded_file = st.file_uploader("**Input Image**")
if uploaded_file:
    # remove rotation (when using smart-phone camera) and resize
    input_img = glib.remove_rotation(uploaded_file.getvalue())
    st.session_state.input_img = input_img
else:
    input_img = st.session_state.input_img

if input_img:
    st.image(input_img)
    input_slider = st.slider("Image strength", 1, 100, 60)

selection = st.selectbox("**Load Prompt**", some_prompts.keys())
prompt_text = st.text_area("**Input Text Prompt**", value = some_prompts[selection], 
                           height=50, key="prompt")

process_button = st.button("Generate", type="primary")

gen_img = ''
st.subheader("Generated Image")
if process_button:
    with st.spinner("Drawing..."):
        image_bytes = input_img.getvalue() if input_img else None
        gen_img = glib.get_altered_image_from_model(prompt_content=prompt_text, image_bytes=image_bytes,
                                                    img_strength=input_slider)
        st.session_state.gen_img = gen_img
else:
    gen_img = st.session_state.gen_img
    
# generated output and download elements
if gen_img:
    st.image(gen_img)
    img_bytes = gen_img.getvalue()
    
    # placeholder for outputs
    outs = st.empty().container()
    checkbox = outs.checkbox("Upload for printing", value=True, key="checkbox")
    filename = outs.text_input("Image name:", key="filename")
    
    # Check if the filename already exists
    if filename and checkbox and glib.checkFileinS3(filename+".jpg"):
        filename = ''
        st.error("File already exists")
    
    outs.download_button(label="Download",
        data=img_bytes, 
        file_name=filename+".jpg",
        mime="image/jpg",
        disabled=not filename,
        on_click=s3_uploader, 
        args=[checkbox, filename+".jpg", img_bytes])
genAI.py
aws-samples/amazon-aiml-genai-streamlit-stable-diffusion
import os
import io
import boto3
import json
import base64
from PIL import Image
from PIL import ImageOps
from io import BytesIO
import sagemaker
from stability_sdk_sagemaker.predictor import StabilityPredictor
from stability_sdk.api import GenerationRequest, GenerationResponse, TextPrompt
import streamlit as st
from resizeimage import resizeimage

# SM Jumpstart Consts
region_name = "us-west-2"
storage_bucket = "your-s3-storage"
sm_sdxl_ep = "sdxl-jumpstart"

# Bedrock Consts
endpoint_url = "bedrock-endpoint-url"
bedrock_model_id = 'stability.stable-diffusion-xl'

def checkFileinS3(filename):
    session = boto3.Session()
    s3 = session.client('s3')
    results = s3.list_objects(Bucket=storage_bucket, Prefix=filename)
    return ('Contents' in results)

def uploadFileToS3(fileName, fileBytes):
    session = boto3.Session()
    s3 = session.client('s3', region_name=region_name)
    s3.upload_fileobj(io.BytesIO(fileBytes), storage_bucket, fileName)

# Request for an image from SDXL given a text-prompt
def smjs_create_img_from_prompt(text):
    session = boto3.session.Session()
    deployed_model = StabilityPredictor(endpoint_name=sm_sdxl_ep, 
                                        sagemaker_session=sagemaker.Session(boto_session=session))
    output_img = deployed_model.predict(GenerationRequest(text_prompts=[TextPrompt(text=text)],
                                        style_preset="cinematic", seed = 3))
    return output_img

def smjs_create_img_from_img_and_prompt(text, image, slider):
    strength = 0.5 + (slider / 100)*(0.45)
    session = boto3.session.Session()
    deployed_model = StabilityPredictor(endpoint_name=sm_sdxl_ep, 
                                        sagemaker_session=sagemaker.Session(boto_session=session))
    img = deployed_model.predict(GenerationRequest(text_prompts=[TextPrompt(text=text)],
                                        init_image=image,
                                        cfg_scale=9,
                                        image_strength=strength,
                                        seed=42))
    return img

# Extract the image out from the response
def decode_image(model_response: GenerationResponse):
    image = model_response.artifacts[0].base64
    image_data = base64.b64decode(image.encode())
    image = Image.open(io.BytesIO(image_data))
    result_io = BytesIO()
    image.save(result_io, format=image.format)
    return result_io

# When taking pictures from a smart-phone, they have rotation metadata
# Remove rotation and resize to 512 to 512. There will be some cropping.
def remove_rotation(image_bytes):
    image_io = BytesIO(image_bytes)
    image = Image.open(image_io)
    new_one = ImageOps.exif_transpose(image)
    new_one = resizeimage.resize_cover(new_one, [512,512])
    new_img_io = BytesIO()
    new_one.save(new_img_io, format="PNG")
    return new_img_io
    
def get_resized_image_io(image_bytes):
    image_io = BytesIO(image_bytes)
    image = Image.open(image_io)
    resized_image = image.resize((512, 512))
    resized_io = BytesIO()
    resized_image.save(resized_io, format=image.format)
    return resized_io

def prepare_image_for_endpoint(image_bytes):
    resized_io = get_resized_image_io(image_bytes)
    img_str = base64.b64encode(resized_io.getvalue()).decode("utf-8")
    return img_str

def get_stability_ai_response_image(response):
    response = json.loads(response.get('body').read())
    images = response.get('artifacts')
    image_data = base64.b64decode(images[0].get('base64'))
    return BytesIO(image_data)

def get_altered_image_from_model(prompt_content, image_bytes, img_strength):
    if image_bytes is None:
        img = smjs_create_img_from_prompt(prompt_content)
        return decode_image(img)
    else:
        img_str = prepare_image_for_endpoint(image_bytes)
        img = smjs_create_img_from_img_and_prompt(prompt_content, img_str, img_strength)
        return decode_image(img)

def get_stability_ai_request_body(prompt, image_str = None):
    # see https://platform.stability.ai/docs/features/api-parameters
    body = {"text_prompts": [ {"text": prompt } ], "cfg_scale": 9, "steps": 50, }
    if image_str:
        body["init_image"] = image_str
    return json.dumps(body)

def bedrock_create_img_from_img_and_prompt (text, image):
    body = get_stability_ai_request_body(text, image)
    session = boto3.session.Session()
    bedrock = session.client(service_name='bedrock', region_name=region_name,
                             endpoint_url=endpoint_url) 
    response = bedrock.invoke_model(body=body, modelId=bedrock_model_id, 
                                    contentType="application/json",
                                    accept="application/json")
    output = get_stability_ai_response_image(response)
    return output

glib.py
aws-samples/amazon-aiml-genai-streamlit-stable-diffusion