Update edit_app.py

edit_app.py

@@ -1,13 +1,15 @@
-from __future__ import annotations
+rom __future__ import annotations
 
 import math
+import os
 import random
-
+import csv
 import gradio as gr
 import torch
 from PIL import Image, ImageOps
 from diffusers import StableDiffusionInstructPix2PixPipeline
 
+header = ['file', 'steps', 'seed', 'image_cfg', 'text_cfg', 'prompt', 'input height', 'input width']
 
 help_text = """
 If you're not getting what you want, there may be a few reasons:
@@ -79,27 +81,61 @@ def main():
         randomize_cfg: bool,
         text_cfg_scale: float,
         image_cfg_scale: float,
+        number_of_images: int, 
+        folder_path: str,
+        file_name: str
     ):
-        seed = random.randint(0, 100000) if randomize_seed else seed
-        text_cfg_scale = round(random.uniform(6.0, 9.0), ndigits=2) if randomize_cfg else text_cfg_scale
-        image_cfg_scale = round(random.uniform(1.2, 1.8), ndigits=2) if randomize_cfg else image_cfg_scale
-
-        width, height = input_image.size
-        factor = 512 / max(width, height)
-        factor = math.ceil(min(width, height) * factor / 64) * 64 / min(width, height)
-        width = int((width * factor) // 64) * 64
-        height = int((height * factor) // 64) * 64
-        input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)
-
-        if instruction == "":
-            return [input_image, seed]
-
-        generator = torch.manual_seed(seed)
-        edited_image = pipe(
-            instruction, image=input_image,
-            guidance_scale=text_cfg_scale, image_guidance_scale=image_cfg_scale,
-            num_inference_steps=steps, generator=generator,
-        ).images[0]
+        for i in range(0, int(number_of_images)):
+            seed = random.randint(0, 100000) if randomize_seed else seed
+            text_cfg_scale = round(random.uniform(6.0, 9.0), ndigits=2) if randomize_cfg else text_cfg_scale
+            image_cfg_scale = round(random.uniform(1.2, 1.8), ndigits=2) if randomize_cfg else image_cfg_scale
+
+            width, height = input_image.size
+            factor = 512 / max(width, height)
+            factor = math.ceil(min(width, height) * factor / 64) * 64 / min(width, height)
+            width = int((width * factor) // 64) * 64
+            height = int((height * factor) // 64) * 64
+            #input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)
+
+            if instruction == "":
+                return [input_image, seed]
+
+            generator = torch.manual_seed(seed)
+            edited_image = pipe(
+                instruction, image=input_image,
+                guidance_scale=text_cfg_scale, image_guidance_scale=image_cfg_scale,
+                num_inference_steps=steps, generator=generator,
+            ).images[0]
+
+            if(os.path.exists(folder_path)):
+                file_path = folder_path + "/" + "records.csv"
+                if(os.path.isfile(file_path)):
+                    with open(file_path, 'a', encoding='UTF8') as f:
+                        writer = csv.writer(f)            
+                        if(i == 0):
+                            data = [file_name + "_" + str(i), steps, seed, image_cfg_scale, text_cfg_scale, instruction, height, width]
+                            print(data)
+                            writer.writerow(data)
+                            f.close()
+                        if(i != 0):
+                            data = [file_name + "_" + str(i), steps, seed, image_cfg_scale, text_cfg_scale, instruction, height, width]
+                            print(data)
+                            writer.writerow(data)
+                            f.close()
+                else:
+                    with open(file_path, 'w', encoding = 'UTF8') as f:
+                          writer = csv.writer(f)
+                          writer.writerow(header)
+                          data = [file_name + "_" + str(i), steps, seed, image_cfg_scale, text_cfg_scale, instruction, height, width]
+                          print(data)
+                          writer.writerow(data)
+                          f.close()
+            else:
+                os.mkdir(folder_path)
+
+
+            parent_path = folder_path
+            edited_image.save(f"{parent_path}/{file_name}_{i}.png")
         return [seed, text_cfg_scale, image_cfg_scale, edited_image]
 
     def reset():
@@ -139,6 +175,7 @@ def main():
                 show_label=False,
                 interactive=True,
             )
+            
             seed = gr.Number(value=1371, precision=0, label="Seed", interactive=True)
             randomize_cfg = gr.Radio(
                 ["Fix CFG", "Randomize CFG"],
@@ -149,6 +186,9 @@ def main():
             )
             text_cfg_scale = gr.Number(value=7.5, label=f"Text CFG", interactive=True)
             image_cfg_scale = gr.Number(value=1.5, label=f"Image CFG", interactive=True)
+            number_of_images = gr.Number(value=0, label=f"Number of Images", interactive=True)
+            folder_path = gr.Textbox(lines=1, label="Folder Path", interactive=True)
+            file_name = gr.Textbox(lines=1, label="File Name", interactive=True)
 
         gr.Markdown(help_text)
 
@@ -175,6 +215,9 @@ def main():
                 randomize_cfg,
                 text_cfg_scale,
                 image_cfg_scale,
+                number_of_images,
+                folder_path,
+                file_name
             ],
             outputs=[seed, text_cfg_scale, image_cfg_scale, edited_image],
         )
@@ -189,4 +232,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()