Update app_merged.py

app_merged.py

@@ -22,6 +22,8 @@ import datetime
 from pathlib import Path
 from io import BytesIO
 
+import openai
+from openai import OpenAI
 
 
 from PIL import Image
@@ -716,8 +718,59 @@ def run_rmbg(image):
     mask = pred_pil.resize(image_size)
     image.putalpha(mask)
     return image
+    
 
+def generate_description(object_description,image, detail="high", max_tokens=250):
+    openai_api_key = os.getenv("OPENAI_API_KEY")
+    client = OpenAI(api_key=openai_api_key)
 
+    if image is not None:
+        try:
+            img = image  # No need to open, directly use the PIL Image object
+
+            buffered = io.BytesIO()
+            img.save(buffered, format=IMAGE_FORMAT)
+            img_base64 = base64.b64encode(buffered.getvalue()).decode()
+
+            prompt = f"As if you were describing the interior design, make a detailed caption of this image in one large paragraph. Highlighting textures, furnitures, locations. This object should be included in the description :{object_description}"
+
+            payload = {
+                "model": "gpt-4o-mini",
+                "messages": [{
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                        {"type": "image_url",
+                        "image_url": {"url": f"data:image/jpeg;base64,{img_base64}", "detail": detail}}
+                    ]
+                }],
+                "max_tokens": max_tokens
+            }
+
+            response = client.chat.completions.create(**payload)
+            return response.choices[0].message.content
+        except Exception as e:
+            print(e)
+    else:
+        try:
+            prompt = f"Description: {object_description}. As if you were designing an interior, improve this sentence in one large paragraph. Highlighting textures, furnitures, locations, such that you create a coherent, visually pleasing setting."
+
+            payload = {
+                "model": "gpt-4o-mini",
+                "messages": [{
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                    ]
+                }],
+                "max_tokens": max_tokens
+            }
+
+            response = client.chat.completions.create(**payload)
+            return response.choices[0].message.content
+        except Exception as e:
+            print(e)    
+    
 
 def preprocess_image(image: Image.Image, height=768, width=768):
     image = np.array(image)
@@ -1693,7 +1746,10 @@ with gr.Blocks() as app:
 
                 with gr.Row():
                     with gr.Group():
-                        prompt = gr.Textbox(label="Prompt")
+                        with gr.Column():
+                            prompt = gr.Textbox(label="Prompt")
+                            augment_prompt = gr.Button(value='Augment Prompt')
+                        
                         bg_source = gr.Radio(choices=[e.value for e in list(BGSource)[2:]],
                                             value=BGSource.LEFT.value,
                                             label="Lighting Preference (Initial Latent)", type='value')
@@ -1746,6 +1802,7 @@ with gr.Blocks() as app:
             #     outputs=[result_gallery, output_bg],
             #     run_on_click=True, examples_per_page=1024
             # )
+        augment_prompt.click(generate_description, inputs=[prompt, dummy_image_for_prompt_augmentation], outputs=[prompt])
         ips = [extracted_fg, prompt, image_width, image_height, num_samples, seed, steps, a_prompt, n_prompt, cfg, highres_scale, highres_denoise, lowres_denoise, bg_source]
         relight_button.click(fn=process_relight, inputs=ips, outputs=[result_gallery])
         example_quick_prompts.click(lambda x, y: ', '.join(y.split(', ')[:2] + [x[0]]), inputs=[example_quick_prompts, prompt], outputs=prompt, show_progress=False, queue=False)