update

app.py

@@ -227,6 +227,18 @@ def set_denoising_steps(denoising_steps, sampling_timesteps_state):
     print("set denoising steps to", worldmem.sampling_timesteps)
     return sampling_timesteps_state
 
+def set_context_length(context_length, sampling_context_length_state):
+    worldmem.n_tokens = context_length
+    sampling_context_length_state = context_length
+    print("set context length to", worldmem.n_tokens)
+    return sampling_context_length_state
+
+def set_memory_length(memory_length, sampling_memory_length_state):
+    worldmem.condition_similar_length = memory_length
+    sampling_memory_length_state = memory_length
+    print("set memory length to", worldmem.condition_similar_length)
+    return sampling_memory_length_state
+
 def generate(keys):
     # print("algo frame:", len(worldmem.frames))
     actions = parse_input_to_tensor(keys)
@@ -336,41 +348,51 @@ with gr.Blocks(css=css) as demo:
         # </a>
         # </div>
 
+    example_actions = ["AAAAAAAAAAAADDDDDDDDDDDD", "AAAAAAAAAAAAAAAAAAAAAAAA", "DDDDDDDDEEEEEEEEEESSSAAAAAAAAWWW", "DDDDDDDDDDDDQQQQQQQQQQQQQQQDDDDDDDDDDDD", 
+    "DDDDWWWDDDDDDDDDDDDDDDDDDDDSSS", "SSUNNWWEEEEEEEEEAAASSUNNWWEEEEEEEEE"]
+
+
     with gr.Row(variant="panel"):
         video_display = gr.Video(autoplay=True, loop=True)
         image_display = gr.Image(value=DEFAULT_IMAGE, interactive=False, label="Last Frame")
 
+
     with gr.Row(variant="panel"):
         with gr.Column(scale=2):
             input_box = gr.Textbox(label="Action Sequence", placeholder="Enter action sequence here...", lines=1, max_lines=1)
             log_output = gr.Textbox(label="History Log", interactive=False)
+            gr.Markdown("### Action sequence examples.")
+            with gr.Row():
+                buttons = []
+                for action in example_actions[:2]:
+                    with gr.Column(scale=len(action)):
+                        buttons.append(gr.Button(action))
+            with gr.Row():
+                for action in example_actions[2:4]:
+                    with gr.Column(scale=len(action)):
+                        buttons.append(gr.Button(action))
+            with gr.Row():
+                for action in example_actions[4:6]:
+                    with gr.Column(scale=len(action)):
+                        buttons.append(gr.Button(action))
+
         with gr.Column(scale=1):
-            slider = gr.Slider(minimum=10, maximum=50, value=worldmem.sampling_timesteps, step=1, label="Denoising Steps")
+            slider_denoising_step = gr.Slider(minimum=10, maximum=50, value=worldmem.sampling_timesteps, step=1, label="Denoising Steps")
+            slider_context_length = gr.Slider(minimum=2, maximum=10, value=worldmem.n_tokens, step=1, label="Context Length")
+            slider_memory_length = gr.Slider(minimum=4, maximum=16, value=worldmem.condition_similar_length, step=1, label="Memory Length")
             submit_button = gr.Button("Generate")
             reset_btn = gr.Button("Reset")
     
     sampling_timesteps_state = gr.State(worldmem.sampling_timesteps)
-    
-    example_actions = ["DDDDDDDDEEEEEEEEEESSSAAAAAAAAWWW", "DDDDDDDDDDDDQQQQQQQQQQQQQQQDDDDDDDDDDDD", 
-    "DDDDWWWDDDDDDDDDDDDDDDDDDDDSSSAAAAAAAAAAAAAAAAAAAAAAAA", "SSUNNWWEEEEEEEEEAAA1NNNNNNNNNSSUNNWW"]
+    sampling_context_length_state = gr.State(worldmem.n_tokens)
+    sampling_memory_length_state = gr.State(worldmem.condition_similar_length)
+
 
     def set_action(action):
         return action
     
-    gr.Markdown("### Action sequence examples.")
-    with gr.Row():
-        buttons = []
-        for action in example_actions[:2]:
-            with gr.Column(scale=len(action)):
-                buttons.append(gr.Button(action))
-    with gr.Row():
-        for action in example_actions[2:4]:
-            with gr.Column(scale=len(action)):
-                buttons.append(gr.Button(action))
-    with gr.Row():
-        for action in example_actions[4:5]:
-            with gr.Column(scale=len(action)):
-                buttons.append(gr.Button(action))
+    # gr.Markdown("### Action sequence examples.")
+
 
     for button, action in zip(buttons, example_actions):
             button.click(set_action, inputs=[gr.State(value=action)], outputs=input_box)
@@ -403,7 +425,7 @@ with gr.Blocks(css=css) as demo:
         - N: no-op (do nothing)  
         - 1: switch to hotbar 1  
         - U: use item  
-        5. Higher denoising steps produce more detailed results but take longer. **20 steps** is a good balance between quality and speed.
+        5. Higher denoising steps produce more detailed results but take longer. **20 steps** is a good balance between quality and speed. The same applies to context and memory length.
         6. If you find this project interesting or useful, please consider giving it a ⭐️ on [GitHub]()!
         7. For feedback or suggestions, feel free to open a GitHub issue or contact me directly at **zeqixiao1@gmail.com**.
         """
@@ -418,6 +440,8 @@ with gr.Blocks(css=css) as demo:
     image_display_5.select(lambda: on_image_click(SUNFLOWERS_RAIN_IMAGE), outputs=image_display)
     image_display_6.select(lambda: on_image_click(PLACE_IMAGE), outputs=image_display)
 
-    slider.change(fn=set_denoising_steps, inputs=[slider, sampling_timesteps_state], outputs=sampling_timesteps_state)
+    slider_denoising_step.change(fn=set_denoising_steps, inputs=[slider_denoising_step, sampling_timesteps_state], outputs=sampling_timesteps_state)
+    slider_context_length.change(fn=set_context_length, inputs=[slider_context_length, sampling_context_length_state], outputs=sampling_context_length_state)
+    slider_memory_length.change(fn=set_memory_length, inputs=[slider_memory_length, sampling_memory_length_state], outputs=sampling_memory_length_state)
 
 demo.launch()