Add Support for Video Media

README.md

@@ -52,6 +52,7 @@ for i in range(10):
 # log more than just metrics (files, text, artifacts, model weights or series thereof)
 # exp["config"] = File("/path/to/config.txt")
 # exp["summary"] = Text("first run")
+# exp["rollout"] = Video("/path/to/preview.mp4")
 # exp["model"] = Model(torch.nn.Module)
 # exp["config-series"].append(File("/path/to/config1.txt"))
 # exp["config-series"].append(File("path/to/config2.txt"))

docs/source/api.rst

@@ -76,6 +76,10 @@ Media Wrappers
    :members:
    :show-inheritance:
 
+.. autoclass:: Video
+   :members:
+   :show-inheritance:
+
 .. autoclass:: Model
    :members:
    :show-inheritance:

docs/source/guide/media.rst

@@ -2,25 +2,28 @@
 Logging Media
 #############
 
-LitLogger supports uploading images and text files that are displayed alongside
-your metrics in the experiment view. For new code, prefer the dict-style
-wrappers :class:`~litlogger.media.Image` and :class:`~litlogger.media.Text`.
+LitLogger supports uploading images, text files, and videos that are displayed
+alongside your metrics in the experiment view. For new code, prefer the
+dict-style wrappers :class:`~litlogger.media.Image`,
+:class:`~litlogger.media.Text`, and :class:`~litlogger.media.Video`.
 
 Using the Experiment API
 ========================
 
 .. code-block:: python
 
    import litlogger
-   from litlogger import Image, Text
+   from litlogger import Image, Text, Video
 
    experiment = litlogger.init(name="media-demo")
 
    experiment["preview"] = Image("generated.png")
    experiment["notes"] = Text("epoch 0 summary")
+   experiment["video_preview"] = Video("preview.mp4")
 
    experiment["samples"].append(Image("sample-0.png"), step=0)
    experiment["captions"].append(Text("reconstruction"), step=0)
+   experiment["clips"].append(Video("sample-0.mp4"), step=0)
 
 Legacy Helper API
 =================
@@ -56,6 +59,13 @@ Logging Text
 
    exp.log_media("predictions", "predictions.txt", kind=MediaType.TEXT, step=5)
 
+Logging Video
+=============
+
+.. code-block:: python
+
+   exp.log_media("rollout", "preview.mp4", kind=MediaType.VIDEO, step=5)
+
 Supported Types
 ===============
 
@@ -72,6 +82,9 @@ Supported Types
    * - ``TEXT``
      - ``.txt``, ``.csv``, ``.json``, ``.log``
      - Displayed as text in the experiment view
+   * - ``VIDEO``
+     - ``.mp4``, ``.mov``, ``.avi``, ``.webm``
+     - Displayed as videos in the experiment view
 
 When ``kind`` is not provided, LitLogger guesses the type from the file's MIME
 type. If the type cannot be determined, a ``ValueError`` is raised.

docs/source/guide/standalone.rst

@@ -69,11 +69,12 @@ Files, Media, and Models
 
 .. code-block:: python
 
-   from litlogger import File, Image, Model, Text
+   from litlogger import File, Image, Model, Text, Video
 
    experiment["config"] = File("config.yaml")
    experiment["preview"] = Image("sample.png")
    experiment["notes"] = Text("training summary")
+   experiment["rollout"] = Video("preview.mp4")
    experiment["checkpoint"] = Model("checkpoint.ckpt")
 
 Resume by Name

docs/source/tutorials/file_media_model.rst

@@ -7,6 +7,7 @@ This tutorial covers the new file-like API surface:
 - :class:`~litlogger.media.File`
 - :class:`~litlogger.media.Image`
 - :class:`~litlogger.media.Text`
+- :class:`~litlogger.media.Video`
 - :class:`~litlogger.media.Model`
 
 Static Files
@@ -55,6 +56,18 @@ Images
    experiment["preview"] = Image("preview.png")
    experiment["samples"].append(Image("sample-0.png"), step=0)
 
+Videos
+======
+
+Use :class:`~litlogger.media.Video` for video files or in-memory frame data.
+
+.. code-block:: python
+
+   from litlogger import Video
+
+   experiment["demo"] = Video("preview.mp4")
+   experiment["clips"].append(Video("sample-0.mp4"), step=0)
+
 Models
 ======
 

src/litlogger/__init__.py

@@ -27,7 +27,7 @@
 
 # Import SDK functions
 from litlogger.init import finish, get_metadata, init
-from litlogger.media import File, Image, Model, Text
+from litlogger.media import File, Image, Model, Text, Video
 
 # Global variables
 experiment: Experiment | None = None
@@ -48,6 +48,7 @@
     "Image",
     "Model",
     "Text",
