update policy

app.py

@@ -7,7 +7,7 @@ import torch
 import re
 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
+from utils import clean_response_gpt, setup_model, init_gpt_api, call_gpt_api, get_refine_msg, clean_cache, get_personalize_message, clean_refined_prompt_response_gpt
 from utils import SCENARIOS, PROMPTS, IMAGES, OPTIONS, T2I_MODELS, INSTRUCTION, IMAGE_OPTIONS
 import spaces #[uncomment to use ZeroGPU]
 import transformers
@@ -108,7 +108,6 @@ def personalize_prompt(prompt, history, feedback, like_image, dislike_image):
     client = init_gpt_api()
     messages = get_personalize_message(prompt, history, feedback, like_image, dislike_image)
     outputs = call_gpt_api(messages, client, "gpt-4o", seed, max_tokens=2000, temperature=0.7, top_p=0.9)
-    print(outputs)
     # prompt_list = clean_response_gpt(outputs)
     # print(prompt_list)
     return outputs
@@ -203,6 +202,10 @@ def display_scenario(participant, choice):
         prompt2: gr.update(value=PROMPTS.get(choice, ""), interactive=False),
         images_method1: initial_images1, 
         images_method2: initial_images2,
+        like_image1: None,
+        dislike_image1: None,
+        like_image2: None,
+        dislike_image2: None,
         history_images1: [],
         history_images2: [],
         next_btn1: gr.update(interactive=False), 
@@ -221,9 +224,16 @@ def generate_image(participant, scenario, prompt, active_tab, like_image, dislik
     
     history_prompts = [v["prompt"] for v in responses_memory[participant][method].values()]
     feedback = [v["sim_radio"] for v in responses_memory[participant][method].values()]
-    print(history_prompts, feedback)
+
     personalized_prompt = personalize_prompt(prompt, history_prompts, feedback, like_image, dislike_image)
     
+    personalized_prompt = clean_refined_prompt_response_gpt(personalized_prompt)
+    print(f"Personalized prompt: {personalized_prompt}, {type(personalized_prompt)}")
+    
+    if "I'm sorry, I can't assist with" in personalized_prompt:
+        print("error in gpt...")
+        personalized_prompt = prompt
+
     gallery_images = []
     if method == METHODS[0]:
         for i in range(NUM_IMAGES): 
@@ -491,7 +501,8 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
     participant.change(fn=set_user, inputs=[participant], outputs=[scenario])
     scenario.change(display_scenario, 
         inputs=[participant, scenario], 
-        outputs=[scenario_content, prompt1, prompt2, images_method1, images_method2, history_images1, history_images2, next_btn1, next_btn2, redesign_btn1, redesign_btn2, submit_btn1, submit_btn2])
+        outputs=[scenario_content, prompt1, prompt2, images_method1, images_method2, like_image1, dislike_image1, like_image2, dislike_image2, history_images1, history_images2, next_btn1, next_btn2, redesign_btn1, redesign_btn2, submit_btn1, submit_btn2])
+    
     # prompt1.change(fn=reset_gallery, inputs=[], outputs=[gallery_state1])
     # prompt2.change(fn=reset_gallery, inputs=[], outputs=[gallery_state2])
     next_btn1.click(fn=generate_image, inputs=[participant, scenario, prompt1, active_tab, like_image1, dislike_image1], outputs=[images_method1])

utils.py

@@ -78,6 +78,16 @@ def clean_response_gpt(res: str):
     return prompts
 
 
+def clean_refined_prompt_response_gpt(res: str):
+    # Using regex to extract the refined prompt
+    match = re.search(r"\*\*Refined Prompt:\*\*\n\n(.+)", res, re.DOTALL)
+    if match:
+        refined_prompt = match.group(1).strip()
+    else:
+        refined_prompt = res.strip()  # Fallback: Use full text if no match found
+    return refined_prompt
+
+
 def get_refine_msg(prompt, num_prompts):
     messages = [{"role": "system", "content": f"You are a helpful, respectful and precise assistant. You will be asked to generate {num_prompts} refined prompts. Only respond with those refined prompts"}]
 
