first commit

app.py

@@ -1,7 +1,517 @@
 import gradio as gr
+from gradio.themes.base import Base
+import numpy as np
+import random
+import spaces
+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 SCENARIOS, PROMPTS, IMAGES, OPTIONS, T2I_MODELS, INSTRUCTION
+import spaces #[uncomment to use ZeroGPU]
+import transformers
+import gspread
+import asyncio
+from datetime import datetime
 
-def greet(name):
-    return "Hello " + name + "!!"
+CLIP_MODEL = "ViT-H-14"
+PRETRAINED_CLIP = "laion2b_s32b_b79k"
+default_t2i_model = "black-forest-labs/FLUX.1-dev" # "black-forest-labs/FLUX.1-dev" 
+default_llm_model = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B" # "meta-llama/Meta-Llama-3-8B-Instruct"
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+NUM_IMAGES=4
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+clean_cache() 
+
+selected_pipe = setup_model(default_t2i_model, torch_dtype, device)
+# clip_model, _, preprocess = open_clip.create_model_and_transforms(CLIP_MODEL, pretrained=PRETRAINED_CLIP, device=device)
+llm_pipe = None
+torch.cuda.empty_cache()
+inverted_prompt = ""
+
+VERBAL_MSG = "Please verbally describe key differences found in the image pair."
+DEFAULT_SCENARIO = "Product advertisement"
+METHODS = ["Method 1", "Method 2"]
+MAX_ROUND = 5
+# intermittent memory
+counter1, counter2 = 1, 1
+responses_memory = {}
+assigned_scenarios = list(SCENARIOS.keys())[:2]
+current_task1, current_task2 = METHODS # current task 1 (tab 1)
+task1_success, task2_success = False, False
+
+########################################################################################################
+# Generating images with two methods
+########################################################################################################
+
+
+@spaces.GPU(duration=65)
+def infer(
+    prompt,
+    negative_prompt="",
+    seed=42,
+    randomize_seed=True,
+    width=256,
+    height=256,
+    guidance_scale=5,
+    num_inference_steps=18,
+    progress=gr.Progress(track_tqdm=True),
+):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    generator = torch.Generator().manual_seed(seed)
+    with torch.no_grad():
+        image = selected_pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+        ).images[0]
+
+    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()
+    messages = get_refine_msg(prompt, num_prompts)
+    outputs = call_gpt_api(messages, client, "gpt-4o", seed, max_tokens, temperature, top_p)
+    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 = {
+        "iter": iter,
+        "lr": lr,
+        "batch_size": batch_size,
+        "prompt_len": prompt_len,
+        "weight_decay": 0.1,
+        "prompt_bs": 1,
+        "loss_weight": 1.0,
+        "print_step": 100,
+        "clip_model": CLIP_MODEL,
+        "clip_pretrain": PRETRAINED_CLIP,
+    }
+    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
+
+
+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
+
+
+########################################################################################################
+# Button-related functions
+########################################################################################################
+
+def reset_gallery():
+    return []
+
+def display_error_message(msg, duration=5):
+    gr.Warning(msg, duration=duration)
+
+def display_info_message(msg, duration=5):
+    gr.Info(msg, duration=duration)
+
+def switch_tab(active_tab):
+    print("switching tab")
+    if active_tab == "Task A":
+        return gr.Tabs(selected="Task B")
+    else:
+        return gr.Tabs(selected="Task A")
+
+def set_user(participant):
+    global responses_memory
+    responses_memory[participant] = {METHODS[0]:{}, METHODS[1]:{}}
+
+    id = re.findall(r'\d+', participant)
+    if len(id) == 0 or int(id[0]) % 2 == 0: # name invalid, assign first half scenarios
+        assigned_scenarios = list(SCENARIOS.keys())[:2]
+    else:
+        assigned_scenarios = list(SCENARIOS.keys())[2:]
+    return assigned_scenarios[0]
+
+def display_scenario(participant, choice):
+    # reset intermittent storage when scenario change
+    global counter1, counter2, responses_memory, current_task1, current_task2, task1_success, task2_success
+    
+    task1_success, task2_success = False, False
+    counter1, counter2 = 1, 1
+    
+    if check_participant(participant):
+        responses_memory[participant] = {METHODS[0]:{}, METHODS[1]:{}}
+
+    [current_task1, current_task2] = random.sample(METHODS, 2)
+    if current_task1 == METHODS[0]:
+        initial_images1 = IMAGES[choice]["baseline"]
+        initial_images2 = IMAGES[choice]["ours"]
+    else:
+        initial_images1 = IMAGES[choice]["ours"]
+        initial_images2 = IMAGES[choice]["baseline"]
+    
+    res = { 
+        scenario_content: SCENARIOS.get(choice, ""), 
+        prompt: PROMPTS.get(choice, ""),
+        prompt1: "", 
+        prompt2: "",
+        images_method1: initial_images1, 
+        images_method2: initial_images2,
+        gallery_state1: initial_images1, 
+        gallery_state2: initial_images2, 
+        sim_radio1: None, 
+        sim_radio2: None, 
+        response1: VERBAL_MSG, 
+        response2: VERBAL_MSG, 
+        next_btn1: gr.update(interactive=False), 
+        next_btn2: gr.update(interactive=False), 
+        redesign_btn1: gr.update(interactive=True), 
+        redesign_btn2: gr.update(interactive=True),
+        submit_btn1: gr.update(interactive=False),
+        submit_btn2: gr.update(interactive=False),
+    }
+    return res
+
+def generate_image(participant, scenario, prompt, gallery_state, active_tab):
+    if not check_participant(participant): return [], []
+    global current_task1, current_task2
+    
+    method = current_task1 if active_tab == "Task A" else current_task2
+    
+    if method == METHODS[0]:
+        for i in range(NUM_IMAGES): 
+            img = infer(prompt)
+            gallery_state.append(img)
+            yield gallery_state
+    else:
+        refined_prompts = refine_prompt(gallery_state, 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 redesign(participant, scenario, prompt, sim_radio, response, images_method, 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 method == METHODS[0]:
+            counter1 += 1
+            counter = counter1
+        else:
+            counter2 += 1
+            counter = counter2
+        
+        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)
+        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
+    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
+    method = current_task1 if active_tab == "Task A" else current_task2
+    
+    if check_evaluation(sim_radio, response) 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) 
+
+        try:
+            gc = gspread.service_account(filename='credentials.json')
+            sheet = gc.open("DiverseGen-phase2").sheet1 
+
+            for i, entry in responses_memory[participant][method].items():
+                sheet.append_row([participant, scenario, method, i, entry["prompt"], entry["sim_radio"],  entry["response"]])
+
+            display_info_message("✅ Your answer is saved!")
+
+            # reset counter and update success indicator
+            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)
+            next_scenario = assigned_scenarios[1] if task1_success and task2_success else assigned_scenarios[0]
+            return [], [], None, VERBAL_MSG, prompt_state, next_state, redesign_state, submit_state, tabs, next_scenario
+        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()}
+    else:
+        return {submit_btn1: gr.skip()} if active_tab == "Task A" else {submit_btn2: gr.skip()}
+
+
+########################################################################################################
+# Interface 
+########################################################################################################
+
+css="""
+#col-container {
+    margin: 0 auto;
+    max-width: 700px;
+}
+
+#col-container2 {
+    margin: 0 auto;
+    max-width: 1000px;
+}
+
+#button-container {
+    display: flex;
+    justify-content: center; /* Centers the buttons horizontally */
+}
+"""
+
+with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "Arial", "sans-serif"]), css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(" # 📌 **Diverse Text-to-Image Generation**")
+
+        with gr.Row():
+            participant = gr.Textbox(
+                label="🧑‍💼 Participant ID", placeholder="Please enter you participant id"
+            )
+            scenario = gr.Dropdown(
+                choices=list(SCENARIOS.keys()),
+                # value=DEFAULT_SCENARIO,
+                value=None,
+                label="���� Scenario",
+                interactive=False,
+            )
+        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)
+
+    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.Column(elem_id="col-container"):
+                gallery_state1 = gr.State(IMAGES[DEFAULT_SCENARIO]["baseline"])
+                images_method1 = gr.Gallery(show_label=False, columns=[4], rows=[1], elem_id="gallery")
+            with gr.Column(elem_id="col-container2"):
+                gr.Markdown("### 📝 Evaluation")               
+                sim_radio1 = gr.Radio(
+                    OPTIONS, 
+                    label="How would you evaluate your satisfaction with the generated images, based on your expectations for the specified scenario?",
+                    type="value",
+                    elem_classes=["gradio-radio"]
+                )
+                response1 = gr.Textbox(
+                    label="Verbally describe key differences found in the image pair.",
+                    max_lines=1,
+                    interactive=False,
+                    container=False,
+                    value=VERBAL_MSG
+                )
+                
+                with gr.Row(elem_id="button-container"):
+                    redesign_btn1 = gr.Button("🎨 Redesign", variant="primary", scale=0)
+                    submit_btn1 = gr.Button("✅ Submit", variant="primary", interactive=False, scale=0)
+
+
+        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.Column(elem_id="col-container"):
+                gallery_state2 = gr.State(IMAGES[DEFAULT_SCENARIO]["ours"])
+                images_method2 = gr.Gallery(show_label=False, columns=[4], rows=[1], elem_id="gallery")
+
+            with gr.Column(elem_id="col-container2"):
+                gr.Markdown("### 📝 Evaluation")
+                sim_radio2 = gr.Radio(
+                    OPTIONS, 
+                    label="How would you evaluate your satisfaction with the generated images, based on your expectations for the specified scenario?",
+                    type="value",
+                    elem_classes=["gradio-radio"]
+                )
+                response2 = gr.Textbox(
+                    label="Verbally describe key differences found in the image pair.",
+                    max_lines=1,
+                    interactive=False,
+                    container=False,
+                    value=VERBAL_MSG
+                )
+                with gr.Row(elem_id="button-container"):
+                    redesign_btn2 = gr.Button("🎨 Redesign", variant="primary", scale=0)
+                    submit_btn2 = gr.Button("✅ Submit", variant="primary", interactive=False, scale=0)
+
+
+########################################################################################################
+# Button Function Setup
+########################################################################################################
+
+    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, 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])
+    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])
+    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]
+    )
+    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]
+    )
+    submit_btn1.click(fn=save_response, 
+        inputs=[participant, scenario, prompt1, sim_radio1, response1, images_method1, active_tab], 
+        outputs=[images_method1, gallery_state1, sim_radio1, prompt1, response1, next_btn1, redesign_btn1, submit_btn1, tabs, scenario])
+    
+    submit_btn2.click(fn=save_response, 
+        inputs=[participant, scenario, prompt2, sim_radio2, response2, images_method2, active_tab], 
+        outputs=[images_method2, gallery_state2, sim_radio2, prompt2, response2, next_btn2, redesign_btn2, submit_btn2, tabs, scenario])
+
+
+if __name__ == "__main__":
+    demo.launch()