+    "Video",
     "experiment",
     "finalize",
     "finish",

src/litlogger/experiment_support.py

@@ -20,7 +20,7 @@
 
 from lightning_sdk.lightning_cloud.openapi import V1MediaType
 
-from litlogger.media import File, Image, Model, Text
+from litlogger.media import File, Image, Model, Text, Video
 from litlogger.series import Series
 from litlogger.types import MediaType, Metrics, MetricValue, PhaseType
 
@@ -143,12 +143,8 @@ def resolve_remote_model(exp: "Experiment", key: str) -> Model | Series | None:
 
     @staticmethod
     def rebuild_state(exp: "Experiment") -> None:
-        """Rebuild state from remote metadata, steps, artifacts, and media.
-
-        TODO: Add backend-supported recovery for model bindings so resumed
-        experiments can reconstruct ``Model`` values without storing them in
-        frontend-visible metadata tags.
-        """
+        """Rebuild state from remote metadata, steps, artifacts, and media."""
+        # TODO: add BE support for restoring model states as well
         exp._update_metrics_store()
         tags = getattr(exp._metrics_store, "tags", None) or []
         for tag in tags:
@@ -294,6 +290,9 @@ def wrap_media_file(exp: "Experiment", media_name: str, media_type: V1MediaType)
             text = Text("")
             text.path = media_name
             return text
+
+        if media_type == V1MediaType.VIDEO:
+            return Video(media_name)
         return File(media_name)
 
     @staticmethod
@@ -316,6 +315,8 @@ def media_type_to_v1(exp: "Experiment", media_type: MediaType) -> V1MediaType:
             return V1MediaType.IMAGE
         if media_type == MediaType.TEXT:
             return V1MediaType.TEXT
+        if media_type == MediaType.VIDEO:
+            return V1MediaType.VIDEO
         raise ValueError(f"Unsupported media type for file upload: {media_type}")
 
     @staticmethod
