Spaces:

sud0301
/

orbis2_test

Sleeping

App Files Files Community

Silvio Galesso commited on 13 days ago

Commit

86d8d8b

1 Parent(s): 380643c

appearance fixes

Browse files

Files changed (2) hide show

app2.py +43 -8
orbis2_app_engine.py +8 -0

app2.py CHANGED Viewed

@@ -192,6 +192,10 @@ if UI_ONLY:
     # auto-seek stay a no-op instead of erroring.
     MIN_CONTEXT_END_FRAME = 0
     MIN_CONTEXT_TIME_S = 0.0
 else:
     MODELS_DIR = resolve_models_dir()
     os.environ["ORBIS2_MODELS_DIR"] = str(MODELS_DIR)   # L1/L2 configs expandvars this
@@ -210,7 +214,22 @@ else:
     MIN_CONTEXT_END_FRAME = ENGINE.min_context_end_frame(L1_FRAME_RATE)
     MIN_CONTEXT_TIME_S = MIN_CONTEXT_END_FRAME / CONTEXT_FPS
 DEFAULT_GEN_STEPS = 10      # rollout steps; each yields model.num_pred_frames frames
 # L1 and L2 are sampled at different step counts. L1 (detail) needs more steps for
 # image fidelity; L2 (abstract, consistency-distilled) is designed for very few.
 DEFAULT_L1_STEPS = 4       # L1 sampler steps (NFE)
@@ -291,8 +310,9 @@ is resized to {FRAME_H}×{FRAME_W}. The **L2 abstract predictor** forecasts a la
 future, and the **L1 detail predictor**, conditioned on it, generates the next frames
 autoregressively.
-**Rollout length.** *Rollout steps* sets how many autoregressive steps to run; each
-step emits several image frames, so the final clip is longer than the step count.
 **Sampling.** The two levels are sampled independently, with their own step counts.
 *L1 sampler steps* controls the detail predictor, which integrates the flow-matching
