UI tweaks

ImageState.py

@@ -78,7 +78,7 @@ class ImageState:
     def clear_transforms(self):
         global num
         self.init_transforms()
-        clear_img_dir()
+        clear_img_dir("./img_history")
         num = 0
         return self._render_all_transformations()
     def _apply_vector(self, src, vector):
@@ -151,7 +151,11 @@ class ImageState:
         return self._render_all_transformations()
     def update_images(self, path1, path2, blend_weight):
         if path1 is None and path2 is None:
+            print("no paths")
             return None
+        if path1 == path2:
+            print("paths are the same")
+            print(path1)
         if path1 is None: path1 = path2
         if path2 is None: path2 = path1
         self.path1, self.path2 = path1, path2
@@ -170,7 +174,7 @@ class ImageState:
         prompt_transform = self.transform_history[-1]
         latent_index = int(index / 100 * (prompt_transform.iterations - 1))
         print(latent_index)
-        self.current_prompt_transforms[-1] = prompt_transform.transforms[latent_index]
+        self.current_prompt_transforms[-1] = prompt_transform.transforms[latent_index].to(self.device)
         return self._render_all_transformations()
     # def rescale_mask(self, mask):
     #     rep = mask.clone()
@@ -196,7 +200,7 @@ class ImageState:
         for i, transform in enumerate(self.prompt_optim.optimize(self.blend_latent,
                                                                 positive_prompts,
                                                                 negative_prompts)):
-            transform_log.transforms.append(transform.clone().detach())
+            transform_log.transforms.append(transform.detach().cpu())
             self.current_prompt_transforms[-1] = transform
             with torch.no_grad():
                 image = self._render_all_transformations(return_twice=False)

animation.py

@@ -9,23 +9,6 @@ def clear_img_dir(img_dir):
         os.remove(filename)
     
 
-def create_gif(total_duration, extend_frames, folder, gif_name="face_edit.gif"):
-    images = []
-    paths = glob.glob(folder + "/*")
-    frame_duration = total_duration / len(paths)
-    print(len(paths), "frame dur", frame_duration)
-    durations = [frame_duration] * len(paths)
-    if extend_frames:
-        durations [0] = 1.5
-        durations [-1] = 3
-    for file_name in os.listdir(folder):
-        if file_name.endswith('.png'):
-            file_path = os.path.join(folder, file_name)
-            images.append(imageio.imread(file_path))
-    # images[0] = images[0].set_meta_data({'duration': 1})
-    # images[-1] = images[-1].set_meta_data({'duration': 1})
-    imageio.mimsave(gif_name, images, duration=durations)
-    return gif_name
 
 if __name__ == "__main__":
     # clear_img_dir()

app.py

@@ -5,7 +5,7 @@ import sys
 import wandb
 import torch
 
-from configs import set_major_global, set_major_local, set_small_local
+from configs import set_major_global, set_major_local, set_preset, set_small_local
 import uuid 
 # print()'
 sys.path.append("taming-transformers")
@@ -18,7 +18,7 @@ import edit
 from backend import ImagePromptOptimizer, ProcessorGradientFlow
 from ImageState import ImageState
 from loaders import load_default
-from animation import create_gif
+# from animation import create_gif
 from prompts import get_random_prompts
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -70,16 +70,40 @@ class StateWrapper:
         return state, *state[0].update_images(*args, **kwargs)
     def update_requant(state, *args, **kwargs):
         return state, *state[0].update_requant(*args, **kwargs)
-def ret_id(id):
-    print(id)
-    return(id)
 with gr.Blocks(css="styles.css") as demo:
     id = gr.State(str(uuid.uuid4()))
     state = gr.State([ImageState(vqgan, promptoptim), str(uuid.uuid4())])
     with gr.Row():
         with gr.Column(scale=1):
-            x = gr.Button(label="asd")
-            x.click(ret_id, inputs=id, outputs=id)
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown(value="""## Image Upload
+                                    For best results, crop the photos like in the example pictures""", show_label=False)
+                    with gr.Row():
+                        base_img = gr.Image(label="Base Image", type="filepath")
+                        blend_img = gr.Image(label="Image for face blending (optional)", type="filepath")
+                with gr.Accordion(label="Add Mask", open=False):
+                    mask = gr.Image(tool="sketch", interactive=True)
+                    gr.Markdown(value="Note: You must clear the mask using the rewind button every time you want to change the mask (this is a gradio issue)")
+                    set_mask = gr.Button(value="Set mask")
+                    gr.Text(value="this image shows the mask passed to the model when you press set mask (debugging purposes)")
+                    testim = gr.Image()
+            with gr.Row():
+                gr.Examples(
+                    examples=glob.glob("test_pics/*"),
+                    inputs=base_img,
+                    outputs=blend_img,
+                    fn=set_img_from_example,
+                    ) 
+        with gr.Column(scale=1):
+            out = gr.Image()
+            rewind = gr.Slider(value=100,
+                                label="Rewind back through a prompt transform: Use this to scroll through the iterations of your prompt transformation.",
+                                minimum=0,
+                                maximum=100)
+
+            apply_prompts = gr.Button(variant="primary", value="🎨 Apply Prompts", elem_id="apply")
+            clear = gr.Button(value="❌ Clear all transformations (irreversible)", elem_id="warning")
             blue_eyes = gr.Slider(
                 label="Blue Eyes",
                 minimum=-.8,
@@ -119,39 +143,6 @@ with gr.Blocks(css="styles.css") as demo:
                 maximum=2.,
                 step=0.07,
             )
