update animation creation

ImageState.py

@@ -1,4 +1,7 @@
 # from align import align_from_path
+import imageio
+import glob
+import uuid
 from animation import clear_img_dir
 from backend import ImagePromptOptimizer, log
 import importlib
@@ -38,6 +41,9 @@ class ImageState:
         self.transform_history = []
         self.attn_mask = None
         self.prompt_optim = prompt_optimizer
+        self.state_id = "./" + str(uuid.uuid4())
+        print("NEW INSTANCE")
+        print(self.state_id)
         self._load_vectors()
         self.init_transforms()
     def _load_vectors(self):
@@ -45,6 +51,24 @@ class ImageState:
         self.red_blue_vector = torch.load("./latent_vectors/2blue_eyes.pt", map_location=self.device)
         self.green_purple_vector = torch.load("./latent_vectors/nose_vector.pt", map_location=self.device)
         self.asian_vector = torch.load("./latent_vectors/asian10.pt", map_location=self.device)
+    def create_gif(self, total_duration, extend_frames, gif_name="face_edit.gif"):
+        images = []
+        folder = self.state_id
+        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
     def init_transforms(self):
         self.blue_eyes = torch.zeros_like(self.lip_vector)
         self.lip_size = torch.zeros_like(self.lip_vector)
@@ -104,10 +128,10 @@ class ImageState:
         if self.quant:
             new_latent, _, _ = self.vqgan.quantize(new_latent.to(self.device))
         image = self._decode_latent_to_pil(new_latent)
-        img_dir = "./img_history"
+        img_dir = self.state_id
         if not os.path.exists(img_dir):
             os.mkdir(img_dir)
-        image.save(f"./img_history/img_{num:06}.png")
+        image.save(f"{img_dir}/img_{num:06}.png")
         num += 1
         return (image, image) if return_twice else image
     def apply_gp_vector(self, weight):
@@ -149,14 +173,12 @@ class ImageState:
         latent_index = int(index / 100 * (prompt_transform.iterations - 1))
         print(latent_index)
         self.current_prompt_transforms[-1] = prompt_transform.transforms[latent_index]
-        # print(self.current_prompt_transform)
-        # print(self.current_prompt_transforms.mean())
         return self._render_all_transformations()
-    def rescale_mask(self, mask):
-        rep = mask.clone()
-        rep[mask < 0.03] = -1000000
-        rep[mask >= 0.03] = 1
-        return rep
+    # def rescale_mask(self, mask):
+    #     rep = mask.clone()
+    #     rep[mask < 0.03] = -1000000
+    #     rep[mask >= 0.03] = 1
+    #     return rep
     def apply_prompts(self, positive_prompts, negative_prompts, lr, iterations, lpips_weight, reconstruction_steps):
         transform_log = PromptTransformHistory(iterations + reconstruction_steps)
         transform_log.transforms.append(torch.zeros_like(self.blend_latent, requires_grad=False))

animation.py

@@ -2,15 +2,14 @@ import imageio
 import glob
 import os
 
-def clear_img_dir():
-    img_dir = "./img_history"
+def clear_img_dir(img_dir):
     if not os.path.exists(img_dir):
         os.mkdir(img_dir)
     for filename in glob.glob(img_dir+"/*"):
         os.remove(filename)
     
 
-def create_gif(total_duration, extend_frames, folder="./img_history", gif_name="face_edit.gif"):
+def create_gif(total_duration, extend_frames, folder, gif_name="face_edit.gif"):
     images = []
     paths = glob.glob(folder + "/*")
     frame_duration = total_duration / len(paths)

app.py

@@ -6,7 +6,8 @@ import wandb
 import torch
 
 from configs import set_major_global, set_major_local, set_small_local
-
+import uuid 
+# print()'
 sys.path.append("taming-transformers")
 
 import gradio as gr
@@ -37,6 +38,8 @@ def get_cleared_mask():
     # mask.clear()
 
 class StateWrapper:
+    def create_gif(state, *args, **kwargs):
+        return state, state[0].create_gif(*args, **kwargs)
     def apply_asian_vector(state, *args, **kwargs):
         return state, *state[0].apply_asian_vector(*args, **kwargs)
     def apply_gp_vector(state, *args, **kwargs):
@@ -141,7 +144,7 @@ with gr.Blocks(css="styles.css") as demo:
                                 minimum=0,
                                 maximum=100)
 
-            apply_prompts = gr.Button(value="🎨 Apply Prompts", elem_id="apply")
+            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)
@@ -149,7 +152,7 @@ with gr.Blocks(css="styles.css") as demo:
                 extend_frames = gr.Checkbox(value=True, label="Make first and last frame longer")
                 gif = gr.File(interactive=False)
                 create_animation = gr.Button(value="Create Animation")
-                create_animation.click(create_gif, inputs=[duration, extend_frames], outputs=gif)
+                create_animation.click(StateWrapper.create_gif, inputs=[state, duration, extend_frames], outputs=[state, gif])
 
         with gr.Column(scale=1):
             gr.Markdown(value="""## Text Prompting
@@ -166,12 +169,12 @@ with gr.Blocks(css="styles.css") as demo:
                     with gr.Row():
                         gr.Markdown(value="## Preset Configs", show_label=False)
                     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)
+                        # 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)
                     iterations = gr.Slider(minimum=10,
                                             maximum=60,
                                             step=1,
@@ -181,14 +184,13 @@ with gr.Blocks(css="styles.css") as demo:
                                             maximum=7e-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")
-                with gr.Accordion(label="Advanced Prompt Editing Options", open=False):
                     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")
                     reconstruction_steps = gr.Slider(minimum=0,
                                             maximum=50,
-                                            value=15,
+                                            value=3,
                                             step=1,
                                             label="Steps to run at the end of the optimization, optimizing only the masked perceptual loss. If the edit is changing the identity too much, this setting will run steps at the end that will 'pull' the image back towards the original identity")
                     # discriminator_steps = gr.Slider(minimum=0,
@@ -196,7 +198,7 @@ with gr.Blocks(css="styles.css") as demo:
                     #                         step=1,
                     #                         value=0,
                     #                         label="Steps to run at the end, optimizing only the discriminator loss. This helps to reduce artefacts, but because the model is trained on CelebA, this will make your generations look more like generic white celebrities")
-    clear.click(state.clear_transforms, inputs=[state], outputs=[state, out, mask])
+    clear.click(StateWrapper.clear_transforms, inputs=[state], outputs=[state, out, mask])
     asian_weight.change(StateWrapper.apply_asian_vector, inputs=[state, asian_weight], outputs=[state, out, mask])
     lip_size.change(StateWrapper.apply_lip_vector, inputs=[state, lip_size], outputs=[state, out, mask])
     # hair_green_purple.change(StateWrapper.apply_gp_vector, inputs=[state, hair_green_purple], outputs=[state, out, mask])

configs.py

@@ -4,4 +4,4 @@ def set_small_local():
 def set_major_local():
     return (gr.Slider.update(value=25), gr.Slider.update(value=0.2), 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=2), gr.Slider.update(value=0.2))
+    return (gr.Slider.update(value=30), gr.Slider.update(value=0.1), gr.Slider.update(value=1), gr.Slider.update(value=1))