Video generation integration and multi-modal data - #294
Open
guillaume-byte wants to merge 5 commits into
Open
Conversation
…ld=), video/audio/image/mask/pointcloud Lets any metadata field (not just the sample's own input) carry an image, mask, video clip, audio track, or point cloud -- e.g. a generated/predicted video attached to an ordinary segmentation sample. wl.save_media() writes a small descriptor into a media:<field> column plus the encoded bytes into an in-memory LRU store; GetMedia and the field-aware GetPointCloud stream those bytes to the studio on demand, independent of the sample's own task_type. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…a kind A dataset's __getitem__ can now return a plain string as its "image" (a prompt, a conversation turn) and wl.save_media(kind="text") can attach a caption/reference string to any sample — completing the modality set (image/mask/video/audio/pointcloud/text) on both sides so every main-sample type can also carry every metadata type. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
A text-generation sample's "image" slot is a plain string, not pixels; _get_image_array_and_metadata() passed it straight through and the caller's .ndim check crashed. Return None instead, matching the function's existing "nothing to preview" convention already handled by its only caller (data_service.py's natural-sort stats computation).
Datasets that combine several main-sample modalities behind one loader (e.g. mixing segmentation/video_generation/detection_pointcloud/ text_generation via one dataset object) have no single dataset-level task_type attribute. load_raw_video()/is_video_sample() previously only looked one up via getattr on the dataset object, so every video sample in such a dataset failed the video-shape check and could not stream. Peek at the sample's own getitem-returned metadata dict for a task_type override before falling back to today's shape-only heuristic.
…ave_media() A dataset's getitem metadata dict can already carry arbitrary content, but turning an array/string into real, renderable media (thumbnails, video player, 3D viewer, text panel) previously required calling wl.save_media() explicitly for every field. Add media_store.classify_kind() to auto-detect a value's kind by shape/type alone (array -> image/mask/video/pointcloud, str -> text), and DataSampleTrackingWrapper(auto_render_metadata=True) to attach every classified value through save_media()'s own encode/store path, once per sample right after registration -- no explicit call needed. A "<field>_audio" entry shaped (samples, sample_rate) is muxed into "<field>"'s video instead of becoming its own field; task_type is never auto-rendered, since it is a control value, not content. Off by default, so every other dataset's metadata keeps landing as plain DataFrame values exactly as before -- opting in is a per-loader choice via wl.watch_or_edit(..., auto_render_metadata=True).
Contributor
There was a problem hiding this comment.
Pull request overview
This PR extends WeightsLab’s backend + proto API to support multi-modal samples—especially generative media—by adding a streaming GetMedia RPC and a process-local media store so the UI can render videos/audio/text (both as primary sample payloads and as metadata-attached outputs).
Changes:
- Added
GetMediaserver-streaming RPC (proto + Python servicer) to stream playable media in chunks (MP4/GIF/WAV) with a small LRU cache. - Introduced
save_media()plus a newmedia_storeto attach encoded media to per-sample metadata fields (media:<field>) and surface posters/descriptors inDataStats. - Added new video/audio utilities and expanded data loading paths to handle text-generation and video poster rendering; added tests covering
GetMedia, video utils, and store-streamed point clouds.
Reviewed changes
Copilot reviewed 13 out of 14 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| weightslab/trainer/trainer_services.py | Exposes new GetMedia RPC by delegating to DataService. |
| weightslab/trainer/services/data_service.py | Implements GetMedia, media chunk sizing, media stats emission, and integrates media columns into metadata handling. |
| weightslab/src.py | Adds save_media() and _encode_one_media() for encoding/storing media attached to metadata fields. |
| weightslab/proto/experiment_service.proto | Adds GetMedia, MediaRequest, MediaChunk, and extends PointCloudRequest with field. |
| weightslab/proto/experiment_service_pb2.py | Regenerated protobuf Python definitions for new messages/RPC. |
| weightslab/proto/experiment_service_pb2_grpc.py | Regenerated gRPC Python bindings with GetMedia. |
| weightslab/data/video_utils.py | New utilities for poster rendering, audio encoding, and MP4/GIF clip encoding. |
| weightslab/data/media_store.py | New in-memory LRU store for encoded media bytes + descriptors/posters. |
| weightslab/data/data_utils.py | Adds text-generation and video-aware loading paths (load_raw_video, poster handling). |
| weightslab/data/data_samples_with_ops.py | Adds optional auto-rendering of metadata into media fields via save_media(). |
| weightslab/init.py | Re-exports save_media. |
| tests/gRPC/test_get_point_cloud.py | Extends tests for point clouds streamed from metadata-attached store entries. |
| tests/gRPC/test_get_media.py | Adds servicer-level tests for GetMedia (streaming, caching, field mode, failure modes). |
| tests/data/test_video_utils.py | Adds unit tests for video/audio utilities and routing heuristics. |
Files not reviewed (1)
- weightslab/proto/experiment_service_pb2.py: Generated file
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+5132
to
+5134
| self._media_cache[key] = value | ||
| while len(self._media_cache) > self._MEDIA_CACHE_ENTRIES: | ||
| self._media_cache.pop(next(iter(self._media_cache))) |
Comment on lines
+2816
to
+2825
| arr = _to_numpy(item) if item is not None else None | ||
|
|
||
| if kind == _ms.KIND_AUDIO: | ||
| data = _vu.encode_audio_wav(arr, sample_rate) | ||
| duration = (len(arr) / float(sample_rate)) if sample_rate else 0.0 | ||
| # Audio has no still; the list shows a "♪" chip instead of a thumbnail. | ||
| return data, "audio/wav", b"", { | ||
| "sample_rate": int(sample_rate or 0), | ||
| "duration_seconds": round(float(duration), 3), | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Include video generation usecase with Weightslab, and subsequent modalities not managed by the UI:
And all those modalities are managed and rendered from the UI as input or metadata now.