clip latents differently

3 files changed, 17 insertions, 7 deletions

diff --git a/cli/app/search/search_class.py b/cli/app/search/search_class.py
index f7b2136..105fc9f 100644
--- a/cli/app/search/search_class.py
+++ b/cli/app/search/search_class.py
@@ -96,6 +96,8 @@ def find_nearest_vector(sess, generator, opt_fp_in, opt_dims, out_images, out_la
     'truncation': input_trunc,
   })
 
+  target = tf.compat.v1.placeholder(tf.float32, shape=(batch_size, img_size, img_size, num_channels))
+
   ## clip the Z encoding
 
   opt_clip = 1.0
@@ -104,10 +106,10 @@ def find_nearest_vector(sess, generator, opt_fp_in, opt_dims, out_images, out_la
     tf.random.uniform([batch_size, z_dim], minval=-opt_clip, maxval=opt_clip), input_z)
   clip_latent = tf.assign(input_z, clipped_encoding)
 
-  target = tf.compat.v1.placeholder(tf.float32, shape=(batch_size, img_size, img_size, num_channels))
-
   ## normalize the Y encoding
-  normalized_labels = tf.nn.l2_normalize(input_y)
+  # normalized_labels = tf.nn.l2_normalize(input_y)
+  # tf.reduce_mean(tf.abs(encoding - gen_encoding))
+  normalized_labels = input_y / tf.reduce_sum(input_y)
   clip_labels = tf.assign(input_y, normalized_labels)
 
   ## if computing Feature loss, use these encoders
diff --git a/cli/app/search/search_dense.py b/cli/app/search/search_dense.py
index ac4dc59..46183c7 100644
--- a/cli/app/search/search_dense.py
+++ b/cli/app/search/search_dense.py
@@ -124,6 +124,8 @@ def find_dense_embedding_for_images(params):
     else:
       gen_img = generator(latent, signature=gen_signature)
 
+  gen_img_orig = gen_img
+  
   # Convert generated image to channels_first.
   gen_img = tf.transpose(gen_img, [0, 3, 1, 2])
 
@@ -393,7 +395,7 @@ def find_dense_embedding_for_images(params):
     ])
 
     # Main optimization loop.
-    print("Total iterations: {}".format(params.inv_it))
+    print("Beginning dense iteration...")
     for _ in range(params.inv_it):
 
       _inv_loss, _mse_loss, _feat_loss,\
@@ -435,15 +437,14 @@ def find_dense_embedding_for_images(params):
     out_labels[out_pos:out_pos+BATCH_SIZE] = label_batch
     out_err[out_pos:out_pos+BATCH_SIZE] = rec_err_batch
 
-    gen_images  = sess.run(gen_img)
+    gen_images  = sess.run(gen_img_orig)
     images = vs.data2img(gen_images)
 
     # write encoding, latent to pkl file
     for i in range(BATCH_SIZE):
       out_i = out_pos + i
       sample_fn, ext = os.path.splitext(sample_fns[out_i])
-      image = Image.fromarray(images[i])
-      image = vs.grid_transform(image)
+      image = Image.fromarray(image)
       fp = BytesIO()
       image.save(fp, format='png')
       data = upload_bytes_to_cortex(params.folder_id, sample_fn + "-inverse.png", fp, "image/png")
diff --git a/inversion/live.py b/inversion/live.py
index ffac0af..e8001ef 100644
--- a/inversion/live.py
+++ b/inversion/live.py
@@ -212,6 +212,13 @@ class Interpolator:
     gen_in['y'] = lerp_label.output
     self.gen_img = generator(gen_in, signature=gen_signature)
 
+    gen_layer_name = 'module_apply_' + gen_signature + '/' + params.inv_layer
+    gen_encoding = tf.get_default_graph().get_tensor_by_name(gen_layer_name)
+    ENC_SHAPE = gen_encoding.get_shape().as_list()[1:]
+    set_encoding = tf.get_variable(name='encoding', dtype=tf.float32, shape=[BATCH_SIZE,] + ENC_SHAPE)
+    encoding_mix = LerpParam('encoding_mix', a_in=gen_encoding, b_in=set_encoding, shape=[BATCH_SIZE,] + ENC_SHAPE, datatype="input")
+    tf.contrib.graph_editor.swap_ts(gen_encoding, encoding_mix)
+
     sys.stderr.write("Sin params: {}\n".format(", ".join(self.sin_params.keys())))
     sys.stderr.write("Lerp params: {}\n".format(", ".join(self.lerp_params.keys())))
     sys.stderr.write("Opts: {}\n".format(", ".join(self.opts.keys())))