-            with gr.Row():
-                with gr.Column():
-                    gr.Markdown(value="""## Image Upload
-                                    For best results, crop the photos like in the example pictures""", show_label=False)
-                    with gr.Row():
-                        base_img = gr.Image(label="Base Image", type="filepath")
-                        blend_img = gr.Image(label="Image for face blending (optional)", type="filepath")
-                    # gr.Markdown("## Image Examples")
-                with gr.Accordion(label="Add Mask", open=False):
-                    mask = gr.Image(tool="sketch", interactive=True)
-                    gr.Markdown(value="Note: You must clear the mask using the rewind button every time you want to change the mask (this is a gradio bug)")
-                    set_mask = gr.Button(value="Set mask")
-                    gr.Text(value="this image shows the mask passed to the model when you press set mask (debugging purposes)")
-                    testim = gr.Image()
-                    # # clear_mask = gr.Button(value="Clear mask")
-                    # clear_mask.click(get_cleared_mask, outputs=mask)
-            with gr.Row():
-                gr.Examples(
-                    examples=glob.glob("test_pics/*"),
-                    inputs=base_img,
-                    outputs=blend_img,
-                    fn=set_img_from_example,
-                    # cache_examples=True,
-                    ) 
-        with gr.Column(scale=1):
-            out = gr.Image()
-            rewind = gr.Slider(value=100,
-                                label="Rewind back through a prompt transform: Use this to scroll through the iterations of your prompt transformation.",
-                                minimum=0,
-                                maximum=100)
-
-            apply_prompts = gr.Button(variant="primary", value="🎨 Apply Prompts", elem_id="apply")
-            clear = gr.Button(value="❌ Clear all transformations (irreversible)", elem_id="warning")
             with gr.Accordion(label="💾 Save Animation", open=False):
                 gr.Text(value="Creates an animation of all the steps in the editing process", show_label=False)
                 duration = gr.Number(value=10, label="Duration of the animation in seconds")
@@ -161,7 +152,7 @@ with gr.Blocks(css="styles.css") as demo:
                 create_animation.click(StateWrapper.create_gif, inputs=[state, duration, extend_frames], outputs=[state, gif])
 
         with gr.Column(scale=1):
