Update app.py

app.py

@@ -4,6 +4,7 @@ from diffusers import DiffusionPipeline
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 from share_btn import community_icon_html, loading_icon_html, share_js
+from uuid import uuid4
 
 import torch
 import json
@@ -13,6 +14,20 @@ import gc
 
 lora_list = hf_hub_download(repo_id="multimodalart/LoraTheExplorer", filename="sdxl_loras.json", repo_type="space")
 
+IMAGE_DATASET_DIR = Path("image_dataset") / f"train-{uuid4()}"
+IMAGE_DATASET_DIR.mkdir(parents=True, exist_ok=True)
+IMAGE_JSONL_PATH = IMAGE_DATASET_DIR / "metadata.jsonl"
+
+scheduler = CommitScheduler(
+    repo_id="multimodalart/lora-fusing-preferences",
+    repo_type="dataset",
+    folder_path=IMAGE_DATASET_DIR,
+    path_in_repo=IMAGE_DATASET_DIR.name,
+    every=10
+)
+
+client = InferenceClient()
+
 with open(lora_list, "r") as file:
     data = json.load(file)
     sdxl_loras = [
@@ -66,7 +81,7 @@ div#share-btn-container > div {flex-direction: row;background: black;align-items
 
 original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
 
-def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lora_2_scale=0.5, progress=gr.Progress(track_tqdm=True)):
+def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lora_2_scale=0.5, seed, progress=gr.Progress(track_tqdm=True)):
   state_dict_1 = copy.deepcopy(shuffled_items[0]['state_dict'])
   state_dict_2 = copy.deepcopy(shuffled_items[1]['state_dict'])
   pipe = copy.deepcopy(original_pipe)
@@ -79,12 +94,15 @@ def merge_and_run(prompt, negative_prompt, shuffled_items, lora_1_scale=0.5, lor
   
   if negative_prompt == "":
     negative_prompt = None
-      
-  image = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=20, width=768, height=768).images[0]
+          
+  if(seed < 0):
+      seed = random.randint(0, 2147483647)
+  generator = torch.Generator(device="cuda").manual_seed(seed)
+  image = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=20, width=768, height=768, generator=generator).images[0]
   del pipe
   gc.collect()
   torch.cuda.empty_cache()
-  return image, gr.update(visible=True)
+  return image, gr.update(visible=True), seed
 
 def get_description(item):
       trigger_word = item["trigger_word"]
@@ -108,6 +126,15 @@ def shuffle_images():
     
     return title_1, prompt_description_1, repo_id_1, title_2, prompt_description_2, repo_id_2, prompt, two_shuffled_items, scale, scale
 
+def save_preferences(lora_1_id, lora_1_scale, lora_2_id, lora_2_scale, prompt, generated_image, thumbs_direction, seed):
+    image_path = IMAGE_DATASET_DIR / f"{uuid4()}.png"
+    with scheduler.lock:
+        generated_image.save(image_path)
+        with IMAGE_JSONL_PATH.open("a") as f:
+            json.dump({"prompt": prompt, "file_name":image_path.name, "lora_1_id": lora_2_id, "lora_1_scale": lora_1_scale, "lora_2_id": lora_2_id, "lora_2_scale": lora_2_scale, "thumbs_direction": thumbs_direction, "seed": seed}, f)
+            f.write("\n")
+    return gr.update(visible=True)
+
 with gr.Blocks(css=css) as demo:
   shuffled_items = gr.State()
   title = gr.HTML(
@@ -147,9 +174,11 @@ with gr.Blocks(css=css) as demo:
                         community_icon = gr.HTML(community_icon_html)
                         loading_icon = gr.HTML(loading_icon_html)
                         share_button = gr.Button("Share to community", elem_id="share-btn")
-
+            post_eval = gr.Markdown("Thanks for evaluating. The dataset with evaluations is [here](#)", visible=False)
   with gr.Accordion("Advanced settings", open=False):
     negative_prompt = gr.Textbox(label="Negative prompt")
+    seed = gr.Slider(label="Seed", info="-1 denotes a random seed", minimum=-1, maximum=2147483647, value=-1)
+    last_used_seed = gr.Slider(label="Last used seed", info="The seed used in the last generation", minimum=0, maximum=2147483647, value=-1, interactive=False)
     with gr.Row():
       lora_1_scale = gr.Slider(label="LoRA 1 Scale", minimum=0, maximum=1, step=0.1, value=0.7)
       lora_2_scale = gr.Slider(label="LoRa 2 Scale", minimum=0, maximum=1, step=0.1, value=0.7)
@@ -158,8 +187,10 @@ with gr.Blocks(css=css) as demo:
   demo.load(shuffle_images, inputs=[], outputs=[lora_1, lora_1_prompt, lora_1_id, lora_2, lora_2_prompt, lora_2_id, prompt, shuffled_items, lora_1_scale, lora_2_scale], queue=False, show_progress="hidden")
   shuffle_button.click(shuffle_images, outputs=[lora_1, lora_1_prompt, lora_1_id, lora_2, lora_2_prompt, lora_2_id, prompt, shuffled_items, lora_1_scale, lora_2_scale], queue=False, show_progress="hidden")
 
-  run_btn.click(merge_and_run, inputs=[prompt, negative_prompt, shuffled_items, lora_1_scale, lora_2_scale], outputs=[output_image, post_gen_info])
-  prompt.submit(merge_and_run, inputs=[prompt, negative_prompt, shuffled_items, lora_1_scale, lora_2_scale], outputs=[output_image, post_gen_info])
+  run_btn.click(merge_and_run, inputs=[prompt, negative_prompt, shuffled_items, lora_1_scale, lora_2_scale], outputs=[output_image, post_gen_info, last_used_seed])
+  prompt.submit(merge_and_run, inputs=[prompt, negative_prompt, shuffled_items, lora_1_scale, lora_2_scale], outputs=[output_image, post_gen_info, last_used_seed])
+  thumbs_up.click(save_preferences, inputs=[lora_1_id, lora_1_scale, lora_2_id, lora_2_scale, prompt, output_image, gr.State("up"), seed], outputs=[post_eval])
+  thumbs_down.click(save_preferences, inputs=[lora_1_id, lora_1_scale, lora_2_id, lora_2_scale, prompt, output_image, gr.State("down"), seed], outputs=[post_eval])
   share_button.click(None, [], [], _js=share_js)
 demo.queue()
 demo.launch()