live_preview_helpers.py

@@ -0,0 +1,172 @@
+import torch
+import numpy as np
+from typing import Any, Dict, List, Optional, Union
+from diffusers import FluxPipeline
+
+# Helper functions
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class FLUXPipelineWithIntermediateOutputs(FluxPipeline):
+    """
+    Extends the FluxPipeline to yield intermediate images during the denoising process 
+    with progressively increasing resolution for faster generation.
+    """
+    # FLUX pipeline function
+    @torch.inference_mode()
+    def generate_images(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        timesteps: List[int] = None,
+        guidance_scale: float = 3.5,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        max_sequence_length: int = 512,
+    ):
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+        device = self._execution_device
+
+        # 3. Encode prompt
+        lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None
+        prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas,
+            mu=mu,
+        )
+        self._num_timesteps = len(timesteps)
+
+        # Handle guidance
+        guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0]) if self.transformer.config.guidance_embeds else None
+
+        # 6. Denoising loop
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+
+            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+            noise_pred = self.transformer(
+                hidden_states=latents,
+                timestep=timestep / 1000,
+                guidance=guidance,
+                pooled_projections=pooled_prompt_embeds,
+                encoder_hidden_states=prompt_embeds,
+                txt_ids=text_ids,
+                img_ids=latent_image_ids,
+                joint_attention_kwargs=self.joint_attention_kwargs,
+                return_dict=False,
+            )[0]
+            
+            # Yield intermediate result
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+            yield self._decode_latents_to_image(latents, height, width, output_type)
+            torch.cuda.empty_cache()
+
+        # Final image
+        self.maybe_free_model_hooks()
+        torch.cuda.empty_cache()
+
+    def _decode_latents_to_image(self, latents, height, width, output_type, vae=None):
+        """Decodes the given latents into an image."""
+        vae = vae or self.vae
+        latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+        latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
+        image = vae.decode(latents, return_dict=False)[0]
+        return self.image_processor.postprocess(image, output_type=output_type)[0]

