add custom params

app.py

@@ -3,11 +3,30 @@ import requests
 
 prefix = "https://smartfeed-custom-tools.hf.space/gradio_api/file="
 
+extract_models = [
+    "o3",
+    "gpt-5",
+    "o4-mini",
+    "o4-mini-high",
+]
+
+
+def change_image_style(
+    image_url,
+    style_image_url,
+    extract_model,
+    extract_prompt,
+    prompt_prefix
+):
+    image_url = "https://replicate.delivery/pbxt/Mjo5VWBA0fm8oq3BefpbdtXeJitBJc6AobPnnMdFxxCKzT3P/linkedin.jpeg"
+    style_image_url = "https://replicate.delivery/pbxt/N55l5TWGh8mSlNzW8usReoaNhGbFwvLeZR3TX1NL4pd2Wtfv/replicate-prediction-f2d25rg6gnrma0cq257vdw2n4c.png"
 
-def change_image_style(image_url, style_image_url):
     data = {
         "image_url": image_url,
-        "style_image_url": style_image_url
+        "style_image_url": style_image_url,
+        "extract_model": extract_model,
+        "extract_prompt": extract_prompt,
+        "prompt_prefix": prompt_prefix
     }
 
     response = requests.post(
@@ -17,12 +36,22 @@ def change_image_style(image_url, style_image_url):
     )
 
     if response.status_code == 200:
-        resultImageUrl = response.json().get("data").get("image_url")
-        return resultImageUrl
+        data = response.json().get("data")
+        resultImageUrl = data.get("image_url")
+        edit_prompt = data.get("edit_prompt")
+        # replace \n to <br>
+        edit_prompt = edit_prompt.replace("\n", "<br>")
+        return resultImageUrl, edit_prompt
     else:
         raise Exception(f"Error: {response.status_code} - {response.text}")
 
-def generate_image(input_image, style_images):
+def generate_image(
+    input_image,
+    style_images,
+    extract_model,
+    extract_prompt,
+    prompt_prefix
+):
 
     if not input_image:
         raise gr.Error(f"Please upload an input image! Refer to step 1️⃣")
@@ -32,15 +61,18 @@ def generate_image(input_image, style_images):
 
     inputImageUrl = prefix + input_image
 
-    result_images = list[str]()
+    result_images = []
+    markdownStr = ""
     for style_image in style_images:
         if not style_image:
             raise gr.Error(f"Cannot find any style image! Please refer to step 1️⃣")
         styleImageUrl = prefix + style_image[0]
-        result_images.append(change_image_style(inputImageUrl, styleImageUrl))
+        result_image, edit_prompt = change_image_style(inputImageUrl, styleImageUrl, extract_model, extract_prompt, prompt_prefix)
+        result_images.append(result_image)
+        markdownStr += f"{edit_prompt}: <img src='{result_image}' style='zoom: 33%;' />\n"
+
+    return result_images, markdownStr
 
-    return result_images
-    
 
 def swap_to_gallery(images):
     return gr.update(value=images, visible=True), gr.update(visible=True), gr.update(visible=False)
@@ -61,18 +93,23 @@ with gr.Blocks() as demo:
             with gr.Column(visible=False) as clear_button:
                 remove_and_reupload = gr.ClearButton(value="Remove and upload new ones", components=files, size="sm")
             input_image = gr.Image(label="Input Image", type="filepath", interactive=True)
+            extract_model = gr.Dropdown(choices=extract_models, label="Extract Model", value="o3")
+            extract_prompt = gr.Textbox(lines=2, label="Extracted Prompt", value="Accurately extract everything from this photo into a prompt especially the hairstyle excluding the backgrounds. I want to replicate it perfectly. need to be a complete paragraph")
+            prompt_prefix = gr.Textbox(lines=1, label="Prompt Prefix", value="change the backgrounds to match this:")
+
             submit = gr.Button("Submit")
 
         with gr.Column():
             gallery = gr.Gallery(label="Generated Images")
+            markdown = gr.Markdown(label="Generated Images")
 
         files.upload(fn=swap_to_gallery, inputs=files, outputs=[uploaded_files, clear_button, files])
         remove_and_reupload.click(fn=remove_back_to_files, outputs=[uploaded_files, clear_button, files])
 
         submit.click(
             fn=generate_image,
-            inputs=[input_image, uploaded_files],
-            outputs=[gallery]
+            inputs=[input_image, uploaded_files, extract_model, extract_prompt, prompt_prefix],
+            outputs=[gallery, markdown]
         )