update config and rm discriminator

app.py

@@ -139,7 +139,7 @@ with gr.Blocks(css="styles.css") as demo:
                         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=300,
+                                            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",)

app_backend.py

@@ -81,7 +81,6 @@ class ImagePromptOptimizer(nn.Module):
         self.make_grid = make_grid
         self.return_val = return_val
         self.quantize = quantize
-        self.disc = load_disc(self.device)
         self.lpips_weight = lpips_weight
         self.perceptual_loss = lpips.LPIPS(net='vgg').to(self.device)
     def disc_loss_fn(self, logits):
@@ -175,19 +174,19 @@ class ImagePromptOptimizer(nn.Module):
             clip_clone = processed_img.clone()
             clip_clone.register_hook(self.attn_masking)
             clip_clone.retain_grad()
-            with torch.autocast("cuda"):
-                clip_loss = self.get_similarity_loss(pos_prompts, neg_prompts, clip_clone)
-                print("CLIP loss", clip_loss)
-                perceptual_loss = self.perceptual_loss(lpips_input, original_img.clone()) * self.lpips_weight
-                print("LPIPS loss: ", perceptual_loss)
-                with torch.no_grad():
-                    disc_logits = self.disc(transformed_img)
-                    disc_loss = self.disc_loss_fn(disc_logits)
-                    print(f"disc_loss = {disc_loss}")
-                    disc_loss2 = self.disc(processed_img)
+            # with torch.autocast("cuda"):
+            clip_loss = self.get_similarity_loss(pos_prompts, neg_prompts, clip_clone)
+            print("CLIP loss", clip_loss)
+            perceptual_loss = self.perceptual_loss(lpips_input, original_img.clone()) * self.lpips_weight
+            print("LPIPS loss: ", perceptual_loss)
+                # with torch.no_grad():
+                #     disc_logits = self.disc(transformed_img)
+                #     disc_loss = self.disc_loss_fn(disc_logits)
+                #     print(f"disc_loss = {disc_loss}")
+                #     disc_loss2 = self.disc(processed_img)
             if log:
                 wandb.log({"Perceptual Loss": perceptual_loss})
-                wandb.log({"Discriminator Loss": disc_loss})
+                # wandb.log({"Discriminator Loss": disc_loss})
                 wandb.log({"CLIP Loss": clip_loss})
             clip_loss.backward(retain_graph=True)
             perceptual_loss.backward(retain_graph=True)
@@ -209,13 +208,13 @@ class ImagePromptOptimizer(nn.Module):
             lpips_input = processed_img.clone()
             lpips_input.register_hook(self.attn_masking2)
             lpips_input.retain_grad()
-            with torch.autocast("cuda"):
-                perceptual_loss = self.perceptual_loss(lpips_input, original_img.clone()) * self.lpips_weight
-                with torch.no_grad():
-                    disc_logits = self.disc(transformed_img)
-                    disc_loss = self.disc_loss_fn(disc_logits)
-                    print(f"disc_loss = {disc_loss}")
-                    disc_loss2 = self.disc(processed_img)
+            # with torch.autocast("cuda"):
+            perceptual_loss = self.perceptual_loss(lpips_input, original_img.clone()) * self.lpips_weight
+                # with torch.no_grad():
+                #     disc_logits = self.disc(transformed_img)
+                #     disc_loss = self.disc_loss_fn(disc_logits)
+                #     print(f"disc_loss = {disc_loss}")
+                #     disc_loss2 = self.disc(processed_img)
             # print(f"disc_loss2 = {disc_loss2}")
             if log:
                 wandb.log({"Perceptual Loss": perceptual_loss})

configs.py

@@ -2,6 +2,6 @@ import gradio as gr
 def set_small_local():
     return (gr.Slider.update(value=25), gr.Slider.update(value=0.15), gr.Slider.update(value=1), gr.Slider.update(value=4))
 def set_major_local():
-    return (gr.Slider.update(value=25), gr.Slider.update(value=0.25), gr.Slider.update(value=35), gr.Slider.update(value=10))
+    return (gr.Slider.update(value=25), gr.Slider.update(value=0.2), gr.Slider.update(value=36.6), gr.Slider.update(value=10))
 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))