app.py

from __future__ import annotations

import boto3
from botocore.config import Config
from dotenv import load_dotenv
import os
import shutil
from typing import List, Tuple, TYPE_CHECKING
import uuid
import argparse
import logging
from enum import Enum
import tempfile
from pathlib import Path

import requests
import banana_dev as banana
import streamlit as st
from PIL import Image
from streamlit_image_select import image_select
import smart_open

if TYPE_CHECKING:
    from io import BytesIO

logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# Looks for .env file in current directory to pull environment variables. Should
# not overwrite already set environment variables. Used for S3 credentials.
load_dotenv()

_S3_PATH_OUTPUT = "s3://gretel-image-synthetics-use2/data/{identifier}/{image_type}_images.zip"
_GRETEL_USERINFO_ENDPOINT = "https://api.gretel.cloud/users/me"



class UxState(str, Enum):
    LOGIN_VIA_API_KEY = "login_via_api_key"
    UPLOAD1 = "upload1"
    UPLOAD2 = "upload2"
    UPLOAD3 = "upload3"
    PROMPT = "prompt"
    TRAIN = "train"
    FINISHED = "finished"

# Command-line arguments to control some stuff for easier local testing.
# Eventually may want to move everything into functions and have a
# if __name__ == "main" setup instead of everything inline.
parser = argparse.ArgumentParser()
parser.add_argument(
    "--dry-run", action="store_true",
    help="Skip sending train request to backend server.",
)
parser.add_argument(
    "--train-endpoint-url", default=None,
    help="URL of backend server to send train request to. If None, use hardcoded banana setup.",
)
cli_args = parser.parse_args()

def setup_session_state():
    if "key" not in st.session_state:
        st.session_state["key"] = uuid.uuid4().hex

    if "ux_state" not in st.session_state:
        st.session_state["ux_state"] = UxState.LOGIN_VIA_API_KEY

    if "model_inputs" not in st.session_state:
        st.session_state["model_inputs"] = None

    if "concepts" not in st.session_state:
        st.session_state["concepts"] = []
      
    if "prompt_keywords" not in st.session_state:
        st.session_state["prompt_keywords"] = None
    
    if "prompt" not in st.session_state:
        st.session_state["prompt"] = None
        
    if "view" not in st.session_state:
        st.session_state["view"] = False

    if "user_email" not in st.session_state:
        st.session_state["user_email"] = None

    if "user_firstname" not in st.session_state:
        st.session_state["user_firstname"] = None

    if "user_verified" not in st.session_state:
        st.session_state["user_verified"] = False


def bucket_parts(s3_path: str) -> Tuple[str, str]:
    """Split an S3 path into bucket and key.

    Args:
        s3_path: path starting with "s3:"

    Returns:
        Tuple of bucket and key for the path
    """
    parts = s3_path.split("/")
    bucket = parts[2]
    key = "/".join(parts[3:])
    return bucket, key


def generate_s3_get_url(s3_path: str, expiration_seconds: int) -> str:
    """Generate a presigned S3 url to read from an S3 path.

    A presigned url allows anyone accessing that url to read the s3 path without
    needing s3 credentials until the url expires.

    Args:
        s3_path: path starting with "s3:"
        expiration_seconds: how long the url will be valid (does not influence
            lifetime of the underlying s3 object, only the presigned url)

    Returns:
        The presigned url
    """
    bucket, key = bucket_parts(s3_path)

    s3_client = boto3.client("s3", config=Config(signature_version="s3v4", s3={"addressing_style": "path"}))
    download_url = s3_client.generate_presigned_url(
        "get_object",
        Params={
            "Bucket": bucket,
            "Key": key
        },
        ExpiresIn=expiration_seconds
    )
    return download_url


def generate_s3_put_url(s3_path: str, expiration_seconds: int) -> str:
    """Generate a presigned S3 url to write to an S3 path.

    A presigned url allows anyone accessing that url to write to the s3 path
    without needing s3 credentials until the url expires.

    Args:
        s3_path: path starting with "s3:"
        expiration_seconds: how long the url will be valid (does not influence
            lifetime of the underlying s3 object, only the presigned url)

    Returns:
        The presigned url
    """
    bucket, key = bucket_parts(s3_path)

    s3_client = boto3.client("s3", config=Config(signature_version="s3v4", s3={"addressing_style": "path"}))
    upload_url = s3_client.generate_presigned_url(
        "put_object",
        Params={
            "Bucket": bucket,
            "Key": key
        },
        ExpiresIn=expiration_seconds
    )
    return upload_url


