Better gradio keyframing interface

app.py

@@ -8,7 +8,7 @@ import numpy as np
 from PIL import Image, ImageOps, ImageDraw, ImageFont, ImageColor
 from urllib.request import urlopen
 
-#os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+# os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
 root = os.path.dirname(os.path.abspath(__file__))
 static = os.path.join(root, "static")
@@ -33,12 +33,13 @@ unet3d_condition_model_forward_copy = UNet3DConditionModel.forward
 UNet3DConditionModel.forward = unet3d_condition_model_forward
 
 
-
 model_id = "cerspense/zeroscope_v2_576w"
 model_path = model_id
 pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-pipe.to('cuda')
+pipe.to("cuda")
+
+MAX_KEYS = 10
 
 @spaces.GPU(duration=120)
 def core(bundle):
@@ -55,25 +56,41 @@ def core(bundle):
 
 
 def clear_btn_fn():
-    return "", "", "", ""
+
+    return 0, *[""] * MAX_KEYS, *[""] * MAX_KEYS, *[""] * MAX_KEYS, ""
 
 
 def gen_btn_fn(
-    prompts,
-    bboxes,
-    frames,
-    word_prompt_indices,
-    trailing_length,
-    n_spatial_steps,
-    n_temporal_steps,
-    spatial_strengthen_scale,
-    spatial_weaken_scale,
-    temporal_strengthen_scale,
-    temporal_weaken_scale,
-    rand_seed,
-    progress = gr.Progress(),
+    *args,
+    # *prompts,
+    # *bboxes,
+    # *frame_indices,
+    # word_prompt_indices_tb,
+    # trailing_length,
+    # n_spatial_steps,
+    # n_temporal_steps,
+    # spatial_strengthen_scale,
+    # spatial_weaken_scale,
+    # temporal_strengthen_scale,
+    # temporal_weaken_scale,
+    # rand_seed,
+    progress=gr.Progress(),
 ):
 
+    # no prompt at all
+    if not args[0]:
+        return
+
+    rand_seed = args[-1]
+    temporal_weaken_scale = args[-2]
+    temporal_strengthen_scale = args[-3]
+    spatial_weaken_scale = args[-4]
+    spatial_strengthen_scale = args[-5]
+    n_temporal_steps = args[-6]
+    n_spatial_steps = args[-7]
+    trailing_length = args[-8]
+    word_prompt_indices = args[-9]
+
     bundle = {}
     bundle["trailing_length"] = trailing_length
     bundle["num_dd_spatial_steps"] = n_spatial_steps
