modified: app.py

app.py

@@ -54,8 +54,12 @@ if os.path.exists("history_gradio/history.safetensors"):
     os.remove("history_gradio/history.safetensors")
 
 out_root = 'src/gradio_utils/gradio_outputs'
+out_root_prompt = 'src/gradio_utils/gradio_prompts'
 if not os.path.exists(out_root):
     os.makedirs(out_root)
+if not os.path.exists(out_root_prompt):
+    os.makedirs(out_root_prompt)
+
 exp_folders = [d for d in os.listdir(out_root) if d.startswith("exp_") and d[4:].isdigit()]
 if exp_folders:
     max_idx = max(int(d[4:]) for d in exp_folders)
@@ -63,9 +67,12 @@ if exp_folders:
 else:
     name_dir = "exp_0"
 output_dir = os.path.join(out_root, name_dir)
+output_prompt = os.path.join(out_root_prompt, name_dir)
 
 if not os.path.exists(output_dir):
     os.makedirs(output_dir)
+if not os.path.exists(output_prompt):
+    os.makedirs(output_prompt)
 if not os.path.exists("heatmap"):
     os.makedirs("heatmap")
 if not os.path.exists("heatmap/average_heatmaps"):
@@ -74,7 +81,19 @@ source_image = None
 history_tensors = {
         "source img": torch.zeros((1, 1, 1)),
         "prev img": torch.zeros((1, 1, 1))}   
-instructions = ['source']
+instructions = ['']
+
+
+def read_sorted_prompts(folder_path):
+    # List all .txt files and sort them
+    files = sorted([f for f in os.listdir(folder_path) if f.endswith('.txt')])
+    prompts = []
+    for filename in files:
+        file_path = os.path.join(folder_path, filename)
+        with open(file_path, 'r') as f:
+            prompt = f.read().strip()
+            prompts.append(prompt)
+    return prompts
 
 
 @torch.inference_mode()
@@ -84,9 +103,11 @@ def reset():
     if os.path.exists("history_gradio/history.safetensors"):
         os.remove("history_gradio/history.safetensors")
 
-    global out_root, output_dir, history_tensors, source_image, instructions
+    global out_root, out_root_prompt, output_dir, output_prompt, history_tensors, source_image, instructions
     if not os.path.exists(out_root):
         os.makedirs(out_root)
+    if not os.path.exists(out_root_prompt):
+        os.makedirs(out_root_prompt)
     exp_folders = [d for d in os.listdir(out_root) if d.startswith("exp_") and d[4:].isdigit()]
     if exp_folders:
         max_idx = max(int(d[4:]) for d in exp_folders)
@@ -94,14 +115,17 @@ def reset():
     else:
         name_dir = "exp_0"
     output_dir = os.path.join(out_root, name_dir)
+    output_prompt = os.path.join(out_root_prompt, name_dir)
 
     if not os.path.exists(output_dir):
         os.makedirs(output_dir)
+    if not os.path.exists(output_prompt):
+        os.makedirs(output_prompt)
     if not os.path.exists("heatmap"):
         os.makedirs("heatmap")
     if not os.path.exists("heatmap/average_heatmaps"):
         os.makedirs("heatmap/average_heatmaps")