optim_utils.py

@@ -0,0 +1,239 @@
+import random
+import numpy as np
+import requests
+from io import BytesIO
+from PIL import Image
+from statistics import mean
+import copy
+import json
+from typing import Any, Mapping
+import open_clip
+import torch
+
+from sentence_transformers.util import (semantic_search, 
+                                        dot_score, 
+                                        normalize_embeddings)
+
+
+def nn_project(curr_embeds, embedding_layer, print_hits=False):
+    with torch.no_grad():
+        bsz,seq_len,emb_dim = curr_embeds.shape
+        
+        # Using the sentence transformers semantic search which is 
+        # a dot product exact kNN search between a set of 
+        # query vectors and a corpus of vectors
+        curr_embeds = curr_embeds.reshape((-1,emb_dim))
+        curr_embeds = normalize_embeddings(curr_embeds) # queries
+
+        embedding_matrix = embedding_layer.weight
+        embedding_matrix = normalize_embeddings(embedding_matrix)
+        
+        hits = semantic_search(curr_embeds, embedding_matrix, 
+                                query_chunk_size=curr_embeds.shape[0], 
+                                top_k=1,
+                                score_function=dot_score)
+
+        if print_hits:
+            all_hits = []
+            for hit in hits:
+                all_hits.append(hit[0]["score"])
+            print(f"mean hits:{mean(all_hits)}")
+        
+        nn_indices = torch.tensor([hit[0]["corpus_id"] for hit in hits], device=curr_embeds.device)
+        nn_indices = nn_indices.reshape((bsz,seq_len))
+
+        projected_embeds = embedding_layer(nn_indices)
+
+    return projected_embeds, nn_indices
+
+def decode_ids(input_ids, tokenizer, by_token=False):
+    input_ids = input_ids.detach().cpu().numpy()
+
+    texts = []
+
+    if by_token:
+        for input_ids_i in input_ids:
+            curr_text = []
+            for tmp in input_ids_i:
+                curr_text.append(tokenizer.decode([tmp]))
+
+            texts.append('|'.join(curr_text))
+    else:
+        for input_ids_i in input_ids:
+            texts.append(tokenizer.decode(input_ids_i))
+
+    return texts
+
+def get_target_feature(model, preprocess, tokenizer_funct, device, target_images=None, target_prompts=None):
+    if target_images is not None:
+        with torch.no_grad():
+            curr_images = [preprocess(i).unsqueeze(0) for i in target_images]
+            curr_images = torch.concatenate(curr_images).to(device)
+            all_target_features = model.encode_image(curr_images)
+    else:
+        texts = tokenizer_funct(target_prompts).to(device)
+        all_target_features = model.encode_text(texts)
+
+    return all_target_features
+
+def encode_text_embedding(model, text_embedding, ids, avg_text=False):
+        cast_dtype = model.transformer.get_cast_dtype()
+
+        x = text_embedding + model.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = model.transformer(x, attn_mask=model.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = model.ln_final(x)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        if avg_text:
+            x = x[torch.arange(x.shape[0]), :ids.argmax(dim=-1)]
+            x[:, 1:-1]
+            x = x.mean(dim=1) @ model.text_projection
+        else:
+            x = x[torch.arange(x.shape[0]), ids.argmax(dim=-1)] @ model.text_projection
+
+        return x
+        
+def forward_text_embedding(model, embeddings, ids, image_features, avg_text=False, return_feature=False):
+    text_features = encode_text_embedding(model, embeddings, ids, avg_text=avg_text)
+
+    if return_feature:
+        return text_features
+
+    image_features = image_features / image_features.norm(dim=1, keepdim=True)
+    text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+    logits_per_image = image_features @ text_features.t()
+    logits_per_text = logits_per_image.t()
+
+    return logits_per_image, logits_per_text
+    
+def initialize_prompt(tokenizer, token_embedding, args, device, original_prompt):
+    prompt_len = args["prompt_len"]
+
+    # randomly optimize prompt embeddings
+    tokens = tokenizer.encode(original_prompt)
+    if len(tokens) > prompt_len:
+        tokens = tokens[:prompt_len]
+    if len(tokens) < prompt_len:
+        tokens += [0] * (prompt_len - len(tokens))
+    
+    prompt_ids = torch.tensor([tokens] * args["prompt_bs"]).to(device)
+    # prompt_ids = torch.randint(len(tokenizer.encoder), (args.prompt_bs, prompt_len)).to(device)
+    prompt_embeds = token_embedding(prompt_ids).detach()
+    prompt_embeds.requires_grad = True
+
+    # initialize the template
+    template_text = "{}"
+    padded_template_text = template_text.format(" ".join(["<start_of_text>"] * prompt_len))
+    dummy_ids = tokenizer.encode(padded_template_text)
+
+    # -1 for optimized tokens
+    dummy_ids = [i if i != 49406 else -1 for i in dummy_ids]
+    dummy_ids = [49406] + dummy_ids + [49407]
+    dummy_ids += [0] * (77 - len(dummy_ids))
+    dummy_ids = torch.tensor([dummy_ids] * args["prompt_bs"]).to(device)
+
+    # for getting dummy embeds; -1 won't work for token_embedding
+    tmp_dummy_ids = copy.deepcopy(dummy_ids)
+    tmp_dummy_ids[tmp_dummy_ids == -1] = 0
+    dummy_embeds = token_embedding(tmp_dummy_ids).detach()
+    dummy_embeds.requires_grad = False
+    
+    return prompt_embeds, dummy_embeds, dummy_ids
+
+def optimize_prompt_loop(model, tokenizer, token_embedding, all_target_features, args, device, original_prompt):
+    opt_iters = args["iter"]
+    lr = args["lr"]
+    weight_decay = args["weight_decay"]
+    print_step = args["print_step"]
+    batch_size = args["batch_size"]
+    print_new_best = True
+    
+    # initialize prompt
+    prompt_embeds, dummy_embeds, dummy_ids = initialize_prompt(tokenizer, token_embedding, args, device, original_prompt)
+    p_bs, p_len, p_dim = prompt_embeds.shape
+
+    # get optimizer
+    input_optimizer = torch.optim.AdamW([prompt_embeds], lr=lr, weight_decay=weight_decay)
+
+    best_sim = -1000 * args["loss_weight"]
+    best_text = ""
+
+    for step in range(opt_iters):
+        # randomly sample sample images and get features
+        if batch_size is None:
+            target_features = all_target_features
+        else:
+            curr_indx = torch.randperm(len(all_target_features))
+            target_features = all_target_features[curr_indx][0:batch_size]
+            
+        universal_target_features = all_target_features
+
+        # forward projection
+        projected_embeds, nn_indices = nn_project(prompt_embeds, token_embedding, print_hits=False)
+
+        # get cosine similarity score with all target features
+        with torch.no_grad():
+            # padded_embeds = copy.deepcopy(dummy_embeds)
+            padded_embeds = dummy_embeds.detach().clone()
+            padded_embeds[dummy_ids == -1] = projected_embeds.reshape(-1, p_dim)
+            logits_per_image, _ = forward_text_embedding(model, padded_embeds, dummy_ids, universal_target_features)
+            scores_per_prompt = logits_per_image.mean(dim=0)
+            universal_cosim_score = scores_per_prompt.max().item()
+            best_indx = scores_per_prompt.argmax().item()
+        
+        # tmp_embeds = copy.deepcopy(prompt_embeds)
+        tmp_embeds = prompt_embeds.detach().clone()
+        tmp_embeds.data = projected_embeds.data
+        tmp_embeds.requires_grad = True
+        
+        # padding
+        # padded_embeds = copy.deepcopy(dummy_embeds)
+        padded_embeds = dummy_embeds.detach().clone()
+        padded_embeds[dummy_ids == -1] = tmp_embeds.reshape(-1, p_dim)
+        
+        logits_per_image, _ = forward_text_embedding(model, padded_embeds, dummy_ids, target_features)
+        cosim_scores = logits_per_image
+        loss = 1 - cosim_scores.mean()
+        loss = loss * args["loss_weight"]
+        
+        prompt_embeds.grad, = torch.autograd.grad(loss, [tmp_embeds])
+        
+        input_optimizer.step()
+        input_optimizer.zero_grad()
+
+        curr_lr = input_optimizer.param_groups[0]["lr"]
+        cosim_scores = cosim_scores.mean().item()
+
+        decoded_text = decode_ids(nn_indices, tokenizer)[best_indx]
+        if print_step is not None and (step % print_step == 0 or step == opt_iters-1):
+            per_step_message = f"step: {step}, lr: {curr_lr}"
+            # if not print_new_best:
+                # per_step_message = f"\n{per_step_message}, cosim: {universal_cosim_score:.3f}, text: {decoded_text}"
+            # print(per_step_message)
+
+        if best_sim * args["loss_weight"] < universal_cosim_score * args["loss_weight"]:
+            best_sim = universal_cosim_score
+            best_text = decoded_text
+            if print_new_best:
+                print(f"step: {step}, new best cosine sim: {best_sim}, new best prompt: {best_text}")
+
+    if print_step is not None:
+        print(f"best cosine sim: {best_sim}, best prompt: {best_text}")
+
+    return best_text
+
+
+def optimize_prompt(model, preprocess, args, device, target_images=None, target_prompts=None):
+    token_embedding = model.token_embedding
+    tokenizer = open_clip.tokenizer._tokenizer
+    tokenizer_funct = open_clip.get_tokenizer(args["clip_model"])
+
+    all_target_features = get_target_feature(model, preprocess, tokenizer_funct, device, target_images=target_images)
+    learned_prompt = optimize_prompt_loop(model, tokenizer, token_embedding, all_target_features, args, device, target_prompts)
+
+    return learned_prompt
+    

requirements.txt

@@ -0,0 +1,13 @@
+accelerate
+diffusers
+torch
+transformers
+git+https://github.com/huggingface/diffusers.git
+sentencepiece
+openai
+huggingface_hub
+sentence-transformers
+ftfy
+mediapy
+open-clip-torch==2.24.0
+gspread

utils.py

@@ -0,0 +1,104 @@
+import re
+from diffusers import DiffusionPipeline, FluxPipeline
+from live_preview_helpers import FLUXPipelineWithIntermediateOutputs 
+import torch
+import os
+from openai import OpenAI
+import subprocess
+
+T2I_MODELS = {
+    "Stable Diffusion v2.1": "stabilityai/stable-diffusion-2-1", 
+    "SDXL-Turbo": "stabilityai/sdxl-turbo",
+    "Stable Diffusion v3.5-medium": "stabilityai/stable-diffusion-3.5-medium", # Default
+    "Flux.1-dev": "black-forest-labs/FLUX.1-dev",
+}
+
+SCENARIOS = {
+  "Product advertisement": "You are designing an advertising campaign for a new line of coffee machines. To ensure the campaign resonates with a wider audience, you use generative models to create marketing images that showcase a variety of users interacting with the product.",
+  "Tourist promotion": "You are creating a travel campaign to attract a diverse range of visitors to a specific destination. To make the promotional materials more engaging and inclusive, you use generative models to design posters that highlight a broader array of experiences.",
+  "Fictional character generation": "You are creating a narrative superhero game where the player often interacts with multiple other non-player characters in the story. To test how different characters would affect the experience of gameplay, you decide to use generative models to help construct characters for (play)testing.",
+  "Interior Design": "You have a one-bedroom apartment and want to arrange your bed, desk, and dresser in the best way possible. You love the color white and want to ensure your space feels bright and open. To make a decision, you need a way to visualize different furniture placements before setting everything up.",
+#   "Education & accessibility": "You are a grade school teacher and the lesson of the day is about teamwork. Some of your students may have a difficult time visualizing what teamwork looks like because they are either (1) too young, (2) English is not their first language, or (3) they may have cognitive impairments that make it difficult for them to visualize concepts (e.g. aphantasia).."
+}
+
+PROMPTS = {
+    "Product advertisement": "Design a marketing advertisement image for a coffee machine.",
+    "Tourist promotion": "Design a travel promotional poster to showcase the beauty and attractions of a tourist destination.",
+    "Fictional character generation": "Generate a character of a superhero.",
+    "Interior Design": "Generate an one-bedroom apartment interior design.",
+    # "Education & accessibility": "Generate an image of grade school students buildind a sandcastle together on the beach."
+}
+
+IMAGES = {
+    "Product advertisement": {"baseline": ["images/scenario1_base1.png","images/scenario1_base2.png","images/scenario1_base3.png","images/scenario1_base4.png"], 
+                                            "ours": ["images/scenario1_our1.png","images/scenario1_our2.png","images/scenario1_our3.png","images/scenario1_our4.png"]},
+    "Tourist promotion": {"baseline": ["images/scenario5_base1.png","images/scenario5_base2.png","images/scenario5_base3.png","images/scenario5_base4.png"], 
+                                            "ours": ["images/scenario5_our1.png","images/scenario5_our2.png","images/scenario5_our3.png","images/scenario5_our4.png"]},
+    "Fictional character generation": {"baseline": ["images/scenario2_base1.png","images/scenario2_base2.png","images/scenario2_base3.png","images/scenario2_base4.png"], 
+                                            "ours": ["images/scenario2_our1.png","images/scenario2_our2.png","images/scenario2_our3.png","images/scenario2_our4.png"]},
+    "Interior Design": {"baseline": ["images/scenario3_base1.png","images/scenario3_base2.png","images/scenario3_base3.png","images/scenario3_base4.png"], 
+                                            "ours": ["images/scenario3_our1.png","images/scenario3_our2.png","images/scenario3_our3.png","images/scenario3_our4.png"]},
+    # "Education & accessibility": {"baseline": ["images/scenario4_base1.png","images/scenario4_base2.png","images/scenario4_base3.png","images/scenario4_base4.png"], 
+                                            # "ours": ["images/scenario4_our1.png","images/scenario4_our2.png","images/scenario4_our3.png","images/scenario4_our4.png"]},
+}
+
+OPTIONS = ["Very Unsatisfied", "Unsatisfied", "Slightly Unsatisfied", "Neutral", "Slightly Satisfied", "Satisfied", "Very Satisfied"]
+
+INSTRUCTION = "📌 **Instruction**: Now, we want to understand your satisfaction with the images generated. <br /> 📌 Step 1: You will start from evaluating the following images based on the given prompt. <br /> 📌 Step 2: Then please modify the prompt according to your expectations for the given scenario background, and answer the evaluation question **until you are satisfied** with at least one of the images generated below. If you are not satisfied with the generated images, you can repeatedly modify the prompts for at most **5 times**."
+def clean_cache():
+    subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+def setup_model(t2i_model_repo, torch_dtype, device):
+    if t2i_model_repo == "stabilityai/sdxl-turbo" or t2i_model_repo == "stabilityai/stable-diffusion-3.5-medium" or t2i_model_repo == "stabilityai/stable-diffusion-2-1":
+        pipe = DiffusionPipeline.from_pretrained(t2i_model_repo, torch_dtype=torch_dtype).to(device)
+    elif t2i_model_repo == "black-forest-labs/FLUX.1-dev":
+        # pipe = FluxPipeline.from_pretrained(t2i_model_repo, torch_dtype=torch_dtype).to(device)
+        pipe = FLUXPipelineWithIntermediateOutputs.from_pretrained(t2i_model_repo, torch_dtype=torch_dtype).to(device)
+    torch.cuda.empty_cache()
+    
+    return pipe
+
+def init_gpt_api():
+    return OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+def call_gpt_api(messages, client, model, seed, max_tokens, temperature, top_p):
+    completion = client.chat.completions.create(
+        model=model,
+        messages=messages,
+        seed=seed,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+    )
+    return completion.choices[0].message.content
+
+def clean_response_gpt(res: str):
+    prompts = re.findall(r'\d+\.\s"?(.*?)"?(?=\n|$)', res)
+    return prompts
+
+
+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"}]
+
+    message = f"""Given a prompt, modify the prompt for me to explore variations in subject attributes, actions, and contextual details, while retaining the semantic consistency of the original description. 
+
+Follow the following refinement instruction: 
+1. Subject: refine broad terms into specific subsets, focusing on but not restricted on ethinity, gender, age of human. 
+2. Object: modify the brand, color of object(s) only if it's not specified in the prompt. 
+3. Setting: add details to the background environment, such as change of temporal or spatial details (e.g., day to night, indoor to outdoor). 
+4. Action: add more details to the action or specify the object or goal of the action. 
+
+For example, given this prompt: a person is drinking a coffee in a coffee shop, the refined prompts could be: 
+'an elderly woman is drinking a coffee in a coffee shop' (subject adjective) 
+'an asian young woman is drinking a coffee in a coffee shop' (subject adjective) 
+'a young woman is drinking a hot coffee with her left hand in a coffee shop' (action details) 
+'a woman is drinking a coffee in an outdoor coffee shop in the garden' (setting details) 
+If there is no human in the sentence, you do not need to add person intentionally. 
+If you use adjectives, they should be visual. So don't use something like 'interesting'. Please also vary the number of modifications but do not change the number of subjects/objects that have been specified in the prompt. Remember not to change the predefined concepts that have been specified in the prompt. e.g. don't change a boy to several boys.
+
+Can you give me {num_prompts} modified prompts for the prompt '{prompt}' please."""
+
+    messages.append({"role": "user", "content": f"{message}"})
+    return messages