Spaces:
Running
Running
| from PIL import Image | |
| from io import BytesIO | |
| import base64 | |
| import time | |
| import streamlit as st | |
| UNIT_COST = { | |
| "user": 0.01, | |
| "assistant": 0.03, | |
| } | |
| # Convert Image to Base64 | |
| def im_2_b64(image): | |
| image = Image.open(image) | |
| image.thumbnail((512, 512), Image.Resampling.LANCZOS) | |
| image = image.convert("RGB") | |
| buff = BytesIO() | |
| image.save(buff, format="JPEG") | |
| img_str = base64.b64encode(buff.getvalue()) | |
| return img_str | |
| def calculate_cost(): | |
| def get_text_cost(text, unit_cost): | |
| num_of_words = len(text.split()) | |
| tokens = max(1000.0 * num_of_words / 750.0, 0.0) | |
| tokens = tokens / 1000.0 | |
| cost = tokens * unit_cost | |
| return cost | |
| def get_image_cost(unit_cost=0.01): | |
| cost = 0.00255 # 512x512 image: https://openai.com/pricing | |
| return cost | |
| messages = st.session_state.messages | |
| total_cost = 0 | |
| for message in messages: | |
| role = message["role"] | |
| for content in message["content"]: | |
| if content["type"] == "image_url": | |
| total_cost += get_image_cost(UNIT_COST[role]) | |
| else: | |
| total_cost += get_text_cost(content["text"], UNIT_COST[role]) | |
| st.session_state.cost.append(total_cost) | |
| def clear_uploader(): | |
| st.session_state["uploader_key"] += 1 | |
| st.rerun() | |
| def undo(): | |
| if len(st.session_state.messages) > 0: | |
| st.query_params.clear() | |
| msg = st.session_state.messages.pop() | |
| if msg["role"] == "assistant": | |
| st.session_state.cost.pop() | |
| time.sleep(0.1) | |
| st.rerun() | |
| def restart(): | |
| st.query_params.clear() | |
| st.session_state.messages = [] | |
| st.session_state.cost = [] | |
| time.sleep(0.2) | |
| clear_uploader() |