From fe34dea84ac674ac3c572f773b131cd8a3b3a932 Mon Sep 17 00:00:00 2001
From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Date: Wed, 3 Jun 2026 20:00:35 -0700
Subject: [PATCH 1/3] feat(generator): add streaming generation API with
 chunked audio decode

---
 README.md                          |  28 +++++++
 examples/stream_demo.py            |  97 +++++++++++++++++++++
 generator.py                       | 130 +++++++++++++++++++++++++++--
 tests/test_streaming_generation.py |  31 +++++++
 4 files changed, 278 insertions(+), 8 deletions(-)
 create mode 100644 examples/stream_demo.py
 create mode 100644 tests/test_streaming_generation.py

diff --git a/README.md b/README.md
index 0423b4f..0eb218c 100644
--- a/README.md
+++ b/README.md
@@ -172,6 +172,34 @@ audio = generator.generate(
 )
 ```
 
+### Streaming generation
+
+Use `generate_stream()` to receive PCM chunks while frames are still being
+generated. The default `chunk_frames=25` yields about two seconds of 24 kHz
+audio per chunk.
+Chunks are watermarked independently; if SilentCipher rejects a very short
+terminal fragment, that fragment is yielded unchanged.
+
+```python
+import torch
+
+from generator import load_miso_8b
+
+generator = load_miso_8b(device="cuda")
+
+chunks = []
+for chunk in generator.generate_stream(
+    text="This sentence is decoded in chunks.",
+    speaker=0,
+    context=[],
+    max_audio_length_ms=10_000,
+    chunk_frames=25,
+):
+    chunks.append(chunk.cpu())
+
+audio = torch.cat(chunks, dim=0)
+```
+
 ---
 
 ## Weights
diff --git a/examples/stream_demo.py b/examples/stream_demo.py
new file mode 100644
index 0000000..97c3529
--- /dev/null
+++ b/examples/stream_demo.py
@@ -0,0 +1,97 @@
+import argparse
+import os
+import time
+import wave
+
+os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "60")
+os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60")
+
+import torch
+
+from generator import DEFAULT_MISO_TTS_REPO_ID, load_miso_8b
+
+# Disable Triton compilation
+os.environ["NO_TORCH_COMPILE"] = "1"
+
+
+def _pcm16_bytes(audio: torch.Tensor) -> bytes:
+    audio = audio.detach().flatten().to(dtype=torch.float32).cpu()
+    pcm = audio.clamp(-1.0, 1.0).mul(32767.0).to(torch.int16)
+    return pcm.numpy().tobytes()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--text", type=str, default="Hello from streamed Miso TTS.")
+    parser.add_argument("--speaker", type=int, default=0)
+    parser.add_argument("--output", type=str, default="streamed_generation.wav")
+    parser.add_argument("--max-audio-length-ms", type=float, default=10_000)
+    parser.add_argument("--chunk-frames", type=int, default=25)
+    parser.add_argument(
+        "--model-path-or-repo-id",
+        type=str,
+        default=os.environ.get("MISO_TTS_8B_MODEL", DEFAULT_MISO_TTS_REPO_ID),
+    )
+    args = parser.parse_args()
+
+    # Select the best available device, skipping MPS due to float64 limitations.
+    if torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+    print(f"Using device: {device}")
+
+    if os.path.exists(args.model_path_or_repo_id):
+        print(f"Loading Miso TTS model from local path: {args.model_path_or_repo_id}")
+    else:
+        print(
+            "Loading Miso TTS model from Hugging Face: "
+            f"https://huggingface.co/{args.model_path_or_repo_id}"
+        )
+        print("The model will be downloaded and cached automatically if it is not already present.")
+
+    generator = load_miso_8b(device, model_path_or_repo_id=args.model_path_or_repo_id)
+
+    start_time = time.perf_counter()
+    first_audio_time = None
+    chunk_count = 0
+    sample_count = 0
+
+    with wave.open(args.output, "wb") as output_file:
+        output_file.setnchannels(1)
+        output_file.setsampwidth(2)
+        output_file.setframerate(generator.sample_rate)
+
+        for chunk in generator.generate_stream(
+            text=args.text,
+            speaker=args.speaker,
+            context=[],
+            max_audio_length_ms=args.max_audio_length_ms,
+            chunk_frames=args.chunk_frames,
+        ):
+            if chunk.numel() == 0:
+                continue
+
+            if first_audio_time is None:
+                first_audio_time = time.perf_counter()
+
+            chunk_count += 1
+            sample_count += chunk.numel()
+            output_file.writeframes(_pcm16_bytes(chunk))
+            duration_s = chunk.numel() / generator.sample_rate
+            print(f"Wrote chunk {chunk_count}: {duration_s:.2f}s")
+
+    total_time = time.perf_counter() - start_time
+    audio_duration_s = sample_count / generator.sample_rate
+    if first_audio_time is None:
+        print("No audio generated.")
+    else:
+        time_to_first_audio = first_audio_time - start_time
+        print(f"Time to first audio: {time_to_first_audio:.2f}s")
+        print(f"Total generation time: {total_time:.2f}s")
+        print(f"Audio duration: {audio_duration_s:.2f}s")
+        print(f"Successfully generated {args.output}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/generator.py b/generator.py
index 8441cee..8b1f566 100644
--- a/generator.py
+++ b/generator.py
@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 import os
-from typing import List, Optional, Tuple
+from typing import Iterator, List, Optional, Tuple
 
 os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "60")
 os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "60")
@@ -27,6 +27,19 @@ class Segment:
     audio: torch.Tensor
 
 
+def _stack_audio_frames(samples: List[torch.Tensor]) -> torch.Tensor:
+    return torch.stack(samples).permute(1, 2, 0)
+
+
+def _match_num_samples(audio: torch.Tensor, num_samples: int) -> torch.Tensor:
+    if audio.size(0) > num_samples:
+        return audio[:num_samples]
+    if audio.size(0) < num_samples:
+        padding = torch.zeros(num_samples - audio.size(0), dtype=audio.dtype, device=audio.device)
+        return torch.cat([audio, padding], dim=0)
+    return audio
+
+
 def load_llama3_tokenizer():
     """
     https://github.com/huggingface/transformers/issues/22794#issuecomment-2092623992
@@ -114,8 +127,7 @@ def _tokenize_segment(self, segment: Segment) -> Tuple[torch.Tensor, torch.Tenso
 
         return torch.cat([text_tokens, audio_tokens], dim=0), torch.cat([text_masks, audio_masks], dim=0)
 
-    @torch.inference_mode()
-    def generate(
+    def _generate_frames(
         self,
         text: str,
         speaker: int,
@@ -123,7 +135,7 @@ def generate(
         max_audio_length_ms: float = 90_000,
         temperature: float = 0.9,
         topk: int = 50,
-    ) -> torch.Tensor:
+    ) -> Iterator[torch.Tensor]:
         self._model.reset_caches()
 
         max_generation_len = int(max_audio_length_ms / 80)
@@ -140,7 +152,6 @@ def generate(
         prompt_tokens = torch.cat(tokens, dim=0).long().to(self.device)
         prompt_tokens_mask = torch.cat(tokens_mask, dim=0).bool().to(self.device)
 
-        samples = []
         curr_tokens = prompt_tokens.unsqueeze(0)
         curr_tokens_mask = prompt_tokens_mask.unsqueeze(0)
         curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
@@ -157,23 +168,126 @@ def generate(
             if torch.all(sample == 0):
                 break  # eos
 
-            samples.append(sample)
-
             curr_tokens = torch.cat([sample, torch.zeros(1, 1).long().to(self.device)], dim=1).unsqueeze(1)
             curr_tokens_mask = torch.cat(
                 [torch.ones_like(sample).bool(), torch.zeros(1, 1).bool().to(self.device)], dim=1
             ).unsqueeze(1)
             curr_pos = curr_pos[:, -1:] + 1
 
-        audio = self._audio_tokenizer.decode(torch.stack(samples).permute(1, 2, 0)).squeeze(0).squeeze(0)
+            yield sample
 
+    def _decode_frames(self, samples: List[torch.Tensor]) -> torch.Tensor:
+        return self._audio_tokenizer.decode(_stack_audio_frames(samples)).squeeze(0).squeeze(0)
+
+    def _watermark_audio(self, audio: torch.Tensor) -> torch.Tensor:
         # This applies an imperceptible watermark to identify audio as AI-generated.
         # If using Miso TTS in another application, use your own private key and keep it secret.
         audio, wm_sample_rate = watermark(self._watermarker, audio, self.sample_rate, MISO_TTS_WATERMARK)
         audio = torchaudio.functional.resample(audio, orig_freq=wm_sample_rate, new_freq=self.sample_rate)
+        return audio
+
+    def _watermark_stream_chunk(self, audio: torch.Tensor, *, is_final: bool) -> torch.Tensor:
+        target_num_samples = audio.size(0)
+        try:
+            audio = self._watermark_audio(audio)
+        except Exception:
+            if not is_final:
+                raise
+            # SilentCipher may reject very short final chunks. Earlier chunks are
+            # watermarked independently; only the terminal fragment falls back.
+            return audio
+        return _match_num_samples(audio, target_num_samples)
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        text: str,
+        speaker: int,
+        context: List[Segment],
+        max_audio_length_ms: float = 90_000,
+        temperature: float = 0.9,
+        topk: int = 50,
+    ) -> torch.Tensor:
+        samples = list(
+            self._generate_frames(
+                text=text,
+                speaker=speaker,
+                context=context,
+                max_audio_length_ms=max_audio_length_ms,
+                temperature=temperature,
+                topk=topk,
+            )
+        )
+
+        audio = self._decode_frames(samples)
+        audio = self._watermark_audio(audio)
 
         return audio
 
+    def generate_stream(
+        self,
+        text: str,
+        speaker: int,
+        context: List[Segment],
+        max_audio_length_ms: float = 90_000,
+        temperature: float = 0.9,
+        topk: int = 50,
+        chunk_frames: int = 25,
+    ) -> Iterator[torch.Tensor]:
+        if chunk_frames <= 0:
+            raise ValueError("chunk_frames must be greater than 0")
+
+        # Keep inference mode inside the generator body. Decorating a generator
+        # function only wraps iterator construction, not iteration.
+        with torch.inference_mode():
+            chunk: List[torch.Tensor] = []
+            all_samples: List[torch.Tensor] = []
+            streamed_num_samples = 0
+            pending_final_audio: Optional[torch.Tensor] = None
+
+            with self._audio_tokenizer.streaming(1):
+                for sample in self._generate_frames(
+                    text=text,
+                    speaker=speaker,
+                    context=context,
+                    max_audio_length_ms=max_audio_length_ms,
+                    temperature=temperature,
+                    topk=topk,
+                ):
+                    chunk.append(sample)
+                    all_samples.append(sample)
+
+                    if len(chunk) == chunk_frames:
+                        audio = self._decode_frames(chunk)
+                        streamed_num_samples += audio.size(0)
+                        chunk = []
+                        if audio.numel() > 0:
+                            yield self._watermark_stream_chunk(audio, is_final=False)
+
+                if chunk:
+                    audio = self._decode_frames(chunk)
+                    streamed_num_samples += audio.size(0)
+                    if audio.numel() > 0:
+                        pending_final_audio = audio
+
+            if all_samples:
+                # Mimi streaming decode does not expose an explicit flush. Decode
+                # the full code sequence once and emit only the deferred tail so
+                # concatenated stream chunks keep the batch decode length.
+                full_audio = self._decode_frames(all_samples)
+                tail = None
+                if streamed_num_samples < full_audio.size(0):
+                    tail = full_audio[streamed_num_samples:]
+
+                if pending_final_audio is not None:
+                    yield self._watermark_stream_chunk(
+                        pending_final_audio,
+                        is_final=tail is None or tail.numel() == 0,
+                    )
+
+                if tail is not None and tail.numel() > 0:
+                    yield self._watermark_stream_chunk(tail, is_final=True)
+
 
 def _state_dict_from_checkpoint(checkpoint: object) -> dict[str, torch.Tensor]:
     if not isinstance(checkpoint, dict):
diff --git a/tests/test_streaming_generation.py b/tests/test_streaming_generation.py
new file mode 100644
index 0000000..4c0e37f
--- /dev/null
+++ b/tests/test_streaming_generation.py
@@ -0,0 +1,31 @@
+import unittest
+
+import torch
+
+from generator import _match_num_samples, _stack_audio_frames
+
+
+class StreamingGenerationHelpersTest(unittest.TestCase):
+    def test_stack_audio_frames_preserves_time_order(self) -> None:
+        frames = [
+            torch.tensor([[1, 2, 3]]),
+            torch.tensor([[4, 5, 6]]),
+        ]
+
+        codes = _stack_audio_frames(frames)
+
+        expected = torch.tensor([[[1, 4], [2, 5], [3, 6]]])
+        self.assertTrue(torch.equal(codes, expected))
+
+    def test_match_num_samples_trims_or_pads(self) -> None:
+        audio = torch.tensor([1.0, 2.0, 3.0])
+
+        trimmed = _match_num_samples(audio, 2)
+        padded = _match_num_samples(audio, 5)
+
+        self.assertTrue(torch.equal(trimmed, torch.tensor([1.0, 2.0])))
+        self.assertTrue(torch.equal(padded, torch.tensor([1.0, 2.0, 3.0, 0.0, 0.0])))
+
+
+if __name__ == "__main__":
+    unittest.main()

From 4667c226cad45b4984764e62e6f3ea4a6ac4e95e Mon Sep 17 00:00:00 2001
From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Date: Wed, 3 Jun 2026 20:08:16 -0700
Subject: [PATCH 2/3] fix: address self-review findings

---
 generator.py | 113 +++++++++++++++++++++++++++++++--------------------
 1 file changed, 70 insertions(+), 43 deletions(-)

diff --git a/generator.py b/generator.py
index 8b1f566..567d972 100644
--- a/generator.py
+++ b/generator.py
@@ -127,15 +127,13 @@ def _tokenize_segment(self, segment: Segment) -> Tuple[torch.Tensor, torch.Tenso
 
         return torch.cat([text_tokens, audio_tokens], dim=0), torch.cat([text_masks, audio_masks], dim=0)
 
-    def _generate_frames(
+    def _prepare_prompt(
         self,
         text: str,
         speaker: int,
         context: List[Segment],
-        max_audio_length_ms: float = 90_000,
-        temperature: float = 0.9,
-        topk: int = 50,
-    ) -> Iterator[torch.Tensor]:
+        max_audio_length_ms: float,
+    ) -> Tuple[torch.Tensor, torch.Tensor, int]:
         self._model.reset_caches()
 
         max_generation_len = int(max_audio_length_ms / 80)
@@ -152,17 +150,27 @@ def _generate_frames(
         prompt_tokens = torch.cat(tokens, dim=0).long().to(self.device)
         prompt_tokens_mask = torch.cat(tokens_mask, dim=0).bool().to(self.device)
 
-        curr_tokens = prompt_tokens.unsqueeze(0)
-        curr_tokens_mask = prompt_tokens_mask.unsqueeze(0)
-        curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
-
         max_seq_len = 2048
         max_context_len = max_seq_len - max_generation_len
-        if curr_tokens.size(1) >= max_context_len:
+        if prompt_tokens.size(0) >= max_context_len:
             raise ValueError(
                 f"Inputs too long, must be below max_seq_len - max_generation_len: {max_context_len}"
             )
 
+        return prompt_tokens, prompt_tokens_mask, max_generation_len
+
+    def _generate_frames(
+        self,
+        prompt_tokens: torch.Tensor,
+        prompt_tokens_mask: torch.Tensor,
+        max_generation_len: int,
+        temperature: float,
+        topk: int,
+    ) -> Iterator[torch.Tensor]:
+        curr_tokens = prompt_tokens.unsqueeze(0)
+        curr_tokens_mask = prompt_tokens_mask.unsqueeze(0)
+        curr_pos = torch.arange(0, prompt_tokens.size(0)).unsqueeze(0).long().to(self.device)
+
         for _ in range(max_generation_len):
             sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk)
             if torch.all(sample == 0):
@@ -208,15 +216,11 @@ def generate(
         temperature: float = 0.9,
         topk: int = 50,
     ) -> torch.Tensor:
+        prompt_tokens, prompt_tokens_mask, max_generation_len = self._prepare_prompt(
+            text, speaker, context, max_audio_length_ms
+        )
         samples = list(
-            self._generate_frames(
-                text=text,
-                speaker=speaker,
-                context=context,
-                max_audio_length_ms=max_audio_length_ms,
-                temperature=temperature,
-                topk=topk,
-            )
+            self._generate_frames(prompt_tokens, prompt_tokens_mask, max_generation_len, temperature, topk)
         )
 
         audio = self._decode_frames(samples)
@@ -237,56 +241,79 @@ def generate_stream(
         if chunk_frames <= 0:
             raise ValueError("chunk_frames must be greater than 0")
 
-        # Keep inference mode inside the generator body. Decorating a generator
-        # function only wraps iterator construction, not iteration.
+        # Tokenize the prompt (including Mimi encode of any context audio) before
+        # entering Mimi streaming mode, so streamed generation conditions on the
+        # same context tokens as generate().
         with torch.inference_mode():
-            chunk: List[torch.Tensor] = []
-            all_samples: List[torch.Tensor] = []
-            streamed_num_samples = 0
-            pending_final_audio: Optional[torch.Tensor] = None
-
-            with self._audio_tokenizer.streaming(1):
-                for sample in self._generate_frames(
-                    text=text,
-                    speaker=speaker,
-                    context=context,
-                    max_audio_length_ms=max_audio_length_ms,
-                    temperature=temperature,
-                    topk=topk,
-                ):
-                    chunk.append(sample)
-                    all_samples.append(sample)
+            prompt_tokens, prompt_tokens_mask, max_generation_len = self._prepare_prompt(
+                text, speaker, context, max_audio_length_ms
+            )
+
+        frames = self._generate_frames(prompt_tokens, prompt_tokens_mask, max_generation_len, temperature, topk)
+
+        chunk: List[torch.Tensor] = []
+        all_samples: List[torch.Tensor] = []
+        streamed_num_samples = 0
+        pending_final_audio: Optional[torch.Tensor] = None
+
+        # Each chunk is computed fully inside torch.inference_mode() and yielded
+        # outside it. A plain `with torch.inference_mode():` around the whole
+        # generator body would stay active while the generator is suspended at
+        # `yield`, silently putting the caller's loop body into inference mode.
+        with self._audio_tokenizer.streaming(1):
+            finished = False
+            while not finished:
+                out: Optional[torch.Tensor] = None
+                with torch.inference_mode():
+                    while len(chunk) < chunk_frames:
+                        sample = next(frames, None)
+                        if sample is None:
+                            finished = True
+                            break
+                        chunk.append(sample)
+                        all_samples.append(sample)
 
                     if len(chunk) == chunk_frames:
                         audio = self._decode_frames(chunk)
                         streamed_num_samples += audio.size(0)
                         chunk = []
                         if audio.numel() > 0:
-                            yield self._watermark_stream_chunk(audio, is_final=False)
+                            out = self._watermark_stream_chunk(audio, is_final=False)
+                if out is not None:
+                    yield out
 
+            with torch.inference_mode():
                 if chunk:
                     audio = self._decode_frames(chunk)
                     streamed_num_samples += audio.size(0)
                     if audio.numel() > 0:
                         pending_final_audio = audio
 
+        final_chunks: List[torch.Tensor] = []
+        with torch.inference_mode():
             if all_samples:
                 # Mimi streaming decode does not expose an explicit flush. Decode
-                # the full code sequence once and emit only the deferred tail so
-                # concatenated stream chunks keep the batch decode length.
+                # the full code sequence once (in batch mode, outside the streaming
+                # context) and emit only the deferred tail so concatenated stream
+                # chunks keep the batch decode length.
                 full_audio = self._decode_frames(all_samples)
                 tail = None
                 if streamed_num_samples < full_audio.size(0):
                     tail = full_audio[streamed_num_samples:]
 
                 if pending_final_audio is not None:
-                    yield self._watermark_stream_chunk(
-                        pending_final_audio,
-                        is_final=tail is None or tail.numel() == 0,
+                    final_chunks.append(
+                        self._watermark_stream_chunk(
+                            pending_final_audio,
+                            is_final=tail is None or tail.numel() == 0,
+                        )
                     )
 
                 if tail is not None and tail.numel() > 0:
-                    yield self._watermark_stream_chunk(tail, is_final=True)
+                    final_chunks.append(self._watermark_stream_chunk(tail, is_final=True))
+
+        for out in final_chunks:
+            yield out
 
 
 def _state_dict_from_checkpoint(checkpoint: object) -> dict[str, torch.Tensor]:

From 7f76f25466ce1be3139cf79f90f9a3505a1fe54b Mon Sep 17 00:00:00 2001
From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Date: Wed, 3 Jun 2026 20:14:23 -0700
Subject: [PATCH 3/3] fix: decode only uniform chunks inside Mimi streaming
 mode (CUDA graph shape capture)

---
 generator.py | 38 +++++++++++++-------------------------
 1 file changed, 13 insertions(+), 25 deletions(-)

diff --git a/generator.py b/generator.py
index 567d972..f20a74d 100644
--- a/generator.py
+++ b/generator.py
@@ -254,12 +254,16 @@ def generate_stream(
         chunk: List[torch.Tensor] = []
         all_samples: List[torch.Tensor] = []
         streamed_num_samples = 0
-        pending_final_audio: Optional[torch.Tensor] = None
 
         # Each chunk is computed fully inside torch.inference_mode() and yielded
         # outside it. A plain `with torch.inference_mode():` around the whole
         # generator body would stay active while the generator is suspended at
         # `yield`, silently putting the caller's loop body into inference mode.
+        #
+        # Only full chunk_frames-sized chunks are decoded inside Mimi streaming
+        # mode: on CUDA the streaming decoder is wrapped in a CUDA graph captured
+        # at the first chunk's shape, so a shorter residual chunk would raise a
+        # shape mismatch. The residual is covered by the batch tail decode below.
         with self._audio_tokenizer.streaming(1):
             finished = False
             while not finished:
@@ -282,38 +286,22 @@ def generate_stream(
                 if out is not None:
                     yield out
 
-            with torch.inference_mode():
-                if chunk:
-                    audio = self._decode_frames(chunk)
-                    streamed_num_samples += audio.size(0)
-                    if audio.numel() > 0:
-                        pending_final_audio = audio
-
-        final_chunks: List[torch.Tensor] = []
+        tail_audio: Optional[torch.Tensor] = None
         with torch.inference_mode():
             if all_samples:
                 # Mimi streaming decode does not expose an explicit flush. Decode
                 # the full code sequence once (in batch mode, outside the streaming
-                # context) and emit only the deferred tail so concatenated stream
-                # chunks keep the batch decode length.
+                # context) and emit only the deferred tail - the residual frames
+                # plus any samples the streamed chunks have not covered - so
+                # concatenated stream chunks keep the batch decode length.
                 full_audio = self._decode_frames(all_samples)
-                tail = None
                 if streamed_num_samples < full_audio.size(0):
                     tail = full_audio[streamed_num_samples:]
+                    if tail.numel() > 0:
+                        tail_audio = self._watermark_stream_chunk(tail, is_final=True)
 
-                if pending_final_audio is not None:
-                    final_chunks.append(
-                        self._watermark_stream_chunk(
-                            pending_final_audio,
-                            is_final=tail is None or tail.numel() == 0,
-                        )
-                    )
-
-                if tail is not None and tail.numel() > 0:
-                    final_chunks.append(self._watermark_stream_chunk(tail, is_final=True))
-
-        for out in final_chunks:
-            yield out
+        if tail_audio is not None:
+            yield tail_audio
 
 
 def _state_dict_from_checkpoint(checkpoint: object) -> dict[str, torch.Tensor]: