update refine policy

app.py

@@ -7,13 +7,11 @@ 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
+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 SCENARIOS, PROMPTS, IMAGES, OPTIONS, T2I_MODELS, INSTRUCTION, IMAGE_OPTIONS
 import spaces #[uncomment to use ZeroGPU]
 import transformers
 import gspread
-import asyncio
-from datetime import datetime
 
 CLIP_MODEL = "ViT-H-14"
 PRETRAINED_CLIP = "laion2b_s32b_b79k"
@@ -33,7 +31,7 @@ llm_pipe = None
 torch.cuda.empty_cache()
 inverted_prompt = ""
 
-VERBAL_MSG = "Please verbally describe key differences found in the image pair."
+VERBAL_MSG = "Please verbally describe why you are satisfied or unsatisfied at the generated images."
 DEFAULT_SCENARIO = "Product advertisement"
 METHODS = ["Method 1", "Method 2"]
 MAX_ROUND = 5
@@ -78,37 +76,6 @@ def infer(
 
     return image
 
-async def infer_async(prompt):
-    return infer(prompt)
-# generate a batch of images in parallel
-async def generate_batch(prompts):
-    tasks = [infer_async(p) for p in prompts]
-    images = await asyncio.gather(*tasks)  # Run all in parallel
-    return images
-
-@spaces.GPU 
-def call_llm_refine_prompt(prompt, num_prompts=5, max_tokens=1000, temperature=0.7, top_p=0.9):
-    print(f"loading {default_llm_model}")
-    global llm_pipe
-    if not llm_pipe:
-        llm_pipe = transformers.pipeline("text-generation", model=default_llm_model, model_kwargs={"torch_dtype": torch_dtype}, device_map="auto")
-    
-    messages = get_refine_msg(prmpt, num_prompts)
-    terminators = [
-        llm_pipe.tokenizer.eos_token_id,
-        llm_pipe.tokenizer.convert_tokens_to_ids("<|eot_id|>")
-    ]
-    outputs = llm_pipe(
-        messages,
-        max_new_tokens=max_tokens,
-        eos_token_id=terminators,
-        do_sample=True,
-        temperature=temperature,
-        top_p=top_p,
-    )
-    prompt_list = clean_response_gpt(outputs[0]["generated_text"][-1]["content"])
-    return prompt_list
-
 def call_gpt_refine_prompt(prompt, num_prompts=5, max_tokens=1000, temperature=0.7, top_p=0.9):
     seed = random.randint(0, MAX_SEED)
     client = init_gpt_api()
@@ -117,12 +84,6 @@ def call_gpt_refine_prompt(prompt, num_prompts=5, max_tokens=1000, temperature=0
     prompt_list = clean_response_gpt(outputs)
     return prompt_list
 
-def refine_prompt(gallery_state, prompt):
-    modified_prompts = call_gpt_refine_prompt(prompt)
-    return modified_prompts
-    
-    # eval(prompt, inverted_prompt, gallery_state, clip_model, preprocess)
-
 @spaces.GPU(duration=100)
 def invert_prompt(prompt, images, prompt_len=15, iter=1000, lr=0.1, batch_size=2):
     text_params = {
@@ -142,25 +103,15 @@ def invert_prompt(prompt, images, prompt_len=15, iter=1000, lr=0.1, batch_size=2
     # eval(prompt, learned_prompt, optimized_images, clip_model, preprocess)
     # return learned_prompt
 
-
-def eval(prompt, optimized_prompt, optimized_images, clip_model, preprocess):
-    torch.cuda.empty_cache()
-    tokenizer = open_clip.get_tokenizer(CLIP_MODEL)
-    images = [preprocess(i).unsqueeze(0) for i in optimized_images]
-    images = torch.concatenate(images).to(device)
-    
-    with torch.no_grad():
-        image_feat = clip_model.encode_image(images)
-        text_feat = clip_model.encode_text(tokenizer([prompt]).to(device))
-        optimized_text_feat = clip_model.encode_text(tokenizer([optimized_prompt]).to(device))
-
-    image_feat /= image_feat.norm(dim=-1, keepdim=True)
-    text_feat /= text_feat.norm(dim=-1, keepdim=True)
-    optimized_text_feat /= optimized_text_feat.norm(dim=-1, keepdim=True)
-
-    similarity = text_feat.cpu().numpy() @ image_feat.cpu().numpy().T
-    similarity_optimized = optimized_text_feat.cpu().numpy() @ image_feat.cpu().numpy().T
-
+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)
+    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
 
 ########################################################################################################
 # Button-related functions
@@ -182,8 +133,43 @@ def switch_tab(active_tab):
     else:
         return gr.Tabs(selected="Task A")
 
+def check_satisfaction(sim_radio, active_tab):
+    global counter1, counter2, current_task1, current_task2
+    method = current_task1 if active_tab == "Task A" else current_task2
+    counter = counter1 if method == METHODS[0] else counter2
+
+    fully_satisfied_option = ["Satisfied", "Very Satisfied"]  # The value to trigger submit
+    enable_submit = sim_radio in fully_satisfied_option or counter >= MAX_ROUND
+   
+    return gr.update(interactive=enable_submit), gr.update(interactive=(not enable_submit)) 
+
+def check_participant(participant):
+    if participant == "":
+        display_error_message("Please fill your participant id!")
+        return False
+    return True
+
+def check_evaluation(sim_radio):
+    if not sim_radio :
+        display_error_message("❌ Please fill all evaluations before change image or submit.")
+        return False
+    
+    return True
+
+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 set_user(participant):
-    global responses_memory
+    global responses_memory, assigned_scenarios
     responses_memory[participant] = {METHODS[0]:{}, METHODS[1]:{}}
 
     id = re.findall(r'\d+', participant)
@@ -213,25 +199,12 @@ def display_scenario(participant, choice):
     
     res = { 
         scenario_content: SCENARIOS.get(choice, ""), 
-        prompt: PROMPTS.get(choice, ""),
-        prompt1: "", 
-        prompt2: "",
+        prompt1: gr.update(value=PROMPTS.get(choice, ""), interactive=False),
+        prompt2: gr.update(value=PROMPTS.get(choice, ""), interactive=False),
         images_method1: initial_images1, 
         images_method2: initial_images2,
-        gallery_state1: initial_images1, 
-        gallery_state2: initial_images2, 
-        sim_radio1: None, 
-        sim_radio2: None, 
-        dislike_radio1: None,
-        like_radio1: None,
-        dislike_radio2: None,
-        like_radio2: None,
-        like_image1: None,
-        dislike_image1: None,
-        like_image2: None,
-        dislike_image2: None,
-        response1: VERBAL_MSG, 
-        response2: VERBAL_MSG, 
+        history_images1: [],
+        history_images2: [],
         next_btn1: gr.update(interactive=False), 
         next_btn2: gr.update(interactive=False), 
         redesign_btn1: gr.update(interactive=True), 
@@ -241,64 +214,34 @@ def display_scenario(participant, choice):
     }
     return res
 
-def generate_image(participant, scenario, prompt, gallery_state, active_tab):
+def generate_image(participant, scenario, prompt, active_tab, like_image, dislike_image):
     if not check_participant(participant): return [], []
     global current_task1, current_task2
-    
     method = current_task1 if active_tab == "Task A" else current_task2
     
+    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)
+    
+    gallery_images = []
     if method == METHODS[0]:
         for i in range(NUM_IMAGES): 
-            img = infer(prompt)
-            gallery_state.append(img)
-            yield gallery_state
+            img = infer(personalized_prompt)
+            gallery_images.append(img)
+            yield gallery_images
     else:
-        refined_prompts = refine_prompt(gallery_state, prompt)
+        refined_prompts = call_gpt_refine_prompt(personalized_prompt)
         for i in range(NUM_IMAGES): 
             img = infer(refined_prompts[i])
-            gallery_state.append(img)
-            yield gallery_state
-
-def check_satisfaction(sim_radio, active_tab):
-    global counter1, counter2, current_task1, current_task2
-    method = current_task1 if active_tab == "Task A" else current_task2
-    counter = counter1 if method == METHODS[0] else counter2
-
-    fully_satisfied_option = ["Satisfied", "Very Satisfied"]  # The value to trigger submit
-    enable_submit = sim_radio in fully_satisfied_option or counter >= MAX_ROUND
-   
-    return gr.update(interactive=enable_submit), gr.update(interactive=(not enable_submit)) 
-
-def check_participant(participant):
-    if participant == "":
-        display_error_message("Please fill your participant id!")
-        return False
-    return True
-
-def check_evaluation(sim_radio, response):
-    if not sim_radio :
-        display_error_message("❌ Please fill all evaluations before change image or submit.")
-        return False
-    
-    return True
-
-def select_dislike(like_radio, images_method):
-    if like_radio == IMAGE_OPTIONS[0]:
-        return images_method[0]
-    elif like_radio == IMAGE_OPTIONS[1]:
-        return images_method[1]
-    elif like_radio == IMAGE_OPTIONS[2]:
-        return images_method[2]
-    elif like_radio == IMAGE_OPTIONS[3]:
-        return images_method[3]
-    else:
-        return None
+            gallery_images.append(img)
+            yield gallery_images
     
-def redesign(participant, scenario, prompt, sim_radio, response, images_method, active_tab):
+def redesign(participant, scenario, prompt, sim_radio, current_images, history_images, active_tab):
     global counter1, counter2, responses_memory, current_task1, current_task2
     method = current_task1 if active_tab == "Task A" else current_task2
 
-    if check_evaluation(sim_radio, response) and check_participant(participant):
+    if check_evaluation(sim_radio) and check_participant(participant):
         if method == METHODS[0]:
             counter1 += 1
             counter = counter1
@@ -309,62 +252,68 @@ def redesign(participant, scenario, prompt, sim_radio, response, images_method,
         responses_memory[participant][method][counter-1] = {}
         responses_memory[participant][method][counter-1]["prompt"] = prompt
         responses_memory[participant][method][counter-1]["sim_radio"] = sim_radio
-        responses_memory[participant][method][counter-1]["response"] = response
-
-        prompt_state = gr.update(visible=True)
+        # responses_memory[participant][method][counter-1]["response"] = response
+        
+        history_prompts = [[v["prompt"]] for v in responses_memory[participant][method].values()]
+        if not history_images: 
+            history_images = current_images
+        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(interactive=False) if counter >= MAX_ROUND else 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 else gr.update(interactive=False)
 
-        return [], None, VERBAL_MSG, prompt_state, next_state, redesign_state, submit_state
+        return None, None, None, current_images, history_images, examples_state, prompt_state, next_state, redesign_state, submit_state
     else:
         return {submit_btn1: gr.skip()} if active_tab == "Task A" else {submit_btn2: gr.skip()}
 
-def show_message(selected_option):
-    if selected_option:
-        return "Click \"Redesign\" and revise your prompt to create images that may more closely match your expectations."
-    return ""  
-  
-def save_response(participant, scenario, prompt, sim_radio, response, images_method, active_tab):
-    global current_task1, current_task2, counter1, counter2, responses_memory, task1_success, task2_success, assigned_scenarios
+
+def save_response(participant, scenario, prompt, sim_radio, active_tab):
+    global current_task1, current_task2 # not change
+    global task1_success, task2_success, counter1, counter2, responses_memory, assigned_scenarios # will change
+
     method = current_task1 if active_tab == "Task A" else current_task2
-    
-    if check_evaluation(sim_radio, response) and check_participant(participant):
+    if check_evaluation(sim_radio) and check_participant(participant):
         counter = counter1 if method == METHODS[0] else counter2
-        # image_paths = [save_image(img, "method", i) for i, img in enumerate(images_method)]
 
         responses_memory[participant][method][counter] = {}
         responses_memory[participant][method][counter]["prompt"] = prompt
         responses_memory[participant][method][counter]["sim_radio"] = sim_radio
-        responses_memory[participant][method][counter]["response"] = response
-        prompt_state = gr.update(visible=False)
-        next_state = gr.update(visible=False, interactive=False)
-        submit_state = gr.update(interactive=False) 
-        redesign_state = gr.update(interactive=False) 
-
+        # responses_memory[participant][method][counter]["response"] = response
+        
         try:
             gc = gspread.service_account(filename='credentials.json')
             sheet = gc.open("DiverseGen-phase3").sheet1 
 
             for i, entry in responses_memory[participant][method].items():
-                sheet.append_row([participant, scenario, method, i, entry["prompt"], entry["sim_radio"],  entry["response"]])
+                sheet.append_row([participant, scenario, method, i, entry["prompt"], entry["sim_radio"]])
 
             display_info_message("✅ Your answer is saved!")
 
-            # reset counter and update success indicator
+            # reset global variables
             if method == METHODS[0]:
                 counter1 = 1
             else:
                 counter2 = 1
-
             if active_tab == "Task A":
                 task1_success = True
             else:
                 task2_success = True
-
-            tabs = switch_tab(active_tab)
+            # decide if change scenario
             next_scenario = assigned_scenarios[1] if task1_success and task2_success else assigned_scenarios[0]
-            return [], [], None, None, None, None, None, VERBAL_MSG, prompt_state, next_state, redesign_state, submit_state, tabs, next_scenario
+            # update buttons
+            example_state = gr.update(samples=[], visible=False)
+            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) 
+            tabs = switch_tab(active_tab)
+
+            return None, None, None, None, None, [], [], example_state, prompt_state, next_state, redesign_state, submit_state, next_scenario, tabs
+        
         except Exception as e:
             display_error_message(f"❌ Error saving response: {str(e)}")
             return {submit_btn1: gr.skip()} if active_tab == "Task A" else {submit_btn2: gr.skip()}
@@ -388,7 +337,7 @@ css="""
 }
 
 #col-container3 {
-    margin: 0 auto;
+    margin: 0 0 auto auto;
     max-width: 300px;
 }
 
@@ -413,7 +362,6 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
             )
             scenario = gr.Dropdown(
                 choices=list(SCENARIOS.keys()),
-                # value=DEFAULT_SCENARIO,
                 value=None,
                 label="📌 Scenario",
                 interactive=False,
@@ -421,13 +369,6 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
         scenario_content = gr.Textbox(
             label="📖 Background", 
             interactive=False, 
-            # value=SCENARIOS[DEFAULT_SCENARIO]
-        )
-        prompt = gr.Textbox(
-                label="🎨 Prompt",
-                max_lines=1,
-                # value=PROMPTS[DEFAULT_SCENARIO],
-                interactive=False
         )
         active_tab = gr.State("Task A")
         instruction = gr.Markdown(INSTRUCTION)
@@ -435,26 +376,25 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
     with gr.Tabs() as tabs:
         with gr.TabItem("Task A", id="Task A") as task1_tab:
             task1_tab.select(lambda: "Task A", outputs=[active_tab])
-            with gr.Column(elem_id="col-container"):        
-                # gr.Markdown("### Step 2: This is the prompt to generate images, you may modify the prompt after first round evaluation")
-                with gr.Row():
-                    prompt1 = gr.Textbox(
-                            label="🎨 Revise Prompt",
-                            max_lines=1,
-                            placeholder="Enter your prompt",
-                            # value=PROMPTS[DEFAULT_SCENARIO],
-                            scale=4, 
-                            visible=False
-                    )
-                    next_btn1 = gr.Button("Generate", variant="primary", scale=1, interactive=False, visible=False)
-
+            with gr.Row(elem_id="compact-row"):
+                prompt1 = gr.Textbox(
+                        label="🎨 Revise Prompt",
+                        max_lines=5,
+                        placeholder="Enter your prompt",
+                        scale=4, 
+                        visible=True,
+                )
+                next_btn1 = gr.Button("Generate", variant="primary", scale=1, interactive=False, visible=False)
+            with gr.Row(elem_id="compact-row"):
+                example1 = gr.Examples([['']], prompt1, label="Revised Prompt History", visible=False)
+        
             with gr.Row(elem_id="compact-row"):
                 with gr.Column(elem_id="col-container"):
-                    gallery_state1 = gr.State([])
-                    images_method1 = gr.Gallery(show_label=False, columns=[4], rows=[1], height=420, elem_id="gallery")
+                    images_method1 = gr.Gallery(label="Images", columns=[4], rows=[1], height=200, elem_id="gallery")
+                    history_images1 = gr.Gallery(label="History Images", columns=[4], rows=[1], elem_id="gallery")
                 with gr.Column(elem_id="col-container3"):
-                    like_image1 = gr.Image(label="Satisfied Image", width=200, height=200, sources='upload')
-                    dislike_image1 = gr.Image(label="Unsatisfied Image", width=200, height=200, sources='upload')
+                    like_image1 = gr.Image(label="Satisfied Image", width=200, height=200, sources='upload', type="pil")
+                    dislike_image1 = gr.Image(label="Unsatisfied Image", width=200, height=200, sources='upload', type="pil")
             with gr.Column(elem_id="col-container2"):
                 gr.Markdown("### 📝 Evaluation")               
                 sim_radio1 = gr.Radio(
@@ -465,13 +405,13 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
                 )
                 like_radio1 = gr.Radio(
                     IMAGE_OPTIONS, 
-                    label="Select the image you are most satisfied.",
+                    label="Select the image that you find MOST satisfactory. You may leave this section blank if you prefer the previous images.",
                     type="value",
                     elem_classes=["gradio-radio"]
                 )
                 dislike_radio1 = gr.Radio(
                     IMAGE_OPTIONS, 
-                    label="Select the image you are most unsatisfied.",
+                    label="Please choose the image that you find LEAST satisfactory. You may leave this section blank if you are more dislike previous images.",
                     type="value",
                     elem_classes=["gradio-radio"]
                 )
@@ -491,26 +431,25 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
 
         with gr.TabItem("Task B", id="Task B") as task2_tab:
             task2_tab.select(lambda: "Task B", outputs=[active_tab])
-            with gr.Column(elem_id="col-container"):        
-                # gr.Markdown("### Step 2: This is the prompt to generate images, you may modify the prompt after first round evaluation")
-                with gr.Row():
-                    prompt2 = gr.Textbox(
-                            label="🎨 Revise Prompt",
-                            max_lines=1,
-                            placeholder="Enter your prompt",
-                            # value=PROMPTS[DEFAULT_SCENARIO],
-                            scale=4,
-                            visible=False
-                    )
-                    next_btn2 = gr.Button("Generate", variant="primary", scale=1, interactive=False, visible=False)
+            with gr.Row(elem_id="compact-row"):
+                prompt2 = gr.Textbox(
+                        label="🎨 Revise Prompt",
+                        max_lines=5,
+                        placeholder="Enter your prompt",
+                        scale=4,
+                        visible=True,
+                )
+                next_btn2 = gr.Button("Generate", variant="primary", scale=1, interactive=False, visible=False)
+            with gr.Row(elem_id="compact-row"):
+                example2 = gr.Examples([['']], prompt2, label="Revised Prompt History", visible=False)
 
             with gr.Row(elem_id="compact-row"):
                 with gr.Column(elem_id="col-container"):
-                    gallery_state2 = gr.State(IMAGES[DEFAULT_SCENARIO]["ours"])
-                    images_method2 = gr.Gallery(height=420, show_label=False, columns=[4], rows=[1], elem_id="gallery")
+                    images_method2 = gr.Gallery(label="Images", columns=[4], rows=[1], height=200, elem_id="gallery")
+                    history_images2 = gr.Gallery(label="History Images", columns=[4], rows=[1], elem_id="gallery")
                 with gr.Column(elem_id="col-container3"):
-                    like_image2 = gr.Image(label="Satisfied Image", width=200, height=200, sources='upload')
-                    dislike_image2 = gr.Image(label="Unsatisfied Image", width=200, height=200, sources='upload')
+                    like_image2 = gr.Image(label="Satisfied Image", width=200, height=200, sources='upload', type="pil")
+                    dislike_image2 = gr.Image(label="Unsatisfied Image", width=200, height=200, sources='upload', type="pil")
 
             with gr.Column(elem_id="col-container2"):
                 gr.Markdown("### 📝 Evaluation")
@@ -522,13 +461,13 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
                 )
                 like_radio2 = gr.Radio(
                     IMAGE_OPTIONS, 
-                    label="Select the image you are most satisfied.",
+                    label="Select the image that you find MOST satisfactory. You may leave this section blank if you prefer the previous images.",
                     type="value",
                     elem_classes=["gradio-radio"]
                 )
                 dislike_radio2 = gr.Radio(
                     IMAGE_OPTIONS, 
-                    label="Select the image you are most unsatisfied.",
+                    label="Please choose the image that you find LEAST satisfactory. You may leave this section blank if you are more dislike previous images.",
                     type="value",
                     elem_classes=["gradio-radio"]
                 )
@@ -550,35 +489,37 @@ 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, prompt, prompt1, prompt2, images_method1, images_method2, gallery_state1, gallery_state2, sim_radio1, sim_radio2, dislike_radio1, like_radio1, dislike_radio2, like_radio2, like_image1, dislike_image1, like_image2, dislike_image2, response1, response2, 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, gallery_state1, active_tab], outputs=[images_method1])
-    next_btn2.click(fn=generate_image, inputs=[participant, scenario, prompt2, gallery_state2, active_tab], outputs=[images_method2])
+    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])
+    # 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])
+    next_btn2.click(fn=generate_image, inputs=[participant, scenario, prompt2, active_tab, like_image2, dislike_image2], outputs=[images_method2])
     sim_radio1.change(fn=check_satisfaction, inputs=[sim_radio1, active_tab], outputs=[submit_btn1, redesign_btn1])
     sim_radio2.change(fn=check_satisfaction, inputs=[sim_radio2, active_tab], outputs=[submit_btn2, redesign_btn2])
-    dislike_radio1.select(fn=select_dislike, inputs=[dislike_radio1, gallery_state1], outputs=[dislike_image1])
-    like_radio1.select(fn=select_dislike, inputs=[like_radio1, gallery_state1], outputs=[like_image1])
-    dislike_radio2.select(fn=select_dislike, inputs=[dislike_radio2, gallery_state2], outputs=[dislike_image2])
-    like_radio2.select(fn=select_dislike, inputs=[like_radio2, gallery_state2], outputs=[like_image2])
+    dislike_radio1.select(fn=select_image, inputs=[dislike_radio1, images_method1], outputs=[dislike_image1])
+    like_radio1.select(fn=select_image, inputs=[like_radio1, images_method1], outputs=[like_image1])
+    dislike_radio2.select(fn=select_image, inputs=[dislike_radio2, images_method2], outputs=[dislike_image2])
+    like_radio2.select(fn=select_image, inputs=[like_radio2, images_method2], outputs=[like_image2])
 
     redesign_btn1.click(
         fn=redesign, 
-        inputs=[participant, scenario, prompt1, sim_radio1, response1, images_method1, active_tab], 
-        outputs=[gallery_state1, sim_radio1, response1, prompt1, next_btn1, redesign_btn1, submit_btn1]
+        inputs=[participant, scenario, prompt1, sim_radio1, images_method1, history_images1, active_tab], 
+        outputs=[sim_radio1, dislike_radio1, like_radio1, images_method1, history_images1, example1.dataset, prompt1, next_btn1, redesign_btn1, submit_btn1]
     )
     redesign_btn2.click(
         fn=redesign, 
-        inputs=[participant, scenario, prompt2, sim_radio2, response2, images_method2, active_tab], 
-        outputs=[gallery_state2, sim_radio2, response2, prompt2, next_btn2, redesign_btn2, submit_btn2]
+        inputs=[participant, scenario, prompt2, sim_radio2, images_method2, history_images2, active_tab], 
+        outputs=[sim_radio2, dislike_radio2, like_radio2, images_method2, history_images2, example2.dataset, prompt2, next_btn2, redesign_btn2, submit_btn2]
     )
     submit_btn1.click(fn=save_response, 
-        inputs=[participant, scenario, prompt1, sim_radio1, response1, images_method1, active_tab], 
-        outputs=[images_method1, gallery_state1, sim_radio1, dislike_radio1, like_radio1, like_image1, dislike_image1, prompt1, response1, next_btn1, redesign_btn1, submit_btn1, tabs, scenario])
+        inputs=[participant, scenario, prompt1, sim_radio1, active_tab], 
+        outputs=[sim_radio1, dislike_radio1, like_radio1, like_image1, dislike_image1, images_method1, history_images1, example1.dataset, prompt1, next_btn1, redesign_btn1, submit_btn1, scenario, tabs])
     
     submit_btn2.click(fn=save_response, 
-        inputs=[participant, scenario, prompt2, sim_radio2, response2, images_method2, active_tab], 
-        outputs=[images_method2, gallery_state2, sim_radio2,  dislike_radio2, like_radio2, like_image2, dislike_image2, prompt2, response2, next_btn2, redesign_btn2, submit_btn2, tabs, scenario])
+        inputs=[participant, scenario, prompt2, sim_radio2, active_tab], 
+        outputs=[sim_radio2,  dislike_radio2, like_radio2, like_image2, dislike_image2, images_method2, history_images2, example2.dataset, prompt2, next_btn2, redesign_btn2, submit_btn2, scenario, tabs])
 
 
 if __name__ == "__main__":