sublinear 0.2.0__tar.gz → 0.2.2__tar.gz

{sublinear-0.2.0 → sublinear-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sublinear
-Version: 0.2.0
+Version: 0.2.2
 Summary: Python SDK for Sublinear foundation reward models for robotics.
 Author: Sublinear
 License-Expression: MIT

{sublinear-0.2.0 → sublinear-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sublinear"
-version = "0.2.0"
+version = "0.2.2"
 description = "Python SDK for Sublinear foundation reward models for robotics."
 readme = "README.md"
 requires-python = ">=3.10"

{sublinear-0.2.0 → sublinear-0.2.2}/src/sublinear/core.py

@@ -397,6 +397,298 @@ class Sublinear:
         verbose: bool = True,
         **input_json: Any,
     ) -> SimpleNamespace | list[SimpleNamespace]:
-        print('Thank you for trying out Sublinear! We are still in development and will be launching the API soon.')
-        return
+        # print('Thank you for trying out Sublinear! We are still in development and will be launching the API soon.')
+        # return
+        """
+        Create one inference job for each video.
+        All keyword arguments captured by **input_json are passed
+        to every inference job. Each job also receives the view_type
+        associated with that video.
+        """
+        if "downsample_to" not in input_json:
+            input_json["downsample_to"] = 10
+        if "task_description" not in input_json:
+            raise ValueError("task_description must be provided")
+        if "verbose" not in input_json:
+            input_json["verbose"] = True
+        # 
+        if isinstance(video_paths, str):
+            video_paths = [video_paths]
+        if isinstance(video_frames, np.ndarray):
+            video_frames = list(video_frames)
+        if isinstance(video_frames_external_view, np.ndarray):
+            video_frames_external_view = list(video_frames_external_view)
+        if isinstance(video_frames_wrist_view, np.ndarray):
+            video_frames_wrist_view = list(video_frames_wrist_view)
+        # 
+        videos, final_video_paths, final_video_paths_external_view, final_video_paths_wrist_view, view_type_per_video, single_video, video_count = get_final_video_frames(
+                model=model,
+                video_paths=video_paths,
+                video_path=video_path,
+                video_paths_external_view=(video_paths_external_view),
+                video_paths_wrist_view=(video_paths_wrist_view),
+                video_path_external_view=(video_path_external_view),
+                video_path_wrist_view=(video_path_wrist_view),
+                video_frames=video_frames,
+                video_frames_external_view=(video_frames_external_view),
+                video_frames_wrist_view=(video_frames_wrist_view),
+                view_type_per_video=view_type_per_video,
+                view_type=view_type,
+                verbose=verbose,
+        )
+        # 
+        if video_count < 1 or video_count > 100:
+            raise ValueError(f"Number of videos must be between 1 and 100, got {video_count}")
+        # 
+        videos_iterable = (enumerate(videos) if videos is not None else repeat(None, video_count))
+        final_video_paths_iterable = (enumerate(final_video_paths) if final_video_paths is not None else repeat(None, video_count))
+        final_video_paths_external_view_iterable = (enumerate(final_video_paths_external_view) if final_video_paths_external_view is not None else repeat(None, video_count))
+        final_video_paths_wrist_view_iterable = (enumerate(final_video_paths_wrist_view) if final_video_paths_wrist_view is not None else repeat(None, video_count))
+        model = model.lower()
+        # if "dm" not in model and model not in ["sole-r1", "robometer", "roboreward", "topreward"]:
+        #     raise ValueError(
+        #         f"Model '{model}' is not supported. "
+        #         "Supported models are: 'sole-r1', 'robometer', 'roboreward', 'topreward'."
+        #     )
+        active_job_ids: set[str] = set()
+        active_job_ids_lock = threading.Lock()
+        cancel_requested = threading.Event()
+        def register_job(job_id: str) -> None:
+            with active_job_ids_lock:
+                active_job_ids.add(job_id)
+            # Handles a race where Ctrl+C occurs just before registration.
+            if cancel_requested.is_set():
+                try:
+                    self.cancel_job(job_id)
+                except requests.RequestException:
+                    pass
+        def unregister_job(job_id: str) -> None:
+            with active_job_ids_lock:
+                active_job_ids.discard(job_id)
+        def submit_video(
+            indexed_video: tuple[int, list[np.ndarray]],
+            indexed_video_path: tuple[int, str] | None = None,
+            indexed_video_path_external_view: tuple[int, str] | None = None,
+            indexed_video_path_wrist_view: tuple[int, str] | None = None,
+        ) -> SimpleNamespace:
+            if indexed_video is None and (indexed_video_path_external_view is None or indexed_video_path_wrist_view is None):
+                if not indexed_video_path is None:
+                    video_index, video_path = indexed_video_path
+                    # frames = load_video_frames(video_path)
+                elif not indexed_video_path_external_view is None and indexed_video_path_wrist_view is None:
+                    video_index, video_path = indexed_video_path_external_view
+                    # frames = load_video_frames(video_path_external_view)
+                elif indexed_video_path_external_view is None and not indexed_video_path_wrist_view is None:
+                    video_index, video_path_wrist_view = indexed_video_path_wrist_view
+                    # frames = load_video_frames(video_path_wrist_view)
+                # 
+                source_path = Path(video_path)
+                job_input_json = dict(input_json)
+                frame_width, frame_height = get_video_dimensions(
+                    source_path
+                )
+                job_input_json["_frame_width"] = frame_width
+                job_input_json["_frame_height"] = frame_height
+                job_input_json["view_type"] = view_type_per_video[video_index]
+                with source_path.open("rb") as video_file:
+                    created_job = self.create_job(
+                        job_input_json,
+                        video=video_file,
+                        filename=source_path.name,
+                        model=model,
+                        on_job_created=register_job,
+                    )
+            else:
+                if not indexed_video is None:
+                    video_index, frames = indexed_video
+                elif not indexed_video_path_external_view is None and not indexed_video_path_wrist_view is None:
+                    video_index, video_path_external_view = indexed_video_path_external_view
+                    video_index, video_path_wrist_view = indexed_video_path_wrist_view
+                    frames_external = load_video_frames(video_path_external_view)
+                    frames_wrist = load_video_frames(video_path_wrist_view)
+                    if not len(frames_external) == len(frames_wrist):
+                        raise ValueError(f"Number of frames in external view video {len(frames_external)} does not match number of frames in wrist view video {len(frames_wrist)} for video_idx {video_index}")
+                    frames = [np.concatenate((frames_external[i], frames_wrist[i]), axis=1) for i in range(len(frames_external))]
+                else:
+                    raise ValueError("No video data provided for submission.")
+                # 
+                if not frames:
+                    raise ValueError(
+                        f"Video {video_index} has no frames."
+                    )
+                first_frame = np.asarray(frames[0])
+                if (
+                    first_frame.ndim != 3
+                    or first_frame.shape[2] != 3
+                ):
+                    raise ValueError(
+                        f"Video {video_index} frames must have "
+                        "shape (height, width, 3)."
+                    )
+                height, width = first_frame.shape[:2]
+                filename = f"video_{video_index}_{time.time_ns()}.mp4"
+                job_input_json = dict(input_json)
+                job_input_json["_frame_width"] = width
+                job_input_json["_frame_height"] = height
+                job_input_json["view_type"] = (view_type_per_video[video_index])
+                with tempfile.TemporaryDirectory() as temp_directory:
+                    encoded_video_path = (
+                        Path(temp_directory) / filename
+                    )
+                    writer = cv2.VideoWriter(
+                        str(encoded_video_path),
+                        cv2.VideoWriter_fourcc(*"mp4v"),
+                        30.0,
+                        (width, height),
+                    )
+                    if not writer.isOpened():
+                        raise RuntimeError(
+                            f"Unable to initialize MP4 encoder "
+                            f"for video {video_index}."
+                        )
+                    try:
+                        for frame_index, frame in enumerate(frames):
+                            frame = np.asarray(frame)
+                            if (
+                                frame.ndim != 3
+                                or frame.shape != first_frame.shape
+                            ):
+                                raise ValueError(
+                                    f"Frame {frame_index} of video "
+                                    f"{video_index} has shape "
+                                    f"{frame.shape}; expected "
+                                    f"{first_frame.shape}."
+                                )
+                            if frame.dtype != np.uint8:
+                                if (
+                                    np.issubdtype(
+                                        frame.dtype,
+                                        np.floating,
+                                    )
+                                    and frame.size > 0
+                                    and frame.min() >= 0
+                                    and frame.max() <= 1
+                                ):
+                                    frame = frame * 255
+                                frame = np.clip(
+                                    frame,
+                                    0,
+                                    255,
+                                ).astype(np.uint8)
+                            # Input frames are RGB; OpenCV expects BGR.
+                            bgr_frame = cv2.cvtColor(
+                                frame,
+                                cv2.COLOR_RGB2BGR,
+                            )
+                            writer.write(bgr_frame)
+                    finally:
+                        writer.release()
+                    if (
+                        not encoded_video_path.is_file()
+                        or encoded_video_path.stat().st_size == 0
+                    ):
+                        raise RuntimeError(
+                            f"Encoding video {video_index} failed."
+                        )
+                    with encoded_video_path.open("rb") as video_file:
+                        created_job = self.create_job(
+                            job_input_json,
+                            video=video_file,
+                            filename=filename,
+                            model=model,
+                            on_job_created=register_job,
+                        )
+            # 
+            job_id = created_job["job_id"]
+            try:
+                finished_job = self.wait_for_job(job_id)
+            finally:
+                # completed, failed, or canceled are all terminal.
+                unregister_job(job_id)
+            result = finished_job.get("result")
+            if not isinstance(result, dict):
+                raise RuntimeError(
+                    f"Video {video_index} returned an "
+                    "invalid inference result."
+                )
+            return SimpleNamespace(**result)
+        # 
+        worker_count = min(
+            self.max_parallel_jobs,
+            # len(videos),
+            video_count,
+        )
+        executor = ThreadPoolExecutor(max_workers=worker_count)
+        try:
+            results = list(executor.map(submit_video, videos_iterable, final_video_paths_iterable, final_video_paths_external_view_iterable, final_video_paths_wrist_view_iterable,))
+        except KeyboardInterrupt:
+            cancel_requested.set()
+            with active_job_ids_lock:
+                jobs_to_cancel = list(active_job_ids)
+            for job_id in jobs_to_cancel:
+                try:
+                    self.cancel_job(job_id)
+                except requests.RequestException as error:
+                    print(
+                        f"Could not cancel job {job_id}: {error}",
+                        flush=True,
+                    )
+            # Do not wait for queued worker functions to start.
+            executor.shutdown(
+                wait=False,
+                cancel_futures=True,
+            )
+            raise
+        except Exception:
+            # One submission failed, so cancel the remaining batch.
+            cancel_requested.set()
+            with active_job_ids_lock:
+                jobs_to_cancel = list(active_job_ids)
+            for job_id in jobs_to_cancel:
+                try:
+                    self.cancel_job(job_id)
+                except requests.RequestException as cancel_error:
+                    print(
+                        f"Could not cancel job {job_id}: {cancel_error}",
+                        flush=True,
+                    )
+            executor.shutdown(wait=False, cancel_futures=True)
+            raise
+        else:
+            executor.shutdown(wait=True)
+        # with ThreadPoolExecutor(
+        #     max_workers=worker_count
+        # ) as executor:
+        #     try:
+        #         results = list(executor.map(submit_video, videos_iterable, final_video_paths_iterable, final_video_paths_external_view_iterable, final_video_paths_wrist_view_iterable,))
+        #     except KeyboardInterrupt:
+        #         for job_id in active_job_ids:
+        #             try:
+        #                 self.cancel_job(job_id)
+        #             except requests.RequestException:
+        #                 pass
+        #         executor.shutdown(wait=False, cancel_futures=True)
+        #         raise
+        if single_video:
+            return results[0]
+        else:
+            if len(results)==1:
+                return results[0]
+            else:
+                results_rewards = [result.rewards[0] for result in results]
+                # results_output_text = [result.output_text[0] for result in results if result.output_text is not None else None]
+                results_output_text = [result.output_text[0] if result.output_text is not None else None for result in results]
+                results_success_probs = [result.success_probs[0] if result.success_probs else None for result in results]
+                results_success = [result.success[0] if result.success else None for result in results]
+                results = SimpleNamespace(
+                    model=model,
+                    rewards=results_rewards,
+                    output_text=results_output_text,
+                    success_probs=results_success_probs,
+                    success=results_success,
+                )
+        return results
+
+
+
 

{sublinear-0.2.0 → sublinear-0.2.2}/src/sublinear.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sublinear
-Version: 0.2.0
+Version: 0.2.2
 Summary: Python SDK for Sublinear foundation reward models for robotics.
 Author: Sublinear
 License-Expression: MIT