-    instructions = ['source']
+    instructions = ['']
     source_image = None
     history_tensors = {
         "source img": torch.zeros((1, 1, 1)),
@@ -111,7 +135,8 @@ def reset():
     traget_prompt = "(Required) Describe the desired content of the edited image."
     gallery = None
     output_image = None
-    return source_prompt, traget_prompt, gallery, output_image
+    init_image = None
+    return source_prompt, traget_prompt, gallery, output_image, init_image
 
 
 @torch.inference_mode()
@@ -145,7 +170,7 @@ def generate_image(
     init_image=None,
     image2image_strength=0.0,
 ):
-    global ae, t5, clip, model, name, is_schnell, output_dir, add_sampling_metadata, offload, history_tensors
+    global ae, t5, clip, model, name, is_schnell, output_dir, output_prompt, add_sampling_metadata, offload, history_tensors
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch.cuda.empty_cache()
     seed = None
@@ -250,9 +275,14 @@ def generate_image(
     img.save(filename, format="jpeg", exif=exif_data, quality=95, subsampling=0)
     instructions = [prompt]
 
+    prompt_path = os.path.join(output_prompt, f"round_0000.txt")
+    with open(prompt_path, "w") as f:
+        f.write(prompt)
+
         #--------------------    6.4 save editing prompt, update gradio component: gallery      ----------------------#
     img_and_prompt = []
     history_imgs = sorted(os.listdir(output_dir))
+    instructions = read_sorted_prompts(output_prompt)
     for img_file, prompt_txt in zip(history_imgs, instructions):
         img_and_prompt.append((os.path.join(output_dir, img_file), prompt_txt))
         history_gallery = gr.Gallery(value=img_and_prompt, label="History Image", interactive=True, columns=3)
@@ -262,7 +292,7 @@ def generate_image(
 @spaces.GPU(duration=200)
 @torch.inference_mode()
 def edit(init_image, source_prompt, target_prompt, editing_strategy, denoise_strategy, num_steps, guidance, attn_guidance_start_block, inject_step, init_image_2=None):
-    global ae, t5, clip, model, name, is_schnell, output_dir, add_sampling_metadata, offload, source_image, history_tensors, instructions
+    global ae, t5, clip, model, name, is_schnell, output_dir, output_prompt, add_sampling_metadata, offload, source_image, history_tensors, instructions
    
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch.cuda.empty_cache()
@@ -276,6 +306,9 @@ def edit(init_image, source_prompt, target_prompt, editing_strategy, denoise_str
 
     if not any("round_0000" in fname for fname in os.listdir(output_dir)):
         Image.fromarray(init_image).save(os.path.join(output_dir,"round_0000_[source].jpg"))
+        prompt_path = os.path.join(output_prompt, f"round_0000.txt")
+        with open(prompt_path, "w") as f:
+            f.write('')
 
     init_image = init_image[:new_h, :new_w, :]
     width, height = init_image.shape[0], init_image.shape[1]
@@ -436,7 +469,7 @@ def edit(init_image, source_prompt, target_prompt, editing_strategy, denoise_str
         else:
             idx = 1
     formatted_idx = str(idx).zfill(4) # Format as a 4-digit string
-
+    os.makedirs(output_prompt, exist_ok=True)
     #--------------------    6.3 output name      ----------------------#
     if denoise_strategy == 'multi_turn_consistent':
         denoise_strategy = 'MTC'
@@ -458,9 +491,14 @@ def edit(init_image, source_prompt, target_prompt, editing_strategy, denoise_str
     instructions.append(target_prompt)
     print("End Edit")
 
+    prompt_path = os.path.join(output_prompt, f"round_{formatted_idx}.txt")
+    with open(prompt_path, "w") as f:
+        f.write(target_prompt)
+
     #--------------------    6.4 save editing prompt, update gradio component: gallery      ----------------------#
     img_and_prompt = []
     history_imgs = sorted(os.listdir(output_dir))
+    instructions = read_sorted_prompts(output_prompt)
     for img_file, prompt_txt in zip(history_imgs, instructions):
         img_and_prompt.append((os.path.join(output_dir, img_file), prompt_txt))
     history_gallery = gr.Gallery(value=img_and_prompt, label="History Image", interactive=True, columns=3)
@@ -486,7 +524,6 @@ def create_demo(model_name: str, device: str = "cuda" if torch.cuda.is_available
     # Pre-defined examples
     examples = [
         ["src/gradio_utils/gradio_examples/000000000011.jpg", "", "an eagle standing on the branch", ['attn_guidance'], 15, 3.5, 11, 0],
-        ["src/gradio_utils/gradio_examples/221000000002.jpg", "", "a cat wearing a hat standing on the fence", ['attn_guidance'], 15, 3.5, 11, 0],
     ]
 
     with gr.Blocks() as demo:
@@ -531,7 +568,7 @@ def create_demo(model_name: str, device: str = "cuda" if torch.cuda.is_available
             inputs=[init_image, source_prompt, target_prompt, editing_strategy, denoise_strategy, num_steps, guidance, attn_guidance_start_block, inject_step, init_image_2],
             outputs=[output_image, gallery]
         )
-        reset_btn.click(fn = reset, outputs=[source_prompt, target_prompt, gallery, output_image])
+        reset_btn.click(fn = reset, outputs=[source_prompt, target_prompt, gallery, output_image, init_image])
         
         # Add examples
         gr.Examples(