initial commit

Dockerfile

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 FROM python:3.9-slim
 
 WORKDIR /app
@@ -9,6 +24,15 @@ RUN apt-get update && apt-get install -y \
     git \
     && rm -rf /var/lib/apt/lists/*
 
+RUN mkdir -p /tmp/huggingface \
+    && chown -R 1000:1000 /tmp/huggingface     
+                             
+ENV HF_HOME=/tmp/huggingface              \
+    HF_HUB_CACHE=/tmp/huggingface         \
+    TRANSFORMERS_CACHE=/tmp/huggingface   \
+    XDG_CACHE_HOME=/tmp                   \
+    TMPDIR=/tmp
+
 COPY requirements.txt ./
 COPY src/ ./src/
 

README.md

@@ -1,14 +1,14 @@
 ---
-title: Streamlit Template Space
+title: Cosmos Embed1
 emoji: 🚀
 colorFrom: red
 colorTo: red
 sdk: docker
 app_port: 8501
 tags:
-  - streamlit
+- streamlit
 pinned: false
-short_description: Streamlit template space
+short_description: Cosmos-Embed1 demo app
 ---
 
 # Welcome to Streamlit!

requirements.txt

@@ -1,3 +1,23 @@
-altair
-pandas
-streamlit
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+altair==5.5.0
+streamlit==1.45.0
+pandas==2.2.3
+plotly==6.0.1
+faiss-cpu==1.11.0
+transformers==4.44.2
+einops==0.8.1
+torchvision==0.21.0

src/kinetics700_val.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c7495ac043e0565b636b2ac964bf630fa0b0ed6c7569cbd6c92f156ea5899bb
+size 15595889

src/opendrive_val.parquet

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d0610bbbba386744efed46272c39d140ba902e1ff47acddb7ec54c44ff7d444
+size 10482753

src/streamlit_app.py

@@ -1,40 +1,229 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import streamlit as st
+from typing import Union, Optional
+import pandas as pd
+import plotly.express as px
+import faiss
+import numpy as np
+from transformers import AutoModel, AutoProcessor
+import torch
+from datetime import datetime
+
+
+class SelectedIndex:
+    def __init__(self, idx) -> None:
+        self.idx = int(idx)
+        self.timestamp = datetime.now()
+
+    def __eq__(self, value: Union["SelectedIndex", int]) -> bool:
+        if isinstance(value, SelectedIndex):
+            return self.idx == value.idx
+        return self.idx == int(value)
+
+    def __ne__(self, value: Union["SelectedIndex", int]) -> bool:
+        return not self.__eq__(value)
+
+    def is_valid(self) -> bool:
+        return self.idx >= 0
+
+
+@st.cache_data
+def load_data(path: str):
+    df = pd.read_parquet(path)
+    embs = np.stack(df["embedding"].tolist()).astype("float32")
+    faiss.normalize_L2(embs)
+    D = embs.shape[1]
+    index = faiss.IndexFlatIP(D)
+    index.add(embs)
+    return df, index, embs
+
+
+def load_model() -> tuple[AutoModel, AutoProcessor]:
+    if "preprocessor" not in st.session_state:
+        st.session_state.preprocessor = AutoProcessor.from_pretrained(
+            "nvidia/Cosmos-Embed1-224p", trust_remote_code=True, token=True,
+        )
+    if "model" not in st.session_state:
+        model = AutoModel.from_pretrained(
+            "nvidia/Cosmos-Embed1-224p", trust_remote_code=True, token=True,
+        )
+        model.eval()
+        st.session_state.model = model
+    return st.session_state.model, st.session_state.preprocessor
+
+
+def preview_video(df, idx, slot, height=420, margin_top=30, autoplay=True, title=None) -> None:
+    if title:
+        slot.markdown(f"### {title}")
+    start = int(df.loc[idx, "span_start"])
+    end = int(df.loc[idx, "span_end"])
+    youtube_id = df.loc[idx, "youtube_id"]
+    url = f"https://www.youtube.com/embed/{youtube_id}?start={start}&end={end}"
+    sep = "?" if "?" not in url else "&"
+    params = f"{sep}mute=1&rel=0"
+    if autoplay:
+        params += "&autoplay=1"
+    slot.markdown(
+        f'''
+        <div style="margin-top:{margin_top}px">
+          <iframe width="100%" height="{height}"
+                  src="{url}{params}"
+                  frameborder="0"
+                  allow="accelerometer; autoplay; encrypted-media; gyroscope; picture-in-picture"
+                  allow="autoplay; fullscreen" allowfullscreen>
+          </iframe>
+        </div>
+        ''',
+        unsafe_allow_html=True
+    )
+
+
+def get_nearest_ids(vec, k=5, ignore_self=True) -> list:
+    q = vec.reshape(1, -1).astype("float32")
+    faiss.normalize_L2(q)
+    topk = k + 1 if ignore_self else k
+    _, I = index.search(q, topk)
+    ids = I[0]
+    return ids[1:].tolist() if ignore_self else ids.tolist()
+
+
+def get_most_recent_selection() -> tuple[Optional[int], str]:
+    if st.session_state.text_selection.is_valid() and st.session_state.click_selection.is_valid():
+        if st.session_state.text_selection.timestamp > st.session_state.click_selection.timestamp:
+            return st.session_state.text_selection.idx, "text"
+        return st.session_state.click_selection.idx, "click"
+    if st.session_state.text_selection.is_valid():
+        return st.session_state.text_selection.idx, "text"
+    if st.session_state.click_selection.is_valid():
+        return st.session_state.click_selection.idx, "click"
+    return None, ""
+
+
+def reset_state() -> None:
+    if "text_selection" not in st.session_state:
+        st.session_state.text_selection = SelectedIndex(-1)
+    if "click_selection" not in st.session_state:
+        st.session_state.click_selection = SelectedIndex(-1)
+    if "text_query" not in st.session_state:
+        st.session_state.text_query = ""
+
+# ─── App setup ────────────────────────────────────────────────────
+
+st.set_page_config(layout="wide")
+reset_state()
+model, preprocessor = load_model()
+file_map = {"kinetics700 (val)": "src/kinetics700_val.parquet", "opendv (val)": "src/opendrive_val.parquet"}
+st.title("🔍 Search with Cosmos-Embed1")
+
+col1, col2 = st.columns([2,2])
+with col1:
+    dataset = st.selectbox("Select dataset", list(file_map.keys()), on_change=reset_state)
+df, index, embs = load_data(file_map[dataset])
+
+# initialize session state
+if "text_selection" not in st.session_state:
+    st.session_state.text_selection = SelectedIndex(-1)
+if "click_selection" not in st.session_state:
+    st.session_state.click_selection = SelectedIndex(-1)
+if "text_query" not in st.session_state:
+    st.session_state.text_query = ""
+
+# ─── Layout ────────────────────────────────────────────────────────
+
+# LEFT: scatter
+with col1:
+    fig = px.scatter(
+        df, x="x", y="y",
+        hover_name="tar_key", hover_data=["cluster_id"],
+        color="cluster_id", color_continuous_scale="Turbo",
+        title="t-SNE projection (click to select)"
+    )
+    fig.update_layout(
+        dragmode="zoom",
+        margin=dict(l=5, r=5, t=40, b=5),
+        xaxis_title=None, yaxis_title=None,
+        coloraxis_colorbar=dict(title="")
+    )
+    fig.update_xaxes(showticklabels=False, showgrid=False, zeroline=False,
+                     showline=True, linecolor="black", mirror=True)
+    fig.update_yaxes(showticklabels=False, showgrid=False, zeroline=False,
+                     showline=True, linecolor="black", mirror=True)
+    fig.update_layout(annotations=[dict(
+        text="k-means cluster", xref="paper", yref="paper",
+        x=1.02, y=0.5, textangle=90, showarrow=False
+    )])
+
+    most_recent_idx, most_recent_method = get_most_recent_selection()
+    if most_recent_idx is not None and most_recent_method == "text":
+        x0, y0 = df.iloc[most_recent_idx][["x", "y"]]
+        span = 6.0
+        fig.update_layout(
+            xaxis_range=[x0 - span, x0 + span],
+            yaxis_range=[y0 - span, y0 + span],
+            transition={"duration": 1},
+        )
+
+    click_event = st.plotly_chart(
+        fig, use_container_width=True,
+        on_select="rerun", selection_mode="points"
+    )
+
+# RIGHT: text input & preview
+with col2:
+    if click_event and click_event.get("selection", {}).get("point_indices"):
+        curr_click = click_event["selection"]["point_indices"][0]
+        if curr_click != st.session_state.click_selection:
+            # new click so update the previous selection and wipe any text query
+            st.session_state.click_selection = SelectedIndex(curr_click)
+            st.session_state.text_query = ""
+
+    # text input (will pick up cleared or existing text)
+    text_query = st.text_input(
+        "Search via text",
+        key="text_query",
+        help="Type a query and press Enter"
+    )
+
+    # if user typed text (and pressed Enter), override selection
+    if text_query:
+        with torch.no_grad():
+            model_input = preprocessor(text=[text_query])
+            emb_out = model.get_text_embeddings(**model_input).text_proj.cpu().numpy()
+        idx_text, = get_nearest_ids(emb_out, k=1, ignore_self=False)
+        if st.session_state.text_selection != idx_text:
+            # new text so update the previous selection and wipe any text query
+            st.session_state.text_selection = SelectedIndex(idx_text)
+            st.rerun()
+
+    # main preview
+    preview_slot = st.empty()
+    most_recent, most_recent_modality = get_most_recent_selection()
+    if most_recent is not None:
+        preview_video(df, most_recent, preview_slot)
+    else:
+        preview_slot.write("⏳ Waiting for selection…")
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# BOTTOM: 5 nearest neighbors
+st.markdown("### 🎬 5 Closest Videos")
+if most_recent is not None:
+    ignore_self = most_recent_modality == "click"
+    nn_ids = get_nearest_ids(embs[most_recent], k=5, ignore_self=ignore_self)
+    cols = st.columns(5)
+    for c, nid in zip(cols, nn_ids):
+        preview_video(df, nid, c, height=180, margin_top=5, autoplay=False)
+else:
+    st.write("Use a click or a text query above to list neighbors.")