def zip_and_upload_images(identifier: str, uploaded_files: List[BytesIO], image_type: str) -> str:
    """Save images as zip file to s3 for use in backend.
    Blocks until images are processed, added to zip file, and uploaded to S3.
    Args:
        identifier: unique identifier for the run, used in s3 link
        uploaded_files: BytesIO or UploadedFile from streamlit fileuploader
        image_type: string to identify different batches of images used in the
            backend model/training. Currently used values: "face", "theme"
    Returns:
        S3 location of zip file containing png images.
    """
    with tempfile.TemporaryDirectory() as temp_dir_name:
        logger.info(f"Working from temp dir to zip and upload images: {temp_dir_name}")
        temp_dir = Path(temp_dir_name)
        if not os.path.exists(temp_dir / identifier):
            os.makedirs(temp_dir / identifier)

        logger.info("Processing uploaded images")
        for num, uploaded_file in enumerate(uploaded_files):
            file_ = Image.open(uploaded_file).convert("RGB")
            file_.save(temp_dir / identifier / f"{num}_test.png")
        local_zip_filestem = str(temp_dir / f"{identifier}_{image_type}_images")
        logger.info("Making zip archive")
        shutil.make_archive(local_zip_filestem, "zip", temp_dir / identifier)
        local_zip_filename = f"{local_zip_filestem}.zip"

        logger.info("Uploading zip file to s3")
        # TODO: can we define expiration when making the s3 path?
        # Probably if we use the boto3 library instead of smart open
        s3_path = _S3_PATH_OUTPUT.format(identifier=identifier, image_type=image_type)

        with open(local_zip_filename, "rb") as fin:
            with smart_open.open(s3_path, "wb") as fout:
                fout.write(fin.read())
        logger.info(f"Completed upload to {s3_path}")

        return s3_path

def train_model(model_inputs):
    # logger.info("Skipping model training since --dry-run is enabled.")
    # logger.info(f"model_inputs: {model_inputs}")
    # st.markdown(str(model_inputs))

    # if cli_args.train_endpoint_url is None:
    # Use banana backend
    api_key = "03cdd72e-5c04-4207-bd6a-fd5712c1740e"
    model_key = "12f50d2a-fc6a-4334-b77f-e97fcabdee65"
    #st.markdown(str(model_inputs))
    #print(model_inputs)
    _ = banana.run(api_key, model_key, model_inputs)
    # else:
    #     # Send request directly to specified url
    #     _ = requests.post(cli_args.train_endpoint_url, json=model_inputs)


def switch_ux_state(new_state: UxState):
    st.session_state['ux_state'] = new_state
    st.experimental_rerun()

def run_enter_api_key():
    api_key_input = st.empty()
    with api_key_input.form(key='user_auth_api_key'):
        api_key_input = st.text_input(label='Please enter your Gretel API Key', type='password')
        st.caption("Don't have a Gretel Cloud account yet? [Sign up](https://gretel.ai/signup) for free now!")
        submit_button = st.form_submit_button(label='Submit')
        if submit_button:
            r = requests.get(_GRETEL_USERINFO_ENDPOINT, headers={'authorization': api_key_input})
            if r.status_code != 200:
                st.error('API key could not be verified')
                return
            me = r.json().get('data', {}).get('me', {})
            email = me.get('email')
            if email is None:
                st.error('No e-mail associated with this API key')
                return
            st.session_state["user_email"] = email
            st.session_state["user_firstname"] = me.get('firstname')
            st.session_state["user_verified"] = True
            
            switch_ux_state(UxState.UPLOAD1)

