Spaces:
Runtime error
Runtime error
| 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}'") | |