imprived speed of mdn sampling

app.py

@@ -30,18 +30,22 @@ for writer_id in [120, 80]:
     all_loaded_data.append(loaded_data)
 
 default_loaded_data = all_loaded_data[-1]
+mdn_words = []
+mdn_mean_Ws = []
+all_word_mdn_Ws = []
+all_word_mdn_Cs = []
 
 # data for writer interpolation
-words = []
+writer_words = []
 writer_mean_Ws = []
 all_word_writer_Ws = []
 all_word_writer_Cs = []
 weight = 0.7
 
 def update_target_word(target_word):
-    words.clear()
+    writer_words.clear()
     for word in target_word.split(" "):
-        words.append(word)
+        writer_words.append(word)
 
     writer_mean_Ws.clear()
     for loaded_data in all_loaded_data:
@@ -50,7 +54,7 @@ def update_target_word(target_word):
 
     all_word_writer_Ws.clear()
     all_word_writer_Cs.clear()
-    for word in words:
+    for word in writer_words:
         all_writer_Ws, all_writer_Cs = convenience.get_DSD(net, word, writer_mean_Ws, all_loaded_data, device)
         all_word_writer_Ws.append(all_writer_Ws)
         all_word_writer_Cs.append(all_writer_Cs)
@@ -62,20 +66,8 @@ def update_target_word(target_word):
 def update_writer_slider(val):
     global weight
     weight = val
-    im = Image.fromarray(np.zeros([160, 750]))
-    dr = ImageDraw.Draw(im)
-    width = 50
-
-    for word, all_writer_Ws, all_writer_Cs in zip(words, all_word_writer_Ws, all_word_writer_Cs):
-        all_W_c = convenience.get_writer_blend_W_c([1 - weight, weight], all_writer_Ws, all_writer_Cs)
-        all_commands = convenience.get_commands(net, word, all_W_c)
-
-        for [x, y, t] in all_commands:
-            if t == 0:
-                dr.line((px+width, py, x+width, y), 255, 1)
-            px, py = x, y
-        width += np.max(all_commands[:, 0]) + 25
-
+    net.clamp_mdn = 0
+    im = convenience.draw_words(writer_words, all_word_writer_Ws, all_word_writer_Cs, [1 - weight, weight], net)
     return im.convert("RGB")
 
 def update_chosen_writers(writer1, writer2):
@@ -87,37 +79,77 @@ def update_chosen_writers(writer1, writer2):
         all_loaded_data.append(loaded_data)
 
     return gr.Slider.update(label=f"{writer1} vs. {writer2}"), update_writer_slider(weight)
-'''
-def choose_writer(writ, c1, c2, c3, c4, val):
-    all_loaded_data.clear()
-    w = int(writ.split(" ")[1])
-    loaded_data = dl.next_batch(TYPE='TRAIN', uid=w, tids=list(range(num_samples)))
-    all_loaded_data.append(loaded_data)
-    return char_grid(c1, c2, c3, c4, val)
-'''
 
 # for character blend
 def interpolate_chars(c1, c2, weight):
+    """Generates an image of handwritten text based on target_sentence"""
+
     net.clamp_mdn = 0
-    image = convenience.sample_blended_chars([1 - weight, weight], [c1, c2], net, [default_loaded_data], device).convert("RGB")
-    return image
+
+    letters = [c1, c2]
+    character_weights = [1 - weight, weight]
+    M = len(letters)
+    mean_global_W = convenience.get_mean_global_W(net, all_loaded_data[0], device)
+
+    all_Cs = torch.zeros(1, M, convenience.L, convenience.L)
+    for i in range(M):  # get corners of grid
+        W_vector, char_matrix =  convenience.get_DSD(net, letters[i], [mean_global_W], [default_loaded_data], device)
+        all_Cs[:, i, :, :] = char_matrix
+
+    all_Ws = mean_global_W.reshape(1, 1, convenience.L)
+
+    all_W_c = convenience.get_character_blend_W_c(character_weights, all_Ws, all_Cs)
+    all_commands = convenience.get_commands(net, letters[0], all_W_c)
+
+    width = 60
+    x_offset = 325
+    im = Image.fromarray(np.zeros([160, 750]))
+    dr = ImageDraw.Draw(im)
+    for [x, y, t] in all_commands:
+        if t == 0:
+            dr.line((
+                px + width/2 + x_offset,
+                py - width/2,  # letters are shifted down for some reason
+                x + width/2 + + x_offset,
+                y - width/2), 255, 1)
+        px, py = x, y
+
+        
+    return im.convert("RGB")
 
 def choose_blend_chars(c1, c2):
    return gr.Slider.update(label=f"'{c1}' vs. '{c2}'")
 
 # for MDN
