Update app.py

app.py

@@ -9,8 +9,6 @@ from typing import List, Tuple, TYPE_CHECKING
 import uuid
 import argparse
 import logging
-import sendgrid
-from sendgrid.helpers.mail import Mail, Email, To, Content
 from enum import Enum
 import tempfile
 from pathlib import Path
@@ -34,60 +32,53 @@ 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"
+_GRETEL_USERINFO_ENDPOINT = "https://api.gretel.cloud/users/me"
+
+
 
 class UxState(str, Enum):
-    LOGIN = "login"
-    VERIFY_EMAIL = "verify_email"
+    LOGIN_VIA_API_KEY = "login_via_api_key"
     UPLOAD1 = "upload1"
     UPLOAD2 = "upload2"
-    CAPTCHA = "captcha"
+    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.
 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
+        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 "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 "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 "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 "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.
@@ -162,15 +153,12 @@ def generate_s3_put_url(s3_path: str, expiration_seconds: int) -> 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: 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.
     """
@@ -201,268 +189,207 @@ def zip_and_upload_images(identifier: str, uploaded_files: List[BytesIO], image_
 
         return s3_path
 
-def send_email(to_email, user_code):
-    sg = sendgrid.SendGridAPIClient(api_key=os.environ.get('SENDGRID_API_KEY'))
-    from_email = Email("santhosh@gretel.ai")
-    to_email = To(to_email)
-    subject = "One Time Code"
-    content = Content("text/plain", f"Here is your one-time code: {user_code}")
-    mail = Mail(from_email, to_email, subject, content)
-    mail_json = mail.get()
-    response = sg.client.mail.send.post(request_body=mail_json)
-
-
-# Create a function to generate a captcha
-def generate_captcha():
-    # Make a GET request to the API endpoint to generate a captcha
-    response = requests.get(CAPTCHA_ENDPOINT)
-
-    # If the request was successful, return the API response
-    if response.status_code == 200:
-        return response.json()
-    else:
-        logger.warn(f"Error from generate captcha request: {response.json()}")
-        # Otherwise, return an error message
-        return {"error": "Failed to generate captcha"}
-
-# Create a function to verify the captcha
-def verify_captcha(captcha_id, captcha_response):
-    # Make a POST request to the API endpoint with the captcha ID and response
-    verify_json = {"uuid": captcha_id, "captcha": captcha_response}
-    response = requests.post(
-        VERIFY_ENDPOINT, json=verify_json,
-    )
-    logger.info(f"Response from captcha verify: {response}")
-
-    # If the request was successful, return the API response
-    if response.status_code == 200:
-        return response.json()
-
-    # Otherwise, return an error message
-    return {"error": "Failed to verify captcha"}
-
 def train_model(model_inputs):
-    api_key = os.environ.get('API_KEY')
-    model_key = os.environ.get('MODEL_KEY')
+    api_key = "03cdd72e-5c04-4207-bd6a-fd5712c1740e"
+    model_key = "bd2c55f5-84bb-40f9-82fb-196ca68b1c1d"
     st.markdown(str(model_inputs))
     _ = banana.run(api_key, model_key, model_inputs)
 
-def run_login():
-    user_email_input = st.empty()
-    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_email"] = text_input
-            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-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["ux_state"] = UxState.UPLOAD1
-                st.experimental_rerun()
-            else:
-                st.markdown("Please Enter Correct Code!")
+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', type='primary')
+        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"]
-    face_images = st.empty()
-    with face_images.form("my_form"):
+    images = st.empty()
+    with images.form("concept_one_form"):
         uploaded_files = st.file_uploader(
-            "Choose image files", accept_multiple_files=True, type=["png", "jpg", "jpeg"]
+            "Choose first concept image files", accept_multiple_files=True, type=["png", "jpg", "jpeg"]
+        )
+        
+        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?',
         )
-        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["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
+                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!')
-            st.session_state["ux_state"] = UxState.UPLOAD2
-            st.experimental_rerun()
-
+            if concept:
+                switch_ux_state(UxState.UPLOAD2)
+            else: 
+                switch_ux_state(UxState.PROMPT)
 
 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(
-            "Choose a Theme!",
-            images=[
-                "https://gretel-image-synthetics.s3.us-west-2.amazonaws.com/theme-images/got.png",
-                "https://gretel-image-synthetics.s3.us-west-2.amazonaws.com/theme-images/ironman.png",
-                "https://ichef.bbci.co.uk/images/ic/640x360/p09t1hg0.jpg",
-            ],
-            captions=["Game of Thrones", "Iron Man", "Thor"],
-            return_value="index",
+    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"]
         )
+        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)
 
-        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:
-                with st.spinner():
-                    dictionary = {
-                        0: [
-                            "https://gretel-image-synthetics.s3.us-west-2.amazonaws.com/data/game-of-thrones.zip",
-                            "game-of-thrones",
-                        ],
-                        1: ["https://gretel-image-synthetics.s3.us-west-2.amazonaws.com/data/iron-man.zip", "iron-man"],
-                        2: ["https://gretel-image-synthetics.s3.us-west-2.amazonaws.com/data/thor.zip", "thor"],
-                    }
-                    st.session_state["model_inputs"] = {
-                        "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["ux_state"] = UxState.CAPTCHA
-                st.experimental_rerun()
-        with col2:
-            submitted_4 = st.form_submit_button(
-                "If none of the themes interest you, click here!"
-            )
-            if submitted_4:
-                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!")
-            uploaded_files_2 = st.file_uploader(
-                "Choose image files",
-                accept_multiple_files=True,
-                type=["png", "jpg", "jpeg"],
-            )
-            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["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
-                        "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["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!")
-
+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"]
+        )
+        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:
-            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.
+            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")
+        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(f"Congratulations, your model is training.")
+    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("You may close this browser window/tab.")
+    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"]
 
-    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()
+    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}'")
+        raise ValueError(f"Internal app error, unknown ux_state='{ux_state}'")
+