@@ -87,26 +104,14 @@ def gen_btn_fn(
     bundle["token_inds"] = [int(v) for v in word_prompt_indices.split(",")]
 
     bundle["keyframe"] = []
-    frames = frames.split(";")
-    bboxes = bboxes.split(";")
-    if ";" in prompts:
-        prompts = prompts.split(";")
-    else:
-        prompts = [prompts for i in range(len(frames))]
-
-    assert (
-        len(frames) == len(bboxes) == len(prompts)
-    ), "Inconsistent number of keyframes in the given inputs."
-
-    frames.pop()
-    bboxes.pop()
-    prompts.pop()
-
-    for i in range(len(frames)):
+
+    for i in range(MAX_KEYS):
         keyframe = {}
-        keyframe["bbox_ratios"] = [float(v) for v in bboxes[i].split(",")]
-        keyframe["frame"] = int(frames[i])
-        keyframe["prompt"] = prompts[i]
+        if not args[i]:
+            break
+        keyframe["prompt"] = args[i]
+        keyframe["bbox_ratios"] = [float(v) for v in args[i + MAX_KEYS].split(",")]
+        keyframe["frame"] = int(args[i + 2 * MAX_KEYS])
         bundle["keyframe"].append(keyframe)
     print(bundle)
     result = core(bundle)
@@ -114,6 +119,15 @@ def gen_btn_fn(
     return path
 
 
+def keyframe_update(num):
+    keyframes = []
+    for i in range(num):
+        keyframes.append(gr.Row(visible=True))
+    for i in range(MAX_KEYS - num):
+        keyframes.append(gr.Row(visible=False))
+    return keyframes
+
+
 def save_mask(inputs):
     layers = inputs["layers"]
     if not layers:
@@ -191,7 +205,7 @@ with gr.Blocks(
 ) as main:
 
     description = """
-    <h1 align="center" style="font-size: 48px">TrailBlazer: Trajectory Control for Diffusion-Based Video Generation</h1>
+    <h1 align="center" style="font-size: 48px">TrailBlazer: Trajectory Control for Diffusion-Based Video Generation (v0.0.3)</h1>
     <h4 align="center" style="margin: 0;">If you like our project, please give us a star ✨ at our Huggingface space, and our Github repository.</h4>
         <br>
         <span align="center" style="font-size: 18px">
@@ -205,8 +219,9 @@ with gr.Blocks(
     <p>
         <strong>Usage:</strong> Our Gradio app is implemented based on our executable script CmdTrailBlazer in our github repository. Please see our general information below for a quick guidance, as well as the hints within the app widgets.
     <ul>
-    <li>Basic: The bounding box (bbox) is the tuple of four floats for the rectangular corners: left, top, right, bottom in the normalized ratio. The Word prompt indices is a list of 1-indexed numbers determining the prompt word.</li>
-    <li>Advanced Options: We also offer some key parameters to adjust the synthesis result. Please see our paper for more information about the ablations.</li>
+    <li><strong>Basic</strong>: Our app workflow is straightforward. First, select the Number of keyframes, then fill all the values in the appeared prompt/frame indice/bounding box(bbox) for each keyframe, as well as the word prompt indices. Finally, hit the Generate button to run the TrailBlazer. It is roughly 90secs to get the result. We also provide the <strong>SketchPadHelper</strong> to visually design our bbox format</li>
+
+    <li><strong>Advanced Options</strong>: We also offer some key parameters to adjust the synthesis result. Please see our paper for more information about the ablations.</li>
     </ul>
 
     For your initial use, it is advisable to select one of the examples provided below and attempt to swap the subject first (e.g., cat -> lion). Subsequently, define the keyframe with the associated bbox/frame/prompt. Please note that our current work is based on the ZeroScope (cerspense/zeroscope_v2_576w) model. Using prompts that are commonly recognized in the ZeroScope model context is recommended.
@@ -214,25 +229,66 @@ with gr.Blocks(
     """
 
     gr.HTML(description)
-    dummy_note = gr.Textbox(
-        interactive=True, label="Note", visible=False
-    )
+    dummy_note = gr.Textbox(interactive=True, label="Note", visible=False)
+
+    keyframes = []
+    prompts = []
+    bboxes = []
+    frame_indices = []
+
     with gr.Row():
         with gr.Column(scale=2):
             with gr.Row():
                 with gr.Tab("Main"):
-                    text_prompt_tb = gr.Textbox(
-                        interactive=True, label="Keyframe: Prompt"
-                    )
-                    bboxes_tb = gr.Textbox(interactive=True, label="Keyframe: Bboxes")
-                    frame_tb = gr.Textbox(
-                        interactive=True, label="Keyframe: frame indices"
+
+                    # text_prompt_tb = gr.Textbox(
+                    #     interactive=True, label="Keyframe: Prompt"
+                    # )
+                    # bboxes_tb = gr.Textbox(interactive=True, label="Keyframe: Bboxes")
+                    # frame_tb = gr.Textbox(
+                    #     interactive=True, label="Keyframe: frame indices"
+                    # )
+
+                    dropdown = gr.Dropdown(
+                        label="Number of keyframes",
+                        choices=range(2, MAX_KEYS),
                     )
+                    for i in range(MAX_KEYS):
+                        with gr.Row(visible=False) as row:
+                            text = f"Keyframe #{i}"
+                            text = gr.HTML(text, visible=True)
+                            prompt = gr.Textbox(
+                                None,
+                                label=f"Prompt #{i}",
+                                visible=True,
+                                interactive=True,
+                                scale=4,
+                            )
+                            frame_ids = gr.Textbox(
+                                None,
+                                label=f"Frame indice #{i}",
+                                visible=True,
+                                interactive=True,
+                                scale=1,
+                            )
+                            bbox = gr.Textbox(
+                                None,
+                                label=f"BBox #{i}",
+                                visible=True,
+                                interactive=True,
+                                scale=3,
+                            )
+                            prompts.append(prompt)
+                            bboxes.append(bbox)
+                            frame_indices.append(frame_ids)
+                        keyframes.append(row)
+                    dropdown.change(keyframe_update, dropdown, keyframes)
+
                     with gr.Row():
                         word_prompt_indices_tb = gr.Textbox(
                             interactive=True, label="Word prompt indices:"
                         )
-                        text = "<strong>Hint</strong>: Each keyframe ends with <strong>SEMICOLON</strong>, and <strong>COMMA</strong> for separating each value in the keyframe. The prompt field can be a single prompt without semicolon, or multiple prompts ended semicolon. One can use the SketchPadHelper tab to help to design the bboxes field."
+                        text = "<strong>Hint</strong>: Each keyframe is associated with a prompt, frame indice, and the corresponding bbox. The bbox is the tuple of the four floats determining the four bbox corners (left, top, right, bottom) in normalized ratio. The word prompt indices is 1-indexed value to indicate the word in prompt. Note that we use <strong>COMMA</strong> to separate the multiple values."
                         gr.HTML(text)
                     with gr.Row():
                         clear_btn = gr.Button(value="Clear")
@@ -343,57 +399,132 @@ with gr.Blocks(
             with gr.Row():
                 out_gen_1 = gr.Video(visible=True, show_label=False)
 
+    with gr.Row():
+        gr.Markdown("## Two keyframes example")
     with gr.Row():
         gr.Examples(
             examples=[
                 [
+                    "assets/gradio/fish-RL.mp4",
+                    "It generates clownfish at right, then move to left",
+                    2,
                     "A clownfish swimming in a coral reef",
-                    "0.5,0.35,1.0,0.65; 0.0,0.35,0.5,0.65;",
-                    "0; 24;",
-                    "1, 2",
+                    "A clownfish swimming in a coral reef",
+                    "0",
+                    "24",
+                    "0.5,0.35,1.0,0.65",
+                    "0.0,0.35,0.5,0.65",
+                    "1,2",
                     "123451232531",
-                    "It generates clownfish at right, then move to left",
-                    "assets/gradio/fish-RL.mp4",
-                ],
-                [
-                    "A cat is running on the grass",
-                    "0.0,0.35,0.4,0.65; 0.6,0.35,1.0,0.65; 0.0,0.35,0.4,0.65;"
-                    "0.6,0.35,1.0,0.65; 0.0,0.35,0.4,0.65;",
-                    "0; 6; 12; 18; 24;",
-                    "1, 2",
-                    "123451232530",
-                    "The cat will run Left/Right/Left/Right",
-                    "assets/gradio/cat-LRLR.mp4",
                 ],
                 [
+                    "assets/gradio/fish-TL2BR.mp4",
+                    "The fish moves from top left to bottom right, from far to near.",
+                    2,
                     "A fish swimming in the ocean",
-                    "0.0,0.0,0.1,0.1; 0.5,0.5,1.0,1.0;",
-                    "0; 24;",
+                    "A fish swimming in the ocean",
+                    "0",
+                    "24",
+                    "0.0,0.0,0.1,0.1",
+                    "0.5,0.5,1.0,1.0",
                     "1, 2",
                     "0",
-                    "The fish moves from top left to bottom right, from far to near.",
-                    "assets/gradio/fish-TL2BR.mp4"
                 ],
                 [
+                    "assets/gradio/tiger-TL2BR.mp4",
+                    "Same with the above but now the prompt associates with tiger",
+                    2,
+                    "A tiger walking alone down the street",
                     "A tiger walking alone down the street",
-                    "0.0,0.0,0.1,0.1; 0.5,0.5,1.0,1.0;",
-                    "0; 24;",
+                    "0",
+                    "24",
+                    "0.0,0.0,0.1,0.1",
+                    "0.5,0.5,1.0,1.0",
                     "1, 2",
                     "0",
-                    "Same with the above but now the prompt associates with tiger",
-                    "assets/gradio/tiger-TL2BR.mp4"
                 ],
                 [
-                    "A white cat walking on the grass; A yellow dog walking on the grass;",
-                    "0.7,0.4,1.0,0.65; 0.0,0.4,0.3,0.65;",
-                    "0; 24;",
+                    "assets/gradio/Cat2Dog.mp4",
+                    "The subject will deformed from cat to dog.",
+                    2,
+                    "A white cat walking on the grass",
+                    "A yellow dog walking on the grass",
+                    "0",
+                    "24",
+                    "0.7,0.4,1.0,0.65",
+                    "0.0,0.4,0.3,0.65",
                     "1,2,3",
                     "123451232531",
-                    "The subject will deformed from cat to dog.",
-                    "assets/gradio/Cat2Dog.mp4",
                 ],
             ],
-            inputs=[text_prompt_tb, bboxes_tb, frame_tb, word_prompt_indices_tb, rand_seed, dummy_note, out_gen_1],
+            inputs=[
+                out_gen_1,
+                dummy_note,
+                dropdown,
+                prompts[0],
+                prompts[1],
+                frame_indices[0],
+                frame_indices[1],
+                bboxes[0],
+                bboxes[1],
+                word_prompt_indices_tb,
+                rand_seed,
+            ],
+            outputs=None,
+            fn=None,
+            cache_examples=False,
+        )
+
+    with gr.Row():
+        gr.Markdown("## Five keyframes example")
+    with gr.Row():
+        gr.Examples(
+            examples=[
+                [
+                    "assets/gradio/cat-LRLR.mp4",
+                    "The poor cat will run Left/Right/Left/Right :(",
+                    5,
+                    "A cat is running on the grass",
+                    "A cat is running on the grass",
+                    "A cat is running on the grass",
+                    "A cat is running on the grass",
+                    "A cat is running on the grass",
+                    "0",
+                    "6",
+                    "12",
+                    "18",
+                    "24",
+                    "0.0,0.35,0.4,0.65",
+                    "0.6,0.35,1.0,0.65",
+                    "0.0,0.35,0.4,0.65",
+                    "0.6,0.35,1.0,0.65",
+                    "0.0,0.35,0.4,0.65",
+                    "1, 2",
+                    "123451232530",
+                ],
+            ],
+            inputs=[
+                out_gen_1,
+                dummy_note,
+                dropdown,
+                prompts[0],
+                prompts[1],
+                prompts[2],
+                prompts[3],
+                prompts[4],
+                frame_indices[0],
+                frame_indices[1],
+                frame_indices[2],
+                frame_indices[3],
+                frame_indices[4],
+                bboxes[0],
+                bboxes[1],
+                bboxes[2],
+                bboxes[3],
+                bboxes[4],
+                word_prompt_indices_tb,
+                rand_seed,
+            ],
             outputs=None,
             fn=None,
             cache_examples=False,
@@ -402,16 +533,16 @@ with gr.Blocks(
     clear_btn.click(
         clear_btn_fn,
         inputs=[],
-        outputs=[text_prompt_tb, bboxes_tb, frame_tb, word_prompt_indices_tb],
+        outputs=[dropdown, *prompts, *bboxes, *frame_indices, word_prompt_indices_tb],
         queue=False,
     )
 
     gen_btn.click(
         gen_btn_fn,
         inputs=[
-            text_prompt_tb,
-            bboxes_tb,
-            frame_tb,
+            *prompts,
+            *bboxes,
+            *frame_indices,
             word_prompt_indices_tb,
             trailing_length,
             n_spatial_steps,