Fixed timing

app.py

@@ -105,7 +105,8 @@ def diffusion_chat(question, noising, max_it, pause_length):
     )
     yield render_html("Iteration 0 (initial noise)",
                       highlight_tokens(current_tokens[answer_start:], answer_start, just_noised_indices, color="red"))
-    time.sleep(pause_length)
+    
+    start = time.perf_counter()
 
     last_tokens = []
     prev_decoded = []
@@ -113,9 +114,10 @@ def diffusion_chat(question, noising, max_it, pause_length):
     unmasked_mask = [False] * len(current_tokens)
 
     for i in range(max_it):
-
+        
         generated_tokens, confidences = generate_diffusion_text(current_tokens, top_p, top_k)
         current_tokens = ori_input_tokens[:answer_start] + generated_tokens[answer_start:]
+        
 
         # GREEN highlighting: compare to previous tokens
         new_decoded = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:])
@@ -125,6 +127,8 @@ def diffusion_chat(question, noising, max_it, pause_length):
         }
         prev_decoded = new_decoded
 
+        time.sleep(max(pause_length - (time.perf_counter() - start), 0))
+
         yield render_html(f"Iteration {i+1}/{max_it} (after generation)",
                           highlight_tokens(current_tokens[answer_start:], answer_start,  diff_indices, color="green"))
         time.sleep(pause_length)
@@ -154,6 +158,7 @@ def diffusion_chat(question, noising, max_it, pause_length):
     
             yield render_html(f"Iteration {i+1}/{max_it} (before noising)",
                               highlight_tokens(current_tokens[answer_start:], answer_start, just_noised_indices, color="red"))
+            start = time.perf_counter()
     
             current_tokens = ori_input_tokens[:answer_start] + noised_answer[answer_start:]