openai

app.py

@@ -5,9 +5,13 @@ Date: 2023-05-19
 """
 import os
 from typing import Union
+from datetime import datetime
+
 import streamlit as st
 import requests
 from PIL import Image
+from langchain.callbacks import get_openai_callback
+import openai
 
 def download_image(image_url:str, 
                    fname:str=None, 
@@ -31,11 +35,11 @@ def download_image(image_url:str,
         return None
 
 def save_rgba_image(uploaded_file_name:str, uploaded_file_buffer:object,
-                    Default_fold="images/"):
+                    default_fold:str="images/"):
     """
     # Save the uploaded file to a specific location
     """
-    file_path = os.path.join(Default_fold[:-1], uploaded_file_name)
+    file_path = os.path.join(default_fold[:-1], uploaded_file_name)
     with open(file_path, "wb") as file:
         file.write(uploaded_file_buffer)
     image_file = Image.open(file_path)
@@ -43,13 +47,34 @@ def save_rgba_image(uploaded_file_name:str, uploaded_file_buffer:object,
     image_file_rgba.save(file_path)
     return file_path
 
+def edit_image(img_name:str,
+    prompt:str="change backgroud to a beautiful sunset over the ocean")-> str:
+    """
+    code is generated by being
+    im_size : one of ['256x256', '512x512', '1024x1024'] - 'size'
+    """
+    response = openai.Image.create_edit(
+        image=open(img_name, "rb"),
+        prompt=prompt,
+        size="256x256",
+        n=1,
+        response_format="url"
+    )
+    img_url = response["data"][0]["url"]
+    return img_url
+
 def main():
     """
     Main function for the streamlit app.
     """
+    default_fold = "images/"
+    # load_dotenv()
     st.set_page_config(page_title="Avatar Generation")
     st.header("Artistic Avatar Generation 💖")
 
+    api_key = st.text_input("Enter your OpenAI API key:")
+    openai.api_key = api_key
+
     mode = st.radio("Choose your option", ("Upload photo by file", "Upload photo by url"))
     photo_filename = None
 
@@ -58,15 +83,15 @@ def main():
         if uploaded_file is not None:
             st.image(uploaded_file, width=256, caption="Your photo")
             # Save the uploaded file to a specific location
-            photo_filename = save_rgba_image(uploaded_file.name, uploaded_file.getbuffer())
+            photo_filename = save_rgba_image(uploaded_file.name, uploaded_file.getbuffer(),
+                                             default_fold=default_fold)
     else:
         photo_url = st.text_input("Enter your photo url:",
                               value="https://pngimg.com/uploads/monkey/monkey_PNG18720.png") 
         if photo_url:
             st.image(photo_url, width=256, caption="Your photo")
-            photo_filename = download_image(photo_url) # e.g., "images/monkey_PNG18720.png"
-
-    st.write("Your photo filename:", photo_filename)
+            photo_filename = download_image(photo_url, default_fold=default_fold)
+            # e.g., photo_url = "images/monkey_PNG18720.png"
 
     if photo_filename:
         image = Image.open(photo_filename)
@@ -75,7 +100,17 @@ def main():
         prompt = st.text_input("Prompt to transform to avatar: ",
                                 value="Change background to a beautiful sunset over the ocean")
         if prompt:
-            st.write("Prompt:", prompt)
+            with get_openai_callback() as oai_call_back:
+                # img_url = gen_image(prompt)
+                img_url = edit_image(photo_filename, prompt)
+                print(oai_call_back)
+            print(f"Avatar image url: {img_url}")
+            st.image(img_url, width=256, caption=f"prompt: {prompt}")
+            fname = f"{prompt}_{datetime.now().strftime('%Y%m%d%H%M%S')}.png"
+            output_filename = download_image(img_url, fname=fname)
+            if output_filename:
+                st.write(f"Your avatar filename: {output_filename}")
+                st.write("Done!")
 
 if __name__ == '__main__':
     main()

requirements.txt

@@ -1,3 +1,5 @@
 streamlit
 requests
-Pillow
+Pillow
+langchain
+openai