-def mdn_(word, maxs, maxr):
-    image = convenience.mdn_single_sample(word, maxs, maxr, net, [default_loaded_data], device).convert("RGB")
-    return image
-"""
-def char_vid(word):
-    #make word list
-    convenience.char_interpolation_video(word_list, 10, net, [default_loaded_data], device).convert('RGB')
-    vid_path = f"/content/drive/MyDrive/Colab Notebooks/Spring22/decoupled-style-descriptors-eb/results/abcdefg_video.mov"
-    return gr.Video.update(value=vid_path)
-"""
+
+def update_mdn_word(target_word):
+    mdn_words.clear()
+    for word in target_word.split(" "):
+        mdn_words.append(word)
+
+    mdn_mean_Ws.clear()
+    mean_global_W = convenience.get_mean_global_W(net, default_loaded_data, device)
+    mdn_mean_Ws.append(mean_global_W)
+
+    all_word_mdn_Ws.clear()
+    all_word_mdn_Cs.clear()
+    for word in mdn_words:
+        all_writer_Ws, all_writer_Cs = convenience.get_DSD(net, word, mdn_mean_Ws, [default_loaded_data], device)
+        all_word_mdn_Ws.append(all_writer_Ws)
+        all_word_mdn_Cs.append(all_writer_Cs)
+
+    return sample_mdn(net.scale_sd, net.clamp_mdn)
+
+
+def sample_mdn(maxs, maxr):
+    net.clamp_mdn = maxr
+    net.scale_sd = maxs
+    im = convenience.draw_words(mdn_words, all_word_mdn_Ws, all_word_mdn_Cs, [1], net)
+    return im.convert("RGB")
+
 
 update_target_word("hello world")
+update_mdn_word("hello world")
 
 with gr.Blocks() as demo:
     with gr.Tabs():
@@ -163,39 +195,6 @@ with gr.Blocks() as demo:
             char1.change(fn=choose_blend_chars, inputs=[char1, char2], outputs=[char_slider])
             char2.change(fn=choose_blend_chars, inputs=[char1, char2], outputs=[char_slider])
         
-        """
-        with gr.TabItem("Character Grid"): #slow
-          with gr.Row():
-              with gr.Column():
-                  char1 = gr.Dropdown(choices=avail_char_list, value="y", label="Character 1")
-              with gr.Column():
-                  char2 = gr.Dropdown(choices=avail_char_list, value="s", label="Character 2")
-              with gr.Column():
-                  char3 = gr.Dropdown(choices=avail_char_list, value="u", label="Character 3")
-              with gr.Column():
-                  char4 = gr.Dropdown(choices=avail_char_list, value="n", label="Character 4")
-          with gr.Row(): 
-              submit_button = gr.Button(value="Blend y, s, u, and n!")
-          '''
-          with gr.Row():
-              radio_options2 = ["Writer " + str(n) for n in writer_options]
-              writer = gr.Radio(radio_options2, value="Writer 80", label="Style for Writer")
-              writer.change(fn=choose_writer, inputs=[writer, char1, char2, char3, char4, slider2], outputs=[output])
-          '''
-          #slider2 = gr.Slider(2, 20, value=10, label="Grid Size", step=1)
-
-          default_image = convenience.sample_character_grid(['y', 's', 'u', 'n'], 10, net, [default_loaded_data], device).convert("RGB")
-          output = gr.Image(default_image)
-          
-          char1.change(fn=choose_grid_chars, inputs=[char1, char2, char3, char4], outputs=[submit_button])
-          char2.change(fn=choose_grid_chars, inputs=[char1, char2, char3, char4], outputs=[submit_button])
-          char3.change(fn=choose_grid_chars, inputs=[char1, char2, char3, char4], outputs=[submit_button])
-          char4.change(fn=choose_grid_chars, inputs=[char1, char2, char3, char4], outputs=[submit_button])
-          
-          
-          #slider2.change(fn=char_grid, inputs=[char1, char2, char3, char4, slider2], outputs=[output])
-          submit_button.click(fn=char_grid, inputs=[char1, char2, char3, char4], outputs=[output])
-        """
         
         with gr.TabItem("Add Randomness"):
           mdn_word = gr.Textbox(label="Target Word", value="hello world", max_lines=1)
@@ -211,14 +210,14 @@ with gr.Blocks() as demo:
             with gr.Column():
               scale_rand = gr.Slider(0, 3, value=0.5, label="Scale of Randomness")
           with gr.Row():
-              _button = gr.Button(value="Re!")
+              mdn_sample_button = gr.Button(value="Resample!")
           with gr.Row():
               default_im = convenience.mdn_single_sample("hello world", 0.5, 1, net, [default_loaded_data], device).convert('RGB')
               mdn_output = gr.Image(default_im)
 
