refactoring optimization loop

ImageState.py

@@ -102,7 +102,7 @@ class ImageState:
         x = Image.fromarray(x, "L")
         return x
 
-    @torch.inference_mode()
+    @torch.no_grad()
     def _render_all_transformations(self, return_twice=True):
         global num
         current_vector_transforms = (
@@ -150,7 +150,7 @@ class ImageState:
             clear_img_dir(self.img_dir)
         return self.blend(blend_weight)
 
-    @torch.inference_mode()
+    @torch.no_grad()
     def blend(self, weight):
         _, latent = blend_paths(
             self.vqgan,
@@ -163,7 +163,7 @@ class ImageState:
         self.blend_latent = latent
         return self._render_all_transformations()
 
-    @torch.inference_mode()
+    @torch.no_grad()
     def rewind(self, index):
         if not self.transform_history:
             print("No history")
@@ -221,7 +221,7 @@ class ImageState:
         ):
             transform_log.transforms.append(transform.detach().cpu())
             self.current_prompt_transforms[-1] = transform
-            with torch.inference_mode():
+            with torch.no_grad():
                 image = self._render_all_transformations(return_twice=False)
             if log:
                 wandb.log({"image": wandb.Image(image)})

animation.py

@@ -4,10 +4,6 @@ import os
 
 
 def clear_img_dir(img_dir):
-    if not os.path.exists("img_history"):
-        os.mkdir("img_history")
-    if not os.path.exists(img_dir):
-        os.mkdir(img_dir)
     for filename in glob.glob(img_dir + "/*"):
         os.remove(filename)
 

backend.py

@@ -140,7 +140,7 @@ class ImagePromptEditor(nn.Module):
         return newgrad
 
     def _get_next_inputs(self, transformed_img):
-        processed_img = loop_post_process(transformed_img)  # * self.attn_mask
+        processed_img = loop_post_process(transformed_img)
         processed_img.retain_grad()
 
         lpips_input = processed_img.clone()
@@ -154,51 +154,53 @@ class ImagePromptEditor(nn.Module):
         return (processed_img, lpips_input, clip_input)
 
     def _optimize_CLIP_LPIPS(self, optim, original_img, vector, pos_prompts, neg_prompts):
-        optim.zero_grad()
-        transformed_img = self(vector)
-        processed_img, lpips_input, clip_input = self._get_next_inputs(
-            transformed_img
-        )
-        with torch.autocast("cuda"):
-            clip_loss = self._get_CLIP_loss(pos_prompts, neg_prompts, clip_input)
-            print("CLIP loss", clip_loss)
-            perceptual_loss = (
-                self.perceptual_loss(lpips_input, original_img.clone())
-                * self.lpips_weight
+        for i in (range(self.iterations)):
+            optim.zero_grad()
+            transformed_img = self(vector)
+            processed_img, lpips_input, clip_input = self._get_next_inputs(
+                transformed_img
             )
-            print("LPIPS loss: ", perceptual_loss)
-            print("Sum Loss", perceptual_loss + clip_loss)
-        if log:
-            wandb.log({"Perceptual Loss": perceptual_loss})
-            wandb.log({"CLIP Loss": clip_loss})
-        
-        # These gradients will be masked if attn_mask has been set
-        clip_loss.backward(retain_graph=True)
-        perceptual_loss.backward(retain_graph=True)
-
-        optim.step()
-        yield vector
+            with torch.autocast("cuda"):
+                clip_loss = self._get_CLIP_loss(pos_prompts, neg_prompts, clip_input)
+                print("CLIP loss", clip_loss)
+                perceptual_loss = (
+                    self.perceptual_loss(lpips_input, original_img.clone())
+                    * self.lpips_weight
+                )
+                print("LPIPS loss: ", perceptual_loss)
+                print("Sum Loss", perceptual_loss + clip_loss)
+            if log:
+                wandb.log({"Perceptual Loss": perceptual_loss})
+                wandb.log({"CLIP Loss": clip_loss})
+            
+            # These gradients will be masked if attn_mask has been set
+            clip_loss.backward(retain_graph=True)
+            perceptual_loss.backward(retain_graph=True)
+
+            optim.step()
+            yield vector
 
     def _optimize_LPIPS(self, vector, original_img, optim):
-        optim.zero_grad()
-        transformed_img = self(vector)
-        processed_img = loop_post_process(transformed_img)  # * self.attn_mask
-        processed_img.retain_grad()
-
-        lpips_input = processed_img.clone()
-        lpips_input.register_hook(self._attn_mask_inverse)
-        lpips_input.retain_grad()
-        with torch.autocast("cuda"):
-            perceptual_loss = (
-                self.perceptual_loss(lpips_input, original_img.clone())
-                * self.lpips_weight
-            )
-        if log:
-            wandb.log({"Perceptual Loss": perceptual_loss})
-        print("LPIPS loss: ", perceptual_loss)
-        perceptual_loss.backward(retain_graph=True)
-        optim.step()
-        yield vector
+        for i in range(self.reconstruction_steps):
+            optim.zero_grad()
+            transformed_img = self(vector)
+            processed_img = loop_post_process(transformed_img)
+            processed_img.retain_grad()
+
+            lpips_input = processed_img.clone()
+            lpips_input.register_hook(self._attn_mask_inverse)
+            lpips_input.retain_grad()
+            with torch.autocast("cuda"):
+                perceptual_loss = (
+                    self.perceptual_loss(lpips_input, original_img.clone())
+                    * self.lpips_weight
+                )
+            if log:
+                wandb.log({"Perceptual Loss": perceptual_loss})
+            print("LPIPS loss: ", perceptual_loss)
+            perceptual_loss.backward(retain_graph=True)
+            optim.step()
+            yield vector
 
     def optimize(self, latent, pos_prompts, neg_prompts):
         self.set_latent(latent)
@@ -209,10 +211,10 @@ class ImagePromptEditor(nn.Module):
         vector = torch.randn_like(self.latent, requires_grad=True, device=self.device)
         optim = torch.optim.Adam([vector], lr=self.lr)
 
-        for i in tqdm(range(self.iterations)):
-            yield self._optimize_CLIP_LPIPS(optim, original_img, vector, pos_prompts, neg_prompts)
+        for transform in self._optimize_CLIP_LPIPS(optim, original_img, vector, pos_prompts, neg_prompts):
+            yield transform
 
         print("Running LPIPS optim only")
-        for i in range(self.reconstruction_steps):
-            yield self._optimize_LPIPS(vector, original_img, transformed_img, optim)
+        for transform in self._optimize_LPIPS(vector, original_img, optim):
+            yield transform
         yield vector if self.return_val == "vector" else self.latent + vector