-            gr.Markdown(value="""## Text Prompting
+            gr.Markdown(value="""## ✍️ Prompt Editing
             See readme for a prompting guide. Use the '|' symbol to separate prompts. Use the "Add mask" section to make local edits. Negative prompts are highly recommended""", show_label=False)
             positive_prompts = gr.Textbox(label="Positive prompts",
                                             value="a picture of a woman with a very big nose | a picture of a woman with a large wide nose | a woman with an extremely prominent nose")
@@ -171,29 +162,35 @@ with gr.Blocks(css="styles.css") as demo:
             gen_prompts.click(get_random_prompts, outputs=[positive_prompts, negative_prompts])
             with gr.Row():
                 with gr.Column():
-                    gr.Text(value="⚙️ Prompt Editing Configuration", show_label=False)
                     with gr.Row():
-                        gr.Markdown(value="## Preset Configs", show_label=False)
+                        gr.Markdown(value="## ⚙ Prompt Editing Config", show_label=False)
+                    with gr.Accordion(label="Config Tutorial", open=False):
+                        gr.Markdown(value="""
+                        - If results are not changing enough, increase the learning rate or decrease the perceptual loss weight
+                        - To make local edits, use the 'Add Mask' section
+                        - If using a mask and the image is changing too much outside of the masked area, try increasing the perceptual loss weight or lowering the learning rate
+                        - Use the rewind slider to scroll through the iterations of your prompt transformation, you can resume editing from any point in the history. 
+                        - I recommend starting prompts with 'a picture of a'
+                        - To avoid shifts in gender, you can use 'a person' instead of 'a man' or 'a woman', especially in the negative prompts. 
+                        - The more 'out-of-domain' the prompts are, the more you need to increase the learning rate and decrease the perceptual loss weight. For example, trying to make a black person have platinum blond hair is more out-of-domain than the same transformation on a caucasian person. 
+                        - 
+                        """)
                     with gr.Row():
                         # with gr.Column():
-                        small_local = gr.Button(value="Small Masked Changes (e.g. add lipstick)", elem_id="small_local").style(full_width=False)
-                        # with gr.Column():
-                        major_local = gr.Button(value="Major Masked Changes (e.g. change hair color or nose size)").style(full_width=False)
-                        # with gr.Column():
-                        major_global = gr.Button(value="Major Global Changes (e.g. change race / gender").style(full_width=False)
+                        presets = gr.Dropdown(default="Select a preset", label="Preset Configs", choices=["Small Masked Changes (e.g. add lipstick)", "Major Masked Changes (e.g. change hair color or nose size)", "Major Global Changes (e.g. change race / gender"])
                     iterations = gr.Slider(minimum=10,
                                             maximum=60,
                                             step=1,
                                             value=20,
                                             label="Iterations: How many steps the model will take to modify the image. Try starting small and seeing how the results turn out, you can always resume with afterwards",)
                     learning_rate = gr.Slider(minimum=4e-3,
-                                            maximum=7e-1,
+                                            maximum=1,
                                             value=1e-1,
                                             label="Learning Rate: How strong the change in each step will be (you should raise this for bigger changes (for example, changing hair color), and lower it for more minor changes. Raise if changes aren't strong enough")
                     lpips_weight = gr.Slider(minimum=0,
                                             maximum=50,
                                             value=1,
-                                            label="Perceptual similarity weight (Keeps areas outside of the mask looking similar to the original. Increase if the rest of the image is changing too much while you're trying to change make a localized edit")
+                                            label="Perceptual Loss weight (Keeps areas outside of the mask looking similar to the original. Increase if the rest of the image is changing too much while you're trying to change make a localized edit")
                     reconstruction_steps = gr.Slider(minimum=0,
                                             maximum=50,
                                             value=3,
@@ -213,11 +210,12 @@ with gr.Blocks(css="styles.css") as demo:
     # requantize.change(StateWrapper.update_requant, inputs=[state, requantize], outputs=[state, out, mask])
     base_img.change(StateWrapper.update_images, inputs=[state, base_img, blend_img, blend_weight], outputs=[state, out, mask])
     blend_img.change(StateWrapper.update_images, inputs=[state, base_img, blend_img, blend_weight], outputs=[state, out, mask])
-    small_local.click(set_small_local, outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
-    major_local.click(set_major_local, outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
-    major_global.click(set_major_global, outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
+    # small_local.click(set_small_local, outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
+    # major_local.click(set_major_local, outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
+    # major_global.click(set_major_global, outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
     apply_prompts.click(StateWrapper.apply_prompts, inputs=[state, positive_prompts, negative_prompts, learning_rate, iterations, lpips_weight, reconstruction_steps], outputs=[state, out, mask])
     rewind.change(StateWrapper.rewind, inputs=[state, rewind], outputs=[state, out, mask])
     set_mask.click(StateWrapper.set_mask, inputs=[state, mask], outputs=[state, testim])
+    presets.change(set_preset, inputs=[presets], outputs=[iterations, learning_rate, lpips_weight, reconstruction_steps])
 demo.queue()
 demo.launch(debug=True, enable_queue=True)

configs.py

@@ -1,7 +1,15 @@
 import gradio as gr
 def set_small_local():
-    return (gr.Slider.update(value=25), gr.Slider.update(value=0.15), gr.Slider.update(value=5), gr.Slider.update(value=4))
+    return (gr.Slider.update(value=18), gr.Slider.update(value=0.15), gr.Slider.update(value=5), gr.Slider.update(value=4))
 def set_major_local():
     return (gr.Slider.update(value=25), gr.Slider.update(value=0.187), gr.Slider.update(value=36.6), gr.Slider.update(value=6))
 def set_major_global():
-    return (gr.Slider.update(value=30), gr.Slider.update(value=0.1), gr.Slider.update(value=1), gr.Slider.update(value=1))
+    return (gr.Slider.update(value=30), gr.Slider.update(value=0.1), gr.Slider.update(value=1), gr.Slider.update(value=1))
+def set_preset(config_str):
+    choices=["Small Masked Changes (e.g. add lipstick)", "Major Masked Changes (e.g. change hair color or nose size)", "Major Global Changes (e.g. change race / gender"]
+    if config_str == choices[0]:
+        return set_small_local()
+    elif config_str == choices[1]:
+        return set_major_local()
+    elif config_str == choices[2]:
+        return set_major_global()