import streamlit as st
import os
from prompts import prompts
from constants import JSON_SCHEMA_FOR_GPT, UPDATED_MODEL_ONLY_SCHEMA, JSON_SCHEMA_FOR_LOC_ONLY
from gpt import runAssistant, checkRunStatus, retrieveThread, createAssistant, saveFileOpenAI, startAssistantThread, \
create_chat_completion_request_open_ai_for_summary, addMessageToThread, create_image_completion_request_gpt
from summarizer import create_brand_html, create_langchain_openai_query
from theme import flux_generated_image, flux_generated_image_seed
import time
from PIL import Image
import io
def process_run(st, thread_id, assistant_id):
run_id = runAssistant(thread_id, assistant_id)
status = 'running'
while status != 'completed':
with st.spinner('. . .'):
time.sleep(20)
status = checkRunStatus(thread_id, run_id)
thread_messages = retrieveThread(thread_id)
for message in thread_messages:
if not message['role'] == 'user':
return message["content"]
else:
pass
def page1():
st.title("Upload Product")
st.markdown("Add a Product
", unsafe_allow_html=True)
st.markdown("Upload your product images, more images you upload better the AI learns
",
unsafe_allow_html=True)
uploaded_files = st.file_uploader("Upload Images", accept_multiple_files=True, key="uploaded_files_key")
product_description = st.text_area("Describe the product", value=st.session_state.get("product_description", ""))
col1, col2 = st.columns([1, 2])
with col1:
if st.button("Save"):
st.session_state['uploaded_files'] = uploaded_files
st.session_state['product_description'] = product_description
st.success("Product information saved!")
with col2:
if st.button("Add product and move to next page"):
if not uploaded_files:
st.warning("Please upload at least one image.")
elif not product_description:
st.warning("Please provide a description for the product.")
else:
st.session_state['uploaded_files'] = uploaded_files
st.session_state['product_description'] = product_description
st.session_state['page'] = "Page 2"
def page2():
st.title("Tell us about your shoot preference")
st.markdown("What are you shooting today?
", unsafe_allow_html=True)
shoot_type = st.radio("Select your shoot type:", ["Editorial", "Catalogue"], index=0)
st.session_state['shoot_type'] = shoot_type
brand_link = st.text_input("Add your brand link:", value=st.session_state.get("brand_link", ""))
st.session_state['brand_link'] = brand_link
if st.button("Get Brand Summary"):
if brand_link:
brand_summary_html = create_brand_html(brand_link)
brand_summary = create_langchain_openai_query(brand_summary_html)
st.session_state['brand_summary'] = brand_summary
st.success("Brand summary fetched!")
else:
st.warning("Please add a brand link.")
brand_summary_value = st.session_state.get('brand_summary', "")
editable_summary = st.text_area("Brand Summary:", value=brand_summary_value, height=100)
st.session_state['brand_summary'] = editable_summary
product_info = st.text_area("Tell us something about your product:", value=st.session_state.get("product_info", ""))
st.session_state['product_info'] = product_info
reference_images = st.file_uploader("Upload Reference Images", accept_multiple_files=True,
key="reference_images_key")
st.session_state['reference_images'] = reference_images
if st.button("Give Me Ideas"):
st.session_state['page'] = "Page 3"
def page3():
st.title("Scene Suggestions")
st.write("Based on your uploaded product and references!")
feedback = st.chat_input("Provide feedback:")
if not st.session_state.get("assistant_initialized", False):
assistant_id = createAssistant("You are a helpful assistant who is an expert in Fashion Shoots.")
updated_prompt = prompts["IDEA_GENERATION_PROMPT"].format(
brand_details=st.session_state["brand_summary"],
product_details=st.session_state["product_info"],
type_of_shoot=st.session_state["shoot_type"],
json_schema=JSON_SCHEMA_FOR_GPT,
product_name=st.session_state["product_description"]
)
file_locations = []
for uploaded_file in st.session_state['uploaded_files']:
bytes_data = uploaded_file.getvalue()
image = Image.open(io.BytesIO(bytes_data))
image.verify()
location = f"temp_image_{uploaded_file.name}"
with open(location, "wb") as f:
f.write(bytes_data)
file_locations.append(location)
image.close()
for uploaded_file in st.session_state['reference_images']:
bytes_data = uploaded_file.getvalue()
image = Image.open(io.BytesIO(bytes_data))
image.verify()
location = f"temp2_image_{uploaded_file.name}"
with open(location, "wb") as f:
f.write(bytes_data)
file_locations.append(location)
image.close()
file_ids = [saveFileOpenAI(location) for location in file_locations]
thread_id = startAssistantThread(file_ids, updated_prompt, "yes", "yes")
st.session_state.assistant_id = assistant_id
st.session_state.thread_id = thread_id
st.session_state.assistant_initialized = True
regenerate_images(thread_id, assistant_id)
if feedback:
if 'images' in st.session_state and 'descriptions' in st.session_state:
for image_path in st.session_state['images']:
os.remove(image_path)
del st.session_state['images']
del st.session_state['descriptions']
del st.session_state["json_descriptions"]
addMessageToThread(st.session_state.thread_id, feedback)
regenerate_images(st.session_state.thread_id, st.session_state.assistant_id)
selected_image_index = None
cols = st.columns(1)
for i in range(len(st.session_state["images"])):
with cols[i]:
st.image(st.session_state.images[i], caption=st.session_state.descriptions[i], use_column_width=True)
if st.radio(f"Select {i + 1}", [f"Select Image {i + 1}"], key=f"radio_{i}"):
selected_image_index = i
if selected_image_index is not None and st.button("Refine"):
st.session_state.selected_image_index = selected_image_index
st.session_state.selected_image = st.session_state.images[selected_image_index]
st.session_state.selected_text = st.session_state.descriptions[selected_image_index]
st.session_state['page'] = "Page 4"
if st.button("Go Back!"):
st.session_state.page = "Page 2"
def regenerate_images(thread_id, assistant_id):
"""Helper function to generate images and descriptions."""
response_from_process_list = []
for _ in range(1): # Assuming you generate 1 set of image/description
response_from_process = process_run(st, thread_id, assistant_id)
response_from_process_list.append(response_from_process)
summary_list = []
for final_response in response_from_process_list:
prompt_for_idea_summary = prompts["IDEA_SUMMARY_PROMPT"].format(
json_schema=str(final_response)
)
summary = create_chat_completion_request_open_ai_for_summary(prompt_for_idea_summary, "No")
summary_list.append(summary)
# Generate images based on the summaries
flux_generated_theme_image = []
for summary in summary_list:
theme_image = flux_generated_image(summary)
flux_generated_theme_image.append(theme_image["file_name"])
# Save the new images and descriptions in session state
st.session_state["images"] = flux_generated_theme_image
st.session_state["descriptions"] = summary_list
st.session_state["json_descriptions"] = response_from_process_list
def page4():
import json
selected_theme_text_by_user = st.session_state.json_descriptions[st.session_state.selected_image_index]
print(selected_theme_text_by_user)
schema_for_model_bg = {"type": "object",
"properties": {
"Model": {
"type": "string",
"description": "The model name or identifier."
},
"Background": {
"type": "string",
"description": "Description or type of the background."
}},
"required": ["Model", "Background"],
"additionalProperties": False
}
session_state_desp = st.session_state["json_descriptions"]
prompt_to_get_details = (f"You are provided with a brief of a Fashion Shoot : "
f"{session_state_desp}.\n Now provide me a JSON which will"
f"have two keys ```Model``` and ```Background```. Provide all detail's"
f"present about model and background in the brief provided by you. Just provide a "
f"natural langauge description. I will use it as description of model and "
f"background needed by the brand Output JSON following the schema")
response_from_open_ai = create_chat_completion_request_open_ai_for_summary(prompt_to_get_details,
schema_name="model_bg",
json_schema=schema_for_model_bg,
json_mode="yes")
json_response_from_open_ai = json.loads(response_from_open_ai)
with (st.sidebar):
st.title(st.session_state["product_info"])
st.write("Product Image")
st.image(st.session_state['uploaded_files'])
st.text("Scene Suggestion:")
st.image(st.session_state.selected_image)
dimensions = st.text_input("Enter Dimensions e.g 3:4, 1:2", key="Dimensions")
seed = st.selectbox(
"Seed Preference",
("Fixed", "Random"),
)
if seed == "Fixed":
seed_number = st.number_input("Enter an integer:", min_value=1, max_value=100000, value=10, step=1)
else:
seed_number = 0
st.text("Thanks will take care")
model_preference = st.selectbox(
"Model Preference",
("Create Own/Edit Pre-filled", "Ideas", "Upload Reference"),
)
if model_preference == "Create Own/Edit Pre-filled":
pre_filled_model_details = st.text_area("Model Idea", value=json_response_from_open_ai["Model"],
key="Model Idea")
elif model_preference == "Ideas":
prompt_to_generate_idea = ("Your task is to create model ideas for shoot of a product of a brand. "
"The details about the brand: ```{brand_details}.\n The product: {product_name},"
"which is: ```{product_details}```.\n Reference images for the product and "
"brands shoot idea is already provided with you. Additionally brand wants to "
"have a ```{type_of_shoot}``` of the model. Now based on all provided details, "
"think step by step and provide your ideas about what type of model the brand"
"should need based on mentioned JSON format. Also provide a combined prompt "
"which the brand will use to create a shoot image. While creating the "
"combined prompt as mentioned in the JSON schema, do not miss any details you"
" mentioned in the JSON.")
updated_model_idea_gen_prompt = prompt_to_generate_idea.format(
brand_details=st.session_state["brand_summary"],
product_details=st.session_state["product_info"],
type_of_shoot=st.session_state["shoot_type"],
product_name=st.session_state["product_description"]
)
response_for_only_model = create_chat_completion_request_open_ai_for_summary(updated_model_idea_gen_prompt
, schema_name="model_only",
json_schema=
UPDATED_MODEL_ONLY_SCHEMA,
json_mode="yes")
pre_filled_model_details = st.text_area("Model Idea", value=response_for_only_model,
key="Model Idea")
else:
uploaded_files = st.file_uploader("Upload one Model Reference Image here",
accept_multiple_files=False, key="uploader")
bytes_data = uploaded_files.getvalue()
image = Image.open(io.BytesIO(bytes_data))
image.verify()
location = f"temp_image_{uploaded_files.name}"
with open(location, "wb") as f:
f.write(bytes_data)
image.close()
prompt_to_generate_idea = ("Follow this JSON Schema : {json_schema_model_only}."
"Your task is to create model ideas for shoot of a product of a brand. "
"The details about the brand: ```{brand_details}.\n The product: {product_name},"
"which is: ```{product_details}```.\n Reference images for the product and "
"brands shoot idea is already provided with you. Additionally brand wants to "
"have a ```{type_of_shoot}``` of the model. Now based on all provided details, "
"think step by step and provide your ideas about what type of model the brand"
"should need based on mentioned JSON format. Also provide a combined prompt "
"which the brand will use to create a shoot image. While creating the "
"combined prompt as mentioned in the JSON schema, do not miss any details you"
" mentioned in the JSON.")
updated_model_idea_gen_prompt = prompt_to_generate_idea.format(
json_schema_model_only=UPDATED_MODEL_ONLY_SCHEMA,
brand_details=st.session_state["brand_summary"],
product_details=st.session_state["product_info"],
type_of_shoot=st.session_state["shoot_type"],
product_name=st.session_state["product_description"]
)
json_response = create_image_completion_request_gpt(location, updated_model_idea_gen_prompt)
pre_filled_model_details = st.text_area("Model Idea", value=json_response,
key="Model Idea")
background_preference = st.selectbox(
"Background Preference",
("Create Own/Edit Pre-filled", "Ideas", "Upload Reference"),
)
if background_preference == "Create Own/Edit Pre-filled":
pre_filled_background_details = st.text_area("Background Idea",
value=json_response_from_open_ai["Background"],
key="Background Idea")
elif background_preference == "Ideas":
prompt_to_generate_idea = ("Follow this JSON Schema : {json_schema_background_only}."
"Your task is to create location/background ideas for shoot of a "
"product of a brand. "
"The details about the brand: ```{brand_details}.\n The product: {product_name},"
"which is: ```{product_details}```.\n Reference images for the product and "
"brands shoot idea is already provided with you. Additionally brand wants to "
"have a ```{type_of_shoot}``` of the model. Now based on all provided details, "
"think step by step and provide your ideas about what type of location the brand"
"should need based on mentioned JSON format. Also provide a combined prompt "
"which the brand will use to create a shoot image. While creating the "
"combined prompt as mentioned in the JSON schema, do not miss any details you"
" mentioned in the JSON.")
updated_bg_idea_gen_prompt = prompt_to_generate_idea.format(
json_schema_background_only=JSON_SCHEMA_FOR_LOC_ONLY,
brand_details=st.session_state["brand_summary"],
product_details=st.session_state["product_info"],
type_of_shoot=st.session_state["shoot_type"],
product_name=st.session_state["product_description"]
)
response_for_only_bg = create_chat_completion_request_open_ai_for_summary(updated_bg_idea_gen_prompt,
schema_name="bg_o",
json_schema=JSON_SCHEMA_FOR_LOC_ONLY,
json_mode="yes")
pre_filled_background_details = st.text_area("Background Idea", value=response_for_only_bg,
key="Background Idea")
else:
uploaded_files = st.file_uploader("Upload one Background Reference Image here",
accept_multiple_files=False, key="uploader")
bytes_data = uploaded_files.getvalue()
image = Image.open(io.BytesIO(bytes_data))
image.verify()
location = f"temp2_image_{uploaded_files.name}"
with open(location, "wb") as f:
f.write(bytes_data)
image.close()
prompt_to_generate_idea = ("Follow this JSON Schema : {json_schema_bg_only}."
"Your task is to create Background/Location ideas for shoot of a "
"product of a brand. "
"The details about the brand: ```{brand_details}.\n The product: {product_name},"
"which is: ```{product_details}```.\n Reference images for the product and "
"brands shoot idea is already provided with you. Additionally brand wants to "
"have a ```{type_of_shoot}``` of the model. Now based on all provided details, "
"think step by step and provide your ideas about what type of location the brand"
"should need based on mentioned JSON format. Also provide a combined prompt "
"which the brand will use to create a shoot image. While creating the "
"combined prompt as mentioned in the JSON schema, do not miss any details you"
" mentioned in the JSON.")
updated_bg_idea_gen_prompt = prompt_to_generate_idea.format(
json_schema_bg_only=JSON_SCHEMA_FOR_LOC_ONLY,
brand_details=st.session_state["brand_summary"],
product_details=st.session_state["product_info"],
type_of_shoot=st.session_state["shoot_type"],
product_name=st.session_state["product_description"]
)
json_response = create_image_completion_request_gpt(location, updated_bg_idea_gen_prompt)
pre_filled_background_details = st.text_area("Background Idea", value=json_response,
key="Background Idea")
start_chat = st.button("Start Chat")
if "mood_chat_messages" not in st.session_state:
st.session_state["mood_chat_messages"] = []
if seed and dimensions and model_preference and background_preference:
if start_chat:
final_mood_board_image_prompt = prompts["FINAL_PROMPT_GENERATION"].format(
brand_details=st.session_state["brand_summary"],
product_details=st.session_state["product_info"],
type_of_shoot=st.session_state["shoot_type"],
product_name=st.session_state["product_description"],
model_details=pre_filled_model_details,
location_details=pre_filled_background_details,
theme_details=str(selected_theme_text_by_user),
chat_history=str(st.session_state["mood_chat_messages"])
)
prompt_for_flux_mood_board = create_chat_completion_request_open_ai_for_summary(
final_mood_board_image_prompt, "No", system_message=prompts["SYSTEM_PROMPT_FOR_MOOD_BOARD"])
if seed == "Fixed":
generated_flux_image = flux_generated_image_seed(prompt_for_flux_mood_board, seed_number, dimensions)
else:
generated_flux_image = flux_generated_image(prompt_for_flux_mood_board)
st.session_state["mood_chat_messages"].append({
"role": "AI",
"message": prompt_for_flux_mood_board,
"image": generated_flux_image["file_name"]
})
# for message in st.session_state["mood_chat_messages"]:
# if message["role"] == "AI":
# st.write(f"Caimera AI: {message['message']}")
# st.image(message['image'])
#else:
# st.write(f"**You**: {message['message']}")
user_input = st.chat_input("Type your message here...")
if user_input:
st.session_state["mood_chat_messages"].append({"role": "User", "message": user_input})
prompt_for_flux_mood_board_n = create_chat_completion_request_open_ai_for_summary(
user_input, "No", system_message=prompts["SYSTEM_PROMPT_FOR_MOOD_BOARD"])
if seed == "Fixed":
generated_flux_image_n = flux_generated_image_seed(prompt_for_flux_mood_board_n, seed_number,
dimensions)
else:
generated_flux_image_n = flux_generated_image(prompt_for_flux_mood_board_n)
st.session_state["mood_chat_messages"].append({
"role": "AI",
"message": prompt_for_flux_mood_board_n,
"image": generated_flux_image_n["file_name"]
})
for message in st.session_state["mood_chat_messages"]:
if message["role"] == "AI":
st.write(f"**AI**: {message['message']}")
st.image(message['image'])
else:
st.write(f"**You**: {message['message']}")
print(seed_number)
if 'page' not in st.session_state:
st.session_state.page = "Page 1"
# Routing between pages
if st.session_state.page == "Page 1":
page1()
elif st.session_state.page == "Page 2":
page2()
elif st.session_state.page == "Page 3":
page3()
elif st.session_state.page == "Page 4":
page4()