Spaces:
Sleeping
Sleeping
| 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("<h2 style='color:#FF5733; font-weight:bold;'>Add a Product</h2>", unsafe_allow_html=True) | |
| st.markdown("<p style='color:#444;'>Upload your product images, more images you upload better the AI learns</p>", | |
| 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("<h3 style='color:#444;'>What are you shooting today?</h3>", 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 | |
| } | |
| prompt_to_get_details = (f"You are provided with a brief of a Fashion Shoot : " | |
| f"{st.session_state[\"json_descriptions\"]}.\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() | |