Update app.py

app.py

@@ -9,26 +9,43 @@ import os
 # OpenAI API Key
 api_key = os.getenv("OPENAI_API_KEY")
 
+
 # Function to encode the image
 def encode_image(image_array):
     # Convert numpy array to an image file and encode it in base64
     img = Image.fromarray(np.uint8(image_array))
-    img_buffer = os.path.join("/tmp", f"temp_image_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg")
+    img_buffer = os.path.join(
+        "/tmp", f"temp_image_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
+    )
     img.save(img_buffer, format="JPEG")
-    
+
     with open(img_buffer, "rb") as image_file:
         return base64.b64encode(image_file.read()).decode("utf-8")
 
+
 # Function to generate product description using OpenAI API
-def generate_product_description(image, text_input=None):
+def generate_product_description(image, description_type, custom_instruction=None):
     # Encode the uploaded image
     base64_image = encode_image(image)
 
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {api_key}"
+    headers = {"Content-Type": "application/json", "Authorization": f"Bearer {api_key}"}
+
+    # Set the description type or custom instruction
+    description_prompts = {
+        "Short Formal": "Create a compelling and succinct product description from the image.",
+        "Bullet Points": "Provide a detailed product description in bullet points based on the image.",
+        "Amazon Optimized": "Write an Amazon-style product description that includes key features, benefits, and a call to action.",
+        "Fashion": "Generate a stylish and trendy product description suitable for a fashion item based on the image.",
+        "Sport": "Create an energetic and engaging product description for a sports-related item based on the image.",
     }
 
+    if description_type == "Other" and custom_instruction:
+        instruction = custom_instruction
+    else:
+        instruction = description_prompts.get(
+            description_type, "Create a product description based on the image."
+        )
+
     # Payload with base64 encoded image as a Data URL
     payload = {
         "model": "gpt-4o-mini",
@@ -36,28 +53,31 @@ def generate_product_description(image, text_input=None):
             {
                 "role": "user",
                 "content": [
-                    {"type": "text", "text": text_input or "What's in this image?"},
+                    {"type": "text", "text": instruction},
                     {
                         "type": "image_url",
-                        "image_url": {
-                            "url": f"data:image/jpeg;base64,{base64_image}"
-                        }
-                    }
-                ]
+                        "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},
+                    },
+                ],
             }
         ],
-        "max_tokens": 300
+        "max_tokens": 300,
     }
 
-    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+    response = requests.post(
+        "https://api.openai.com/v1/chat/completions", headers=headers, json=payload
+    )
     response_data = response.json()
 
     # Handle errors
     if response.status_code != 200:
-        raise ValueError(f"OpenAI API Error: {response_data.get('error', {}).get('message', 'Unknown Error')}")
+        raise ValueError(
+            f"OpenAI API Error: {response_data.get('error', {}).get('message', 'Unknown Error')}"
+        )
+
+    # Extract and return only the generated message content
+    return response_data["choices"][0]["message"]["content"]
 
-    # Extract and return the generated message
-    return response_data["choices"][0]["message"]
 
 css = """
   #output {
@@ -68,18 +88,44 @@ css = """
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.Markdown("WordLift Product Description Generation - [GPT-4o-mini Demo]")
+    gr.Markdown("WordLift Product Description Generation - [FREE]")
     with gr.Tab(label="WordLift Product Description Generation"):
         with gr.Row():
             with gr.Column():
                 input_img = gr.Image(label="Input Picture")
-                text_input = gr.Textbox(label="Additional Instructions (Optional)")
+                description_type = gr.Dropdown(
+                    label="Select Description Type",
+                    choices=[
+                        "Short Formal",
+                        "Bullet Points",
+                        "Amazon Optimized",
+                        "Fashion",
+                        "Sport",
+                        "Other",
+                    ],
+                    value="Short Formal",
+                )
+                custom_instruction = gr.Textbox(
+                    label="Custom Instruction (Only for 'Other')", visible=False
+                )
                 submit_btn = gr.Button(value="Submit")
             with gr.Column():
                 output_text = gr.Textbox(label="Output Text")
 
+        # Toggle visibility of custom instruction based on selected type
+        def toggle_custom_instruction(type_selection):
+            return gr.update(visible=(type_selection == "Other"))
+
+        description_type.change(
+            toggle_custom_instruction,
+            inputs=[description_type],
+            outputs=[custom_instruction],
+        )
+
         submit_btn.click(
-            generate_product_description, [input_img, text_input], [output_text]
+            generate_product_description,
+            [input_img, description_type, custom_instruction],
+            [output_text],
         )
 
 # Launch Gradio app