def run_upload_initial():
    identifier = st.session_state["key"]
    images = st.empty()
    with images.form("concept_one_form"):
        uploaded_files = st.file_uploader(
            "Choose first concept image files", accept_multiple_files=True, type=["png", "jpg", "jpeg"]
        )
        st.caption(
            """
            Files containing your training images (JPG, PNG, etc. size not restricted). These images contain your 'subject' that you want the trained model to embed in the output domain for later generating customized scenes beyond the training images. For best results, use images without noise or unrelated objects in the background.
            """
        )
        
        token = st.text_input("Token Name")
        st.caption(
            """
            The `token name` you use to describe your training images should be in the format: `a [identifier] [class noun]`, where the `[identifier]` should be a rare token. Relatively short sequences with 1-3 letters work the best (e.g. `sks`, `xjy`). `[class noun]` is a coarse class descriptor of the subject (e.g. cat, dog, watch, etc.). For example, your `token` can be: `a sks dog`, or with some extra description `a photo of a sks dog`. The trained model will learn to bind a unique identifier with your specific subject in the `instance_data`.
            """
        )
        class_token = st.text_input("Token Class")
        st.caption(
            """
            The `token class` is a description of the coarse class of your training images, in the format of `a [class noun]`, optionally with some extra description. `token_class` is used to alleviate overfitting to your customised images (the trained model should still keep the learnt prior so that it can still generate different dogs when the `[identifier]` is not in the prompt). Corresponding to the examples of the `token` above, the `token_class` can be `a dog` or `a photo of a dog`.
            """
        )
        concept = st.checkbox(
            'Would you like to fine-tune on a second concept?',
        )
        submitted = st.form_submit_button(f"Upload")
        if submitted:
            with st.spinner('Uploading...'):
                concept_information_dictionary = {
                    "file_path": generate_s3_get_url(zip_and_upload_images(
                    identifier, uploaded_files, "concept_one"), expiration_seconds=3600),
                    "token": token,
                    "class_token": class_token
                }
                st.session_state["concepts"].append(concept_information_dictionary)
            st.success(f'Uploading {len(uploaded_files)} files done!')
            if concept:
                switch_ux_state(UxState.UPLOAD2)
            else: 
                switch_ux_state(UxState.PROMPT)

def run_upload_secondary():
    identifier = st.session_state["key"]
    images = st.empty()
    with images.form("concept_two_form"):
        uploaded_files = st.file_uploader(
            "Choose second concept image files", accept_multiple_files=True, type=["png", "jpg", "jpeg"]
        )
        st.caption(
            """
            Files containing your training images (JPG, PNG, etc. size not restricted). These images contain your 'subject' that you want the trained model to embed in the output domain for later generating customized scenes beyond the training images. For best results, use images without noise or unrelated objects in the background.
            """
        )
        token = st.text_input("Token Name")
        st.caption(
            """
            The `token name` you use to describe your training images should be in the format: `a [identifier] [class noun]`, where the `[identifier]` should be a rare token. Relatively short sequences with 1-3 letters work the best (e.g. `sks`, `xjy`). `[class noun]` is a coarse class descriptor of the subject (e.g. cat, dog, watch, etc.). For example, your `token` can be: `a sks dog`, or with some extra description `a photo of a sks dog`. The trained model will learn to bind a unique identifier with your specific subject in the `instance_data`.
            """
        )
        class_token = st.text_input("Token Class")
        st.caption(
            """
            The `token class` is a description of the coarse class of your training images, in the format of `a [class noun]`, optionally with some extra description. `token_class` is used to alleviate overfitting to your customised images (the trained model should still keep the learnt prior so that it can still generate different dogs when the `[identifier]` is not in the prompt). Corresponding to the examples of the `token` above, the `token_class` can be `a dog` or `a photo of a dog`.
            """
        )
        next_concept = st.checkbox(
            'Would you like to fine-tune on a third concept?',
        )
        submitted = st.form_submit_button(f"Upload")
        if submitted:
            with st.spinner('Uploading...'):
                concept_information_dictionary = {
                    "file_path": generate_s3_get_url(zip_and_upload_images(
                    identifier, uploaded_files, "concept_two"), expiration_seconds=3600),
                    "token": token,
                    "class_token": class_token
                }
                st.session_state["concepts"].append(concept_information_dictionary)
            st.success(f'Uploading {len(uploaded_files)} files done!')
            if next_concept:
                switch_ux_state(UxState.UPLOAD3)
            else: 
                switch_ux_state(UxState.PROMPT)

