Adjusting inference params

app.py

@@ -34,32 +34,14 @@ from utils.sr_utils import *
 
 device = torch.device('cuda')
 
-# torch.hub.download_url_to_file('https://images.pexels.com/photos/68767/divers-underwater-ocean-swim-68767.jpeg', 'coralreef.jpeg')
-
-# def fetch(url_or_path):
-#     if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):
-#         r = requests.get(url_or_path)
-#         r.raise_for_status()
-#         fd = io.BytesIO()
-#         fd.write(r.content)
-#         fd.seek(0)
-#         return fd
-#     return open(url_or_path, 'rb')
-
-# def parse_prompt(prompt):
-#     if prompt.startswith('http://') or prompt.startswith('https://'):
-#         vals = prompt.rsplit(':', 2)
-#         vals = [vals[0] + ':' + vals[1], *vals[2:]]
-#     else:
-#         vals = prompt.rsplit(':', 1)
-#     vals = vals + ['', '1'][len(vals):]
-#     return vals[0], float(vals[1])
 clip_model_vit_b_32 = clip.load('ViT-B/32', device=device)[0].eval().requires_grad_(False)
 clip_model_vit_b_16 = clip.load('ViT-B/16', device=device)[0].eval().requires_grad_(False)
 clip_models = {'ViT-B/32': clip_model_vit_b_32, 'ViT-B/16': clip_model_vit_b_16}
 
 clip_normalize = T.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])
-
+clip_model = 'ViT-B/16'
+sideX, sideY = 512, 512  # Resolution
+inv_color_scale = 1.6
 
 class MakeCutouts(torch.nn.Module):
     def __init__(self, cut_size, cutn):
@@ -149,7 +131,18 @@ class CLIPActivationLoss(nn.Module):
         return -loss if self.maximize else loss
 
 
-def optimize_network(seed, num_iterations, optimizer_type, lr):
+def optimize_network(
+    seed, 
+    opt_type, 
+    lr, 
+    num_iterations, 
+    cutn, 
+    layer,
+    neuron,
+    class_token,
+    maximize,
+    display_rate = 20
+):
     global itt
     itt = 0
 
@@ -183,9 +176,9 @@ def optimize_network(seed, num_iterations, optimizer_type, lr):
     # Initialize input noise
     net_input = torch.zeros([1, input_depth, sideY, sideX], device=device).normal_().div(10).detach()
 
-    if optimizer_type == 'Adam':
+    if opt_type == 'Adam':
         optimizer = torch.optim.Adam(net.parameters(), lr)
-    elif optimizer_type == 'MADGRAD':
+    elif opt_type == 'MADGRAD':
         optimizer = MADGRAD(net.parameters(), lr, momentum=0.9)
     scaler = torch.cuda.amp.GradScaler()
 
@@ -234,7 +227,6 @@ def inference(
     lr, 
     num_iterations, 
     cutn, 
-    clip_model,
     layer,
     neuron,
     class_token,
@@ -249,14 +241,26 @@ def inference(
     # Begin optimization / generation
     gc.collect()
     torch.cuda.empty_cache()
-    out = optimize_network(seed, num_iterations, opt_type, lr)
+    out = optimize_network(
+            seed, 
+            opt_type, 
+            lr, 
+            num_iterations, 
+            cutn, 
+            layer,
+            neuron,
+            class_token,
+            maximize,
+            display_rate
+        )
+
     out.save(f'dip_{timestring}.png', quality=100)
     if save_progress_video:
         video_writer.close()
     return out 
 
 iface = gr.Interface(fn=inference, 
-                     inputs=["number", "text", "number", "number", "number", "text", "number", "number", 
+                     inputs=["number", "text", "number", "number", "number", "number", "number", 
                                 gr.inputs.Checkbox(default=False, label="class_token"),
                                 gr.inputs.Checkbox(default=True, label="maximise"),
                                 "number"], 

requirements.txt

@@ -8,4 +8,5 @@ lpips
 numpy
 imageio
 einops
-madgrad
+madgrad
+cv2