Spaces:
Runtime error
Runtime error
| import base64 | |
| import boto3 | |
| from botocore.config import Config | |
| from dotenv import load_dotenv | |
| import os | |
| import shutil | |
| from typing import List, Tuple | |
| 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 | |
| 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 | |
| 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" | |
| # 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() | |
| if "key" not in st.session_state: | |
| st.session_state["key"] = uuid.uuid4().hex | |
| if "captcha_bool" not in st.session_state: | |
| st.session_state["captcha_bool"] = False | |
| if "model_inputs" not in st.session_state: | |
| st.session_state["model_inputs"] = None | |
| if ( | |
| "s3_face_file_path" not in st.session_state | |
| and "s3_theme_file_path" not in st.session_state | |
| ): | |
| st.session_state["s3_face_file_path"] = None | |
| st.session_state["s3_theme_file_path"] = None | |
| if "view" not in st.session_state: | |
| st.session_state["view"] = False | |
| if "train_view" not in st.session_state: | |
| st.session_state["train_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 "login" not in st.session_state: | |
| st.session_state["login"] = None | |
| if "user_auth_sess" not in st.session_state: | |
| st.session_state["user_auth_sess"] = False | |
| if "user_email" not in st.session_state: | |
| st.session_state["email_provided"] = True | |
| def callback(): | |
| st.session_state["button_clicked"] = True | |
| 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[str], 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 | |
| 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 | |
| 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) | |
| 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 | |
| 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): | |
| 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) | |
| if st.session_state["email_provided"]: | |
| 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_auth_sess"] = True | |
| st.session_state["email_provided"] = False | |
| send_email(text_input, str(st.session_state["key"])) | |
| if st.session_state["user_auth_sess"]: | |
| user_auth = st.empty() | |
| user_email_input.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["login"] = True | |
| else: | |
| st.markdown("Please Enter Correct Code!") | |
| if st.session_state["login"]: | |
| identifier = st.session_state["key"] | |
| user_auth.empty() | |
| user_email_input.empty() | |
| face_images = st.empty() | |
| with face_images.form("my_form"): | |
| uploaded_files = st.file_uploader( | |
| "Choose image files", accept_multiple_files=True, type=["png", "jpg", "jpeg"] | |
| ) | |
| submitted = st.form_submit_button(f"Upload") | |
| if submitted: | |
| with st.spinner('Uploading...'): | |
| st.session_state["s3_face_file_path"] = zip_and_upload_images( | |
| identifier, uploaded_files, "face" | |
| ) | |
| st.success(f'Uploading {len(uploaded_files)} files done!') | |
| 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", | |
| ) | |
| col1, col2 = st.columns([0.17, 1]) | |
| 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"] = { | |
| "superhero_file_path": dictionary[img][0], | |
| # Use presigned url since backend does not have credentials | |
| "person_file_path": generate_s3_get_url(st.session_state["s3_face_file_path"], expiration_seconds=3600), | |
| "superhero_prompt": dictionary[img][1], | |
| "num_images": 50, | |
| } | |
| st.success("Success!") | |
| 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"]: | |
| 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"], | |
| ) | |
| title = st.text_input("Theme Name") | |
| 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( | |
| identifier, uploaded_files_2, "theme" | |
| ) | |
| 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, | |
| "num_images": 50, | |
| } | |
| st.success('Done!') | |
| train = st.empty() | |
| with train.form("training"): | |
| col1, col3, col2 = st.columns(3) | |
| with col1: | |
| email = st.text_input("Enter Email") | |
| with col2: | |
| submitted = st.form_submit_button("Train Model!") | |
| if submitted: | |
| if not email: | |
| st.markdown('Please input an email!') | |
| else: | |
| 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"] = 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. |