@@ -108,37 +118,37 @@ def encode_image(image):
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 def get_personalize_message(prompt, history_prompts, history_feedback, like_image, dislike_image):
-    messages = [{"role": "system", "content": f"You will act as a prompt optimization assistant that helps refine an original prompt based on user feedback over multiple rounds of image generation. The goal is to dynamically adjust the prompt to better align with user preferences while preserving the original intent."}]
-
-    message = f"""The process consists of a maximum of 5 rounds. 
-    Users start with an initial prompt and generate 4 images per round. After reviewing the images, users will modify the prompt based on their preferences. Then we will generate new images based on the modified prompt. 
-    Users will rate the generated images on a scale from ["Very Unsatisfied", "Unsatisfied", "Slightly Unsatisfied", "Neutral", "Slightly Satisfied", "Satisfied", "Very Satisfied"], indicating how satisfied they are with the results.
-    
-    Your task is to analyze the sequence of modified prompts and corresponding ratings to refine the prompt dynamically, ensuring improved results in the next rounds. For each new round, you should:
-    Incorporate the user's modifications: Use the latest user-revised prompt as a reference but retain essential details from previous rounds if they contributed positively.
-    
-    Analyze user ratings:
-    If the rating is high ("Satisfied", "Very Satisfied") → Maintain key aspects of the most recent prompt since it aligns well with user preferences.
-    If the rating is medium ("Slightly Unsatisfied", "Neutral", "Slightly Satisfied") → Adjust minor details that could improve alignment with the user’s preferences, considering the changes from previous rounds.
-    If the rating is low ("Very Unsatisfied", "Unsatisfied") → Identify aspects that might be causing dissatisfaction (e.g., unwanted elements, style mismatch) and rework the prompt while keeping the user’s core intent intact.
-
-    Refine the prompt intelligently and ensure the following: 
-    - The updated prompt reflects user feedback without unnecessary repetition.
-    - Unwanted elements (if any) from previous rounds are removed.
-    - Preferred elements are retained and enhanced.
-    - The modifications remain subtle but progressive to ensure smooth refinement over multiple rounds.
-    - Maintain coherence: Avoid drastic changes that might deviate from the original intent unless the user explicitly requests them.
-    
-
-    Now given the following revised prompts and ratings from user\n:
-"""
-    
+    messages = [
+    {"role": "system", "content": f"You are a prompt refinement assistant. Your task is to improve a user’s text prompt based on their prompt revision history, satisfaction ratings, and preferences inferred from selected images. Your goal is to refine the prompt while maintaining its original meaning, improving clarity, specificity, and alignment with user preferences."}
+]
+
+    message = f"""The refinement should preserve the core meaning of the current prompt while improving its clarity, specificity, and style based on user feedback.
+
+    ### **Input Format:**
+    1. **Prompt History**: A list of previously revised prompts and their corresponding satisfaction ratings.
+    2. **Rating Scale**: Very Unsatisfied, Unsatisfied, Slightly Unsatisfied, Neutral, Slightly Satisfied, Satisfied, Very Satisfied
+    3. **User-Selected Image Preferences**:
+    - **Preferred Image**: The image the user found most satisfactory.
+    - **Disliked Image**: The image the user found least satisfactory.
+    *Note: These images are for reference only and should be used to infer stylistic preferences rather than directly modifying prompt content.*
+    4. **Current Prompt**: The latest prompt from the user, which requires refinement.
+
+    ### **Refinement Guidelines:**
+    - Identify and retain/expand patterns/elements in past revisions and correlate them with satisfaction ratings.
+    - Avoid or adjust features that led to lower ratings.
+    - Improve clarity, specificity, and descriptive quality while ensuring the prompt remains faithful to its current prompt's meaning.
+    - The preferred image reflects desirable attributes; the disliked image indicates elements to avoid. Use these for reference but **do not describe them.**
+    - Output only the refined prompt, no explanations, disclaimers, or formatting.
+
+    The first provided image is the user's preferred image, and the second is the disliked image.
+    Now, refine the following current prompt based on the given user history and preferences:\n"""
+
+    message += "Prompt History\n"
     for his_prompt, feedback in zip(history_prompts, history_feedback):
-        message += f"Revised prompt: {his_prompt}; Rating: {feedback}\n"
+        message += f"{his_prompt}: {feedback}\n"
+    
+    message += f"Current Prompt: '{prompt}'\n Refined Prompt:"
     
-    message += f"\nWe also provide the user's preferred image during this process as the first image provided and the disliked image as the second image\n"
-    message += "Now, please optimize current prompt and only output the modified prompt: '{prompt}'"""
-
     messages.append({
             "role": "user",
             "content": [
@@ -150,7 +160,7 @@ def get_personalize_message(prompt, history_prompts, history_feedback, like_imag
         messages[-1]["content"].append({
                     "type": "image_url",
                     "image_url": {
-                        "url": f"data:image/jpeg;base64,{like_image_base64}",
+                        "url": f"data:image/png;base64,{like_image_base64}",
                     },
                 })
     if dislike_image:
@@ -158,11 +168,9 @@ def get_personalize_message(prompt, history_prompts, history_feedback, like_imag
         messages[-1]["content"].append({
                     "type": "image_url",
                     "image_url": {
-                        "url": f"data:image/jpeg;base64,{dislike_image_base64}",
+                        "url": f"data:image/png;base64,{dislike_image_base64}",
                     },
                 })
-        
-    print(messages)
     
     return messages