@@ -504,7 +524,7 @@ def render_trajectory_plot(raw_json: str):
 # because it's inserted as real page <head> content.
 # --------------------------------------------------------------------------
 TRAJ_CANVAS_HTML_JS = f"""
-<div style="position:relative; width:{TRAJ_CANVAS_SIZE_PX}px;">
   <canvas id="traj-canvas" width="{TRAJ_CANVAS_SIZE_PX}" height="{TRAJ_CANVAS_SIZE_PX}"
           style="display:block; border:1px solid #999; border-radius:8px; cursor:crosshair; background:#fafafa;">
   </canvas>
@@ -843,7 +863,7 @@ def _parse_cursor_time_s(raw) -> float | None:
 # GPU inference
 # ----------------------------------------------------------------------------
 @spaces.GPU(duration=300)
-def run_rollout(video_path: str, num_gen_steps: int, l1_steps: int, l2_steps: int,
                 seed: int, trajectory, cursor_time_s: str, progress=gr.Progress()):
     if video_path is None:
         raise gr.Error("Please upload a short mp4 clip first.")
@@ -867,6 +887,10 @@ def run_rollout(video_path: str, num_gen_steps: int, l1_steps: int, l2_steps: in
         # eta is not exposed: 0.0 keeps sampling on the deterministic ODE.
         eta = DEFAULT_ETA
         job = Path(tempfile.gettempdir()) / f"orbis2_{uuid.uuid4().hex[:8]}"
         say(0.02, "⏳ Loading…")
@@ -979,9 +1003,10 @@ def build_demo():
         with gr.Row():
             with gr.Column(scale=1):
-                n_steps_gen = gr.Slider(4, 30, value=DEFAULT_GEN_STEPS, step=2,
-                                        label="Rollout steps",
-                                        info="Each step emits several frames")
                 l1_steps = gr.Slider(5, 20, value=DEFAULT_L1_STEPS, step=1,
                                      label="L1 sampler steps",
                                      info="Detail predictor — more steps, sharper frames")
@@ -997,6 +1022,10 @@ def build_demo():
                 inp = gr.Video(label="Upload a short mp4 clip", sources=["upload"],
                                elem_id="input_video")
                 cursor_time_s = gr.Textbox(elem_id="cursor_time_s", value="-1", visible=False)
             with gr.Column(scale=1):
                 gr.Markdown(
@@ -1005,6 +1034,12 @@ def build_demo():
                     "it untouched, or hit Clear, to run the rollout unconditionally."
                 )
                 gr.HTML(TRAJ_CANVAS_HTML_JS)
                 raw_points_json = gr.Textbox(elem_id="raw_points_json", visible=False)
                 smoothed_points_json = gr.Textbox(elem_id="smoothed_points_json", visible=False)
                 trajectory_state = gr.State()
@@ -1027,7 +1062,7 @@ def build_demo():
         btn.click(
             run_rollout,
-            inputs=[inp, n_steps_gen, l1_steps, l2_steps, seed_in, trajectory_state, cursor_time_s],
             outputs=outs,
             concurrency_limit=1,
         )

     # auto-seek stay a no-op instead of erroring.
     MIN_CONTEXT_END_FRAME = 0
     MIN_CONTEXT_TIME_S = 0.0
+    # No model to read model.num_pred_frames from -- 1 is just a placeholder so the
+    # UI-only slider below has *some* seconds-per-step to show; it's never used for
+    # an actual rollout (run_rollout raises before reaching the model in this mode).
+    FRAMES_PER_GEN_STEP = 1
 else:
     MODELS_DIR = resolve_models_dir()
     os.environ["ORBIS2_MODELS_DIR"] = str(MODELS_DIR)   # L1/L2 configs expandvars this
     MIN_CONTEXT_END_FRAME = ENGINE.min_context_end_frame(L1_FRAME_RATE)
     MIN_CONTEXT_TIME_S = MIN_CONTEXT_END_FRAME / CONTEXT_FPS
+    FRAMES_PER_GEN_STEP = ENGINE.frames_per_rollout_step
+# Seconds of output video one rollout step yields, at L1_FRAME_RATE fps -- lets the
+# UI expose rollout length in seconds while the model itself still steps in whole
+# rollout steps (see GEN_SECONDS_STEP / run_rollout's seconds -> steps conversion below).
+SECONDS_PER_GEN_STEP = FRAMES_PER_GEN_STEP / L1_FRAME_RATE
 DEFAULT_GEN_STEPS = 10      # rollout steps; each yields model.num_pred_frames frames
+GEN_STEPS_MIN, GEN_STEPS_MAX, GEN_STEPS_STEP = 4, 30, 2   # bounds on the underlying step count
+# The slider itself is in seconds (see build_demo()); these are its bounds, derived
+# from the step-count bounds above via SECONDS_PER_GEN_STEP.
+GEN_SECONDS_MIN = GEN_STEPS_MIN * SECONDS_PER_GEN_STEP
+GEN_SECONDS_MAX = GEN_STEPS_MAX * SECONDS_PER_GEN_STEP
+GEN_SECONDS_STEP = GEN_STEPS_STEP * SECONDS_PER_GEN_STEP
+DEFAULT_GEN_SECONDS = DEFAULT_GEN_STEPS * SECONDS_PER_GEN_STEP
 # L1 and L2 are sampled at different step counts. L1 (detail) needs more steps for
 # image fidelity; L2 (abstract, consistency-distilled) is designed for very few.
 DEFAULT_L1_STEPS = 4       # L1 sampler steps (NFE)
 future, and the **L1 detail predictor**, conditioned on it, generates the next frames
 autoregressively.
+**Rollout length.** *Generated video length (s)* sets how much video to generate; under
+the hood this is rounded to a whole number of autoregressive steps, each of which emits
+{FRAMES_PER_GEN_STEP} image frame{"s" if FRAMES_PER_GEN_STEP != 1 else ""} ({SECONDS_PER_GEN_STEP:.2g}s of video per step).
 **Sampling.** The two levels are sampled independently, with their own step counts.
 *L1 sampler steps* controls the detail predictor, which integrates the flow-matching
 # because it's inserted as real page <head> content.
 # --------------------------------------------------------------------------
 TRAJ_CANVAS_HTML_JS = f"""
