Update app.py

app.py

@@ -1,29 +1,30 @@
-import base64
+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
+from typing import List, Tuple, TYPE_CHECKING
 import uuid
-import zipfile
 import argparse
 import logging
 import sendgrid
 from sendgrid.helpers.mail import Mail, Email, To, Content
-
-from glob import glob
-from io import BytesIO
-from itertools import cycle
+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 st_btn_select import st_btn_select
 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)
@@ -33,29 +34,74 @@ logger.setLevel(logging.INFO)
 load_dotenv()
 
 _S3_PATH_OUTPUT = "s3://gretel-image-synthetics-use2/data/{identifier}/{image_type}_images.zip"
+CAPTCHA_ENDPOINT = "https://captcha-api.akshit.me/v2/generate"
+VERIFY_ENDPOINT = "https://captcha-api.akshit.me/v2/verify"
+
+class UxState(str, Enum):
+    LOGIN = "login"
+    VERIFY_EMAIL = "verify_email"
+    UPLOAD1 = "upload1"
+    UPLOAD2 = "upload2"
+    CAPTCHA = "captcha"
+    TRAIN = "train"
 
 # 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
+
+    if "model_inputs" not in st.session_state:
+        st.session_state["model_inputs"] = None
+
+    if "initial_concept_file_path" not in st.session_state:
+        st.session_state["initial_concept_file_path"] = None
+       
+    if "initial_token" not in st.session_state:
+        st.session_state["initial_token"] = None
+    
+    if "initial_class_token" not in st.session_state:
+        st.session_state["initial_class_token"] = None
+     
+    if "secondary_concept_file_path" not in st.session_state:
+        st.session_state["secondary_concept_file_path"] = None
+    
+    if "secondary_token" not in st.session_state:
+        st.session_state["secondary_token"] = None
+        
+    if "secondary_class_token" not in st.session_state:
+        st.session_state["secondary_class_token"] = None
+    
+    if "prompt_keywords" not in st.session_state:
+        st.session_state["prompt_keywords"] = None
+        
+    if "view" not in st.session_state:
+        st.session_state["view"] = False
 
-bool_session_variables = ['captcha_bool', 'view', 'train_view', 'login', 'user_auth_sess', 'face_images_uploaded', 'theme_images_uploaded']
-for var in bool_session_variables:
-    if var not in st.session_state:
-        st.session_state[var] = False
-
-obj_session_variables = ['s3_face_file_path', 's3_theme_file_path', 'captcha_response', 'user_email']
-for var in obj_session_variables:
-    if var not in st.session_state:
-        st.session_state[var] = None
-
-if "key" not in st.session_state:
-    st.session_state["key"] = uuid.uuid4().hex
+    if "captcha_response" not in st.session_state:
+        st.session_state["captcha_response"] = None
 
-if "captcha" not in st.session_state:
-    st.session_state["captcha"] = {}
+    if "captcha" not in st.session_state:
+        st.session_state["captcha"] = {}
 
-def callback():
-    st.session_state["button_clicked"] = True
+    if "user_email" not in st.session_state:
+        st.session_state["user_email"] = None
 
 
 def bucket_parts(s3_path: str) -> Tuple[str, str]:
@@ -129,43 +175,46 @@ def generate_s3_put_url(s3_path: str, expiration_seconds: int) -> str:
     return upload_url
 
 