-          max_rand.change(fn=mdn_, inputs=[mdn_word, scale_rand, max_rand], outputs=[mdn_output])
-          scale_rand.change(fn=mdn_, inputs=[mdn_word, scale_rand, max_rand], outputs=[mdn_output])        
-          _button.click(fn=mdn_, inputs=[mdn_word, scale_rand, max_rand], outputs=[mdn_output])
-          mdn_word.submit(fn=mdn_, inputs=[mdn_word, scale_rand, max_rand], outputs=[mdn_output])
+          max_rand.change(fn=sample_mdn, inputs=[scale_rand, max_rand], outputs=[mdn_output])
+          scale_rand.change(fn=sample_mdn, inputs=[scale_rand, max_rand], outputs=[mdn_output])        
+          mdn_sample_button.click(fn=sample_mdn, inputs=[scale_rand, max_rand], outputs=[mdn_output])
+          mdn_word.submit(fn=update_mdn_word, inputs=[mdn_word], outputs=[mdn_output])
 
 demo.launch()

convenience.py

@@ -295,26 +295,19 @@ def mdn_video(target_word, num_samples, scale_sd, clamp_mdn, net, all_loaded_dat
     us_target_word = re.sub(r"\s+", '_', target_word)
     os.makedirs(f"./results/{us_target_word}_mdn_samples", exist_ok=True)
     for i in range(num_samples):
-        im = Image.fromarray(np.zeros([160, 750]))
-        dr = ImageDraw.Draw(im)
-        width = 50
-
         net.scale_sd = scale_sd
         net.clamp_mdn = clamp_mdn
 
         mean_global_W = get_mean_global_W(net, all_loaded_data[0], device)
 
+        word_Ws = []
+        word_Cs = []
         for word in words:
             writer_Ws, writer_Cs = get_DSD(net, word, [mean_global_W], [all_loaded_data[0]], device)
-            all_W_c = get_writer_blend_W_c([1], writer_Ws, writer_Cs)
-            all_commands = get_commands(net, word, all_W_c)
-
-            for [x, y, t] in all_commands:
-                if t == 0:
-                    dr.line((px+width, py, x+width, y), 255, 1)
-                px, py = x, y
-            width += np.max(all_commands[:, 0]) + 25
-
+            word_Ws.append(writer_Ws)
+            word_Cs.append(writer_Cs)
+            
+        im = draw_words(words, word_Ws, word_Cs, [1], net)
         im.convert("RGB").save(f'results/{us_target_word}_mdn_samples/sample_{i}.png')
     # Convert fromes to video using ffmpeg
     photos = ffmpeg.input(f'results/{us_target_word}_mdn_samples/sample_*.png', pattern_type='glob', framerate=10)
@@ -325,27 +318,19 @@ def sample_blended_writers(writer_weights, target_sentence, net, all_loaded_data
     """Generates an image of handwritten text based on target_sentence"""
     words = target_sentence.split(' ')
 
-    im = Image.fromarray(np.zeros([160, 750]))
-    dr = ImageDraw.Draw(im)
-    width = 50
-
     writer_mean_Ws = []
     for loaded_data in all_loaded_data:
         mean_global_W = get_mean_global_W(net, loaded_data, device)
         writer_mean_Ws.append(mean_global_W)
 
+    word_Ws = []
+    word_Cs = []
     for word in words:
-        all_writer_Ws, all_writer_Cs = get_DSD(net, word, writer_mean_Ws, all_loaded_data, device)
-        all_W_c = get_writer_blend_W_c(writer_weights, all_writer_Ws, all_writer_Cs)
-        all_commands = get_commands(net, word, all_W_c)
-
-        for [x, y, t] in all_commands:
-            if t == 0:
-                dr.line((px+width, py, x+width, y), 255, 1)
-            px, py = x, y
-        width += np.max(all_commands[:, 0]) + 25
-
-    return im
+        writer_Ws, writer_Cs = get_DSD(net, word, writer_mean_Ws, all_loaded_data, device)
+        word_Ws.append(writer_Ws)
+        word_Cs.append(writer_Cs)
+        
+    return draw_words(words, word_Ws, word_Cs, writer_weights, net)
 
 
 def sample_character_grid(letters, grid_size, net, all_loaded_data, device="cpu"):
@@ -418,26 +403,12 @@ def writer_interpolation_video(target_sentence, transition_time, net, all_loaded
 
     for i in range(n - 1):
         for j in range(transition_time):
-            im = Image.fromarray(np.zeros([160, 750]))
-            dr = ImageDraw.Draw(im)
-            width = 50
-
             completion = j/(transition_time)
 
             individual_weights = [1 - completion, completion]
             writer_weights = [0] * i + individual_weights + [0] * (n - 2 - i)
 
-            for k, word in enumerate(words):
-                all_writer_Ws, all_writer_Cs = word_Ws[k], word_Cs[k]
-                all_W_c = get_writer_blend_W_c(writer_weights, all_writer_Ws, all_writer_Cs)
-                all_commands = get_commands(net, word, all_W_c)
-
-                for [x, y, t] in all_commands:
-                    if t == 0:
-                        dr.line((px+width, py, x+width, y), 255, 1)
-                    px, py = x, y
-                width += np.max(all_commands[:, 0]) + 25
-
+            im = draw_words(words, word_Ws, word_Cs, writer_weights, net)
             im.convert("RGB").save(f"./results/{target_sentence}_blend_frames/frame_{str(i * transition_time + j).zfill(3)}.png")
 
     # Convert fromes to video using ffmpeg
@@ -452,35 +423,25 @@ def mdn_single_sample(target_word, scale_sd, clamp_mdn, net, all_loaded_data, de
     max_scale: the maximum value used to scale SD while sampling (increment is based on num samples)
     '''
     words = target_word.split(' ')
-    im = Image.fromarray(np.zeros([160, 750]))
-    dr = ImageDraw.Draw(im)
-    width = 50
 
     net.scale_sd = scale_sd
     net.clamp_mdn = clamp_mdn
 
     mean_global_W = get_mean_global_W(net, all_loaded_data[0], device)
 
+    word_Ws = []
+    word_Cs = []
     for word in words:
         writer_Ws, writer_Cs = get_DSD(net, word, [mean_global_W], [all_loaded_data[0]], device)
-        all_W_c = get_writer_blend_W_c([1], writer_Ws, writer_Cs)
-        all_commands = get_commands(net, word, all_W_c)
+        word_Ws.append(writer_Ws)
+        word_Cs.append(writer_Cs)
 
-        for [x, y, t] in all_commands:
-            if t == 0:
-                dr.line((px+width, py, x+width, y), 255, 1)
-            px, py = x, y
-        width += np.max(all_commands[:, 0]) + 25
+    return draw_words(words, word_Ws, word_Cs, [1], net)
 
-    return im
 
 def sample_blended_chars(character_weights, letters, net, all_loaded_data, device="cpu"):
     """Generates an image of handwritten text based on target_sentence"""
 
-    width = 60
-    im = Image.fromarray(np.zeros([100, 100]))
-    dr = ImageDraw.Draw(im)
-
     M = len(letters)
     mean_global_W = get_mean_global_W(net, all_loaded_data[0], device)
 
@@ -494,6 +455,9 @@ def sample_blended_chars(character_weights, letters, net, all_loaded_data, devic
     all_W_c = get_character_blend_W_c(character_weights, all_Ws, all_Cs)
     all_commands = get_commands(net, letters[0], all_W_c)
 
+    width = 60
+    im = Image.fromarray(np.zeros([100, 100]))
+    dr = ImageDraw.Draw(im)
     for [x, y, t] in all_commands:
         if t == 0:
             dr.line((
@@ -512,8 +476,6 @@ def char_interpolation_video(letters, transition_time, net, all_loaded_data, dev
 
     os.makedirs(f"./results/{''.join(letters)}_frames", exist_ok=True) # make a folder for the frames
 
-    width = 50
-
     M = len(letters)
     mean_global_W = get_mean_global_W(net, all_loaded_data[0], device)
 
@@ -530,11 +492,11 @@ def char_interpolation_video(letters, transition_time, net, all_loaded_data, dev
             individual_weights = [1 - completion, completion]
             character_weights = [0] * i + individual_weights + [0] * (M - 2 - i)
             all_W_c = get_character_blend_W_c(character_weights, all_Ws, all_Cs)
-            all_commands = get_commands(net, "change this later!", all_W_c)
+            all_commands = get_commands(net, letters[i], all_W_c)
 
             im = Image.fromarray(np.zeros([100, 100]))
             dr = ImageDraw.Draw(im)
-
+            width = 50
             for [x, y, t] in all_commands:
                 if t == 0:
                     dr.line((
@@ -553,3 +515,18 @@ def char_interpolation_video(letters, transition_time, net, all_loaded_data, dev
     videos.run(overwrite_output=True)
 
 
+def draw_words(words, word_Ws, word_Cs, writer_weights, net):
+    im = Image.fromarray(np.zeros([160, 750]))
+    dr = ImageDraw.Draw(im)
+    width = 50
+    for word, all_writer_Ws, all_writer_Cs in zip(words, word_Ws, word_Cs):
+        all_W_c = get_writer_blend_W_c(writer_weights, all_writer_Ws, all_writer_Cs)
+        all_commands = get_commands(net, word, all_W_c)
+
+        for [x, y, t] in all_commands:
+            if t == 0:
+                dr.line((px+width, py, x+width, y), 255, 1)
+            px, py = x, y
+        width += np.max(all_commands[:, 0]) + 25
+
+    return im