213 lines
7.2 KiB
Python
213 lines
7.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from vllm_omni.experimental.fullduplex.joyvl.bridges.backend import ModelBackend
|
|
from vllm_omni.experimental.fullduplex.joyvl.decision import prompts
|
|
|
|
|
|
@dataclass
|
|
class WorkingChunk:
|
|
messages: list[dict[str, Any]] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class MidTermSummary:
|
|
chunk_index: int
|
|
frame_range: str
|
|
summary_text: str
|
|
|
|
|
|
@dataclass
|
|
class QAEntry:
|
|
query: str
|
|
query_time: str
|
|
responses: list[tuple[str, str]] = field(default_factory=list)
|
|
archived_in_chunk: int = 0
|
|
|
|
|
|
@dataclass
|
|
class SessionMemory:
|
|
long_term_memory: str = ""
|
|
long_term_blocks: list[str] = field(default_factory=list)
|
|
mid_term_summaries: list[MidTermSummary] = field(default_factory=list)
|
|
qa_history: list[QAEntry] = field(default_factory=list)
|
|
|
|
|
|
def build_memory_prefix(
|
|
memory: SessionMemory,
|
|
*,
|
|
current_query: str | None,
|
|
query_in_current_chunk: bool,
|
|
keep_qa_history: bool,
|
|
current_chunk_index: int,
|
|
) -> str:
|
|
has_query = bool(current_query)
|
|
sections: list[str] = []
|
|
|
|
history_parts: list[str] = []
|
|
if has_query and memory.long_term_memory:
|
|
history_parts.append(memory.long_term_memory)
|
|
for entry in memory.mid_term_summaries:
|
|
history_parts.append(f"<{entry.frame_range}>\n{entry.summary_text}")
|
|
if history_parts:
|
|
sections.append(prompts.VIDEO_HISTORY_HEADER + "\n\n".join(history_parts))
|
|
|
|
if keep_qa_history and has_query:
|
|
qa_entries = [e for e in memory.qa_history if e.archived_in_chunk < current_chunk_index]
|
|
if qa_entries:
|
|
lines = []
|
|
for idx, entry in enumerate(qa_entries, 1):
|
|
parts = [f"#{idx} [{prompts.QA_QUERY_LABEL}@{entry.query_time or 'N/A'}] {entry.query}"]
|
|
for resp_time, payload in entry.responses:
|
|
parts.append(f"[{prompts.QA_RESPONSE_LABEL}@{resp_time}] {payload}")
|
|
lines.append("\n".join(parts))
|
|
sections.append(prompts.QA_HISTORY_HEADER + "\n".join(lines))
|
|
|
|
if has_query and not query_in_current_chunk:
|
|
sections.append(prompts.USER_QUERY_HEADER + "\n" + current_query.strip())
|
|
|
|
return "\n\n".join(sections)
|
|
|
|
|
|
def _sample_indices(n: int, budget: int) -> list[int]:
|
|
if n == 0:
|
|
return []
|
|
if budget <= 0 or n <= budget:
|
|
return list(range(n))
|
|
if budget == 1:
|
|
return [n // 2]
|
|
selected: list[int] = []
|
|
last_pos = -1
|
|
for slot in range(budget):
|
|
target = round(slot * (n - 1) / (budget - 1))
|
|
target = max(target, last_pos + 1)
|
|
target = min(target, n - (budget - slot))
|
|
selected.append(target)
|
|
last_pos = target
|
|
return selected
|
|
|
|
|
|
def _resize_data_url(data_url: str, max_pixels: int) -> str:
|
|
"""Downscale a base64 image data URL so width*height <= max_pixels (preserving
|
|
aspect ratio), mirroring the reference summarizer. No-op if already within budget,
|
|
not a data URL, or PIL is unavailable."""
|
|
if max_pixels <= 0 or not data_url.startswith("data:"):
|
|
return data_url
|
|
import base64
|
|
import io
|
|
import re
|
|
|
|
match = re.match(r"data:image/\w+;base64,(.+)", data_url)
|
|
if not match:
|
|
return data_url
|
|
try:
|
|
from PIL import Image
|
|
|
|
img = Image.open(io.BytesIO(base64.b64decode(match.group(1))))
|
|
w, h = img.size
|
|
if w * h <= max_pixels:
|
|
return data_url
|
|
scale = (max_pixels / (w * h)) ** 0.5
|
|
img = img.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS)
|
|
buf = io.BytesIO()
|
|
img.convert("RGB").save(buf, format="JPEG")
|
|
return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()
|
|
except Exception:
|
|
return data_url
|
|
|
|
|
|
class Summarizer:
|
|
def __init__(
|
|
self,
|
|
backend: ModelBackend,
|
|
*,
|
|
key_frames_per_chunk: int = 0,
|
|
mid_term_max_tokens: int = 4000,
|
|
long_term_max_tokens: int = 4000,
|
|
preferred_time_span: float = 10.0,
|
|
max_pixels: int = 262144,
|
|
) -> None:
|
|
self._backend = backend
|
|
self._key_frames = key_frames_per_chunk
|
|
self._mid_max_tokens = mid_term_max_tokens
|
|
self._long_max_tokens = long_term_max_tokens
|
|
self._preferred_time_span = preferred_time_span
|
|
self._max_pixels = max_pixels
|
|
|
|
async def summarize_chunk(
|
|
self,
|
|
chunk_index: int,
|
|
frame_range: str,
|
|
frames: list[tuple[str, str]],
|
|
) -> str:
|
|
if not frames:
|
|
return prompts.EMPTY_CHUNK_SUMMARY.format(frame_range=frame_range)
|
|
|
|
picked = [frames[i] for i in _sample_indices(len(frames), self._key_frames)]
|
|
prompt = prompts.MID_TERM_SUMMARY_PROMPT.format(
|
|
chunk_index=chunk_index,
|
|
frame_range=frame_range,
|
|
length_instruction="",
|
|
preferred_time_span=f"{self._preferred_time_span:g} seconds",
|
|
)
|
|
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
|
|
for time_range, data_url in picked:
|
|
content.append({"type": "text", "text": f"<{time_range}>"})
|
|
content.append({"type": "image_url", "image_url": {"url": _resize_data_url(data_url, self._max_pixels)}})
|
|
|
|
text, _ = await self._backend.generate(
|
|
[{"role": "user", "content": content}],
|
|
max_tokens=self._mid_max_tokens,
|
|
temperature=0.8,
|
|
top_p=0.9,
|
|
extra_body={"top_k": 40, "repetition_penalty": 1.1, "presence_penalty": 1.0, "greedy": False},
|
|
)
|
|
return text.strip()
|
|
|
|
async def compress_to_long_term(self, mid_terms: list[MidTermSummary]) -> str:
|
|
"""Compress a batch of mid-term summaries into a single long-term block.
|
|
|
|
Returns one block; the caller (InteractionBrain) accumulates blocks and applies
|
|
the sliding window, matching the reference long-term memory behavior.
|
|
"""
|
|
if not mid_terms:
|
|
return ""
|
|
|
|
merged_range = f"{_range_start(mid_terms[0].frame_range)}-{_range_end(mid_terms[-1].frame_range)}"
|
|
summaries_text = "\n\n".join(f"<{m.frame_range}>\n{m.summary_text}" for m in mid_terms)
|
|
prompt = prompts.LONG_TERM_COMPRESS_PROMPT.format(
|
|
merged_range=merged_range,
|
|
summaries_text=summaries_text,
|
|
length_instruction="",
|
|
)
|
|
compressed, _ = await self._backend.generate(
|
|
[{"role": "user", "content": prompt}],
|
|
max_tokens=self._long_max_tokens,
|
|
temperature=0.3,
|
|
top_p=0.7,
|
|
extra_body={"top_k": 30, "repetition_penalty": 1.1, "presence_penalty": 0.5, "greedy": False},
|
|
)
|
|
return f"<{merged_range}>\n{compressed.strip()}"
|
|
|
|
async def aclose(self) -> None:
|
|
await self._backend.aclose()
|
|
|
|
|
|
def _range_start(frame_range: str) -> str:
|
|
for sep in (" ~ ", "-"):
|
|
if sep in frame_range:
|
|
return frame_range.split(sep, 1)[0].strip()
|
|
return frame_range.strip()
|
|
|
|
|
|
def _range_end(frame_range: str) -> str:
|
|
for sep in (" ~ ", "-"):
|
|
if sep in frame_range:
|
|
return frame_range.split(sep, 1)[-1].strip()
|
|
return frame_range.strip()
|