update personalization

app.py

@@ -1,4 +1,3 @@
-
 import gradio as gr
 import numpy as np
 import random
@@ -11,7 +10,7 @@ import open_clip
 from optim_utils import optimize_prompt
 from utils import (
     clean_response_gpt, setup_model, init_gpt_api, call_gpt_api,
-    get_refine_msg, clean_cache, get_personalize_message,
+    get_refine_msg, clean_cache, get_personalize_message, get_personalized_simplified,
     clean_refined_prompt_response_gpt, IMAGES, OPTIONS, T2I_MODELS,
     INSTRUCTION, IMAGE_OPTIONS, PROMPTS, SCENARIOS  
 )
@@ -41,6 +40,7 @@ torch.cuda.empty_cache()
 METHOD = "Experimental" 
 counter = 1
 enable_submit = False
+redesign_flag = False
 responses_memory = {METHOD: {}}
 example_data = [
     [
@@ -100,7 +100,8 @@ def call_gpt_refine_prompt(prompt, num_prompts=5, max_tokens=1000, temperature=0
 def personalize_prompt(prompt, history, feedback, like_image, dislike_image):
     seed = random.randint(0, MAX_SEED)
     client = init_gpt_api()
-    messages = get_personalize_message(prompt, history, feedback, like_image, dislike_image)
+    # messages = get_personalize_message(prompt, history, feedback, like_image, dislike_image)
+    messages = get_personalized_simplified(prompt, like_image, dislike_image)
     outputs = call_gpt_api(messages, client, "gpt-4o", seed, max_tokens=2000, temperature=0.7, top_p=0.9)
     return outputs
 
@@ -121,12 +122,12 @@ def invert_prompt(prompt, images, prompt_len=15, iter=500, lr=0.1, batch_size=2)
     }
     inverted_prompt = optimize_prompt(clip_model, preprocess, text_params, device, target_images=images, target_prompts=prompt)
 
-    # eval(prompt, learned_prompt, optimized_images, clip_model, preprocess)
-    # return learned_prompt
-
 # =========================
 # UI Helper Functions
 # =========================
+# Store generated images for selection
+current_generated_images = []
+
 def reset_gallery():
     return []
 
@@ -136,105 +137,106 @@ def display_error_message(msg, duration=5):
 def display_info_message(msg, duration=5):
     gr.Info(msg, duration=duration)
 
-def check_satisfaction(sim_radio):
-    global enable_submit, counter
-    fully_satisfied_option = ["Satisfied", "Very Satisfied"]
-    if_submit = (sim_radio in fully_satisfied_option) or enable_submit or (counter > MAX_ROUND)
-    return gr.update(interactive=if_submit)
-
-def select_image(like_radio, images_method):
-    if like_radio == IMAGE_OPTIONS[0]:
-        return images_method[0][0]
-    elif like_radio == IMAGE_OPTIONS[1]:
-        return images_method[1][0]
-    elif like_radio == IMAGE_OPTIONS[2]:
-        return images_method[2][0]
-    elif like_radio == IMAGE_OPTIONS[3]:
-        return images_method[3][0]
-    else:
-        return None
-
-def check_evaluation(sim_radio):
-    if not sim_radio:
+def check_evaluation(sim_radio, like_image, dislike_image):
+    if not sim_radio or not like_image or not dislike_image:
         display_error_message("❌ Please fill all evaluations before changing image or submitting.")
         return False
     return True
 
 def generate_image(prompt, like_image, dislike_image):
-    global responses_memory
+    global responses_memory, current_generated_images
     history_prompts = [v["prompt"] for v in responses_memory[METHOD].values()]
     feedback = [v["sim_radio"] for v in responses_memory[METHOD].values()]
-    personalized = prompt
-    # personalized = personalize_prompt(prompt, history_prompts, feedback, like_image, dislike_image)
-    # personalized = clean_refined_prompt_response_gpt(personalized)
-    # if "I'm sorry, I can't assist with" in personalized:
-    #     personalized = prompt
+    print(feedback, like_image, dislike_image)
+    if like_image and dislike_image and feedback:
+        personalized = personalize_prompt(prompt, history_prompts, feedback, like_image, dislike_image)
+    else:
+        personalized = prompt
     gallery_images = []
+    current_generated_images = []  # Reset the stored images
     refined_prompts = call_gpt_refine_prompt(personalized)
     for i in range(NUM_IMAGES):
         img = infer(refined_prompts[i])
         gallery_images.append(img)
+        current_generated_images.append(img)  # Store for selection
         yield gallery_images
 
-def redesign(prompt, sim_radio, like_radio, dislike_radio, current_images, history_images, like_image, dislike_image):
-    global counter, enable_submit, responses_memory
-    if check_evaluation(sim_radio):
+def on_gallery_select(evt: gr.SelectData):
+    """Handle gallery image selection and return the selected image"""
+    global current_generated_images
+    if current_generated_images and evt.index < len(current_generated_images):
+        return current_generated_images[evt.index]
+    return None
+
+def handle_like_drag(selected_image):
+    """Handle setting an image as liked"""
+    return selected_image
+
+def handle_dislike_drag(selected_image):
+    """Handle setting an image as disliked"""
+    return selected_image
+
+def redesign(prompt, sim_radio, current_images, history_images, like_image, dislike_image):
+    global counter, responses_memory, redesign_flag
+    
+    if check_evaluation(sim_radio, like_image, dislike_image):
         responses_memory[METHOD][counter] = {
             "prompt": prompt,
             "sim_radio": sim_radio,
             "response": "",
-            "satisfied_img": f"round {counter}, {like_radio}",
-            "unsatisfied_img": f"round {counter}, {dislike_radio}",
+            "satisfied_img": f"round {counter}, liked image",
+            "unsatisfied_img": f"round {counter}, disliked image",
         }
 
-        enable_submit = True if sim_radio in ["Satisfied", "Very Satisfied"] or enable_submit else False
-
         history_prompts = [[v["prompt"]] for v in responses_memory[METHOD].values()]
+        
+        # Update history images
         if not history_images:
-            history_images = current_images
+            history_images = current_images.copy() if current_images else []
         elif current_images:
             history_images.extend(current_images)
+        
         current_images = []
 
         examples_state = gr.update(samples=history_prompts, visible=True)
         prompt_state = gr.update(interactive=True)
         next_state = gr.update(visible=True, interactive=True)
         redesign_state = gr.update(interactive=False) if counter >= MAX_ROUND else gr.update(interactive=True)
-        submit_state = gr.update(interactive=True) if counter >= MAX_ROUND or enable_submit else gr.update(interactive=False)
-
+        
         counter += 1
+        redesign_flag = True
 
-        return None, None, None, current_images, history_images, examples_state, prompt_state, next_state, redesign_state, submit_state
-    else:
-        return {submit_btn: gr.skip()}
+        display_info_message(f"✅ Round {counter-1} feedback saved! You can continue redesigning or restart.")
 
-def save_response(prompt, sim_radio, like_radio, dislike_radio, like_image, dislike_image):
-    global counter, enable_submit, responses_memory
-
-    if check_evaluation(sim_radio):
-        # Save the final round entry
-        responses_memory[METHOD][counter] = {
-            "prompt": prompt,
-            "sim_radio": sim_radio,
-            "response": "",
-            "satisfied_img": f"round {counter}, {like_radio}",
-            "unsatisfied_img": f"round {counter}, {dislike_radio}",
-        }
-
-        # Reset states
-        counter = 1
-        enable_submit = False
-
-        # Reset buttons
-        prompt_state = gr.update(interactive=False)
-        next_state = gr.update(visible=False, interactive=False)
-        submit_state = gr.update(interactive=False)
-        redesign_state = gr.update(interactive=False)
-
-        display_info_message("✅ Your answer is saved!")
-        return None, None, None, prompt_state, next_state, redesign_state, submit_state
+        return None, current_images, history_images, examples_state, prompt_state, next_state, redesign_state
     else:
-        return {submit_btn: gr.skip()}
+        return gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip()
+
+def save_response(prompt, sim_radio, like_image, dislike_image):
+    global counter, responses_memory, redesign_flag, current_generated_images
+    
+    # Reset all global variables
+    responses_memory[METHOD] = {}
+    counter = 1
+    redesign_flag = False
+    current_generated_images = []
+    
+    # Reset UI states
+    prompt_state = gr.update(value="", interactive=True)
+    next_state = gr.update(visible=True, interactive=True)
+    redesign_state = gr.update(interactive=False)
+    submit_state = gr.update(interactive=False)
+    sim_radio_state = gr.update(value=None)
+    like_image_state = gr.update(value=None)
+    dislike_image_state = gr.update(value=None)
+    gallery_state = []
+    history_gallery_state = []
+    examples_state = gr.update(samples=[['']], visible=True)
+
+    display_info_message("🔄 Session restarted! You can begin with a new prompt.")
+
+    return (sim_radio_state, prompt_state, next_state, redesign_state, 
+            like_image_state, dislike_image_state, gallery_state, history_gallery_state, examples_state)
 
 # =========================
 # Interface (single tab, no participant/scenario/background)
@@ -256,6 +258,7 @@ css = """
 #button-container {
     display: flex;
     justify-content: center;
+    gap: 10px;
 }
 #compact-compact-row {
     width:100%;
@@ -315,9 +318,56 @@ css = """
     max-width: 150px;
     display: inline-block;
 }
+.instruction-box {
+    background: linear-gradient(135deg, #e8f4fd 0%, #f0f8ff 100%);
+    border: 2px solid #3498db;
+    border-radius: 12px;
+    padding: 20px;
+    margin: 15px 0;
+    color: #2c3e50;
+}
+.instruction-title {
+    font-size: 1.2em;
+    font-weight: bold;
+    margin-bottom: 15px;
+    color: #2c3e50;
+    display: flex;
+    align-items: center;
+    gap: 8px;
+}
+.step-list {
+    list-style: none;
+    padding: 0;
+    margin: 0;
+}
+.step-item {
+    background: rgba(52, 152, 219, 0.1);
+    border-radius: 8px;
+    padding: 12px 16px;
+    margin: 8px 0;
+    border-left: 4px solid #3498db;
+}
+.step-number {
+    font-weight: bold;
+    color: #3498db;
+    margin-right: 8px;
+}
+.personalization-header {
+    background: linear-gradient(135deg, #ff6b6b, #ee5a24);
+    color: white;
+    padding: 15px;
+    border-radius: 10px 10px 0 0;
+    margin: -10px -10px 15px -10px;
+    text-align: center;
+    font-weight: bold;
+    font-size: 1.1em;
+}
 """
 
 with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "Arial", "sans-serif"]), css=css) as demo:
+    # State variable to hold selected image
+    selected_image = gr.State(None)
+    
     with gr.Column(elem_id="col-container", elem_classes=["header-section"]):
         gr.HTML('<div class="logo-container"><img src="https://huggingface.co/spaces/PAI-GEN/POET/resolve/main/images/icon.png" alt="POET Logo"></div>')
         gr.Markdown("### Supporting Prompting Creativity with Automated Expansion of Text-to-Image Generation")
@@ -325,7 +375,7 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
         gr.HTML("""
         <div style="text-align: center;">
             <a href="https://arxiv.org/pdf/2504.13392" target="_blank" class="paper-link">
-                📄 Read the Full Paper .
+                📄 Read the Full Paper
             </a>
         </div>
         """)
@@ -337,13 +387,15 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
         
         gr.Markdown("""
         <div class="authors-section">
-            <a href="https://scholar.google.com/citations?user=HXED4kIAAAAJ&hl=en">Evans Han</a>, <a href"https://www.aliceqian.com/">Alice Qian Zhang</a>, 
-            <a href="https://haiyizhu.com/">Haiyi Zhu</a>, <a href="https://www.andrew.cmu.edu/user/hongs/">Hong Shen</a>, 
-            <a href="https://pliang279.github.io/">Paul Pu Liang</a>, <a href="https://janeon.github.io/">Jane Hsieh</a>
+            <a href="https://scholar.google.com/citations?user=HXED4kIAAAAJ&hl=en">Evans Han</a>, 
+            <a href="https://www.aliceqian.com/">Alice Qian Zhang</a>, 
+            <a href="https://haiyizhu.com/">Haiyi Zhu</a>, 
+            <a href="https://www.andrew.cmu.edu/user/hongs/">Hong Shen</a>, 
+            <a href="https://pliang279.github.io/">Paul Pu Liang</a>, 
+            <a href="https://janeon.github.io/">Jane Hsieh</a>
         </div>
         """, elem_classes=["authors-section"])
         
-        # gr.Markdown("---")
 
     with gr.Tab(""):
         with gr.Row(elem_id="compact-row"):
@@ -360,47 +412,99 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
 
         with gr.Row(elem_id="compact-row"):
             with gr.Column(elem_id="col-container"):
-                images_method = gr.Gallery(label="Images", columns=[4], rows=[1], height=400, elem_id="gallery", format="png")
+                images_method = gr.Gallery(
+                    label="Generated Images (Click to select, then set to Like/Dislike image)", 
+                    columns=[4], 
+                    rows=[1], 
+                    height=400, 
+                    interactive=False,
+                    elem_id="gallery", 
+                    format="png"
+                )
 
             with gr.Column(elem_id="col-container3"):
-                like_image = gr.Image(label="Satisfied Image", width=200, height=200, sources='upload', format="png", type="filepath", visible=False)
-                dislike_image = gr.Image(label="Unsatisfied Image", width=200, height=200, sources='upload', format="png", type="filepath", visible=False)
-
-        with gr.Column(elem_id="col-container2", visible=False):
-            gr.Markdown("### 📝 Evaluation")
-            sim_radio = gr.Radio(
-                OPTIONS,
-                label="How would you rate your satisfaction with the generated images?",
-                type="value",
-                elem_classes=["gradio-radio"]
-            )
-            like_radio = gr.Radio(
-                IMAGE_OPTIONS,
-                label="Select your all-time favorite image (optional).",
-                type="value",
-                elem_classes=["gradio-radio"]
-            )
-            dislike_radio = gr.Radio(
-                IMAGE_OPTIONS,
-                label="Select your all-time least satisfactory image (optional).",
-                type="value",
-                elem_classes=["gradio-radio"]
-            )
-
-        with gr.Column(elem_id="col-container2", visible=False):
-            example = gr.Examples([['']], prompt, label="Revised Prompt History", visible=False)
-            history_images = gr.Gallery(label="History Images", columns=[4], rows=[1], elem_id="gallery", format="png")
-
-            with gr.Row(elem_id="button-container"):
-                redesign_btn = gr.Button("🎨 Redesign", variant="primary", scale=0)
-                submit_btn = gr.Button("✅ Submit", variant="primary", interactive=False, scale=0)
+                like_btn = gr.Button("👍 Set as Liked (Optional for personalization)", size="sm", variant="secondary")
+                like_image = gr.Image(
+                    label="Satisfied Image", 
+                    width=150, 
+                    height=150, 
+                    interactive=False,
+                    format="png", 
+                    type="filepath"
+                )
+                dislike_btn = gr.Button("👎 Set as Disliked (Optional for personalization)", size="sm", variant="secondary")
+                dislike_image = gr.Image(
+                    label="Unsatisfied Image", 
+                    width=150, 
+                    height=150, 
+                    interactive=False,
+                    format="png", 
+                    type="filepath"
+                )
+                
+        with gr.Accordion("🎯 Advanced: Personalized Image Redesign", open=False, elem_id="col-container2"):
+            gr.HTML("""
+            <div class="instruction-box">
+                <div class="instruction-title">
+                    📋 How to Use Personalized Redesign
+                </div>
+                <div class="step-list">
+                    <div class="step-item">
+                        <span class="step-number">1.</span>
+                        <strong>Rate Your Satisfaction:</strong> Provide a satisfaction score for the current generated images
+                    </div>
+                    <div class="step-item">
+                        <span class="step-number">2.</span>
+                        <strong>Select Preferences:</strong> Choose your most liked and disliked images
+                    </div>
+                    <div class="step-item">
+                        <span class="step-number">3.</span>
+                        <strong>Save & Iterate:</strong> Click "Save Personalized Data" before redesgining your prompt and clicking "Generate" 
+                    </div>
+                    <div class="step-item">
+                        <span class="step-number">4.</span>
+                        <strong>Restart Anytime:</strong> Use the "Restart" button to begin a fresh session
+                    </div>
+                </div>
+            </div>
+            """)
+
+            with gr.Column(elem_id="col-container2"):
+                gr.Markdown("### 📊 Rate Current Generation")
+                with gr.Row():
+                    sim_radio = gr.Radio(
+                        OPTIONS,
+                        label="How satisfied are you with the current generated images?",
+                        type="value",
+                        show_label=True,
+                        container=True,
+                        scale=1
+                    )
+                
+                with gr.Row(elem_id="button-container"):
+                    with gr.Column(scale=1):
+                        redesign_btn = gr.Button("💾 Save Personalization Data", variant="primary", size="lg")
+                    with gr.Column(scale=1):
+                        submit_btn = gr.Button("🔄 Restart Session", variant="secondary", size="lg")
+
+
+            with gr.Column(elem_id="col-container2"):
+                example = gr.Examples([['']], prompt, label="📝 Prompt History", visible=True)
+                history_images = gr.Gallery(
+                    label="🗃️ Generation History", 
+                    columns=[4], 
+                    rows=[1], 
+                    elem_id="gallery", 
+                    format="png", 
+                    interactive=False,
+                )
 
         with gr.Column(elem_id="col-container2"):
             gr.Markdown("### 🌟 Examples")
-            ex1 = gr.Image(label="Image 1", width=200, height=200, sources='upload', format="png", type="filepath", visible=False)
-            ex2 = gr.Image(label="Image 2", width=200, height=200, sources='upload', format="png", type="filepath", visible=False)
-            ex3 = gr.Image(label="Image 3", width=200, height=200, sources='upload', format="png", type="filepath", visible=False)
-            ex4 = gr.Image(label="Image 4", width=200, height=200, sources='upload', format="png", type="filepath", visible=False)
+            ex1 = gr.Image(label="Image 1", width=200, height=200, format="png", type="filepath", visible=False)
+            ex2 = gr.Image(label="Image 2", width=200, height=200, format="png", type="filepath", visible=False)
+            ex3 = gr.Image(label="Image 3", width=200, height=200, format="png", type="filepath", visible=False)
+            ex4 = gr.Image(label="Image 4", width=200, height=200, format="png", type="filepath", visible=False)
 
             gr.Examples(
                 examples=[[ex[0], ex[1][0], ex[1][1], ex[1][2], ex[1][3]] for ex in example_data],
@@ -410,28 +514,44 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
 # =========================
 # Wiring
 # =========================
-    sim_radio.change(fn=check_satisfaction, inputs=[sim_radio], outputs=[submit_btn])
-
-    dislike_radio.select(fn=select_image, inputs=[dislike_radio, images_method], outputs=[dislike_image])
-    like_radio.select(fn=select_image, inputs=[like_radio, images_method], outputs=[like_image])
+    # Handle gallery selection
+    images_method.select(
+        fn=on_gallery_select,
+        inputs=[],
+        outputs=[selected_image]
+    )
+    
+    # Handle like/dislike button clicks
+    like_btn.click(
+        fn=handle_like_drag,
+        inputs=[selected_image],
+        outputs=[like_image]
+    )
+    
+    dislike_btn.click(
+        fn=handle_dislike_drag,
+        inputs=[selected_image],
+        outputs=[dislike_image]
+    )
 
     next_btn.click(
         fn=generate_image,
         inputs=[prompt, like_image, dislike_image],
         outputs=[images_method]
-    ).success(lambda: [gr.update(interactive=True), gr.update(interactive=True)], outputs=[next_btn, prompt])
+    ).success(lambda: [gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True)], 
+    outputs=[next_btn, prompt, redesign_btn, submit_btn])
 
     redesign_btn.click(
         fn=redesign,
-        inputs=[prompt, sim_radio, like_radio, dislike_radio, images_method, history_images, like_image, dislike_image],
-        outputs=[sim_radio, dislike_radio, like_radio, images_method, history_images, example.dataset, prompt, next_btn, redesign_btn, submit_btn]
+        inputs=[prompt, sim_radio, images_method, history_images, like_image, dislike_image],
+        outputs=[sim_radio, images_method, history_images, example.dataset, prompt, next_btn, redesign_btn]
     )
 
     submit_btn.click(
         fn=save_response,
-        inputs=[prompt, sim_radio, like_radio, dislike_radio, like_image, dislike_image],
-        outputs=[sim_radio, dislike_radio, like_radio, prompt, next_btn, redesign_btn, submit_btn]
+        inputs=[prompt, sim_radio, like_image, dislike_image],
+        outputs=[sim_radio, prompt, next_btn, redesign_btn, like_image, dislike_image, images_method, history_images, example.dataset]
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

utils.py

@@ -171,7 +171,37 @@ def get_personalize_message(prompt, history_prompts, history_feedback, like_imag
                         "url": f"data:image/png;base64,{dislike_image_base64}",
                     },
                 })
-    print(messages)
+    return messages
+
+def get_personalized_simplified(prompt, like_image, dislike_image):
+    messages = [
+    {"role": "system", "content": f"You are a prompt refinement assistant. Your task is to improve a user’s text prompt based on his liked and disliked images. Your goal is to refine the prompt while maintaining its original meaning, improving clarity, specificity, and alignment with user preferences."}
+]
+
+    message = f"""The first given image is user's liked image, refine prompt with style, and content user likes. The second given image is user's disliked image, refine prompt to avoid those elements or style of this image."""
+    
+    messages.append({
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"{message}"},
+            ],
+        })
+    if like_image:
+        like_image_base64 = encode_image(like_image)
+        messages[-1]["content"].append({
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/png;base64,{like_image_base64}",
+                    },
+                })
+    if dislike_image:
+        dislike_image_base64 = encode_image(dislike_image)
+        messages[-1]["content"].append({
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/png;base64,{dislike_image_base64}",
+                    },
+                })
     return messages
 
 @spaces.GPU