@@ -358,15 +359,14 @@ def upload_media_value(
 
     @staticmethod
     def upload_model_value(exp: "Experiment", key: str, value: Model) -> None:
-        """Upload a model through litmodels and bind the remote wrapper.
-
-        TODO: Persist model recovery data via backend-supported experiment
-        bindings so resumed experiments can rebuild these wrappers.
-        """
+        """Upload a model through litmodels and bind the remote wrapper."""
+        # TODO: Persist model recovery data via backend-supported experiment
+        # bindings so resumed experiments can rebuild these wrappers.
         cloud_account = exp._metrics_store.cluster_id
         model_name = value._log_model(
             experiment_name=exp.name,
             teamspace=exp._teamspace,
+            key=exp._model_experiment_name(key),
             experiment=exp,
             cloud_account=cloud_account if isinstance(cloud_account, str) else None,
         )

src/litlogger/media.py

@@ -282,6 +282,181 @@ def _media_type(self) -> MediaType:
         return MediaType.IMAGE
 
 
+class Video(File):
+    DEFAULT_FPS = 24
+
+    """Represents a video to be logged.
+
+    Accepts:
+    - a file path string
+    - a MoviePy clip
+    - a numpy array / torch tensor of frames
+
+    Supported array shapes:
+    - (T, H, W)            grayscale frames
+    - (T, H, W, C)         frames in HWC format
+    - (T, C, H, W)         frames in CHW format, where C is 1/3/4
+
+    Args:
+        data: Video data or a path to a video file.
+        format: Output container/extension, default "mp4".
+        description: Optional description.
+        fps: Frames per second used when rendering arrays or clips that do
+            not already carry fps metadata.
+    """
+
+    def __init__(
+        self,
+        data: Any,
+        format: str = "mp4",  # noqa: A002
+        description: str = "",
+        fps: float | None = None,
+    ) -> None:
+        self._data = data
+        self._format = format
+        self._fps = fps
+        self._temp_path: str | None = None
+
+        if isinstance(data, str):
+            super().__init__(data, description=description)
+        else:
+            super().__init__("", description=description)
+
+    def _get_upload_path(self) -> str:
+        if isinstance(self._data, str):
+            return super()._get_upload_path()
+        return self._render_to_temp()
+
+    def _render_to_temp(self) -> str:
+        suffix = f".{self._format.lower()}"
+        fd, path = tempfile.mkstemp(suffix=suffix)
+        os.close(fd)
+        self._temp_path = path
+
+        data = self._data
+
+        # torch.Tensor -> numpy
+        try:
+            import torch
+
+            if isinstance(data, torch.Tensor):
+                data = data.detach().cpu().numpy()
+        except ImportError:
+            pass
+
+        # 1) MoviePy clip
+        clip = self._maybe_as_moviepy_clip(data)
+        if clip is not None:
+            fps = self._fps or getattr(clip, "fps", None) or self.DEFAULT_FPS
+            self._write_moviepy_clip(clip, path, fps)
+            return path
+
+        # 2) numpy array of frames
+        try:
+            np = import_module("numpy")
+
+            if isinstance(data, np.ndarray):
+                clip = self._moviepy_clip_from_array(data, fps=self._fps or self.DEFAULT_FPS)
+                self._write_moviepy_clip(clip, path, self._fps or self.DEFAULT_FPS)
+                return path
+
+        except ImportError:
+            pass
+
+        raise TypeError(f"Unsupported video type: {type(data).__name__}")
+
+    def _maybe_as_moviepy_clip(self, data: Any) -> None | Any:
+        """Return data if it looks like a MoviePy clip, else None."""
+        try:
+            # MoviePy 2.x layout
+            video_clip_mod = import_module("moviepy.video.VideoClip")
+            video_clip = video_clip_mod.VideoClip
+            if isinstance(data, video_clip):
+                return data
+        except Exception:
+            pass
+
+        try:
+            # Older common import path
+            editor_mod = import_module("moviepy.editor")
+            video_clip = editor_mod.VideoClip
+            if isinstance(data, video_clip):
+                return data
+        except Exception:
+            pass
+
+        return None
+
+    def _moviepy_clip_from_array(self, data: Any, fps: float) -> Any:
+        np = import_module("numpy")
+
+        if data.dtype != np.uint8:
+            if np.issubdtype(data.dtype, np.floating):
+                # common case: [0,1] floats
+                if data.size and data.max() <= 1.0:
+                    data = (data * 255).clip(0, 255).astype(np.uint8)
+                else:
+                    data = data.clip(0, 255).astype(np.uint8)
+            else:
+                data = data.clip(0, 255).astype(np.uint8)
+
+        if data.ndim not in (3, 4):
+            raise ValueError(
+                f"Unsupported array shape for video: {data.shape}. Expected (T,H,W), (T,H,W,C), or (T,C,H,W)."
+            )
+
+        # (T, H, W) -> grayscale => expand to (T, H, W, 1)
+        if data.ndim == 3:
+            data = data[..., None]
+
+        # Handle TCHW -> THWC when channel axis is second
+        if data.ndim == 4 and data.shape[1] in (1, 3, 4) and data.shape[-1] not in (1, 3, 4):
+            data = data.transpose(0, 2, 3, 1)
+
+        if data.shape[-1] == 1:
+            # MoviePy ImageSequenceClip works best with 2D or RGB arrays;
+            # convert grayscale to RGB for consistency.
+            data = np.repeat(data, 3, axis=-1)
+
+        if data.shape[-1] not in (3, 4):
+            raise ValueError(f"Unsupported channel count for video frames: {data.shape[-1]}")
+
+        # Drop alpha for now unless you explicitly want to preserve/use masks.
+        if data.shape[-1] == 4:
+            data = data[..., :3]
+
+        try:
+            # MoviePy 2.x
+            image_sequence_clip = import_module("moviepy.video.io.ImageSequenceClip").ImageSequenceClip
+        except Exception:
+            # Older common path
+            image_sequence_clip = import_module("moviepy.editor").ImageSequenceClip
+
+        # list(...) avoids some ndarray edge cases in callers and matches
+        # common usage for frame sequences.
+        return image_sequence_clip(list(data), fps=fps)
+
+    def _write_moviepy_clip(self, clip: Any, path: str, fps: float) -> None:
+        # For MP4, libx264 is the usual sensible default.
+        kwargs: dict[str, Any] = {
+            "fps": fps,
+            "logger": None,
+        }
+
+        ext = os.path.splitext(path)[1].lower()
+        if ext == ".mp4":
+            kwargs["codec"] = "libx264"
+            # Better for browser progressive playback.
+            kwargs["ffmpeg_params"] = ["-movflags", "+faststart"]
+
+        clip.write_videofile(path, **kwargs)
+
+    @property
+    @override
+    def _media_type(self) -> MediaType:
+        return MediaType.VIDEO
+
+
 class Text(File):
     """Represents text content to be logged.
 
@@ -413,12 +588,13 @@ def _log_model(
         *,
         experiment_name: str,
         teamspace: Teamspace,
+        key: str | None = None,
         experiment: Any = None,
         cloud_account: str | None = None,
         verbose: bool = False,
     ) -> str:
         """Upload this model to the registry and return its registry name."""
-        model_name = self._registry_name(self.registry_name or experiment_name, teamspace)
+        model_name = self._registry_name(self.registry_name or key or experiment_name, teamspace)
 
         if self._model_kind == "artifact":
             upload_model(