-def zip_and_upload_images(identifier: str, uploaded_files: List[str], image_type: str) -> str:
+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: list of file names
+        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.
     """
-    if not os.path.exists(identifier):
-        os.makedirs(identifier)
-
-    logger.info("Processing uploaded images")
-    for num, uploaded_file in enumerate(uploaded_files):
-        file_ = Image.open(uploaded_file).convert("RGB")
-        file_.save(f"{identifier}/{num}_test.png")
-    local_zip_filestem = f"{identifier}_{image_type}_images"
-    logger.info("Making zip archive")
-    shutil.make_archive(local_zip_filestem, "zip", 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
+    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 send_email(to_email, user_code):
     sg = sendgrid.SendGridAPIClient(api_key=os.environ.get('SENDGRID_API_KEY'))
@@ -178,9 +227,6 @@ def send_email(to_email, user_code):
     response = sg.client.mail.send.post(request_body=mail_json)
 
 
-CAPTCHA_ENDPOINT = "https://captcha-api.akshit.me/v2/generate"
-VERIFY_ENDPOINT = "https://captcha-api.akshit.me/v2/verify"
-
 # Create a function to generate a captcha
 def generate_captcha():
     # Make a GET request to the API endpoint to generate a captcha
@@ -211,55 +257,73 @@ def verify_captcha(captcha_id, captcha_response):
     return {"error": "Failed to verify captcha"}
 
 def train_model(model_inputs):
-    api_key = "03cdd72e-5c04-4207-bd6a-fd5712c1740e"
-    model_key = "bd2c55f5-84bb-40f9-82fb-196ca68b1c1d"
-    _ = banana.run(api_key, model_key, model_inputs)
+    if cli_args.dry_run:
+        logger.info("Skipping model training since --dry-run is enabled.")
+        logger.info(f"model_inputs: {model_inputs}")
+        return
+
+    if cli_args.train_endpoint_url is None:
+        # Use banana backend
+        api_key = "03cdd72e-5c04-4207-bd6a-fd5712c1740e"
+        model_key = "1a3b4ce5-164f-4efb-9f4a-c2ad3a930d0b"
+        st.markdown(str(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)
 
-#first page - user provides email
-if st.session_state["user_email"] is None:
+
+def run_login():
     user_email_input = st.empty()
-    with user_email_input.form(key='user_email_provided'):
+    with user_email_input.form(key='user_auth'):
         text_input = st.text_input(label='Please Enter Your Email')
         submit_button = st.form_submit_button(label='Submit')
-        if submit_button: 
-            st.session_state["user_auth_sess"] = True
+        if submit_button:
             st.session_state["user_email"] = text_input
-            send_email(st.session_state["user_email"], str(st.session_state["key"]))
-            user_email_input.empty()
-    
-#second page - user inputs one time code
-if st.session_state["user_auth_sess"]:
+            send_email(text_input, str(st.session_state["key"]))
+            st.session_state["ux_state"] = UxState.VERIFY_EMAIL
+            # TODO: alternately run this submit log in a callback to the input?
+            # or otherwise ensure we execute the runner for the new state
+            st.experimental_rerun()
+
+
+def run_verify_email():
     user_auth = st.empty()
-    with user_auth.form("one-time-code"):
+    with user_auth.form("one-code"):
         text_input = st.text_input(label='Please Input One Time Code')
         submit_button = st.form_submit_button(label='Submit')
         if submit_button:
             if text_input == st.session_state["key"]:
-                st.session_state["login"] = True
+                st.session_state["ux_state"] = UxState.UPLOAD1
+                st.experimental_rerun()
             else:
                 st.markdown("Please Enter Correct Code!")
 
-#third page - user inputs face images
-if st.session_state["login"]:
+
+def run_upload_initial():
     identifier = st.session_state["key"]
-    user_auth.empty()
     face_images = st.empty()
-    with face_images.form("face_images_upload"):
+    with face_images.form("my_form"):
         uploaded_files = st.file_uploader(
             "Choose image files", accept_multiple_files=True, type=["png", "jpg", "jpeg"]
         )
+        initial_concept_token = st.text_input("Token Name")
+        initial_concept_class_token = st.text_input("Token Class")
         submitted = st.form_submit_button(f"Upload")
         if submitted:
             with st.spinner('Uploading...'):
-                st.session_state["s3_face_file_path"] = zip_and_upload_images(
+                st.session_state["initial_concept_file_path"] = zip_and_upload_images(
                     identifier, uploaded_files, "face"
                 )
+                st.session_state["initial_token"] = initial_concept_token
+                st.session_state["initial_class_token"] = initial_concept_class_token
             st.success(f'Uploading {len(uploaded_files)} files done!')
-            st.session_state["face_images_uploaded"] = True
+            st.session_state["ux_state"] = UxState.UPLOAD2
+            st.experimental_rerun()
+
 
-#fourth page - user inputs theme images
-if st.session_state["face_images_uploaded"]:
-    face_images.empty()
+def run_upload_secondary():
+    identifier = st.session_state["key"]
     preset_theme_images = st.empty()
     with preset_theme_images.form("choose-preset-theme"):
         img = image_select(
@@ -274,6 +338,7 @@ if st.session_state["face_images_uploaded"]:
         )
 
         col1, col2 = st.columns([0.17, 1])
+        prompt_keywords = st.text_input("Prompt Keywords")
         with col1:
             submitted_3 = st.form_submit_button("Submit!")
             if submitted_3:
@@ -287,13 +352,19 @@ if st.session_state["face_images_uploaded"]:
                         2: ["https://gretel-image-synthetics.s3.us-west-2.amazonaws.com/data/thor.zip", "thor"],
                     }
                     st.session_state["model_inputs"] = {
-                        "superhero_file_path": dictionary[img][0],
-                        "person_file_path": generate_s3_get_url(st.session_state["s3_face_file_path"], expiration_seconds=3600),
-                        "superhero_prompt": dictionary[img][1],
+                        "secondary_concept_file_path": dictionary[img][0],
+                        # Use presigned url since backend does not have credentials
+                        "initial_token": st.session_state["initial_token"],
+                        "secondary_token": dictionary[img][1],
+                        "initial_class_token": st.session_state["initial_class_token"],
+                        "secondary_class_token": 'superhero',
+                        "initial_concept_file_path": generate_s3_get_url(st.session_state["initial_concept_file_path"], expiration_seconds=3600),
                         "num_images": 50,
+                        "prompt_keywords": prompt_keywords
                     }
                 st.success("Success!")
-                st.session_state["theme_images_uploaded"] = True
+                st.session_state["ux_state"] = UxState.CAPTCHA
+                st.experimental_rerun()
         with col2:
             submitted_4 = st.form_submit_button(
                 "If none of the themes interest you, click here!"
@@ -302,6 +373,7 @@ if st.session_state["face_images_uploaded"]:
                 st.session_state["view"] = True
 
     if st.session_state["view"]:
+        # TODO: split into it's own ux state and function?
         custom_theme_images = st.empty()
         with custom_theme_images.form("input_custom_themes"):
             st.markdown("If none of the themes interest you, please input your own!")
@@ -310,86 +382,113 @@ if st.session_state["face_images_uploaded"]:
                 accept_multiple_files=True,
                 type=["png", "jpg", "jpeg"],
             )
-            title = st.text_input("Theme Name")
+            secondary_concept_token = st.text_input("Token Name")
+            secondary_concept_class_token = st.text_input("Token Class")
+            prompt_keywords = st.text_input("Prompt Keywords")
             submitted_3 = st.form_submit_button("Submit!")
             if submitted_3:
                 with st.spinner('Uploading...'):
-                    st.session_state["s3_theme_file_path"] = zip_and_upload_images(
+                    st.session_state["secondary_concept_file_path"] = zip_and_upload_images(
                         identifier, uploaded_files_2, "theme"
                     )
+                    #st.markdown(secondary_concept_file_path)
                     st.session_state["model_inputs"] = {
                         # Use presigned urls since backend does not have credentials
-                        "superhero_file_path": generate_s3_get_url(st.session_state["s3_theme_file_path"], expiration_seconds=3600),
-                        "person_file_path": generate_s3_get_url(st.session_state["s3_face_file_path"], expiration_seconds=3600),
-                        "superhero_prompt": title,
+                        "initial_concept_file_path": generate_s3_get_url(st.session_state["initial_concept_file_path"], expiration_seconds=3600),
+                        "secondary_concept_file_path": generate_s3_get_url(st.session_state["secondary_concept_file_path"], expiration_seconds=3600),
+                        "initial_token": st.session_state["initial_token"],
+                        "secondary_token": secondary_concept_token,
+                        "initial_class_token": st.session_state["initial_class_token"],
+                        "secondary_class_token": secondary_concept_class_token,
                         "num_images": 50,
+                        "prompt_keywords": prompt_keywords
                     }
                 st.success('Done!')
-                st.session_state["theme_images_uploaded"] = True
+                st.session_state["ux_state"] = UxState.CAPTCHA
+                st.experimental_rerun()
+                
+
+
+def run_captcha():
+    captcha_form = st.empty()
+    with captcha_form.form("captcha_form", clear_on_submit=True):
+        # Create container to create image/text input out of order from the
+        # format submit button. Needed since we need to know the status of the
+        # form submit to know what the captcha should do.
+        captcha_container = st.container()
+        display_captcha = True
+        # TODO: Submit button renders first, then drops down once the image is
+        # fetched leading to page reflow. Would be nice to not have reflow, but
+        # we need to know if the submit button was previously pressed and if the
+        # captcha was solved to generate and display a new captcha or not.
+        # Possible solution is use an on_click callback to set a session_state
+        # variable to access whether the button was pushed or not instead of the
+        # return value here.
+        submitted = st.form_submit_button("Submit Captcha!")
 
-#fifth page - user enters captcha and trains model
-if st.session_state["theme_images_uploaded"]:
-    if st.session_state["view"]:
-        custom_theme_images.empty()
-    preset_theme_images.empty()
-    train = st.empty()
-    with train.form("training"):
-        submitted = st.form_submit_button("Train Model!")
         if submitted:
-            st.session_state["captcha_bool"] = True
-
-    if st.session_state["captcha_bool"]:
-        captcha_form = st.empty()
-        with captcha_form.form("captcha_form", clear_on_submit=True):
-            # Create container to create image/text input out of order from the
-            # format submit button. Needed since we need to know the status of the
-            # form submit to know what the captcha should do.
-            captcha_container = st.container()
-            display_captcha = True
-            # TODO: Submit button renders first, then drops down once the image is
-            # fetched leading to page reflow. Would be nice to not have reflow, but
-            # we need to know if the submit button was previously pressed and if the
-            # captcha was solved to generate and display a new captcha or not.
-            # Possible solution is use an on_click callback to set a session_state
-            # variable to access whether the button was pushed or not instead of the
-            # return value here.
-            submitted = st.form_submit_button("Submit Captcha!")
-
-            if submitted:
-                result = verify_captcha(st.session_state['captcha']['uuid'], st.session_state["captcha_response"])
-                del st.session_state["captcha_response"]
-                if 'message' in result and result['message'] == 'CAPTCHA_SOLVED':
-                    st.session_state['captcha'] = {}
-                    display_captcha = False
-                    with st.spinner("Model Fine Tuning..."):
-                        st.session_state["model_inputs"]["identifier"] = st.session_state["key"]
-                        st.session_state["model_inputs"]["email"] = st.session_state["user_email"]
-                        s3_output_path = _S3_PATH_OUTPUT.format(identifier=st.session_state["key"], image_type="generated")
-                        # 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"])
-                        st.session_state["train_view"] = True
-                else:
-                    st.error(result['error'])
-
-            if display_captcha:
-                # Generate new captcha and display. Occurs on first load with the
-                # captcha_bool=True, or after previously failed captcha attempts.
-                result = generate_captcha()
-                captcha_id = result['uuid']
-                captcha_image = result['captcha']
-
-                st.session_state['captcha']['uuid'] = captcha_id
-                st.session_state['captcha']['captcha'] = captcha_image
-
-                captcha_container.image(captcha_image, width=300)
-
-                captcha_container.text_input("Enter the captcha response", key="captcha_response")
-                # Submit button already setup previously.
+            result = verify_captcha(st.session_state['captcha']['uuid'], st.session_state["captcha_response"])
+            del st.session_state["captcha_response"]
+            if 'message' in result and result['message'] == 'CAPTCHA_SOLVED':
+                st.session_state['captcha'] = {}
+                display_captcha = False
+                with st.spinner("Model Fine Tuning..."):
+                    st.session_state["model_inputs"]["identifier"] = st.session_state["key"]
+                    st.session_state["model_inputs"]["email"] = st.session_state["user_email"]
+                    s3_output_path = _S3_PATH_OUTPUT.format(identifier=st.session_state["key"], image_type="generated")
+                    # 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"])
+                    st.session_state["ux_state"] = UxState.TRAIN
+                    st.experimental_rerun()
+            else:
+                st.error(result['error'])
+
+        if display_captcha:
+            # Generate new captcha and display. Occurs on first run of the
+            # captcha state, or after previously failed captcha attempts.
+            result = generate_captcha()
+            captcha_id = result['uuid']
+            captcha_image = result['captcha']
+
+            st.session_state['captcha']['uuid'] = captcha_id
+            st.session_state['captcha']['captcha'] = captcha_image
+
+            captcha_container.image(captcha_image, width=300)
+
+            captcha_container.text_input("Enter the captcha response", key="captcha_response")
+            # Submit button already setup previously.
+
+
+def run_train():
+    st.write(f"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("You may close this browser window/tab.")
+
+
+if __name__ == "__main__":
+    setup_session_state()
+
+    ux_state = st.session_state["ux_state"]
+
+    if ux_state == UxState.LOGIN:
+        run_login()
+    elif ux_state == UxState.VERIFY_EMAIL:
+        run_verify_email()
+    elif ux_state == UxState.UPLOAD1:
+        run_upload_initial()
+    elif ux_state == UxState.UPLOAD2:
+        run_upload_secondary()
+    elif ux_state == UxState.CAPTCHA:
+        run_captcha()
+    elif ux_state == UxState.TRAIN:
+        run_train()
+    else:
+        raise ValueError(f"Internal app error, unknown ux_state='{ux_state}'")