+<div style="position:relative; width:{TRAJ_CANVAS_SIZE_PX}px; margin:0 auto;">
   <canvas id="traj-canvas" width="{TRAJ_CANVAS_SIZE_PX}" height="{TRAJ_CANVAS_SIZE_PX}"
           style="display:block; border:1px solid #999; border-radius:8px; cursor:crosshair; background:#fafafa;">
   </canvas>
 # GPU inference
 # ----------------------------------------------------------------------------
 @spaces.GPU(duration=300)
+def run_rollout(video_path: str, gen_seconds: float, l1_steps: int, l2_steps: int,
                 seed: int, trajectory, cursor_time_s: str, progress=gr.Progress()):
     if video_path is None:
         raise gr.Error("Please upload a short mp4 clip first.")
         # eta is not exposed: 0.0 keeps sampling on the deterministic ODE.
         eta = DEFAULT_ETA
+        # The model itself only understands whole rollout steps -- convert the
+        # seconds the user picked back into a step count (see SECONDS_PER_GEN_STEP).
+        num_gen_steps = max(1, round(gen_seconds / SECONDS_PER_GEN_STEP))
         job = Path(tempfile.gettempdir()) / f"orbis2_{uuid.uuid4().hex[:8]}"
         say(0.02, "⏳ Loading…")
         with gr.Row():
             with gr.Column(scale=1):
+                n_seconds_gen = gr.Slider(GEN_SECONDS_MIN, GEN_SECONDS_MAX, value=DEFAULT_GEN_SECONDS,
+                                          step=GEN_SECONDS_STEP,
+                                          label="Generated video length (s)",
+                                          info=f"Generated in chunks of {SECONDS_PER_GEN_STEP:.2g}s each")
                 l1_steps = gr.Slider(5, 20, value=DEFAULT_L1_STEPS, step=1,
                                      label="L1 sampler steps",
                                      info="Detail predictor — more steps, sharper frames")
                 inp = gr.Video(label="Upload a short mp4 clip", sources=["upload"],
                                elem_id="input_video")
                 cursor_time_s = gr.Textbox(elem_id="cursor_time_s", value="-1", visible=False)
+                gr.Markdown(
+                    "Use the video cursor to pick where the model generation should "
+                    "start (i.e. the last context frame)."
+                )
             with gr.Column(scale=1):
                 gr.Markdown(
                     "it untouched, or hit Clear, to run the rollout unconditionally."
                 )
                 gr.HTML(TRAJ_CANVAS_HTML_JS)
+                gr.HTML(
+                    '<p style="font-size:11px; color:#555; margin-top:6px;">'
+                    "Disclaimer: hand-drawn trajectories are likely to be "
+                    "out-of-distribution for the model and can generate unexpected "
+                    "outputs. Happy crashing!</p>"
+                )
                 raw_points_json = gr.Textbox(elem_id="raw_points_json", visible=False)
                 smoothed_points_json = gr.Textbox(elem_id="smoothed_points_json", visible=False)
                 trajectory_state = gr.State()
         btn.click(
             run_rollout,
+            inputs=[inp, n_seconds_gen, l1_steps, l2_steps, seed_in, trajectory_state, cursor_time_s],
             outputs=outs,
             concurrency_limit=1,
         )

orbis2_app_engine.py CHANGED Viewed

@@ -168,6 +168,14 @@ class RolloutEngine:
         )
         return indexer.get_required_l1_start_offset() + l1_context_frames - 1
     def ensure_ready(self, device="cuda", compile=True, compile_mode="reduce-overhead",
                       compile_artifacts=None, speed_scale=1.0, yaw_rate_scale=1.0):
         """Move the model to `device` and wrap it with torch.compile, exactly once.

         )
         return indexer.get_required_l1_start_offset() + l1_context_frames - 1
+    @property
+    def frames_per_rollout_step(self):
+        """Number of output frames each autoregressive rollout step yields
+        (model.num_pred_frames). CPU-only (model metadata), callable right after
+        load(). Lets callers convert the "rollout steps" the model actually takes
+        into seconds of generated video (frames_per_rollout_step / l1_frame_rate)."""
+        return int(self.model.num_pred_frames)
     def ensure_ready(self, device="cuda", compile=True, compile_mode="reduce-overhead",
                       compile_artifacts=None, speed_scale=1.0, yaw_rate_scale=1.0):
         """Move the model to `device` and wrap it with torch.compile, exactly once.