def run_upload_third():
    identifier = st.session_state["key"]
    images = st.empty()
    with images.form("concept_three_form"):
        uploaded_files = st.file_uploader(
            "Choose third concept image files", accept_multiple_files=True, type=["png", "jpg", "jpeg"]
        )
        st.caption(
            """
            Files containing your training images (JPG, PNG, etc. size not restricted). These images contain your 'subject' that you want the trained model to embed in the output domain for later generating customized scenes beyond the training images. For best results, use images without noise or unrelated objects in the background.
            """
        )
        token = st.text_input("Token Name")
        st.caption(
            """
            The `token name` you use to describe your training images should be in the format: `a [identifier] [class noun]`, where the `[identifier]` should be a rare token. Relatively short sequences with 1-3 letters work the best (e.g. `sks`, `xjy`). `[class noun]` is a coarse class descriptor of the subject (e.g. cat, dog, watch, etc.). For example, your `token` can be: `a sks dog`, or with some extra description `a photo of a sks dog`. The trained model will learn to bind a unique identifier with your specific subject in the `instance_data`.
            """
        )
        class_token = st.text_input(f"Token Class")
        st.caption(
            """
            The `token class` is a description of the coarse class of your training images, in the format of `a [class noun]`, optionally with some extra description. `token_class` is used to alleviate overfitting to your customised images (the trained model should still keep the learnt prior so that it can still generate different dogs when the `[identifier]` is not in the prompt). Corresponding to the examples of the `token` above, the `token_class` can be `a dog` or `a photo of a dog`.
            """
        )
        submitted = st.form_submit_button(f"Upload")
        if submitted:
            with st.spinner('Uploading...'):
                concept_information_dictionary = {
                    "file_path": generate_s3_get_url(zip_and_upload_images(
                    identifier, uploaded_files, "concept_three"), expiration_seconds=3600),
                    "token": token,
                    "class_token": class_token
                }
                st.session_state["concepts"].append(concept_information_dictionary)
            st.success(f'Uploading {len(uploaded_files)} files done!')
            switch_ux_state(UxState.PROMPT)

def run_prompts():
    identifier = st.session_state["key"]
    prompt_form = st.empty()
    with prompt_form.form("prompt_form"):
        #prompt = st.text_input("Token Name")
        st.caption(
            """
            The `[Prompt]` and `[Prompt Keywords]` are descriptions of what you would like the model to generate.

            We recommend using a simple sentence covering 1 or 2 concepts for the Prompt. Nouns or adjectives are preferred, as verbs can be more challenging for the model. Examples of good `[prompts]`:
            - A charlock plant in the snow
            - A beaver wearing a suit
            - A cat samurai with a pet pug
            
            You may also add 1 to 3 comma-separated `[Prompt Keywords]` to describe the desired "mood" of the generated images. Examples of good `[Prompt Keywords]` include: concept art, steampunk, trending in ArtStation, good composition, hyper realistic, vivid colors, oil on canvas, Vincent van Gogh.
            """
        )
        st.text("")
        full_prompt = st.text_input("Prompt")
        prompt_keywords = st.text_input(f"Prompt Keywords")
        submitted = st.form_submit_button(f"Submit")
        if submitted:
            st.session_state["prompt_keywords"] = prompt_keywords
            st.session_state["prompt"] = full_prompt
            st.session_state["ux_state"] = UxState.TRAIN

def run_train():
    st.write("Congratulations, your model is training.")
    st.write(f"We'll send an email to {st.session_state['user_email']} when it's finished, usually about 20-30 minutes.")
    st.write("Closing this tab will not affect the ongoing image generation.")
    with st.spinner("Training in progress..."):
        st.session_state["model_inputs"] = {
                        "concepts": st.session_state["concepts"],
                        "num_images": 50,
                        "prompt": st.session_state["prompt"],
                        "prompt_keywords": st.session_state["prompt_keywords"]
        }
        s3_output_path = _S3_PATH_OUTPUT.format(identifier=st.session_state["key"], image_type="generated")
        st.session_state['model_inputs']['identifier'] = st.session_state["key"]
        st.session_state['model_inputs']['email'] = st.session_state["user_email"]
        # The backend does not have s3 credentials, so generate
        # presigned urls for the backend to use to write and read
        # the generated images.
        st.session_state['model_inputs']['output_s3_url_get'] = generate_s3_get_url(
            s3_output_path, expiration_seconds=60 * 60 * 24,
        )
        st.session_state['model_inputs']['output_s3_url_put'] = generate_s3_put_url(
            s3_output_path, expiration_seconds=3600,
        )
        train_model(st.session_state['model_inputs'])
        switch_ux_state(UxState.FINISHED)


def run_finished():
    st.success('Image generation completed!')
    st.write(f"We've sent an email to {st.session_state['user_email']} with a link to your generated images. Check it out!")

if __name__ == "__main__":
    setup_session_state()

    ux_state = st.session_state["ux_state"]

    runners = {
        UxState.LOGIN_VIA_API_KEY: run_enter_api_key,
        UxState.UPLOAD1: run_upload_initial,
        UxState.UPLOAD2: run_upload_secondary,
        UxState.UPLOAD3: run_upload_third,
        UxState.PROMPT: run_prompts,
        UxState.TRAIN: run_train,
        UxState.FINISHED: run_finished,
    }
    if (runner := runners.get(ux_state)) is not None:
        runner()
    else:
        raise ValueError(f"Internal app error, unknown ux_state='{ux_state}'")