Upload evaluated Quasar weights with portable inference and pinned codecs

README.md

@@ -31,10 +31,12 @@ hybrid recurrence/attention backbone. Text, visual codes, audio codes, triangle
 meshes and tool-call proposals share one vocabulary and model. Frozen codecs
 convert between media and discrete tokens.
 
-**Release status: documentation and evaluation evidence only. No model weights,
-training data, codec weights or runnable inference service are included in this
-repository.** This card describes the preserved v5 model and a completed
-curriculum-research branch as of **21 October 2026**. Separate architecture
+**Release status: private, runnable research checkpoint release.** The default
+root weights are the completed curriculum-research checkpoint shown below; the
+preserved v5 comparison weights are in `variants/v5`. Matching source code,
+text tokenizers and pinned frozen visual/audio codecs are included. Training
+datasets, optimizer states and an inference service are not included.
+This card describes these two snapshots as of **21 October 2026**. Separate architecture
 experiments are still being evaluated and are excluded from these results.
 
 **Quality status:** a focused research branch produces recognizable images and
@@ -71,13 +73,14 @@ excludes that auxiliary term.
 
 ## Checkpoint identity
 
-The measurements below refer to these local artifacts, not to downloadable Hub
-weights. SHA-256 values identify the exact `model.safetensors` files.
+The downloadable weights match the exact artifacts measured below. SHA-256 values
+identify each unchanged `model.safetensors` file; packaging only makes metadata
+paths portable and supplies the matching runtime.
 
 | Artifact | Role | SHA-256 |
 |---|---|---|
-| v5 full image continuation | Preserved comparison model | `e20242e9971b6a7cef7ca23a65c25faa70a007807a55a2c0cfe35f4f5682a643` |
-| `joint_repair/final` | Isolated, unpromoted curriculum-research model | `de4846e5e63c0004cf0b6331127ab421d6843f6efde04243c6d5448ab4545240` |
+| `variants/v5/model.safetensors` | Preserved v5 comparison model | `e20242e9971b6a7cef7ca23a65c25faa70a007807a55a2c0cfe35f4f5682a643` |
+| Root `model.safetensors` | Unpromoted `joint_repair/final` research model | `de4846e5e63c0004cf0b6331127ab421d6843f6efde04243c6d5448ab4545240` |
 
 ## Intended use and limitations
 
@@ -218,23 +221,39 @@ research session and is **not generated-image quality**. The earlier 64px
 real-photo probe measured 18.1176 dB; its original aspirational 20 dB gate remained
 false and was not relabeled as passed.
 
-## Local usage requirements
-
-There is no `from_pretrained("Atlas-AI-research/Quasar")` weights release here.
-Running the model requires the matching local Quasar implementation, a compatible
-checkpoint directory with `config.json`, `model.safetensors` and tokenizer metadata,
-and independently available pinned codecs. The project is not a self-contained
-standard Transformers model, and local codec references must be resolved explicitly.
+## Download and run
+
+This is a custom PyTorch checkpoint package, not a Transformers `AutoModel`.
+Use the included `quasar` implementation rather than
+`AutoModel.from_pretrained("Atlas-AI-research/Quasar")`. Authentication is needed
+while the repository is private. Preserve a working CUDA-compatible PyTorch
+installation; the measured environment used PyTorch 2.11.0+cu128. Install a suitable
+PyTorch build separately when needed, then install the remaining inference dependencies:
+
+```bash
+hf auth login
+hf download Atlas-AI-research/Quasar --local-dir Quasar
+cd Quasar
+python -m pip install -r requirements-inference.txt
+python verify_release.py --codecs
+python infer.py --task image --prompt "make a picture of a cat" --output cat.png
+python infer.py --task video --prompt "a small cyan square moving left on a gray background" --output motion.gif
+```
 
-With those prerequisites present, the existing local API is:
+The default inference device is CUDA; the CLI refuses a busy GPU. Verification
+uses CPU only. CPU inference is available with `--device cpu`; it does not train
+the model. Select the preserved baseline with `--checkpoint variants/v5`.
+CPU-only entry points hide CUDA and exclude the unused optional TorchAO quantizer
+within their process, avoiding GPU probing by an unrelated installed backend.
+The included package can also be imported directly from the downloaded directory:
 
 ```python
 from PIL import Image
 from quasar.inference.extended import ExtendedGenerator
 from quasar.inference.decoding import Sampling
 
-# Point to a verified local checkpoint; this documentation repo has no weights.
-generator = ExtendedGenerator("PATH_TO_LOCAL_QUASAR_CHECKPOINT", "cuda")
+# Root contains the exact research model shown in this card.
+generator = ExtendedGenerator(".", "cuda")
 sampling = Sampling(temperature=0.5, top_k=64, top_p=0.9, seed=20261021)
 pixels = generator.generate_image("make a picture of a cat", sampling, resolution=128)
 Image.fromarray(pixels).save("quasar_cat.png")
@@ -243,15 +262,30 @@ frames = generator.generate_video(
 )
 ```
 
-This is a usage example, not a claim of acceptable output for arbitrary prompts.
-GPU use must be exclusive with local training/evaluation workers. No inference
-server, model weights or codec binaries are uploaded by creating this card.
+These examples do not establish acceptable output for arbitrary prompts. GPU use
+must be exclusive with local training/evaluation workers. No inference server is
+deployed. About 1.96 GiB of model/package files include both 361 MB Quasar weight
+files and one shared copy of each frozen codec; training data and optimizer states
+are omitted. The image/code mappings and original trained weights are unchanged.
+
+The runtime is frozen from source commit
+`4afbdccdbd3b0967b161a68658cfd3b92c42636f`, before the separate architecture
+experiments. Two packaging-only edits resolve visual/audio codec paths relative
+to the downloaded checkpoint. The included pinned visual codec Python is loaded
+locally with `trust_remote_code=True`; review it as executable source. Inference
+does not fetch additional model code or weights. Source identities and file hashes
+are in [release.json](release.json) and [files.sha256.json](files.sha256.json).
+The packaged release passed CPU finite-forward/parameter/hash checks for both
+checkpoints and CPU encode/decode checks for both pinned codecs; see
+[release_validation.json](release_validation.json). This verification does not
+claim new semantic quality or replace the held-out results above.
 
 ## Licensing and provenance
 
-The local Quasar project does not currently supply a project/weight distribution
-license. This card does not assign one or grant a weights license. No permissive
-license badge is implied by using PyTorch or permissively licensed components.
+This private upload does not select a distribution license for Quasar's code or
+weights. No permissive license badge is implied by using PyTorch or permissively
+licensed components. Component notices remain separate and are included with
+their codec files; this upload does not make the repository public.
 
 The existing frozen Emu3 visual tokenizer is pinned at revision
 `c81f916ad371289e205310a7539255e8a9396488` under Apache-2.0. The local Descript DAC
@@ -261,6 +295,10 @@ Quasar weights license. Dataset, underlying page, photograph and per-object term
 remain separate; for example, the FineWeb database license does not replace
 underlying web-page rights.
 
+The visual component retains upstream copyright headers and Apache-2.0 terms in
+`codecs/vision/LICENSE`. The audio component retains the official MIT license and
+the pinned upstream weights statement under `codecs/audio`.
+
 ## Reproducibility and next work
 
 The completed R&D snapshot passed 146 CPU tests with 1 skip; the legacy integration

codecs/audio/OFFICIAL_LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023-present, Descript
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

codecs/audio/OFFICIAL_WEIGHTS_TERMS.md

@@ -0,0 +1,189 @@
+# Descript Audio Codec (.dac): High-Fidelity Audio Compression with Improved RVQGAN
+
+This repository contains training and inference scripts
+for the Descript Audio Codec (.dac), a high fidelity general
+neural audio codec, introduced in the paper titled **High-Fidelity Audio Compression with Improved RVQGAN**.
+
+![](https://static.arxiv.org/static/browse/0.3.4/images/icons/favicon-16x16.png) [arXiv Paper: High-Fidelity Audio Compression with Improved RVQGAN
+](http://arxiv.org/abs/2306.06546) <br>
+📈 [Demo Site](https://descript.notion.site/Descript-Audio-Codec-11389fce0ce2419891d6591a68f814d5)<br>
+⚙ [Model Weights](https://github.com/descriptinc/descript-audio-codec/releases/download/0.0.1/weights.pth)
+
+👉 With Descript Audio Codec, you can compress **44.1 KHz audio** into discrete codes at a **low 8 kbps bitrate**.  <br>
+🤌 That's approximately **90x compression** while maintaining exceptional fidelity and minimizing artifacts.  <br>
+💪 Our universal model works on all domains (speech, environment, music, etc.), making it widely applicable to generative modeling of all audio.  <br>
+👌 It can be used as a drop-in replacement for EnCodec for all audio language modeling applications (such as AudioLMs, MusicLMs, MusicGen, etc.) <br>
+
+<p align="center">
+<img src="./assets/comparsion_stats.png" alt="Comparison of compressions approaches. Our model achieves a higher compression factor compared to all baseline methods. Our model has a ~90x compression factor compared to 32x compression factor of EnCodec and 64x of SoundStream. Note that we operate at a target bitrate of 8 kbps, whereas EnCodec operates at 24 kbps and SoundStream at 6 kbps. We also operate at 44.1 kHz, whereas EnCodec operates at 48 kHz and SoundStream operates at 24 kHz." width=35%></p>
+
+
+## Usage
+
+### Installation
+```
+pip install descript-audio-codec
+```
+OR
+
+```
+pip install git+https://github.com/descriptinc/descript-audio-codec
+```
+
+### Weights
+Weights are released as part of this repo under MIT license.
+We release weights for models that can natively support 16 kHz, 24kHz, and 44.1kHz sampling rates.
+Weights are automatically downloaded when you first run `encode` or `decode` command. You can cache them using one of the following commands
+```bash
+python3 -m dac download # downloads the default 44kHz variant
+python3 -m dac download --model_type 44khz # downloads the 44kHz variant
+python3 -m dac download --model_type 24khz # downloads the 24kHz variant
+python3 -m dac download --model_type 16khz # downloads the 16kHz variant
+```
+We provide a Dockerfile that installs all required dependencies for encoding and decoding. The build process caches the default model weights inside the image. This allows the image to be used without an internet connection. [Please refer to instructions below.](#docker-image)
+
+
+### Compress audio
+```
+python3 -m dac encode /path/to/input --output /path/to/output/codes
+```
+
+This command will create `.dac` files with the same name as the input files.
+It will also preserve the directory structure relative to input root and
+re-create it in the output directory. Please use `python -m dac encode --help`
+for more options.
+
+### Reconstruct audio from compressed codes
+```
+python3 -m dac decode /path/to/output/codes --output /path/to/reconstructed_input
+```
+
+This command will create `.wav` files with the same name as the input files.
+It will also preserve the directory structure relative to input root and
+re-create it in the output directory. Please use `python -m dac decode --help`
+for more options.
+
+### Programmatic Usage
+```py
+import dac
+from audiotools import AudioSignal
+
+# Download a model
+model_path = dac.utils.download(model_type="44khz")
+model = dac.DAC.load(model_path)
+
+model.to('cuda')
+
+# Load audio signal file
+signal = AudioSignal('input.wav')
+
+# Encode audio signal as one long file
+# (may run out of GPU memory on long files)
+signal.to(model.device)
+
+x = model.preprocess(signal.audio_data, signal.sample_rate)
+z, codes, latents, _, _ = model.encode(x)
+
+# Decode audio signal
+y = model.decode(z)
+
+# Alternatively, use the `compress` and `decompress` functions
+# to compress long files.
+
+signal = signal.cpu()
+x = model.compress(signal)
+
+# Save and load to and from disk
+x.save("compressed.dac")
+x = dac.DACFile.load("compressed.dac")
+
+# Decompress it back to an AudioSignal
+y = model.decompress(x)
+
+# Write to file
+y.write('output.wav')
+```
+
+### Docker image
+We provide a dockerfile to build a docker image with all the necessary
+dependencies.
+1. Building the image.
+    ```
+    docker build -t dac .
+    ```
+2. Using the image.
+
+    Usage on CPU:
+    ```
+    docker run dac <command>
+    ```
+
+    Usage on GPU:
+    ```
+    docker run --gpus=all dac <command>
+    ```
+
+    `<command>` can be one of the compression and reconstruction commands listed
+    above. For example, if you want to run compression,
+
+    ```
+    docker run --gpus=all dac python3 -m dac encode ...
+    ```
+
+
+## Training
+The baseline model configuration can be trained using the following commands.
+
+### Pre-requisites
+Please install the correct dependencies
+```
+pip install -e ".[dev]"
+```
+
+## Environment setup
+
+We have provided a Dockerfile and docker compose setup that makes running experiments easy.
+
+To build the docker image do:
+
+```
+docker compose build
+```
+
+Then, to launch a container, do:
+
+```
+docker compose run -p 8888:8888 -p 6006:6006 dev
+```
+
+The port arguments (`-p`) are optional, but useful if you want to launch a Jupyter and Tensorboard instances within the container. The
+default password for Jupyter is `password`, and the current directory
+is mounted to `/u/home/src`, which also becomes the working directory.
+
+Then, run your training command.
+
+
+### Single GPU training
+```
+export CUDA_VISIBLE_DEVICES=0
+python scripts/train.py --args.load conf/ablations/baseline.yml --save_path runs/baseline/
+```
+
+### Multi GPU training
+```
+export CUDA_VISIBLE_DEVICES=0,1
+torchrun --nproc_per_node gpu scripts/train.py --args.load conf/ablations/baseline.yml --save_path runs/baseline/
+```
+
+## Testing
+We provide two test scripts to test CLI + training functionality. Please
+make sure that the trainig pre-requisites are satisfied before launching these
+tests. To launch these tests please run
+```
+python -m pytest tests
+```
+
+## Results
+
+<p align="left">
+<img src="./assets/objective_comparisons.png" width=75%></p>

codecs/audio/README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+pipeline_tag: feature-extraction
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

codecs/audio/config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "DacModel"
+  ],
+  "codebook_dim": 8,
+  "codebook_loss_weight": 1.0,
+  "codebook_size": 1024,
+  "commitment_loss_weight": 0.25,
+  "decoder_hidden_size": 1536,
+  "downsampling_ratios": [
+    2,
+    4,
+    5,
+    8
+  ],
+  "encoder_hidden_size": 64,
+  "hidden_size": 1024,
+  "hop_length": 512,
+  "model_type": "dac",
+  "n_codebooks": 32,
+  "quantizer_dropout": 0.0,
+  "sampling_rate": 24000,
+  "torch_dtype": "float32",
+  "transformers_version": "4.42.0.dev0",
+  "upsampling_ratios": [
+    8,
+    5,
+    4,
+    2
+  ]
+}

codecs/audio/model.safetensors

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

codecs/vision/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

codecs/vision/README.md

@@ -0,0 +1,92 @@
+---
+license: apache-2.0
+library_name: transformers
+---
+
+<div align='center'>
+<h1>Emu3: Next-Token Prediction is All You Need</h1h1>
+<h3></h3>
+
+[Emu3 Team, BAAI](https://www.baai.ac.cn/english.html)
+
+| [Project Page](https://emu.baai.ac.cn) | [Paper](https://huggingface.co/papers/2409.18869) | [🤗HF Models](https://huggingface.co/collections/BAAI/emu3-66f4e64f70850ff358a2e60f) | [github](https://github.com/baaivision/Emu3) | [Demo](https://huggingface.co/spaces/BAAI/Emu3) |
+
+
+</div>
+
+<div align='center'>
+<img src="https://github.com/baaivision/Emu3/blob/main/assets/arch.png?raw=True" class="interpolation-image" alt="arch." height="80%" width="70%" />
+</div>
+
+We introduce **Emu3**, a new suite of state-of-the-art multimodal models trained solely with **<i>next-token prediction</i>**! By tokenizing images, text, and videos into a discrete space, we train a single transformer from scratch on a mixture of multimodal sequences.
+
+### Emu3 excels in both generation and perception
+**Emu3** outperforms several well-established task-specific models in both generation and perception tasks, surpassing flagship open models such as SDXL, LLaVA-1.6 and OpenSora-1.2, while eliminating the need for diffusion or compositional architectures.
+
+<div align='center'>
+<img src="https://github.com/baaivision/Emu3/blob/main//assets/comparison.png?raw=True" class="interpolation-image" alt="comparison." height="80%" width="80%" />
+</div>
+
+### Highlights
+
+- **Emu3** is capable of generating high-quality images following the text input, by simply predicting the next vision token. The model naturally supports flexible resolutions and styles.
+- **Emu3** shows strong vision-language understanding capabilities to see the physical world and provides coherent text responses. Notably, this capability is achieved without depending on a CLIP and a pretrained LLM.
+- **Emu3** simply generates a video causally by predicting the next token in a video sequence, unlike the video diffusion model as in Sora. With a video in context, Emu3 can also naturally extend the video and predict what will happen next. 
+
+### Quickstart for Autoencoding
+```python
+import os
+import os.path as osp
+
+from PIL import Image
+import torch
+from transformers import AutoModel, AutoImageProcessor
+
+MODEL_HUB = "BAAI/Emu3-VisionTokenizer"
+
+model = AutoModel.from_pretrained(MODEL_HUB, trust_remote_code=True).eval().cuda()
+processor = AutoImageProcessor.from_pretrained(MODEL_HUB, trust_remote_code=True)
+
+# TODO: you need to modify the path here
+VIDEO_FRAMES_PATH = "YOUR_VIDEO_FRAMES_PATH"
+
+video = os.listdir(VIDEO_FRAMES_PATH)
+video.sort()
+video = [Image.open(osp.join(VIDEO_FRAMES_PATH, v)) for v in video]
+
+images = processor(video, return_tensors="pt")["pixel_values"]
+images = images.unsqueeze(0).cuda()
+
+# image autoencode
+image = images[:, 0]
+print(image.shape)
+with torch.no_grad():
+    # encode
+    codes = model.encode(image)
+    # decode
+    recon = model.decode(codes)
+
+recon = recon.view(-1, *recon.shape[2:])
+recon_image = processor.postprocess(recon)["pixel_values"][0]
+recon_image.save("recon_image.png")
+
+# video autoencode
+images = images.view(
+    -1,
+    model.config.temporal_downsample_factor,
+    *images.shape[2:],
+)
+
+print(images.shape)
+with torch.no_grad():
+    # encode
+    codes = model.encode(images)
+    # decode
+    recon = model.decode(codes)
+
+recon = recon.view(-1, *recon.shape[2:])
+recon_images = processor.postprocess(recon)["pixel_values"]
+for idx, im in enumerate(recon_images):
+    im.save(f"recon_video_{idx}.png")
+
+```

codecs/vision/config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "Emu3VisionVQModel"
+  ],
+  "attn_resolutions": [
+    3
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_emu3visionvq.Emu3VisionVQConfig",
+    "AutoModel": "modeling_emu3visionvq.Emu3VisionVQModel"
+  },
+  "ch": 256,
+  "ch_mult": [
+    1,
+    2,
+    2,
+    4
+  ],
+  "codebook_size": 32768,
+  "double_z": false,
+  "dropout": 0.0,
+  "embed_dim": 4,
+  "in_channels": 3,
+  "model_type": "Emu3VisionVQ",
+  "num_res_blocks": 2,
+  "out_channels": 3,
+  "temporal_downsample_factor": 4,
+  "torch_dtype": "float32",
+  "transformers_version": "4.44.0",
+  "z_channels": 4
+}

codecs/vision/configuration_emu3visionvq.py

@@ -0,0 +1,106 @@
+# coding=utf-8
+# Copyright 2024 The Emu team, BAAI and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+""" Emu3VisionVQ model configuration """
+
+from typing import List
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class Emu3VisionVQConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Emu3VisionVQ`]. It is used to instantiate an video movq 
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a configuration to the VQ model presented in Emu3 paper.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        codebook_size (`int`, *optional*, defaults to 32768):
+            Codebook size of the VQ model.
+        embed_dim (`int`, *optional*, defaults to 4):
+            Dimension of the quantized vector in codebook.
+        z_channels (`int`, *optional*, defaults to 4):
+            Dimension of the output channel of encoder and the input channel of decoder
+        double_z (`bool`, *optional*, defaults to False):
+            Whether double the output dim of the encoder.
+        in_channels (`int`, *optional*, defaults to 3):
+            Input channel of encoder.
+        out_channels (`int`, *optional*, defaults to 3):
+            Output channel of decoder.
+        temporal_downsample_factor (`int`, *optional*, defaults to 4):
+            Temporal downsample factor.
+        ch (`int`, *optional*, defaults to 256):
+            Basic channel number of the intermediate blocks.
+        ch_mult (`List[int]`, *optional*, defaults to `[1, 2, 2, 4]`):
+            Channel scaling factor of the intermediate blocks.
+        num_res_blocks (`int`, *optional*, defaults to 2):
+            Residual block number in each stage.
+        attn_resolutions (`List[int]`, *optional*, defaults to 3):
+            Stage indices to apply attention.
+        dropout (`float`, *optional*, defaults to 0.0):
+            Dropout probability.
+
+    ```python
+    >>> from transformers import Emu3VisionVQ, Emu3VisionVQConfig
+
+    >>> # Initializing a video VQ model of Emu3 configuration
+    >>> configuration = Emu3VisionVQConfig()
+
+    >>> # Initializing a model from the Emu3 VQ model style configuration
+    >>> model = Emu3VisionVQModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "Emu3VisionVQ"
+
+    def __init__(
+        self,
+        codebook_size: int = 32768,
+        embed_dim: int = 4,
+        z_channels: int = 4,
+        double_z: bool = False,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        temporal_downsample_factor: int = 4,
+        ch: int = 256,
+        ch_mult: List[int] = [1, 2, 2, 4],
+        num_res_blocks: int = 2,
+        attn_resolutions: List[int] = [3],
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.codebook_size = codebook_size
+        self.embed_dim = embed_dim
+        self.z_channels = z_channels
+        self.double_z = double_z
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.temporal_downsample_factor = temporal_downsample_factor
+        self.ch = ch
+        self.ch_mult = ch_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_resolutions = attn_resolutions
+        self.dropout = dropout

codecs/vision/image_processing_emu3visionvq.py

@@ -0,0 +1,442 @@
+# coding=utf-8
+# Copyright 2024 The Emu team, BAAI and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for Emu3VisionVQ."""
+
+
+import math
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers.image_transforms import (
+    convert_to_rgb,
+    resize,
+    to_channel_dimension_format,
+)
+from transformers.image_utils import (
+    IMAGENET_STANDARD_MEAN,
+    IMAGENET_STANDARD_STD,
+    ChannelDimension,
+    ImageInput,
+    PILImageResampling,
+    get_image_size,
+    infer_channel_dimension_format,
+    is_scaled_image,
+    make_list_of_images,
+    to_numpy_array,
+    valid_images,
+    validate_preprocess_arguments,
+)
+from transformers.utils import TensorType, is_vision_available, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+if is_vision_available():
+    from PIL import Image
+
+
+def smart_resize(
+    height: int, width: int, factor: int = 8, min_pixels: int = 512 * 512, max_pixels: int = 1024 * 1024
+):
+    """Rescales the image so that the following conditions are met:
+
+    1. Both dimensions (height and width) are divisible by 'factor'.
+
+    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+    3. The aspect ratio of the image is maintained as closely as possible.
+
+    """
+    if height < factor or width < factor:
+        raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")
+    elif max(height, width) / min(height, width) > 5:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than 5, got {max(height, width) / min(height, width)}"
+        )
+
+    h_bar = round(height / factor) * factor
+    w_bar = round(width / factor) * factor
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = math.floor(height / beta / factor) * factor
+        w_bar = math.floor(width / beta / factor) * factor
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = math.ceil(height * beta / factor) * factor
+        w_bar = math.ceil(width * beta / factor) * factor
+
+    return h_bar, w_bar
+
+
+class Emu3VisionVQImageProcessor(BaseImageProcessor):
+    r"""
+    Constructs a Emu3VisionVQ image processor that dynamically resizes images based on the original images.
+
+    Args:
+        do_resize (`bool`, *optional*, defaults to `True`):
+            Whether to resize the image's (height, width) dimensions.
+        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
+            Resampling filter to use when resizing the image.
+        do_rescale (`bool`, *optional*, defaults to `True`):
+            Whether to rescale the image by the specified scale `rescale_factor`.
+        rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
+            Scale factor to use if rescaling the image.
+        do_normalize (`bool`, *optional*, defaults to `True`):
+            Whether to normalize the image.
+        image_mean (`float` or `List[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
+            Mean to use if normalizing the image. This is a float or list of floats for each channel in the image.
+        image_std (`float` or `List[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
+            Standard deviation to use if normalizing the image. This is a float or list of floats for each channel in the image.
+        do_convert_rgb (`bool`, *optional*, defaults to `True`):
+            Whether to convert the image to RGB.
+        min_pixels (`int`, *optional*, defaults to `512 * 512`):
+            The min pixels of the image to resize the image.
+        max_pixels (`int`, *optional*, defaults to `1024 * 1024`):
+            The max pixels of the image to resize the image.
+        spatial_factor (`int`, *optional*, defautls to 8):
+            The spatial downsample factor the image will be downsampled in feature extracting phase
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        resample: PILImageResampling = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: Union[int, float] = 1 / 255,
+        do_normalize: bool = True,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = True,
+        min_pixels: int = 512 * 512,
+        max_pixels: int = 1024 * 1024,
+        spatial_factor: int = 8,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.do_resize = do_resize
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
+        self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
+        self.min_pixels = min_pixels
+        self.max_pixels = max_pixels
+        self.size = {"min_pixels": min_pixels, "max_pixels": max_pixels}
+        self.do_convert_rgb = do_convert_rgb
+        self.spatial_factor = spatial_factor
+
+    def _preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        resample: PILImageResampling = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        spatial_factor: Optional[int] = None,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        output_data_format: Optional[Union[str, ChannelDimension]] = ChannelDimension.FIRST,
+    ):
+        """
+        Preprocess an image or batch of images. Copy of the `preprocess` method from `CLIPImageProcessor`.
+
+        Args:
+            images (`ImageInput`):
+                Image or batch of images to preprocess. Expects pixel values ranging from 0 to 255. If pixel values range from 0 to 1, set `do_rescale=False`.
+            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+                Whether to resize the image.
+            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
+                Resampling filter to use if resizing the image. This can be one of the `PILImageResampling` enums.
+            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+                Whether to rescale the image.
+            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+                Scale factor to use if rescaling the image.
+            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+                Whether to normalize the image.
+            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+                Mean to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.
+            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+                Standard deviation to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.
+            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+                Whether to convert the image to RGB.
+            spatial_factor (`int`, *optional*, defaults to `self.spatial_factor`):
+                The spatial downsample factor the image will be downsampled in feature extracting phase
+            input_data_format (`ChannelDimension` or `str`, *optional*):
+                The channel dimension format for the input image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.   - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+            output_data_format (`ChannelDimension`, *optional*, defaults to `ChannelDimension.FIRST`):
+                The channel dimension format for the output image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - Unset: Use the channel dimension format of the input image.
+        """
+        spatial_factor = spatial_factor if spatial_factor is not None else self.spatial_factor
+
+        images = make_list_of_images(images)
+        if do_convert_rgb:
+            images = [convert_to_rgb(image) for image in images]
+
+        # All transformations expect numpy arrays.
+        images = [to_numpy_array(image) for image in images]
+
+        if is_scaled_image(images[0]) and do_rescale:
+            logger.warning_once(
+                "It looks like you are trying to rescale already rescaled images. If the input"
+                "pixel_values.append()images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
+            )
+
+        if input_data_format is None:
+            # We assume that all images have the same channel dimension format.
+            input_data_format = infer_channel_dimension_format(images[0])
+
+        height, width = get_image_size(images[0], channel_dim=input_data_format)
+        resized_height, resized_width = height, width
+        processed_images = []
+        for image in images:
+            if do_resize:
+                resized_height, resized_width = smart_resize(
+                    height,
+                    width,
+                    factor=spatial_factor,
+                    min_pixels=self.min_pixels,
+                    max_pixels=self.max_pixels,
+                )
+                image = resize(
+                    image, size=(resized_height, resized_width), resample=resample, input_data_format=input_data_format
+                )
+
+            if do_rescale:
+                image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format)
+
+            if do_normalize:
+                image = self.normalize(
+                    image=image, mean=image_mean, std=image_std, input_data_format=input_data_format
+                )
+
+            image = to_channel_dimension_format(image, output_data_format, input_channel_dim=input_data_format)
+            processed_images.append(image)
+
+        image = np.array(processed_images)
+        return image
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        resample: PILImageResampling = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        spatial_factor: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+        output_data_format: Optional[Union[str, ChannelDimension]] = ChannelDimension.FIRST,
+    ):
+        """
+        Args:
+            images (`ImageInput`):
+                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
+                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
+            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+                Whether to resize the image.
+            resample (`int`, *optional*, defaults to `self.resample`):
+                Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only
+                has an effect if `do_resize` is set to `True`.
+            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+                Whether to rescale the image.
+            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+                Rescale factor to rescale the image by if `do_rescale` is set to `True`.
+            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+                Whether to normalize the image.
+            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+                Whether to convert the image to RGB.
+            spatial_factor (`int`, *optional*, defaults to `self.spatial_factor`):
+                The spatial downsample factor the image will be downsampled in feature extracting phase
+            return_tensors (`str` or `TensorType`, *optional*):
+                The type of tensors to return. Can be one of:
+                - Unset: Return a list of `np.ndarray`.
+                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+            input_data_format (`ChannelDimension` or `str`, *optional*):
+                The channel dimension format for the input image. If unset, the channel dimension format is inferred
+                from the input image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+            output_data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
+                The channel dimension format for the output image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - Unset: Use the channel dimension format of the input image.
+        """
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        spatial_factor = spatial_factor if spatial_factor is not None else self.spatial_factor
+
+        images = make_list_of_images(images)
+        if images is None or not valid_images(images):
+            raise ValueError(
+                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
+                "torch.Tensor, tf.Tensor or jax.ndarray."
+            )
+
+        validate_preprocess_arguments(
+            rescale_factor=rescale_factor,
+            do_normalize=do_normalize,
+            image_mean=image_mean,
+            image_std=image_std,
+            do_resize=do_resize,
+            size=self.size,
+            resample=resample,
+        )
+
+        pixel_values = []
+        for image in images:
+            norm_image = self._preprocess(
+                image,
+                do_resize=do_resize,
+                resample=resample,
+                do_rescale=do_rescale,
+                rescale_factor=rescale_factor,
+                do_normalize=do_normalize,
+                image_mean=image_mean,
+                image_std=image_std,
+                do_convert_rgb=do_convert_rgb,
+                spatial_factor=spatial_factor,
+                input_data_format=input_data_format,
+                output_data_format=output_data_format,
+            )
+            pixel_values.extend(norm_image)
+        pixel_values = np.array(pixel_values)
+        data = {"pixel_values": pixel_values}
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    def postprocess(
+        self,
+        images: ImageInput,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        return_tensors: str | TensorType = "PIL.Image.Image",
+        input_data_format: Optional[Union[str, ChannelDimension]] = None,
+    ):
+        """
+        Postprocess an image or batch of images tensor. Postprocess is the reverse process of preprocess.
+        The parameters should be same as in preprocess.
+
+        Args:
+            images (`ImageInput`):
+                Image to postprocess. Expects a single or batch of images with pixel values ranging from -1 to 1.
+            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+                Whether to rescale the image.
+            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+                Rescale factor to rescale the image by if `do_rescale` is set to `True`.
+            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+                Whether to normalize the image.
+            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+            return_tensors (`str` or `TensorType`, *optional*):
+                The type of tensors to return. Can be one of:
+                - Unset: Return a list of `np.ndarray`.
+                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+            input_data_format (`ChannelDimension` or `str`, *optional*):
+                The channel dimension format for the input image. If unset, the channel dimension format is inferred
+                from the input image. Can be one of:
+                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+        """
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+        rescale_factor = 1 / rescale_factor
+
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        image_mean, image_std = self.inverse_meanstd(image_mean, image_std)
+
+        images = make_list_of_images(images)
+        if isinstance(images[0], Image.Image):
+            return images if len(images) > 1 else images[0]
+
+        if input_data_format is None:
+            # We assume that all images have the same channel dimension format.
+            input_data_format = infer_channel_dimension_format(images[0])
+
+        pixel_values = []
+        for image in images:
+            image = to_numpy_array(image)
+            if do_normalize:
+                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
+
+            if do_rescale:
+                image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format)
+                image = image.clip(0, 255).astype(np.uint8)
+
+            if do_normalize and do_rescale and return_tensors == "PIL.Image.Image":
+                image = to_channel_dimension_format(image, ChannelDimension.LAST, input_channel_dim=input_data_format)
+                pixel_values.append(Image.fromarray(image))
+            else:
+                pixel_values.extend(image)
+
+        data = {"pixel_values": pixel_values}
+        return_tensors = return_tensors if return_tensors != "PIL.Image.Image" else None
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    def inverse_meanstd(self, image_mean, image_std):
+        image_mean = self.to_tuple(image_mean)
+        image_std = self.to_tuple(image_std)
+
+        rev_image_mean = tuple(-m / s for m, s in zip(image_mean, image_std))
+        rev_image_std = tuple(1 / s for s in image_std)
+
+        return rev_image_mean, rev_image_std
+
+    def to_tuple(self, value, dim=3):
+        if isinstance(value, int | float):
+            return (value,) * dim
+
+        return tuple(value)

codecs/vision/model.safetensors

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

codecs/vision/modeling_emu3visionvq.py

@@ -0,0 +1,822 @@
+# coding=utf-8
+# Copyright 2024 The Emu team, BAAI and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+""" Emu3VisionVQ model """
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from transformers.modeling_utils import PreTrainedModel
+
+from .configuration_emu3visionvq import Emu3VisionVQConfig
+
+
+class Emu3VisionVQActivation(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, x: torch.Tensor):
+        return x * torch.sigmoid(x)
+
+
+class Emu3VisionVQUpsample(nn.Module):
+
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+    def forward(self, x: torch.Tensor):
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class Emu3VisionVQDownsample(nn.Module):
+
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=3,
+            stride=2,
+            padding=0,
+        )
+
+    def forward(self, x: torch.Tensor):
+        pad = (0, 1, 0, 1)
+        x = F.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Emu3VisionVQCausalConv3d(nn.Module):
+
+    def __init__(
+        self,
+        in_channel: int,
+        out_channel: int,
+        kernel_size: Union[int, Tuple[int, ...]] = (3, 1, 1),
+        stride: Union[int, Tuple[int, ...]] = (1, 1, 1),
+    ):
+        super().__init__()
+
+        if isinstance(kernel_size, int):
+            kernel_size = (kernel_size,) * 3
+        if isinstance(stride, int):
+            stride = (stride,) * 3
+        
+        hw_pad = [k - s for k, s in zip(kernel_size[1:], stride[1:])]
+        self.padding = tuple()
+        for p in hw_pad[::-1]:
+            self.padding += (p // 2 + p % 2, p // 2)
+        self.padding += (2, 0)
+
+        self.conv = nn.Conv3d(
+            in_channel,
+            out_channel,
+            kernel_size,
+            stride=stride,
+        )
+
+    def forward(self, x: torch.Tensor):
+        x = F.pad(x, self.padding)
+        x = self.conv(x)
+        return x
+
+
+class Emu3VisionVQResnetTemporalBlock(nn.Module):
+
+    def __init__(
+        self, 
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        stride = (1, 1, 1)
+        kernel_size = (3, 3, 3)
+
+        self.norm1 = nn.BatchNorm3d(in_channels)
+        self.conv1 = Emu3VisionVQCausalConv3d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+        )
+        self.norm2 = nn.BatchNorm3d(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = Emu3VisionVQCausalConv3d(
+            out_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+        )
+        self.act = Emu3VisionVQActivation()
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = Emu3VisionVQCausalConv3d(
+                    in_channels,
+                    out_channels,
+                    kernel_size=kernel_size,
+                    stride=stride,
+                )
+            else:
+                self.nin_shortcut = nn.Conv3d(
+                    in_channels,
+                    out_channels,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                )
+
+    def forward(self, x: torch.Tensor):
+        h = self.norm1(x)
+        h = self.act(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Emu3VisionVQSpatialNorm(nn.Module):
+
+    def __init__(
+        self,
+        f_channels: int,
+        zq_channels: int,
+        norm_layer: nn.Module = nn.GroupNorm,
+        add_conv: bool = False,
+        num_groups: int = 32,
+        eps: float = 1e-6,
+        affine: bool = True,
+    ):
+        super().__init__()
+        self.norm_layer = norm_layer(
+            num_channels=f_channels,
+            num_groups=num_groups,
+            eps=eps,
+            affine=affine,
+        )
+
+        self.add_conv = add_conv
+        if self.add_conv:
+            self.conv = nn.Conv2d(
+                zq_channels,
+                zq_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+            )
+
+        self.conv_y = nn.Conv2d(
+            zq_channels,
+            f_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.conv_b = nn.Conv2d(
+            zq_channels,
+            f_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+
+    def forward(self, x: torch.Tensor, zq: torch.Tensor):
+        zq = F.interpolate(zq, size=x.shape[-2:], mode="nearest")
+
+        if self.add_conv:
+            zq = self.conv(zq)
+
+        x = self.norm_layer(x)
+        x = x * self.conv_y(zq) + self.conv_b(zq)
+        return x
+
+
+class Emu3VisionVQResnetBlock(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+        zq_ch: Optional[int] = None,
+        add_conv: bool = False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+        self.zq_ch = zq_ch
+
+        if zq_ch is None:
+            norm_kwargs = dict(num_groups=32, eps=1e-6, affine=True)
+            self.norm1 = nn.GroupNorm(num_channels=in_channels, **norm_kwargs)
+            self.norm2 = nn.GroupNorm(num_channels=out_channels, **norm_kwargs)
+        else:
+            self.norm1 = Emu3VisionVQSpatialNorm(in_channels, zq_ch, add_conv=add_conv)
+            self.norm2 = Emu3VisionVQSpatialNorm(out_channels, zq_ch, add_conv=add_conv)
+
+        self.conv1 = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(
+            out_channels,
+            out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        self.act = Emu3VisionVQActivation()
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = nn.Conv2d(
+                    in_channels,
+                    out_channels,
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                )
+            else:
+                self.nin_shortcut = nn.Conv2d(
+                    in_channels,
+                    out_channels,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                )
+
+    def forward(self, x: torch.Tensor, zq: Optional[torch.Tensor] = None):
+        norm_args = tuple() if self.zq_ch is None else (zq, )
+
+        h = self.norm1(x, *norm_args)
+        h = self.act(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h, *norm_args)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Emu3VisionVQAttnBlock(nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        zq_ch: Optional[int] = None,
+        add_conv: bool = False
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.zq_ch = zq_ch
+
+        if zq_ch is None:
+            norm_kwargs = dict(num_groups=32, eps=1e-6, affine=True)
+            self.norm = nn.GroupNorm(num_channels=in_channels, **norm_kwargs)
+        else:
+            self.norm = Emu3VisionVQSpatialNorm(in_channels, zq_ch, add_conv=add_conv)
+
+        self.q = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.k = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.v = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.proj_out = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+
+    def forward(self, x: torch.Tensor, zq: Optional[torch.Tensor] = None):
+        norm_args = tuple() if self.zq_ch is None else (zq, )
+
+        nx = self.norm(x, *norm_args)
+        q = self.q(nx)
+        k = self.k(nx)
+        v = self.v(nx)
+
+        # compute attention
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w)
+        k = k.reshape(b, c, h * w)
+        score = torch.bmm(q.permute(0, 2, 1), k)
+        score = score / (c ** 0.5)
+        score = F.softmax(score, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h * w)
+        v = torch.bmm(v, score.permute(0, 2, 1))
+        v = v.reshape(b, c, h, w)
+
+        v = self.proj_out(v)
+
+        return x + v
+
+
+class Emu3VisionVQTemporalUpsample(nn.Module):
+
+    def __init__(
+        self,
+        in_channel: int,
+        out_channel: int,
+        kernel_size: Tuple[int, ...] = (3, 3, 3),
+        stride: Tuple[int, ...] = (1, 1, 1)
+    ):
+        super().__init__()
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+        self.conv = Emu3VisionVQCausalConv3d(
+            in_channel,
+            out_channel,
+            kernel_size,
+            stride=stride,
+        )
+        
+    def forward(self, x: torch.Tensor):
+        b, c, t, h, w = x.shape
+        x = x.permute(0, 1, 3, 4, 2).contiguous().view(b, -1, t)
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = x.view(b, c, h, w, -1).permute(0, 1, 4, 2, 3).contiguous()
+        x = self.conv(x)
+        return x
+
+
+class Emu3VisionVQTemporalDownsample(nn.Module):
+
+    def __init__(
+        self,
+        in_channel: int,
+        out_channel: int,
+        kernel_size: Tuple[int, ...] = (4, 3, 3),
+        stride: Tuple[int, ...] = (2, 1, 1),
+    ):
+        super().__init__()
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+        self.kernel_size = kernel_size
+
+        self.conv = Emu3VisionVQCausalConv3d(
+            in_channel,
+            out_channel,
+            kernel_size=kernel_size,
+            stride=stride,
+        )
+        
+    def forward(self, x: torch.Tensor):
+        x = self.conv(x)
+        return x
+
+
+class Emu3VisionVQVectorQuantizer(nn.Module):
+
+    def __init__(self, config: Emu3VisionVQConfig):
+        super().__init__()
+        self.embedding = nn.Embedding(config.codebook_size, config.embed_dim)
+        self.embedding.weight.data.uniform_(-1.0 / config.codebook_size, 1.0 / config.codebook_size)
+
+    def forward(self, x: torch.Tensor):
+        # b t c h w -> b t h w c
+        b, t, c, h, w = x.shape
+        x = x.permute(0, 1, 3, 4, 2).contiguous()
+        x_flattened = x.view(-1, c)
+
+        codebook = self.embedding.weight
+
+        d = torch.sum(x_flattened ** 2, dim=1, keepdim=True) + \
+            torch.sum(codebook ** 2, dim=1) - 2 * \
+            torch.einsum('bd,dn->bn', x_flattened, codebook.permute(1, 0))
+
+        indices = torch.argmin(d, dim=1)
+        indices = indices.view(b, t, h, w)
+        return indices
+
+
+class Emu3VisionVQEncoder(nn.Module):
+
+    def __init__(self, config: Emu3VisionVQConfig):
+        super().__init__()
+        self.ch = config.ch
+        self.num_resolutions = len(config.ch_mult)
+        self.num_res_blocks = config.num_res_blocks
+        self.in_channels = config.in_channels
+
+        # downsampling
+        self.conv_in = nn.Conv2d(
+            self.in_channels,
+            self.ch,
+            kernel_size=3,
+            stride=1,
+            padding=1
+        )
+
+        in_ch_mult = (1,) + tuple(config.ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = config.ch * in_ch_mult[i_level]
+            block_out = config.ch * config.ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(
+                    Emu3VisionVQResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        dropout=config.dropout,
+                    )
+                )
+                block_in = block_out
+                if i_level in config.attn_resolutions:
+                    attn.append(Emu3VisionVQAttnBlock(block_in))
+
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Emu3VisionVQDownsample(block_in)
+
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = Emu3VisionVQResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=config.dropout,
+        )
+        self.mid.attn_1 = Emu3VisionVQAttnBlock(block_in)
+        self.mid.block_2 = Emu3VisionVQResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=config.dropout,
+        )
+
+        # end
+        self.norm_out = nn.GroupNorm(num_channels=block_in, num_groups=32, eps=1e-6, affine=True)
+
+        out_z_channels = 2 * config.z_channels if config.double_z else config.z_channels
+        self.conv_out = nn.Conv2d(
+            block_in,
+            out_z_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+        
+        temporal_down_blocks = int(math.log2(config.temporal_downsample_factor))
+        self.time_conv = nn.ModuleList()
+
+        for i in range(temporal_down_blocks):
+            conv = Emu3VisionVQTemporalDownsample(out_z_channels, out_z_channels)
+            self.time_conv.append(conv)
+
+        self.time_res_stack = nn.Sequential(*[
+            Emu3VisionVQResnetTemporalBlock(
+                in_channels=out_z_channels,
+                out_channels=out_z_channels,
+                dropout=config.dropout,
+            ) for _ in range(self.num_res_blocks)
+        ])
+
+        self.act = Emu3VisionVQActivation()
+
+    def forward(self, x: torch.Tensor):
+        t = x.shape[1]
+        x = x.reshape(-1, *x.shape[2:])
+
+        # downsampling
+        h = self.conv_in(x)
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](h)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+
+            if i_level != self.num_resolutions - 1:
+                h = self.down[i_level].downsample(h)
+
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        
+        # end
+        h = self.norm_out(h)
+        h = self.act(h)
+
+        h = self.conv_out(h)
+
+        h = h.reshape(-1, t, *h.shape[1:])
+        h = h.permute(0, 2, 1, 3, 4)
+
+        for conv in self.time_conv:
+            h = self.act(conv(h))
+
+        h = self.time_res_stack(h)
+        h = h.permute(0, 2, 1, 3, 4)
+
+        return h
+
+
+class Emu3VisionVQDecoder(nn.Module):
+
+    def __init__(self, config: Emu3VisionVQConfig):
+        super().__init__()
+        self.ch = config.ch
+        self.num_resolutions = len(config.ch_mult)
+        self.num_res_blocks = config.num_res_blocks
+
+        in_ch_mult = (1,) + tuple(config.ch_mult)
+        zq_ch = config.embed_dim
+
+        block_in = config.ch * config.ch_mult[-1]
+        self.time_res_stack = nn.Sequential(*[
+            Emu3VisionVQResnetTemporalBlock(
+                in_channels=config.z_channels,
+                out_channels=config.z_channels,
+                dropout=config.dropout,
+            ) for _ in range(config.num_res_blocks)
+        ])
+
+        tempo_upsample_block_num = int(math.log2(config.temporal_downsample_factor))
+        self.time_conv = nn.ModuleList()
+        for i in range(tempo_upsample_block_num):
+            conv = Emu3VisionVQTemporalUpsample(config.z_channels, config.z_channels)
+            self.time_conv.append(conv)
+            
+        self.conv_in = nn.Conv2d(
+            config.z_channels,
+            block_in,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = Emu3VisionVQResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=config.dropout,
+            zq_ch=zq_ch,
+        )
+        self.mid.attn_1 = Emu3VisionVQAttnBlock(block_in, zq_ch)
+        self.mid.block_2 = Emu3VisionVQResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            dropout=config.dropout,
+            zq_ch=zq_ch,
+        )
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = config.ch * config.ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                block.append(
+                    Emu3VisionVQResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        dropout=config.dropout,
+                        zq_ch=zq_ch,
+                    )
+                )
+                block_in = block_out
+                if i_level in config.attn_resolutions:
+                    attn.append(Emu3VisionVQAttnBlock(block_in, zq_ch))
+
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Emu3VisionVQUpsample(block_in)
+
+            self.up.insert(0, up)
+
+        self.act = Emu3VisionVQActivation()
+
+        self.norm_out = Emu3VisionVQSpatialNorm(block_in, zq_ch)
+        self.conv_out = nn.Conv2d(
+            block_in,
+            config.out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+    def forward(self, z: torch.Tensor, zq: torch.Tensor):
+        z_zq = torch.cat((z, zq), dim=0)
+        z_zq = z_zq.permute(0, 2, 1, 3, 4)
+        z_zq = self.time_res_stack(z_zq)
+
+        for conv in self.time_conv:
+            z_zq = self.act(conv(z_zq))
+
+        z_zq = z_zq.permute(0, 2, 1, 3, 4)
+
+        h, zq = torch.chunk(z_zq, 2, dim=0)
+
+        h = h.reshape(-1, *h.shape[2:])
+        zq = zq.reshape(-1, *zq.shape[2:])
+        
+        h = self.conv_in(h)
+        
+        # middle
+        h = self.mid.block_1(h, zq)
+        h = self.mid.attn_1(h, zq)
+        h = self.mid.block_2(h, zq)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](h, zq)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h, zq)
+
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        h = self.norm_out(h, zq)
+        h = self.act(h)
+        h = self.conv_out(h)
+
+        return h
+
+
+class Emu3VisionVQPretrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = Emu3VisionVQConfig
+    base_model_prefix = "emuvideovq"
+    main_input_name = "pixel_values"
+    _no_split_modules = ["Emu3VisionVQResnetBlock", "Emu3VisionVQAttnBlock", "Emu3VisionVQResnetTemporalBlock"]
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Conv2d, nn.Conv3d)):
+            nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
+        # copied from the `reset_parameters` method of `class Linear(Module)` in `torch`.
+        elif isinstance(module, nn.Linear):
+            nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+            if module.bias is not None:
+                fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)
+                bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+                nn.init.uniform_(module.bias, -bound, bound)
+        elif isinstance(module, (nn.BatchNorm2d, nn.BatchNorm3d, nn.GroupNorm)):
+            nn.init.constant_(module.weight, 1)
+            nn.init.constant_(module.bias, 0)
+
+
+class Emu3VisionVQModel(Emu3VisionVQPretrainedModel):
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+
+        self.encoder = Emu3VisionVQEncoder(config)
+        self.decoder = Emu3VisionVQDecoder(config)
+        self.quantize = Emu3VisionVQVectorQuantizer(config)
+
+        self.quant_conv = Emu3VisionVQCausalConv3d(config.z_channels, config.embed_dim)
+        self.post_quant_conv = Emu3VisionVQCausalConv3d(config.embed_dim, config.z_channels)
+
+        self.spatial_scale_factor = 2 ** (len(config.ch_mult) - 1)
+
+        self.post_init()
+
+    def encode(self, x: torch.Tensor):
+        ndim = x.ndim
+        if ndim == 4:
+            t = self.config.temporal_downsample_factor
+            b, c, h, w = x.shape
+            x = x.unsqueeze(1).repeat(1, t, 1, 1, 1)
+        elif ndim == 5:
+            b, t, c, h, w = x.shape
+
+        h = self.encoder(x)
+
+        # b t c h w -> b c t h w
+        h = h.permute(0, 2, 1, 3, 4)
+        h = self.quant_conv(h)
+        # b c t h w -> b t c h w
+        h = h.permute(0, 2, 1, 3, 4)
+
+        codes = self.quantize(h)
+
+        if ndim == 4:
+            codes = codes.squeeze(1)
+
+        return codes
+
+    def decode(self, x: torch.Tensor):
+        ndim = x.ndim
+        if ndim == 3:
+            x = x.unsqueeze(1)
+
+        b, t, h, w = x.shape
+        quant = self.quantize.embedding(x.flatten())
+        c = quant.shape[-1]
+        quant = quant.view(b, t, h, w, c).permute(0, 4, 1, 2, 3).contiguous()
+        quant2 = self.post_quant_conv(quant)
+
+        quant = quant.permute(0, 2, 1, 3, 4)
+        quant2 = quant2.permute(0, 2, 1, 3, 4)
+
+        video = self.decoder(quant2, quant)
+        video = video.reshape(
+            b,
+            t * self.config.temporal_downsample_factor,
+            self.config.out_channels,
+            h * self.spatial_scale_factor,
+            w * self.spatial_scale_factor,
+        )
+        if ndim == 3:
+            return video[:, 0]
+        return video
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype

codecs/vision/preprocessor_config.json

@@ -0,0 +1,29 @@
+{
+  "auto_map": {
+      "AutoImageProcessor": "image_processing_emu3visionvq.Emu3VisionVQImageProcessor"
+  },
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "Emu3VisionVQImageProcessor",
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "max_pixels": 1048576,
+  "min_pixels": 262144,
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "max_pixels": 1048576,
+    "min_pixels": 262144
+  },
+  "spatial_factor": 8
+}

config.json

@@ -0,0 +1,174 @@
+{
+  "model": {
+    "architecture": "axiom1_1",
+    "vocab_size": 54144,
+    "d_model": 640,
+    "n_layers": 12,
+    "n_heads": 10,
+    "n_kv_heads": 10,
+    "ffn_mult": 2.6666666666666665,
+    "max_seq_len": 1024,
+    "rope_theta": 10000.0,
+    "norm_eps": 1e-06,
+    "vis_positions": 256,
+    "visual_offset": 16640,
+    "visual_size": 32768,
+    "tie_embeddings": true,
+    "init_std": 0.02,
+    "grad_checkpoint": true,
+    "attention_layers": [
+      3,
+      7,
+      11
+    ],
+    "d_cell": 1280,
+    "n_cells": 2,
+    "top_k": 2,
+    "scan_chunk": 64,
+    "balance_weight": 0.01
+  },
+  "vocab": {
+    "text_size": 16384,
+    "visual_size": 32768,
+    "audio_size": 1024,
+    "audio_books": 4,
+    "coordinate_bins": 512,
+    "max_vertices": 128,
+    "max_faces": 128,
+    "frame_tokens": 64,
+    "max_frames": 8,
+    "max_audio_frames": 160,
+    "schema_version": 3,
+    "special_tokens": [
+      "pad",
+      "bos",
+      "eos",
+      "input",
+      "output",
+      "text",
+      "/text",
+      "image",
+      "/image",
+      "video",
+      "/video",
+      "frame",
+      "/frame",
+      "audio",
+      "/audio",
+      "mesh",
+      "/mesh",
+      "vertices",
+      "faces",
+      "tools",
+      "/tools",
+      "calls",
+      "/calls",
+      "task:text",
+      "task:chat_text",
+      "task:t2i",
+      "task:i2t",
+      "task:t2v",
+      "task:v2t",
+      "task:vcont",
+      "task:t2a",
+      "task:a2t",
+      "task:t2m",
+      "task:m2t",
+      "task:tools",
+      "task:mixed",
+      "task:coding",
+      "time:0",
+      "time:1",
+      "time:2",
+      "time:3",
+      "time:4",
+      "time:5",
+      "time:6",
+      "time:7",
+      "resolution:128"
+    ]
+  },
+  "meta": {
+    "stage": "joint_repair",
+    "step": 512,
+    "tokens_seen": 807174144,
+    "task_target_tokens": {
+      "coding": 373488,
+      "chat_text": 24092354,
+      "text": 257928814,
+      "i2t": 11852250,
+      "t2i": 39056667,
+      "vcont": 6105359,
+      "v2t": 252329,
+      "t2v": 19027778,
+      "t2m": 5739237,
+      "a2t": 207450,
+      "tools": 1430210,
+      "m2t": 477100,
+      "t2a": 10952145,
+      "replay/chat_text": 66428,
+      "replay/text": 225433,
+      "coverage/vcont": 52303,
+      "coverage/t2v": 384846
+    },
+    "stage_task_target_tokens": {
+      "coverage/vcont": 52303,
+      "t2i": 481000,
+      "coverage/t2v": 384846,
+      "replay/chat_text": 23241,
+      "replay/text": 74112
+    },
+    "stages_completed": [
+      "pilot",
+      "text",
+      "image_pilot",
+      "image",
+      "video",
+      "audio_mesh_pilot",
+      "audio_mesh",
+      "unified",
+      "pilot",
+      "image",
+      "video_coverage",
+      "joint_repair"
+    ],
+    "num_params": 90247794,
+    "precision": "bf16",
+    "image_128_trained": true,
+    "image_128_target_tokens": 16775415,
+    "env": {
+      "torch": "2.11.0+cu128",
+      "cuda": "12.8",
+      "python": "3.13.15",
+      "gpu": "NVIDIA GeForce RTX 5070"
+    },
+    "research_only": true,
+    "research_trial": "joint_repair",
+    "research_note": "Experimental branch; not a production promotion or validated semantic capability.",
+    "data": {
+      "preparation": {
+        "audio_codec": {
+          "path": "codecs/audio",
+          "revision": "6ba020b5ba7d9d8076fb90db7e67f27e31980f6e",
+          "sha256": {
+            "config.json": "3af86549c509e5491367dab3f7093014bd559f441ba31fa7427b5f94232368d7",
+            "model.safetensors": "7452e3fc6972991da871ae1a1c9d3e8e219aa247cc0f757018867d6a7fe59aaf",
+            "README.md": "6b1db50f4974b1544fb795491b6fd57a20320ff17686565c50a08407f06075b1"
+          },
+          "codebooks": 4,
+          "boundary_guard_samples": 320,
+          "alignment_samples": -8,
+          "admission_sha256": "7968bc0f22c01b433bd516acc4562ed9ef7543a6e93b8c7c7d7d697a7157bf62",
+          "implementation": "quasar.tokenizers.audio_aligned.AlignedAudioTokenizer"
+        }
+      }
+    },
+    "release_provenance": {
+      "source_weights_sha256": "de4846e5e63c0004cf0b6331127ab421d6843f6efde04243c6d5448ab4545240",
+      "original_config_sha256": "1d7254fcaaca82c66d2eab32ba39cabf673c9c63ed16aabede5d6b1092aeabb1",
+      "runtime_source_commit": "4afbdccdbd3b0967b161a68658cfd3b92c42636f",
+      "metadata_changes": "Portable codec paths; runtime-required metadata retained; workspace paths and preparation-only metadata omitted.",
+      "weights_changed": false
+    }
+  }
+}

files.sha256.json

@@ -0,0 +1,302 @@
+{
+  "config.json": {
+    "bytes": 4547,
+    "sha256": "0a5a60a1af2eb406e8195beed8fb8602aade8252dd3c5ef5bb9c8d4c7a20083c"
+  },
+  "evaluation.json": {
+    "bytes": 3456,
+    "sha256": "70681a6e629a7731f74d3dfa8c0b5692ba61a7cfb520ba307119d5dd20237bd6"
+  },
+  "infer.py": {
+    "bytes": 2730,
+    "sha256": "cdf04813be3d8505cefbaeeee604ff53d022e09fbec7fc321af3544d295b95e6"
+  },
+  "model.safetensors": {
+    "bytes": 361008568,
+    "sha256": "de4846e5e63c0004cf0b6331127ab421d6843f6efde04243c6d5448ab4545240"
+  },
+  "pyproject.toml": {
+    "bytes": 784,
+    "sha256": "320d43b1c10d23a898b64a6fd66253c78b3bbb7a74fd2c4d0f60b40fe3d9943a"
+  },
+  "README.md": {
+    "bytes": 16894,
+    "sha256": "ed5b08fa100594fc8c795a6b8188d99d0563ca736ddcc2cb071c67c909a266c2"
+  },
+  "release.json": {
+    "bytes": 783,
+    "sha256": "2734d29b5d10e09c5de8c6f2cbcc27f205ae1bedc8d46638030262b5aae84896"
+  },
+  "release_validation.json": {
+    "bytes": 660,
+    "sha256": "ffeb7f7738e459e1c9132d284072d39d367b04c751079f7d67d4746ca34c005e"
+  },
+  "requirements-inference.txt": {
+    "bytes": 256,
+    "sha256": "4d38e10bd30904cec18c0025edbe57299ff8034243ed9b670e7ace7657e97b8d"
+  },
+  "verify_release.py": {
+    "bytes": 4126,
+    "sha256": "af41074f5c74ff4090be59bb1d656f5238bd1d31d2c1b2d4612990f1eef7285a"
+  },
+  "assets/attributions.json": {
+    "bytes": 2391,
+    "sha256": "afd033dfe7cc8822899d9946ebec9bb872314127e4ec864e99a03b3fed5e62e3"
+  },
+  "assets/image_comparison.png": {
+    "bytes": 210993,
+    "sha256": "49f61b6696e34dc3589c3cf24aa4da55a5475a369bfbf5ccdefb602a0b41a216"
+  },
+  "assets/motion_comparison.gif": {
+    "bytes": 281900,
+    "sha256": "9184db2e2abe62c33c57de75d08d5c8bf6273600dfadc80842c5ecd5fcfce436"
+  },
+  "quasar/__init__.py": {
+    "bytes": 0,
+    "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+  },
+  "tokenizers/text_tokenizer.json": {
+    "bytes": 1136009,
+    "sha256": "bf620532636ec536d19191eddde091645856e1d6f262d52fdc9f7e9b43ee843a"
+  },
+  "variants/v5/config.json": {
+    "bytes": 4351,
+    "sha256": "5e10caa317356974650e2fa9ecc5bc39db73237ed08a0bcbf7f7a61fedc66c63"
+  },
+  "variants/v5/model.safetensors": {
+    "bytes": 361008568,
+    "sha256": "e20242e9971b6a7cef7ca23a65c25faa70a007807a55a2c0cfe35f4f5682a643"
+  },
+  "variants/v5/tokenizers/text_tokenizer.json": {
+    "bytes": 1136009,
+    "sha256": "bf620532636ec536d19191eddde091645856e1d6f262d52fdc9f7e9b43ee843a"
+  },
+  "variants/v5/tokenizers/vq/vq.safetensors": {
+    "bytes": 524448,
+    "sha256": "0ed59aaface626b762214262572da4d9b0eb841175859fe26c925e606a82069b"
+  },
+  "variants/v5/tokenizers/vq/vq_config.json": {
+    "bytes": 2186,
+    "sha256": "4eae4e47af92536dacc30559498eed3c1318a80fbf603ab6a3a5da6e500e8448"
+  },
+  "tokenizers/vq/vq.safetensors": {
+    "bytes": 524448,
+    "sha256": "0ed59aaface626b762214262572da4d9b0eb841175859fe26c925e606a82069b"
+  },
+  "tokenizers/vq/vq_config.json": {
+    "bytes": 2180,
+    "sha256": "a720078321fd0c30cdedc241ca47c05474d5e91945bb4f1feb100452dffe03de"
+  },
+  "quasar/data/adapters.py": {
+    "bytes": 5338,
+    "sha256": "93dcfca754d0d5d58fefef7c23eb49f91538ad822231422ede0ad2aafe36a1ab"
+  },
+  "quasar/data/collator.py": {
+    "bytes": 1444,
+    "sha256": "dbc6b411df217460ef6578394d4ca3f1a4f6e8be672cdc71771eaf28b1a4a675"
+  },
+  "quasar/data/curated.py": {
+    "bytes": 11234,
+    "sha256": "036bcb4f95a90c3bc338deae2491fc5c416acd3454a862494347712e08c7737e"
+  },
+  "quasar/data/datasets.py": {
+    "bytes": 13870,
+    "sha256": "2a0fccbb576c69351e12e9a60f1c74348cce385e6d06a1786ec87f8bb82f3785"
+  },
+  "quasar/data/download.py": {
+    "bytes": 9387,
+    "sha256": "e98607b44e557c6b9a86ed1d06d487ac810ca2f7458c29f08abbb6fd5a15f19f"
+  },
+  "quasar/data/preprocessing.py": {
+    "bytes": 11794,
+    "sha256": "5a802f8fc5da21699136ffd9941e556474200a3f5038667a1e90293b3af0de25"
+  },
+  "quasar/data/synthetic.py": {
+    "bytes": 18794,
+    "sha256": "4f6046fa0fa8fb426820c8970628ff723d84c6b9539cbf1e02f3bc80b8f3bd05"
+  },
+  "quasar/data/text_preparation.py": {
+    "bytes": 29675,
+    "sha256": "777fb661fa86ef7602e0fca0a707efbc3e1eaa168de3fbda2177506462051fac"
+  },
+  "quasar/data/__init__.py": {
+    "bytes": 0,
+    "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+  },
+  "quasar/evaluation/__init__.py": {
+    "bytes": 0,
+    "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+  },
+  "quasar/inference/decoding.py": {
+    "bytes": 7170,
+    "sha256": "936f834f079dfb2763091897095d4cc65a5e6b6ff16fcbe844e60cdf95bd4152"
+  },
+  "quasar/inference/extended.py": {
+    "bytes": 8322,
+    "sha256": "5ce1f4bd1023f3fb77b1853e2f36d5410ae0595349bf5f9d03aae99d52f1ebab"
+  },
+  "quasar/inference/generate.py": {
+    "bytes": 16454,
+    "sha256": "00a31f84a6917d58d3ebbc09bb04e796b98aa8fc8aef8e32d278652e5aa24e24"
+  },
+  "quasar/inference/universal.py": {
+    "bytes": 6151,
+    "sha256": "50569a27c60c91d5b5a998e018f2e18d63b9483b082faa7225e4ec6a808120eb"
+  },
+  "quasar/inference/__init__.py": {
+    "bytes": 0,
+    "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+  },
+  "quasar/model/attention.py": {
+    "bytes": 2905,
+    "sha256": "22bfd305f5dba981438913d0a8418e4ec9d5e62672740e458db1ce2e8c11585c"
+  },
+  "quasar/model/axiom.py": {
+    "bytes": 6040,
+    "sha256": "1919f42a57cf3f4f0601cb79846f7fe91af826b802a188c63cf0c1a02a4e539d"
+  },
+  "quasar/model/config.py": {
+    "bytes": 3031,
+    "sha256": "be2218d6e051120af3be8ee53aa0ccdd052a316588e69c7d8684540f029d0591"
+  },
+  "quasar/model/layers.py": {
+    "bytes": 1704,
+    "sha256": "859edfce4cce98453da2bd6bc3e24b45ad188d908d3351e0074e90cf55ab4b9f"
+  },
+  "quasar/model/transformer.py": {
+    "bytes": 6531,
+    "sha256": "43ff8614eeb1e04d0d69e20cf235e1e0b5c1a4b8c022554ec21e2a4033440c51"
+  },
+  "quasar/model/__init__.py": {
+    "bytes": 0,
+    "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+  },
+  "quasar/tokenizers/audio.py": {
+    "bytes": 2726,
+    "sha256": "e587422428a3925841754b7a8a928abd0c3bb0fa82bd259b8a17534075f93a55"
+  },
+  "quasar/tokenizers/audio_aligned.py": {
+    "bytes": 1852,
+    "sha256": "d32dbd18faa8eb75b40f7ae1cc64d7e5fffa8f634bb2998f399ed1c0f24adfb2"
+  },
+  "quasar/tokenizers/audio_complete.py": {
+    "bytes": 1068,
+    "sha256": "d50959abc6130c58d4935b4846bc2916475cd51cd26a35a4bcad498d7f51917d"
+  },
+  "quasar/tokenizers/mesh.py": {
+    "bytes": 2144,
+    "sha256": "148cc80bcb900eed289926b3130dcd1d35f1eae2c000511d13a80f3de056ed26"
+  },
+  "quasar/tokenizers/multimodal.py": {
+    "bytes": 9485,
+    "sha256": "5c895527a02624489e7c4ced5912eb2cf0d41912945071ff6795770d740ca5a8"
+  },
+  "quasar/tokenizers/multires_emu.py": {
+    "bytes": 2840,
+    "sha256": "764a04ff63106146cbc0d910b5ff293baf54335ffa93305531d2b971c77da353"
+  },
+  "quasar/tokenizers/restricted_emu.py": {
+    "bytes": 2511,
+    "sha256": "add63d1042dc72dadfd79ec0ad3796e1a6cac2a9023ebf498a54a1d6d3ab8bf8"
+  },
+  "quasar/tokenizers/text.py": {
+    "bytes": 1781,
+    "sha256": "d97bc0c634dab8318fc7528def8d9fb1eb17e21da0c92e0812368c02d313135b"
+  },
+  "quasar/tokenizers/universal.py": {
+    "bytes": 12299,
+    "sha256": "1607e4c9b7eeefe8ee7eacc70bbf53a2bdf934f7636823c5abed2d0a0151dc8c"
+  },
+  "quasar/tokenizers/vision.py": {
+    "bytes": 7566,
+    "sha256": "10f239e83334a06dc49b47fdab7bde9188ade3e5a4612ee1175e0680ba329ca2"
+  },
+  "quasar/tokenizers/__init__.py": {
+    "bytes": 0,
+    "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+  },
+  "quasar/training/checkpoint.py": {
+    "bytes": 3473,
+    "sha256": "39c02d1b30da98b8f3861d304f9db5eea0660cec701d8c58cdb4ef14f1beaefd"
+  },
+  "quasar/training/compilation.py": {
+    "bytes": 2062,
+    "sha256": "812443fb70939fd06d3f8c96e3d75dc8f9beabcbcaef11623cc571291b437877"
+  },
+  "quasar/training/hardware.py": {
+    "bytes": 2776,
+    "sha256": "d9e47b5246c61d09e54b699d9de09d3e42cda29992894b1a58b5b4042e956c2a"
+  },
+  "quasar/training/metrics.py": {
+    "bytes": 2505,
+    "sha256": "19189212aa5513def43109142213f9d1ef8c759162395d26d02d1104b5bbc947"
+  },
+  "quasar/training/migration.py": {
+    "bytes": 1381,
+    "sha256": "0db3c61f34fe0cc202d6bee9a7d4e4a5f58c6da0f23049d58adcf794239b4f00"
+  },
+  "quasar/training/trainer.py": {
+    "bytes": 21181,
+    "sha256": "5a00a53173ca9ee21cce763f9345af1db045b1f0898c073c692aeeb2dbf1c8f2"
+  },
+  "quasar/training/validation_guard.py": {
+    "bytes": 2101,
+    "sha256": "81b2916c4879e06025447fba99111cdbacdcdf4065dc033e1e509b434fdacbdd"
+  },
+  "quasar/training/__init__.py": {
+    "bytes": 0,
+    "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
+  },
+  "codecs/audio/config.json": {
+    "bytes": 542,
+    "sha256": "3af86549c509e5491367dab3f7093014bd559f441ba31fa7427b5f94232368d7"
+  },
+  "codecs/audio/model.safetensors": {
+    "bytes": 298652268,
+    "sha256": "7452e3fc6972991da871ae1a1c9d3e8e219aa247cc0f757018867d6a7fe59aaf"
+  },
+  "codecs/audio/OFFICIAL_LICENSE": {
+    "bytes": 1095,
+    "sha256": "32277f65d28ee0be285c9e2d12970d92bd1331ae9fad490801c3a79243fdb21b"
+  },
+  "codecs/audio/OFFICIAL_WEIGHTS_TERMS.md": {
+    "bytes": 6557,
+    "sha256": "3178619f84b8524105b10518d0527600ca74c01386885995f58bc633da2fcbd9"
+  },
+  "codecs/audio/README.md": {
+    "bytes": 5198,
+    "sha256": "6b1db50f4974b1544fb795491b6fd57a20320ff17686565c50a08407f06075b1"
+  },
+  "codecs/vision/config.json": {
+    "bytes": 592,
+    "sha256": "6717e1bd03d980c6b5e88aa0a47a557d193b0c33afbdb9d24e53a92593ee7473"
+  },
+  "codecs/vision/configuration_emu3visionvq.py": {
+    "bytes": 4176,
+    "sha256": "e4a9c0ee78c5196219a5ab89223ba82d2d20844687d19f877b261e5026ef4c11"
+  },
+  "codecs/vision/image_processing_emu3visionvq.py": {
+    "bytes": 22487,
+    "sha256": "506392861a6842fa3c9a0bb320b3be0228309c93a1f3b45e4aa2caf501141ea4"
+  },
+  "codecs/vision/LICENSE": {
+    "bytes": 11558,
+    "sha256": "1eb85fc97224598dad1852b5d6483bbcf0aa8608790dcc657a5a2a761ae9c8c6"
+  },
+  "codecs/vision/model.safetensors": {
+    "bytes": 1083015124,
+    "sha256": "89536431c69b08b10b449ec309f52dcea22f14b7647317f30f5715273392bbf1"
+  },
+  "codecs/vision/modeling_emu3visionvq.py": {
+    "bytes": 24362,
+    "sha256": "07558b036f40ec751de174487b15fe5c03823ba2edf29449bfa29be6f7f01afc"
+  },
+  "codecs/vision/preprocessor_config.json": {
+    "bytes": 556,
+    "sha256": "98a019cd8820290dbdc53e7366f7e447ee55d085824b63fdb68b2d642dd9edbf"
+  },
+  "codecs/vision/README.md": {
+    "bytes": 3531,
+    "sha256": "5763e0d54ed458232a8a4edb94cbf8b3b7fcf7d76372c2c0ea2c4c7b6e6a4e78"
+  }
+}

infer.py

@@ -0,0 +1,60 @@
+"""Explicit local Quasar inference. No model/codec download or tool execution."""
+from pathlib import Path
+import argparse
+import os
+import sys
+
+ROOT = Path(__file__).resolve().parent
+sys.path.insert(0, str(ROOT))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--checkpoint', type=Path, default=ROOT)
+    parser.add_argument('--task', choices=('text', 'image', 'video'), required=True)
+    parser.add_argument('--prompt', required=True)
+    parser.add_argument('--device', choices=('cpu', 'cuda'), default='cuda')
+    parser.add_argument('--output', type=Path)
+    parser.add_argument('--resolution', type=int, choices=(64, 128), default=128)
+    parser.add_argument('--frames', type=int, default=8)
+    parser.add_argument('--seed', type=int, default=20261021)
+    parser.add_argument('--max-new-tokens', type=int, default=64)
+    args = parser.parse_args()
+    if args.device == 'cpu':
+        os.environ['CUDA_VISIBLE_DEVICES'] = ''
+        os.environ.setdefault('TORCHINDUCTOR_TRITON_DISABLE_DEVICE_DETECTION', '1')
+        # These FP32/BF16 weights do not use the optional GPU quantizer.
+        sys.modules['torchao'] = None
+    import torch
+    torch.set_num_threads(2)
+    if args.device == 'cuda':
+        from quasar.training.hardware import check_gpu
+        check_gpu({'require_cuda': True, 'exclusive_gpu': True, 'min_free_vram_gb': 4})
+    from quasar.inference.extended import ExtendedGenerator
+    from quasar.inference.decoding import Sampling
+    generator = ExtendedGenerator(args.checkpoint.resolve(), args.device)
+    sampling = Sampling(temperature=.5, top_k=64, top_p=.9, seed=args.seed, max_new_tokens=args.max_new_tokens)
+    if args.task == 'text':
+        result = generator.generate_text(args.prompt, sampling)
+        if args.output:
+            args.output.parent.mkdir(parents=True, exist_ok=True)
+            args.output.write_text(result, encoding='utf-8')
+        print(result)
+        return
+    from PIL import Image
+    output = args.output or Path('generated.png' if args.task == 'image' else 'generated.gif')
+    output.parent.mkdir(parents=True, exist_ok=True)
+    if args.task == 'image':
+        pixels = generator.generate_image(args.prompt, sampling, resolution=args.resolution)
+        Image.fromarray(pixels).save(output)
+    else:
+        if output.suffix.lower() != '.gif':
+            parser.error('Video preview output must use .gif')
+        frames = generator.generate_video(args.prompt, args.frames, sampling)
+        images = [Image.fromarray(frame) for frame in frames]
+        images[0].save(output, save_all=True, append_images=images[1:], duration=250, loop=0, disposal=2)
+    print(output.resolve())
+
+
+if __name__ == '__main__':
+    main()

model.safetensors

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

pyproject.toml

@@ -0,0 +1,24 @@
+[project]
+name = "quasar"
+version = "0.1.0"
+description = "Quasar: a unified autoregressive text/image/video foundation-model prototype"
+requires-python = ">=3.10"
+dependencies = [
+  "torch>=2.6", "numpy>=1.26", "pillow>=10.0", "pyyaml>=6.0", "safetensors>=0.4",
+  "tokenizers>=0.19", "huggingface_hub>=0.24", "pyarrow>=14", "imageio-ffmpeg>=0.5", "requests>=2.31",
+]
+
+[project.optional-dependencies]
+app = ["gradio==6.30.0", "psutil>=7.0"]
+dev = ["pytest>=8.0"]
+multimodal = ["jsonschema>=4.20", "transformers==5.5.0", "soundfile>=0.12", "trimesh>=4.0", "scipy>=1.14"]
+
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+include = ["quasar*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

quasar/data/adapters.py

@@ -0,0 +1,87 @@
+"""Offline adapters for the selected native records.
+
+Media bindings are prepared separately with real codecs and audited rights.
+An adapter never downloads an image URL or assumes a dataset card licenses it.
+All reversals of a pair keep the same source group and media assets.
+"""
+from __future__ import annotations
+
+import json
+import ast
+from .curated import digest, validate_calls
+
+
+def _json(x):
+    return json.loads(x) if isinstance(x, str) else x
+
+
+def normalize(dataset: str, row: dict, provenance: dict, binding: dict | None = None) -> list[dict]:
+    """`provenance` supplies id/revision/license/attribution/source_url/split.
+
+    binding: {packet: encoded modality packet, assets: [...rights/checksums...],
+              group: source object/video/speaker/uploader ID, caption: aligned caption}
+    For media, `group` is mandatory: no clip-level or view-level random split.
+    """
+    def record(task, ins, outs, group):
+        return {**provenance, "dataset": dataset, "task": task, "group": group, "inputs": ins, "outputs": outs,
+                **({"assets": binding["assets"], "codec": binding.get("codec", {})} if binding else {})}
+    text = lambda s: {"type": "text", "text": str(s).strip()}
+    if dataset == "mbpp":
+        source_split = provenance.get("source_split")
+        expected = {"train": "train", "validation": "val", "test": "test"}
+        if source_split not in expected or provenance.get("split") != expected[source_split]:
+            raise ValueError("MBPP official train/validation/test provenance must be preserved")
+        code = row["code"].strip()
+        ast.parse(code)  # syntax validation only; dataset code is never executed by this adapter
+        examples = "\n".join(row.get("test_list", []))
+        prompt = row["text"].strip() + ("\n\nExamples:\n" + examples if examples else "")
+        return [record("coding", [text(prompt)], [text("```python\n" + code + "\n```")], str(row["task_id"]))]
+    if dataset in ("builtin", "user"):
+        if not row.get("group"): raise ValueError("local records require a stable source group")
+        return [record(row["task"], row["inputs"], row["outputs"], row["group"])]
+    if dataset == "fineweb_edu":
+        return [record("text", [], [text(row["text"])], row.get("url") or digest(row["text"]))]
+    if dataset == "smoltalk":
+        messages = row["messages"]
+        if not messages or messages[-1]["role"] != "assistant":
+            raise ValueError("SmolTalk row must end with a complete assistant response")
+        context = "\n".join(f"{m['role']}: {m['content']}" for m in messages[:-1])
+        if not context.strip(): raise ValueError("conversation has no conditioning")
+        return [record("chat_text", [text(context)], [text(messages[-1]["content"])], digest(messages))]
+    if dataset == "xlam":
+        tools = _json(row["tools"])
+        calls = _json(row["answers"])
+        validate_calls(tools, calls)
+        # Whole tool schemas (including function descriptions/parameters) define a heldout family.
+        group = digest(sorted(tools, key=lambda x: x.get("function", x)["name"]))
+        return [record("tools", [{"type": "tools", "value": tools}, text(row["query"])], [{"type": "calls", "value": calls}], group)]
+    if binding is None or not binding.get("group") or not binding.get("assets"):
+        raise ValueError("media adapter requires an encoded packet, group ID and audited asset binding")
+    packet, group = binding["packet"], binding["group"]
+    if dataset == "openimages_narratives":
+        if not str(row["dataset_id"]).startswith("open_images"):
+            raise ValueError("only the Open Images subset is selected; COCO/Flickr/ADE rights differ")
+        cap, a, b, kind = row["caption"], "t2i", "i2t", "image"
+    elif dataset == "finevideo":
+        # Required exact-window caption is selected from scene/activity metadata, not a whole-video title.
+        if not binding.get("caption_aligned") or not binding.get("window_seconds"):
+            raise ValueError("FineVideo needs scene-aligned window caption and timestamps")
+        cap, a, b, kind = binding["caption"], "t2v", "v2t", "video"
+    elif dataset == "libritts_r":
+        if not binding.get("complete_utterance"):
+            raise ValueError("speech pair must retain the complete aligned utterance")
+        cap, a, b, kind = row["text_normalized"], "t2a", "a2t", "audio"
+    elif dataset == "fsd50k":
+        # Weak class-label conditioning; never report this as free-form human captioning.
+        cap, a, b, kind = "sound of " + ", ".join(row["labels"]), "t2a", "a2t", "audio"
+    elif dataset == "objaverse_cap3d":
+        cap, a, b, kind = row["caption"], "t2m", "m2t", "mesh"
+    else:
+        raise ValueError(f"no enabled adapter for {dataset}")
+    if packet["type"] != kind or not str(cap).strip(): raise ValueError("packet kind or caption mismatch")
+    out = [record(a, [text(cap)], [packet], group), record(b, [packet], [text(cap)], group)]
+    if dataset == "finevideo" and len(packet["codes"]) >= 4:
+        mid = len(packet["codes"]) // 2
+        out.append(record("vcont", [{"type":"video", "codes":packet["codes"][:mid]}],
+                          [{"type":"video", "codes":packet["codes"][mid:]}], group))
+    return out

quasar/data/collator.py

@@ -0,0 +1,35 @@
+"""Dynamic padding + next-token label construction."""
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+from ..tokenizers.multimodal import Sample
+
+
+def collate(samples: list[Sample], pad_id: int, multiple: int = 8, *, target_length: int | None = None,
+            batch_size: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+    """Right-pad to the batch max (rounded up to `multiple`) and build shifted labels.
+
+    labels[j] = tokens[j+1] where j+1 >= target_start, else -100 (prompt and padding are not
+    trained on).  Right padding is safe under causal attention: real tokens never attend to it.
+    """
+    L = max(len(s.tokens) for s in samples) - 1
+    if target_length is not None:
+        if target_length < L:
+            raise ValueError("fixed padding cannot truncate a sample")
+        L = target_length
+    else:
+        L = -(-L // multiple) * multiple
+    B = len(samples) if batch_size is None else batch_size
+    if B < len(samples):
+        raise ValueError("fixed batch size cannot discard samples")
+    inp = np.full((B, L), pad_id, dtype=np.int64)
+    lab = np.full((B, L), -100, dtype=np.int64)
+    for i, s in enumerate(samples):
+        t = s.tokens.astype(np.int64)
+        n = len(t) - 1
+        inp[i, :n] = t[:-1]
+        lab[i, :n] = t[1:]
+        lab[i, : max(0, s.target_start - 1)] = -100
+    return torch.from_numpy(inp), torch.from_numpy(lab)

quasar/data/curated.py

@@ -0,0 +1,195 @@
+"""Offline normalized records -> validated, deduplicated, tokenized v2 caches.
+
+No network access, model updates, tool execution or implicit license assumptions.
+Records retain dataset, revision, group, split, attribution, asset checksums and
+packet inputs/outputs. Paired generation/understanding rows share one split key.
+"""
+from __future__ import annotations
+
+from collections import Counter, defaultdict
+import hashlib
+import json
+from pathlib import Path
+
+import numpy as np
+
+from ..tokenizers.multimodal import Sample
+from ..tokenizers.universal import UniversalVocab
+from .preprocessing import sha256_file
+
+
+def canonical(value) -> str:
+    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, allow_nan=False)
+
+
+def digest(value) -> str:
+    return hashlib.sha256(canonical(value).encode()).hexdigest()
+
+
+def validate_calls(tools: list[dict], calls: list[dict]) -> None:
+    """Validate tool names and arguments with full JSON Schema; never execute tools."""
+    from jsonschema import Draft202012Validator
+    names = {}
+    for tool in tools:
+        f = tool.get("function", tool)
+        name = f["name"]
+        if name in names:
+            raise ValueError(f"duplicate tool name: {name}")
+        schema = f["parameters"]
+        Draft202012Validator.check_schema(schema)
+        names[name] = Draft202012Validator(schema)
+    for call in calls:
+        if call["name"] not in names:
+            raise ValueError(f"unknown tool: {call['name']}")
+        if not isinstance(call["arguments"], dict):
+            raise ValueError("tool arguments must be a JSON object")
+        names[call["name"]].validate(call["arguments"])
+
+
+def verify_record(r: dict, catalog: dict):
+    for key in ("id", "dataset", "revision", "group", "license", "attribution", "source_url", "task", "inputs", "outputs"):
+        if key not in r or r[key] in (None, "", []):
+            # Text pretraining legitimately has no input conditioning.
+            if key == "inputs" and r.get(key) == []:
+                continue
+            raise ValueError(f"missing provenance/content field: {key}")
+    ds = catalog["datasets"].get(r["dataset"])
+    if ds is None or ds["role"] != "core":
+        raise ValueError("dataset is unknown or disabled/reserve")
+    if r["task"] not in ds["tasks"]:
+        raise ValueError("task not registered for dataset")
+    if r["dataset"] == "mbpp":
+        expected = {"train": "train", "validation": "val", "test": "test"}
+        if r.get("source_split") not in expected or r.get("split") != expected[r["source_split"]]:
+            raise ValueError("MBPP official splits must be retained; held-out solutions cannot be relabeled train")
+    if r["revision"] != ds.get("revision", ds.get("release")):
+        raise ValueError("record revision differs from the pinned catalog")
+    if ds.get("caption_revision") and r.get("caption_revision") != ds["caption_revision"]:
+        raise ValueError("record must pin its caption source revision as well as its asset source")
+    allowed = catalog["policy"]["allowed_asset_licenses"]
+    licenses = r["license"] if isinstance(r["license"], list) else [r["license"]]
+    for lic in licenses:
+        if lic not in allowed:
+            raise ValueError(f"unapproved or ambiguous license: {lic}")
+    declared = ds.get("license")
+    if declared and declared not in licenses and not (declared == "CC-BY" and any(x.startswith("CC-BY-") for x in licenses)):
+        raise ValueError("record license differs from pinned dataset terms")
+    media = [s for s in r["inputs"] + r["outputs"] if s["type"] in ("image", "video", "audio", "mesh")]
+    if media:
+        for asset in r.get("assets", []):
+            for key in ("sha256", "license", "attribution", "source_url", "rights_evidence"):
+                if not asset.get(key):
+                    raise ValueError(f"media asset missing {key}")
+            if asset["license"] not in allowed:
+                raise ValueError("unapproved media asset license")
+            if len(asset["sha256"]) != 64 or any(c not in "0123456789abcdef" for c in asset["sha256"]):
+                raise ValueError("invalid asset checksum")
+        if not r.get("assets"):
+            raise ValueError("media record needs asset rights evidence and checksums")
+    if r["task"] == "tools":
+        tools = next(s["value"] for s in r["inputs"] if s["type"] == "tools")
+        calls = next(s["value"] for s in r["outputs"] if s["type"] == "calls")
+        validate_calls(tools, calls)
+    inputs = {s["type"] for s in r["inputs"]}
+    outputs = {s["type"] for s in r["outputs"]}
+    expected = {"t2i": ("text", "image"), "i2t": ("image", "text"), "t2v": ("text", "video"),
+                "v2t": ("video", "text"), "vcont": ("video", "video"), "t2a": ("text", "audio"),
+                "a2t": ("audio", "text"), "t2m": ("text", "mesh"), "m2t": ("mesh", "text"),
+                "chat_text": ("text", "text"), "coding": ("text", "text"), "tools": ("tools", "calls")}
+    if r["task"] in expected:
+        a, b = expected[r["task"]]
+        if a not in inputs or b not in outputs: raise ValueError("packet types do not match the registered task")
+    if r["task"] == "text" and outputs != {"text"}: raise ValueError("text task needs text targets")
+
+
+def compile_manifest(manifest: str | Path, out: str | Path, vocab: UniversalVocab, tok, catalog: dict,
+                     context: int = 1024, max_disk_bytes: int = 20_000_000_000) -> dict:
+    """Stream into uint16 token files and per-row indices; reject overlong samples without truncation.
+
+    Existing output directories are refused. Failure preserves an INCOMPLETE marker;
+    a loader refuses that directory. Provenance is saved separately from compact indices.
+    """
+    manifest, out = Path(manifest), Path(out)
+    if out.exists():
+        raise FileExistsError(f"refusing to overwrite cache {out}; use a new versioned directory")
+    if tok.vocab_size != vocab.text_size or vocab.size > 65536:
+        raise ValueError("tokenizer size/vocabulary mismatch for uint16 cache")
+    out.mkdir(parents=True)
+    (out / "INCOMPLETE").write_text("preparation in progress")
+    counts, offsets = Counter(), defaultdict(int)
+    groups, assets, seen, seen_split = {}, {}, set(), {}
+    report = {"training_started": False, "records": 0, "target_tokens": 0}
+    handles = {}
+    try:
+        for split in ("train", "val", "test"):
+            handles[split] = ((out / f"{split}.bin").open("wb"), (out / f"{split}.index.jsonl").open("w", encoding="utf-8"))
+        with manifest.open(encoding="utf-8") as src, (out / "provenance.jsonl").open("w", encoding="utf-8") as prov:
+            for lineno, line in enumerate(src, 1):
+                if not line.strip(): continue
+                try:
+                    r = json.loads(line)
+                    verify_record(r, catalog)
+                    group = r["dataset"] + ":" + str(r["group"])
+                    # Explicit official splits win. Otherwise hash source groups, never windows/captions.
+                    split = r.get("split") or ("val" if int(digest(group)[:8], 16) % 100 < 2 else "train")
+                    if split not in handles: raise ValueError("split must be train, val or test")
+                    if group in groups and groups[group] != split: raise ValueError("source group leaks across splits")
+                    groups[group] = split
+                    for a in r.get("assets", []):
+                        h = a["sha256"]
+                        if h in assets and assets[h] != split: raise ValueError("asset content leaks across splits")
+                        assets[h] = split
+                    key = digest([r["task"], r["inputs"], r["outputs"]])
+                    if key in seen:
+                        if seen_split[key] != split: raise ValueError("duplicate example leaks across splits")
+                        counts["duplicates_skipped"] += 1
+                        continue
+                    seen.add(key)
+                    seen_split[key] = split
+                    sample = vocab.build(r["task"], r["inputs"], r["outputs"], tok.encode, context)
+                    if (sum(offsets.values()) + len(sample.tokens)) * 2 > max_disk_bytes:
+                        raise ValueError("prepared token disk budget exceeded")
+                    sample.tokens.astype("<u2").tofile(handles[split][0])
+                    idx = {"offset": offsets[split], "length": len(sample.tokens), "target_start": sample.target_start,
+                           "task": r["task"], "id": r["id"], "dataset": r["dataset"], "group": r["group"]}
+                    handles[split][1].write(canonical(idx) + "\n")
+                    prov.write(canonical({k: r[k] for k in r if k not in ("inputs", "outputs")}) + "\n")
+                    offsets[split] += len(sample.tokens)
+                    counts[split + "/" + r["task"]] += 1
+                    report["records"] += 1
+                    report["target_tokens"] += len(sample.tokens) - sample.target_start
+                except Exception as e:
+                    raise ValueError(f"{manifest}:{lineno}: {e}") from e
+    finally:
+        for pair in handles.values():
+            for f in pair: f.close()
+    tok.save(out / "text_tokenizer.json")
+    (out / "vocab_v2.json").write_text(json.dumps(vocab.to_dict(), indent=2))
+    (out / "catalog.json").write_text(json.dumps(catalog, indent=2))
+    report |= {"counts": dict(counts), "manifest_sha256": sha256_file(manifest), "catalog_sha256": digest(catalog),
+               "token_bytes": sum(offsets.values()) * 2, "context": context, "tokenizer_sha256": sha256_file(out / "text_tokenizer.json")}
+    (out / "stats.json").write_text(json.dumps(report, indent=2))
+    (out / "INCOMPLETE").unlink()
+    return report
+
+
+class UniversalSources:
+    """Same sample/available interface as TaskSources; consumed by existing MixtureLoader."""
+    def __init__(self, cache: str | Path, split: str, vocab: UniversalVocab, seq_len: int = 1024):
+        cache = Path(cache)
+        if (cache / "INCOMPLETE").exists(): raise ValueError("incomplete v2 cache")
+        if json.loads((cache / "vocab_v2.json").read_text()) != vocab.to_dict(): raise ValueError("v2 vocabulary mismatch")
+        self.vocab, self.seq_len, self.rows = vocab, seq_len, defaultdict(list)
+        path = cache / f"{split}.bin"
+        self.tokens = np.memmap(path, dtype="<u2", mode="r") if path.stat().st_size else np.array([], dtype=np.uint16)
+        for line in (cache / f"{split}.index.jsonl").read_text().splitlines():
+            r = json.loads(line)
+            if r["length"] - 1 > seq_len: raise ValueError("cached sample exceeds requested context")
+            self.rows[r["task"]].append(r)
+
+    def available(self): return list(self.rows)
+
+    def sample(self, task: str, rng: np.random.Generator) -> Sample:
+        rows = self.rows[task]
+        r = rows[int(rng.integers(len(rows)))]
+        return Sample(np.array(self.tokens[r["offset"]:r["offset"] + r["length"]], dtype=np.int32), r["target_start"])

quasar/data/datasets.py

@@ -0,0 +1,249 @@
+"""Task sources over the pretokenized caches and a seeded mixture loader.
+
+Every micro-batch holds a single task (so lengths are similar and padding is small) and is
+drawn from an RNG seeded by (seed, step, micro_index): sampling is a pure function of the step,
+which makes runs reproducible and resumable without storing data-loader state.
+"""
+from __future__ import annotations
+
+import json
+import queue
+import threading
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from ..tokenizers.multimodal import Sample, UnifiedVocab
+from ..tokenizers.text import TextTokenizer
+from .collator import collate
+
+BASE_TASKS = ("text", "t2i", "i2t", "i2i", "t2v", "v2t", "vcont")
+CHAT_TASKS = ("chat_text", "chat_mm")
+TASK_MODALITY = {"text": "text", "chat_text": "text", "t2i": "image", "i2t": "image", "i2i": "image",
+                 "t2v": "video", "v2t": "video", "vcont": "video", "chat_mm": "multimodal"}
+
+CHAT_PROMPTS = {
+    "t2i": ["Generate an image of {c}.", "Draw {c}.", "Create a picture of {c}.", "Show me {c}.", "Can you make an image of {c}?"],
+    "t2v": ["Make a video of {c}.", "Generate a video of {c}.", "Create a short clip of {c}.", "Animate {c}."],
+    "i2t": ["Describe this image.", "What do you see in this picture?", "What is in this image?"],
+    "v2t": ["Describe this video.", "What is happening in this video?", "What happens in this clip?"],
+    "vcont": ["Continue the movement in this video.", "Continue this video.", "What happens next? Show me."],
+    "i2v": ["Animate this image: {c}.", "Bring this picture to life: {c}."],
+}
+CHAT_KINDS = ("t2i", "t2v", "i2t", "v2t", "vcont", "i2i", "i2v")
+
+
+def _load_media(cache_dir: Path, name: str, split: str, vocab: UnifiedVocab, keys: tuple[str, ...]) -> tuple[dict, dict] | None:
+    """Load <name>_<split>.npz/.json, dropping rows whose codes are out of range (corruption guard)."""
+    npz, js = cache_dir / f"{name}_{split}.npz", cache_dir / f"{name}_{split}.json"
+    if not npz.exists():
+        return None
+    try:
+        with np.load(npz) as z:
+            arrs = {k: z[k] for k in keys}
+        meta = json.loads(js.read_text())
+    except (OSError, ValueError, KeyError) as e:
+        raise RuntimeError(f"cache {npz} is unreadable ({e}); regenerate with scripts/generate_dataset.py") from e
+    ok = np.ones(len(arrs[keys[0]]), dtype=bool)
+    for a in arrs.values():
+        ok &= (a.reshape(len(a), -1) < vocab.visual_size).all(1)
+    if not ok.all():
+        print(f"[data] {npz.name}: dropping {int((~ok).sum())} corrupted rows")
+    keep = np.nonzero(ok)[0]
+    return {k: a[keep] for k, a in arrs.items()}, {k: [v[i] for i in keep] for k, v in meta.items()}
+
+
+class TaskSources:
+    """Builds one training Sample for any task from the caches of a given split."""
+
+    def __init__(self, cache_dir: str | Path, split: str, vocab: UnifiedVocab, tok: TextTokenizer, seq_len: int, caption_dropout: float = 0.1):
+        cache_dir = Path(cache_dir)
+        self.vocab, self.tok, self.seq_len = vocab, tok, seq_len
+        # Blank the t2i/t2v caption this often so the model also learns p(media) -> classifier-free guidance.
+        self.caption_dropout = caption_dropout if split == "train" else 0.0
+        tp = cache_dir / f"text_{split}.bin"
+        self.text = np.memmap(tp, dtype=np.uint16, mode="r") if tp.exists() else None
+        self.images = _load_media(cache_dir, "images", split, vocab, ("codes",))
+        self.edits = _load_media(cache_dir, "edits", split, vocab, ("src", "tgt"))
+        self.videos = _load_media(cache_dir, "videos", split, vocab, ("codes",))
+        cp = cache_dir / f"chat_{split}.json"
+        self.chat = json.loads(cp.read_text()) if cp.exists() else None
+        # indices with a caption (user videos may be uncaptioned: continuation-only)
+        cap_idx = lambda m: np.array([i for i, c in enumerate(m[1]["captions"]) if c], dtype=np.int64) if m else np.zeros(0, np.int64)
+        self.img_cap, self.vid_cap = cap_idx(self.images), cap_idx(self.videos)
+        need = {"t2i": len(self.img_cap), "i2t": len(self.img_cap), "t2v": len(self.vid_cap), "v2t": len(self.vid_cap),
+                "vcont": self.videos is not None, "i2i": self.edits is not None, "i2v": len(self.vid_cap)}
+        self.chat_kinds = [k for k in CHAT_KINDS if need[k]]
+
+    def available(self) -> list[str]:
+        have = {"text": self.text is not None and len(self.text) > self.seq_len + 1, "chat_text": bool(self.chat),
+                "t2i": len(self.img_cap) > 0, "i2t": len(self.img_cap) > 0, "i2i": self.edits is not None,
+                "t2v": len(self.vid_cap) > 0, "v2t": len(self.vid_cap) > 0, "vcont": self.videos is not None,
+                "chat_mm": bool(self.chat_kinds)}
+        return [t for t, ok in have.items() if ok]
+
+    def enc(self, s: str) -> list[int]:
+        return self.tok.encode(s)
+
+    def _prompt(self, caps: list[str], rng: np.random.Generator) -> list[int]:
+        return [] if rng.random() < self.caption_dropout else self.enc(str(rng.choice(caps)))
+
+    def sample(self, task: str, rng: np.random.Generator) -> Sample:
+        v = self.vocab
+        if task == "text":
+            off = int(rng.integers(0, len(self.text) - self.seq_len - 1))
+            return Sample(np.asarray(self.text[off : off + self.seq_len + 1], dtype=np.int32), 0)
+        if task in ("t2i", "i2t"):
+            (arr, meta), i = self.images, int(self.img_cap[rng.integers(len(self.img_cap))])
+            caps = meta["captions"][i]
+            if task == "t2i":
+                return v.build("t2i", self._prompt(caps, rng), image=arr["codes"][i])
+            return v.build("i2t", self.enc(caps[0]), image=arr["codes"][i])
+        if task == "i2i":
+            (arr, meta), i = self.edits, int(rng.integers(len(self.edits[0]["src"])))
+            return v.build("i2i", self.enc(meta["instructions"][i]), image=arr["src"][i], image2=arr["tgt"][i])
+        if task in ("t2v", "v2t", "vcont"):
+            pool = self.vid_cap if task != "vcont" else np.arange(len(self.videos[0]["codes"]))
+            (arr, meta), i = self.videos, int(pool[rng.integers(len(pool))])
+            caps = meta["captions"][i]
+            if task == "t2v":
+                return v.build("t2v", self._prompt(caps, rng), video=arr["codes"][i])
+            if task == "v2t":
+                return v.build("v2t", self.enc(caps[0]), video=arr["codes"][i])
+            return v.build("vcont", video=arr["codes"][i])
+        if task == "chat_text":
+            r = self.chat[int(rng.integers(len(self.chat)))]
+            s = v.chat(self.enc(r["user"]), self.enc(r["assistant"]))
+            return Sample(s.tokens[: self.seq_len + 1], s.target_start)
+        if task == "chat_mm":
+            return self.chat_mm(str(rng.choice(self.chat_kinds)), rng)
+        raise KeyError(task)
+
+    def chat_mm(self, kind: str, rng: np.random.Generator) -> Sample:
+        """Instruction-formatted multimodal conversation built from the media caches."""
+        v, enc = self.vocab, self.enc
+        pick = lambda xs: str(rng.choice(xs))
+        if kind in ("t2i", "i2t"):
+            (arr, meta), i = self.images, int(self.img_cap[rng.integers(len(self.img_cap))])
+            cap, img = meta["captions"][i][0], arr["codes"][i]
+            if kind == "t2i":
+                return v.chat(enc(pick(CHAT_PROMPTS["t2i"]).format(c=pick(meta["captions"][i]))), enc(f"Here is an image of {cap}.") + v.image_tokens(img))
+            return v.chat(v.image_tokens(img) + enc(" " + pick(CHAT_PROMPTS["i2t"])), enc(f"The image shows {cap}."))
+        if kind == "i2i":
+            (arr, meta), i = self.edits, int(rng.integers(len(self.edits[0]["src"])))
+            instr = meta["instructions"][i]
+            return v.chat(v.image_tokens(arr["src"][i]) + enc(" " + instr[0].upper() + instr[1:] + "."), enc("Here is the edited image.") + v.image_tokens(arr["tgt"][i]))
+        pool = self.vid_cap if kind != "vcont" else np.arange(len(self.videos[0]["codes"]))
+        (arr, meta), i = self.videos, int(pool[rng.integers(len(pool))])
+        cap, vid = (meta["captions"][i] or [""])[0], arr["codes"][i]
+        if kind == "t2v":
+            return v.chat(enc(pick(CHAT_PROMPTS["t2v"]).format(c=pick(meta["captions"][i]))), enc(f"Here is a video of {cap}.") + v.video_tokens(vid))
+        if kind == "v2t":
+            return v.chat(v.video_tokens(vid) + enc(" " + pick(CHAT_PROMPTS["v2t"])), enc(f"The video shows {cap}."))
+        if kind == "vcont":
+            k = len(vid) // 2
+            return v.chat(v.video_tokens(vid[:k]) + enc(" " + pick(CHAT_PROMPTS["vcont"])), enc("Here is the continuation.") + v.video_tokens(vid[k:], first_index=k))
+        # i2v: the image is frame 0; the answer is the rest of the clip
+        return v.chat(v.image_tokens(vid[0]) + enc(" " + pick(CHAT_PROMPTS["i2v"]).format(c=cap)), enc("Here is the animation.") + v.video_tokens(vid[1:], first_index=1))
+
+
+def split_task(name: str) -> tuple[str, str]:
+    """'user/t2i' -> ('user', 't2i'); 't2i' -> ('', 't2i')  (source prefix selects the cache)."""
+    src, _, task = name.rpartition("/")
+    return src, task
+
+
+def available_tasks(sources: dict[str, TaskSources]) -> list[str]:
+    return [f"{k}/{t}" if k else t for k, src in sources.items() for t in src.available()]
+
+
+def compile_bucket_length(length: int, context: int) -> int:
+    """At most four fixed sequence sizes at context 1024; no content truncation."""
+    if not 0 < length <= context:
+        raise ValueError("sample length is outside the model context")
+    return min(context, max(128, 1 << (length - 1).bit_length()))
+
+
+class MixtureLoader:
+    """Yields (task, input_ids, labels) micro-batches; one task per micro-batch, ~`micro_tokens` tokens each.
+
+    `sources` maps a source name to its TaskSources ('' = built-in cache, 'user' = imported data);
+    mixture keys address them as 't2i' or 'user/t2i', so every source/task share is an explicit weight."""
+
+    def __init__(self, sources: dict[str, TaskSources] | TaskSources, weights: dict[str, float], micro_tokens: int, seed: int, prefetch: int = 6,
+                 compile_buckets: bool = False):
+        sources = sources if isinstance(sources, dict) else {"": sources}
+        missing = [t for t, w in weights.items() if w > 0 and t not in available_tasks(sources)]
+        if missing:
+            raise ValueError(f"mixture uses tasks with no cache: {missing}; available: {available_tasks(sources)}")
+        self.sources, self.seed, self.micro_tokens = sources, seed, micro_tokens
+        self.vocab = next(iter(sources.values())).vocab
+        self.tasks = [t for t, w in weights.items() if w > 0]
+        p = np.array([weights[t] for t in self.tasks], dtype=np.float64)
+        self.p = p / p.sum()
+        self.prefetch = prefetch
+        self.compile_buckets = compile_buckets
+
+    def batch(self, step: int, micro: int) -> tuple[str, torch.Tensor, torch.Tensor]:
+        rng = np.random.default_rng([self.seed, step, micro])
+        task = self.tasks[int(rng.choice(len(self.tasks), p=self.p))]
+        samples: list[Sample] = []
+        longest = 0
+        src, base = split_task(task)
+        while True:
+            s = self.sources[src].sample(base, rng)
+            proposed = max(longest, len(s.tokens) - 1 if self.compile_buckets else len(s.tokens))
+            padded = compile_bucket_length(proposed, self.sources[src].seq_len) if self.compile_buckets else proposed
+            if self.compile_buckets and padded > self.micro_tokens:
+                raise ValueError("compiled batch token budget is smaller than its padded context")
+            if samples and (len(samples) + 1) * padded > self.micro_tokens:
+                break
+            longest = proposed
+            samples.append(s)
+        if self.compile_buckets:
+            length = compile_bucket_length(longest, self.sources[src].seq_len)
+            # Empty rows have no target labels and do not change the supervised objective.
+            inp, lab = collate(samples, self.vocab.pad, target_length=length,
+                               batch_size=self.micro_tokens // length)
+        else:
+            inp, lab = collate(samples, self.vocab.pad)
+        return task, inp, lab
+
+    def iterate(self, start_step: int, grad_accum: int):
+        """Background-thread prefetching iterator of (step, micro, task, inp, lab)."""
+        q: queue.Queue = queue.Queue(maxsize=self.prefetch)
+        stop = threading.Event()
+
+        def work():
+            step = start_step
+            while not stop.is_set():
+                for m in range(grad_accum):
+                    item = (step, m, *self.batch(step, m))
+                    while not stop.is_set():
+                        try:
+                            q.put(item, timeout=0.5)
+                            break
+                        except queue.Full:
+                            continue
+                step += 1
+
+        th = threading.Thread(target=work, daemon=True)
+        th.start()
+        try:
+            while True:
+                yield q.get()
+        finally:
+            stop.set()
+
+
+def fixed_eval_batches(sources: dict[str, TaskSources] | TaskSources, tasks: list[str], micro_tokens: int, n_batches: int, seed: int = 999) -> dict[str, list]:
+    """Deterministic held-out batches per task (same every evaluation)."""
+    sources = sources if isinstance(sources, dict) else {"": sources}
+    out = {}
+    for ti, task in enumerate(tasks):
+        if task not in available_tasks(sources):
+            continue
+        ml = MixtureLoader(sources, {task: 1.0}, micro_tokens, seed)
+        out[task] = [ml.batch(10_000 + ti, i)[1:] for i in range(n_batches)]
+    return out

quasar/data/download.py

@@ -0,0 +1,180 @@
+"""Resumable HTTP acquisition with disk reserve, hashes and atomic progress snapshots.
+
+This module never extracts archives, evaluates remote code or initializes CUDA.
+Dataset payloads remain separate from licensed/validated training caches.
+"""
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import hashlib
+import json
+import os
+from pathlib import Path
+import shutil
+import threading
+import time
+
+import requests
+
+
+class AccessBlocked(RuntimeError): pass
+class DiskBudget(RuntimeError): pass
+class FileTooLarge(ValueError): pass
+
+
+def _comparable_path(path: Path) -> Path:
+    """Normalize Windows namespace aliases after resolving symlinks/junctions.
+
+    Non-strict realpath can retain the extended prefix when a parent directory
+    appears concurrently. Compare both spellings consistently; do not skip the
+    resolved-path containment check.
+    """
+    if os.name == "nt":
+        value = str(path)
+        if value.startswith("\\\\?\\UNC\\"):
+            return Path("\\\\" + value[8:])
+        if value.startswith("\\\\?\\") and len(value) >= 7 and value[5:7] == ":\\":
+            return Path(value[4:])
+    return path
+
+
+def exception_summary(exc: Exception, limit: int = 180) -> str:
+    lines = str(exc).splitlines()
+    return type(exc).__name__ + ": " + (lines[0][:limit] if lines else "no additional detail")
+
+
+class Downloads:
+    def __init__(self, root: str | Path, max_gib: float = 70, reserve_gib: float = 25):
+        self.root = Path(root).resolve()
+        self.root.mkdir(parents=True, exist_ok=True)
+        self.maximum, self.reserve = int(max_gib * 2**30), int(reserve_gib * 2**30)
+        self.lock = threading.RLock()
+        self.last_save = 0.0
+        self.bytes = sum(p.stat().st_size for p in self.root.rglob("*") if p.is_file())
+        self.state = {"pid": os.getpid(), "started": time.time(), "training_started": False,
+                      "max_gib": max_gib, "reserve_gib": reserve_gib, "sources": {}}
+        self.current: dict = {}
+
+    def safe_path(self, rel: str) -> Path:
+        p = (self.root / rel).resolve()
+        if not _comparable_path(p).is_relative_to(_comparable_path(self.root)):
+            raise ValueError("download path escapes acquisition directory")
+        return p
+
+    def save(self, force=False):
+        with self.lock:
+            if not force and time.time() - self.last_save < 3: return
+            self.last_save = time.time()
+            self.state |= {"updated": time.time(), "stored_bytes": self.bytes,
+                           "free_bytes": shutil.disk_usage(self.root).free, "current_files": dict(self.current)}
+            tmp = self.root / "download_status.json.tmp"
+            tmp.write_text(json.dumps(self.state, indent=2), encoding="utf-8")
+            # Windows readers briefly deny replacement while their file handle
+            # is open (for example, the live dashboard). Keep the last complete
+            # snapshot visible and retry this atomic swap for a bounded period.
+            for attempt in range(8):
+                try:
+                    os.replace(tmp, self.root / "download_status.json")
+                    break
+                except PermissionError:
+                    if attempt == 7:
+                        raise
+                    time.sleep(0.01 * 2**attempt)
+
+    def event(self, source, status, **kwargs):
+        with self.lock:
+            old = self.state["sources"].setdefault(source, {})
+            old.update({"status": status, "updated": time.time(), **kwargs})
+            self.save(force=True)
+        print(json.dumps({"source": source, "status": status, **kwargs}), flush=True)
+
+    def _claim(self, amount):
+        with self.lock:
+            if self.bytes + amount > self.maximum: raise DiskBudget(f"{self.maximum / 2**30:g} GiB acquisition cap reached")
+            if shutil.disk_usage(self.root).free - amount < self.reserve: raise DiskBudget("25 GiB free-space reserve reached")
+            self.bytes += amount
+
+    def _record(self, record):
+        with self.lock:
+            with (self.root / "files.jsonl").open("a", encoding="utf-8") as f:
+                f.write(json.dumps(record) + "\n")
+
+    def get(self, url: str, rel: str, expected_size: int | None = None, checksum: str | None = None,
+            headers: dict | None = None, max_file_bytes: int = 12 * 2**30, retries: int = 3) -> Path:
+        """Resume .part with validated Content-Range. HTTP credentials are never logged."""
+        path = self.safe_path(rel)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        if path.exists():
+            if expected_size is not None and path.stat().st_size != expected_size: raise ValueError(f"existing size mismatch: {rel}")
+            if checksum:
+                alg, value = checksum.split(":", 1)
+                if self.hash(path, alg) != value: raise ValueError(f"existing checksum mismatch: {rel}")
+            return path
+        part = path.with_name(path.name + ".part")
+        for attempt in range(retries):
+            try:
+                offset = part.stat().st_size if part.exists() else 0
+                h = {**(headers or {}), "Accept-Encoding": "identity"}
+                if offset: h["Range"] = f"bytes={offset}-"
+                with requests.get(url, headers=h, stream=True, timeout=(25, 90)) as response:
+                    if response.status_code in (401, 403): raise AccessBlocked(f"access approval required for {rel} (HTTP {response.status_code})")
+                    response.raise_for_status()
+                    if offset and response.status_code != 206:
+                        # Server ignored range. Restart this owned partial file, accounting for the old bytes.
+                        with self.lock: self.bytes -= offset
+                        offset = 0
+                    if response.status_code == 206 and not response.headers.get("Content-Range", "").startswith(f"bytes {offset}-"):
+                        raise ValueError("server returned a mismatched byte range")
+                    length = int(response.headers.get("Content-Length", "0"))
+                    if offset + length > max_file_bytes or (expected_size is not None and expected_size > max_file_bytes):
+                        raise FileTooLarge(f"file exceeds per-file limit: {rel}")
+                    if expected_size is not None and offset + length and length and expected_size != offset + length:
+                        raise ValueError(f"unexpected remote size for {rel}")
+                    with part.open("ab" if offset else "wb") as f:
+                        for chunk in response.iter_content(1024 * 1024):
+                            if not chunk: continue
+                            if f.tell() + len(chunk) > max_file_bytes: raise FileTooLarge("stream exceeded per-file limit")
+                            self._claim(len(chunk))
+                            try: f.write(chunk)
+                            except Exception:
+                                with self.lock: self.bytes -= len(chunk)
+                                raise
+                            with self.lock: self.current[rel] = {"bytes": f.tell(), "expected": expected_size or offset + length or None}
+                            self.save()
+                if expected_size is not None and part.stat().st_size != expected_size: raise ValueError(f"truncated file: {rel}")
+                if length and part.stat().st_size != offset + length: raise ValueError(f"truncated HTTP body: {rel}")
+                sha = self.hash(part)
+                if checksum:
+                    alg, value = checksum.split(":", 1)
+                    actual = sha if alg == "sha256" else self.hash(part, alg)
+                    if actual != value: raise ValueError(f"checksum mismatch: {rel}; partial retained for inspection")
+                os.replace(part, path)
+                self._record({"path": rel, "bytes": path.stat().st_size, "sha256": sha, "source_url": url,
+                              "publisher_checksum": checksum, "completed": time.time()})
+                with self.lock: self.current.pop(rel, None)
+                self.save(force=True)
+                return path
+            except (AccessBlocked, DiskBudget, FileTooLarge): raise
+            except Exception:
+                if attempt + 1 == retries: raise
+                time.sleep(min(2**attempt, 8))
+        raise AssertionError("unreachable")
+
+    @staticmethod
+    def hash(path, alg="sha256"):
+        h = hashlib.new(alg)
+        with Path(path).open("rb") as f:
+            for chunk in iter(lambda: f.read(4 * 1024 * 1024), b""): h.update(chunk)
+        return h.hexdigest()
+
+    def batch(self, jobs: list[dict], workers=3):
+        errors = []
+        with ThreadPoolExecutor(workers) as pool:
+            futures = {pool.submit(self.get, **job): job["rel"] for job in jobs}
+            for future in as_completed(futures):
+                try: future.result()
+                except Exception as e:
+                    errors.append({"path": futures[future], "error": exception_summary(e)})
+                    if isinstance(e, DiskBudget):
+                        for f in futures: f.cancel()
+        return errors

quasar/data/preprocessing.py

@@ -0,0 +1,250 @@
+"""Corpus download (pinned + licensed), text pretokenization, dataset validation and statistics.
+
+Every external corpus is listed in SOURCES with a pinned revision and its license, and the
+provenance record is written next to the downloaded files so caches can be audited later.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+
+SOURCES: dict[str, dict] = {
+    "tinystories": {
+        "repo_id": "roneneldan/TinyStories",
+        "repo_type": "dataset",
+        "revision": "f54c09fd23315a6f9c86f9dc80f725de7d8f9c64",
+        "files": [
+            "data/train-00000-of-00004-2d5a1467fff1081b.parquet",
+            "data/validation-00000-of-00001-869c898b519ad725.parquet",
+        ],
+        "license": "CDLA-Sharing-1.0",
+        "url": "https://huggingface.co/datasets/roneneldan/TinyStories",
+        "notes": "Synthetic children's stories (GPT-3.5/4 generated). Only 1 of 4 train shards is used.",
+    },
+    "dolly": {
+        "repo_id": "databricks/databricks-dolly-15k",
+        "repo_type": "dataset",
+        "revision": "bdd27f4d94b9c1f951818a7da7fd7aeea5dbff1a",
+        "files": ["databricks-dolly-15k.jsonl"],
+        "license": "CC-BY-SA-3.0",
+        "url": "https://huggingface.co/datasets/databricks/databricks-dolly-15k",
+        "notes": "Human-written instruction/response pairs; used for instruction tuning only.",
+    },
+}
+
+
+def sha256_file(path: str | Path) -> str:
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def download_sources(raw_dir: str | Path, names: list[str] | None = None) -> dict[str, list[Path]]:
+    """Download the pinned corpus files and write a provenance.json record."""
+    from huggingface_hub import hf_hub_download
+
+    raw_dir = Path(raw_dir)
+    out: dict[str, list[Path]] = {}
+    for name in names or list(SOURCES):
+        src = SOURCES[name]
+        dest = raw_dir / name
+        paths = [
+            Path(hf_hub_download(src["repo_id"], f, repo_type=src["repo_type"], revision=src["revision"], local_dir=dest))
+            for f in src["files"]
+        ]
+        record = {k: v for k, v in src.items()} | {"sha256": {p.name: sha256_file(p) for p in paths}}
+        (dest / "provenance.json").write_text(json.dumps(record, indent=2))
+        out[name] = paths
+    return out
+
+
+def load_tinystories(raw_dir: str | Path, split: str) -> list[str]:
+    import pyarrow.parquet as pq
+
+    files = sorted((Path(raw_dir) / "tinystories" / "data").glob(f"{split}-*.parquet"))
+    if not files:
+        raise FileNotFoundError(f"No TinyStories {split} parquet under {raw_dir}; run scripts/generate_dataset.py download")
+    texts: list[str] = []
+    for f in files:
+        texts += [t.strip() for t in pq.read_table(f, columns=["text"]).column("text").to_pylist() if t and t.strip()]
+    return texts
+
+
+def load_dolly(raw_dir: str | Path) -> list[dict]:
+    rows = []
+    with open(Path(raw_dir) / "dolly" / "databricks-dolly-15k.jsonl", encoding="utf-8") as f:
+        for line in f:
+            r = json.loads(line)
+            if r["instruction"].strip() and r["response"].strip():
+                rows.append(r)
+    return rows
+
+
+def tokenize_stream(texts: list[str], encode_batch, prefix: list[int], eos_id: int, out_path: str | Path) -> int:
+    """Concatenate documents as <prefix> doc [EOS] ... into one uint16 token stream (for packing)."""
+    out_path = Path(out_path)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    n = 0
+    with open(out_path, "wb") as f:
+        for i in range(0, len(texts), 4096):
+            arrs = [np.array([*prefix, *ids, eos_id], dtype=np.uint16) for ids in encode_batch(texts[i : i + 4096])]
+            buf = np.concatenate(arrs)
+            buf.tofile(f)
+            n += len(buf)
+    return n
+
+
+def validate_sequence(tokens: np.ndarray, vocab) -> None:
+    """Raise ValueError if a token sequence violates the unified-vocabulary structure rules."""
+    vocab.check_structure(tokens)
+
+
+def dataset_stats(cache_dir: str | Path) -> dict:
+    """Summarise every cache file (token counts / sample counts) for docs and the dashboard."""
+    cache_dir = Path(cache_dir)
+    stats: dict[str, dict] = {}
+    for p in sorted(cache_dir.glob("*.bin")):
+        stats[p.name] = {"tokens": p.stat().st_size // 2}
+    for p in sorted(cache_dir.glob("*.npz")):
+        with np.load(p) as z:
+            stats[p.name] = {k: list(z[k].shape) for k in z.files}
+    return stats
+
+
+# ----------------------------------------------------------------------------- media loading
+def prepare_image(img, res: int = 64) -> np.ndarray:
+    """PIL / array / path -> uint8 [res,res,3] (centre square crop, area resize)."""
+    if isinstance(img, (str, Path)):
+        img = Image.open(img)
+    if isinstance(img, np.ndarray):
+        img = Image.fromarray(img.astype(np.uint8))
+    img = img.convert("RGB")
+    w, h = img.size
+    s = min(w, h)
+    img = img.crop(((w - s) // 2, (h - s) // 2, (w - s) // 2 + s, (h - s) // 2 + s))
+    return np.asarray(img.resize((res, res), Image.BOX if s >= res else Image.BICUBIC))
+
+
+def load_video(src, max_frames: int = 8, fps: float = 4.0, res: int = 64) -> np.ndarray:
+    """Video file (mp4/gif/webm...) or frame list -> uint8 [T,res,res,3], resampled to ~`fps`."""
+    if not isinstance(src, (str, Path)):
+        return np.stack([prepare_image(f, res) for f in list(src)[:max_frames]])
+    import imageio_ffmpeg
+
+    reader = imageio_ffmpeg.read_frames(str(src))
+    meta = next(reader)
+    w, h = meta["size"]
+    stride = max(1, round(float(meta.get("fps") or fps) / fps))
+    frames = []
+    for i, raw in enumerate(reader):
+        if i % stride == 0:
+            frames.append(prepare_image(np.frombuffer(raw, np.uint8).reshape(h, w, 3), res))
+        if len(frames) == max_frames:
+            break
+    reader.close()
+    if not frames:
+        raise ValueError(f"could not decode any frames from {src}")
+    return np.stack(frames)
+
+
+def video_windows(path: str | Path, n_frames: int = 8, fps: float = 4.0, max_windows: int = 4, res: int = 64) -> list[np.ndarray]:
+    """Cut a video file into up to `max_windows` consecutive `n_frames`-frame clips at ~`fps`.
+    Clips too short for one window are sampled evenly at native frame rate; < n_frames frames -> []."""
+    clip = load_video(path, max_frames=n_frames * max_windows, fps=fps, res=res)
+    if len(clip) >= n_frames:
+        return [clip[i : i + n_frames] for i in range(0, len(clip) - n_frames + 1, n_frames)]
+    native = load_video(path, max_frames=10_000, fps=1e9, res=res)  # every frame
+    if len(native) < n_frames:
+        return []
+    return [native[np.linspace(0, len(native) - 1, n_frames).round().astype(int)]]
+
+
+# ----------------------------------------------------------------------------- user datasets
+IMAGE_EXT = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
+VIDEO_EXT = {".mp4", ".mov", ".webm", ".mkv", ".avi", ".gif"}
+USER_KEYS = ("user", "instruction", "prompt", "question", "input")
+ASSISTANT_KEYS = ("assistant", "response", "output", "answer", "completion")
+
+
+def scan_user_dataset(root: str | Path) -> dict[str, list]:
+    """Collect samples from a user-provided folder or .jsonl file, routing records by their keys.
+
+    JSONL records (paths relative to the .jsonl file):
+        {"text": ...}                                             -> text corpus
+        {"image": p, "caption": ...}                              -> image-caption (t2i + i2t)
+        {"image": p, "target_image": p, "instruction": ...}       -> image editing (i2i)
+        {"video": p, "caption": ...}  (caption optional)          -> video-caption (t2v, v2t) / continuation
+        {"instruction"|"prompt"|"user": ..., "response"|"output"|"assistant": ...}  -> instruction tuning
+        {"messages": [{"role": "user", ...}, {"role": "assistant", ...}, ...]}       -> instruction tuning
+    Folders: *.txt -> text documents; image/video files with a same-name .txt sidecar -> captioned
+    media (images without a caption are skipped; videos without one are kept for continuation).
+    """
+    root = Path(root)
+    out: dict[str, list] = {"texts": [], "images": [], "edits": [], "videos": [], "chats": [], "skipped": []}
+    files = [root] if root.is_file() else sorted(p for p in root.rglob("*") if p.is_file())
+    sidecar = {p.with_suffix(".txt") for p in files if p.suffix.lower() in IMAGE_EXT | VIDEO_EXT}
+    for p in files:
+        if p.suffix.lower() == ".jsonl":
+            for ln, line in enumerate(p.read_text(encoding="utf-8", errors="replace").splitlines()):
+                if not line.strip():
+                    continue
+                try:
+                    _route(json.loads(line), p.parent, out)
+                except (ValueError, KeyError, TypeError) as e:
+                    out["skipped"].append(f"{p.name}:{ln + 1}: {e}")
+    # media referenced by a manifest record is not also scanned as a loose file
+    referenced = {Path(x).resolve() for k in ("images", "videos") for x, _ in out[k]} | {Path(x).resolve() for e in out["edits"] for x in e[:2]}
+    for p in files:
+        ext = p.suffix.lower()
+        if ext == ".jsonl" or p.resolve() in referenced:
+            continue
+        if ext == ".txt" and p not in sidecar:
+            t = p.read_text(encoding="utf-8", errors="replace").strip()
+            if t:
+                out["texts"].append(t)
+        elif ext in IMAGE_EXT | VIDEO_EXT:
+            cap_p = p.with_suffix(".txt")
+            cap = cap_p.read_text(encoding="utf-8", errors="replace").strip() if cap_p.exists() else ""
+            if ext in VIDEO_EXT:
+                out["videos"].append((p, cap))
+            elif cap:
+                out["images"].append((p, cap))
+            else:
+                out["skipped"].append(f"{p.name}: image without caption")
+    return out
+
+
+def _route(r: dict, base: Path, out: dict) -> None:
+    get = lambda keys: next((str(r[k]).strip() for k in keys if r.get(k)), "")
+    if "messages" in r:
+        msgs = r["messages"]
+        for a, b in zip(msgs, msgs[1:]):
+            if a.get("role") == "user" and b.get("role") == "assistant" and a.get("content") and b.get("content"):
+                out["chats"].append((str(a["content"]).strip(), str(b["content"]).strip()))
+    elif "video" in r:
+        out["videos"].append((base / r["video"], get(("caption", "text", "prompt"))))
+    elif "image" in r and ("target_image" in r or "edited_image" in r):
+        out["edits"].append((base / r["image"], base / (r.get("target_image") or r["edited_image"]), get(("instruction", "edit", "prompt"))))
+    elif "image" in r:
+        cap = get(("caption", "text", "prompt"))
+        if not cap:
+            raise ValueError("image record without caption")
+        out["images"].append((base / r["image"], cap))
+    elif get(USER_KEYS) and get(ASSISTANT_KEYS):
+        user = get(USER_KEYS) + (("\n\n" + str(r["context"]).strip()) if r.get("context") else "")
+        out["chats"].append((user, get(ASSISTANT_KEYS)))
+    elif get(("text", "content", "story")):
+        out["texts"].append(get(("text", "content", "story")))
+    else:
+        raise ValueError(f"unrecognised record keys {sorted(r)}")
+
+
+def split_of(key: str, val_pct: float) -> str:
+    """Deterministic train/val assignment from a content hash (stable across runs and machines)."""
+    return "val" if int(hashlib.sha1(key.encode("utf-8", "replace")).hexdigest()[:8], 16) % 10_000 < val_pct * 100 else "train"

quasar/data/synthetic.py

@@ -0,0 +1,383 @@
+"""Procedural images and videos with exact ground truth, plus an analyzer that measures them.
+
+Everything here is generated locally (no external data, no license questions).  The same
+attribute vocabulary is used for rendering, captioning and for *measuring* model outputs, so
+prompt-following and motion can be scored quantitatively instead of by eyeballing samples.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field, replace
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+RES = 64
+SS = 4  # supersampling factor for anti-aliased edges
+
+COLORS: dict[str, tuple[int, int, int]] = {
+    "red": (220, 40, 40), "green": (40, 180, 60), "blue": (50, 90, 235), "yellow": (240, 220, 50),
+    "purple": (150, 60, 200), "orange": (245, 140, 30), "white": (240, 240, 240), "cyan": (40, 210, 220),
+}
+BACKGROUNDS: dict[str, tuple[int, int, int]] = {
+    "black": (15, 15, 15), "gray": (105, 105, 105), "navy": (20, 30, 85), "brown": (85, 55, 30),
+}
+SHAPES = ("circle", "square", "triangle")
+SIZES = {"small": 5.0, "medium": 9.0, "large": 13.0}  # half-extent in pixels
+GRID = {  # name -> (col, row) in a 3x3 layout
+    "top left": (0, 0), "top": (1, 0), "top right": (2, 0),
+    "left": (0, 1), "center": (1, 1), "right": (2, 1),
+    "bottom left": (0, 2), "bottom": (1, 2), "bottom right": (2, 2),
+}
+CELL_CENTERS = (16.0, 32.0, 48.0)
+POS_PHRASE = {
+    "top left": "in the top left", "top": "at the top", "top right": "in the top right",
+    "left": "on the left", "center": "in the center", "right": "on the right",
+    "bottom left": "in the bottom left", "bottom": "at the bottom", "bottom right": "in the bottom right",
+}
+# (shape, color) pairs never shown during training: compositional-generalisation probe.
+HOLDOUT_COMBOS = {("triangle", "purple"), ("circle", "orange"), ("square", "cyan")}
+
+MOTIONS = ("left", "right", "up", "down", "bounce", "cw", "ccw", "grow", "shrink", "still", "collide")
+MOTION_PHRASE = {
+    "left": "moving left", "right": "moving right", "up": "moving up", "down": "moving down",
+    "cw": "rotating clockwise", "ccw": "rotating counterclockwise", "grow": "growing",
+    "shrink": "shrinking", "still": "standing still",
+}
+DIR_VEC = {"left": (-1, 0), "right": (1, 0), "up": (0, -1), "down": (0, 1)}
+SHAPE_WORDS = {"circle": ["circle", "ball", "dot"], "square": ["square", "box"], "triangle": ["triangle"]}
+
+
+@dataclass
+class Obj:
+    shape: str
+    color: str
+    half: float
+    cx: float
+    cy: float
+    angle: float = 0.0  # degrees, clockwise on screen
+
+    @property
+    def size_name(self) -> str:
+        return min(SIZES, key=lambda k: abs(SIZES[k] - self.half))
+
+
+@dataclass
+class ImageSpec:
+    obj: Obj
+    bg: str
+    pos: str
+
+
+@dataclass
+class VideoSpec:
+    frames: list[list[Obj]]  # per frame, list of objects
+    bg: str
+    motion: str
+    hdir: str = ""  # horizontal drift for "bounce"
+    meta: dict = field(default_factory=dict)
+
+
+# ----------------------------------------------------------------------------- rendering
+def _polygon(o: Obj) -> list[tuple[float, float]]:
+    h = o.half
+    if o.shape == "square":
+        pts = [(-h, -h), (h, -h), (h, h), (-h, h)]
+    else:  # triangle: apex up, base at bottom, fills half of its bounding square
+        pts = [(0.0, -h), (h, h), (-h, h)]
+    a = math.radians(o.angle)
+    c, s = math.cos(a), math.sin(a)
+    # y points down, so this standard rotation appears clockwise on screen for positive angles
+    return [((x * c - y * s + o.cx) * SS, (x * s + y * c + o.cy) * SS) for x, y in pts]
+
+
+def render(objs: list[Obj], bg: str) -> np.ndarray:
+    """Render objects onto a background; returns uint8 [64, 64, 3]."""
+    im = Image.new("RGB", (RES * SS, RES * SS), BACKGROUNDS[bg])
+    d = ImageDraw.Draw(im)
+    for o in objs:
+        if o.shape == "circle":
+            d.ellipse([(o.cx - o.half) * SS, (o.cy - o.half) * SS, (o.cx + o.half) * SS, (o.cy + o.half) * SS], fill=COLORS[o.color])
+        else:
+            d.polygon(_polygon(o), fill=COLORS[o.color])
+    return np.asarray(im.resize((RES, RES), Image.BOX))
+
+
+def render_video(spec: VideoSpec) -> np.ndarray:
+    return np.stack([render(objs, spec.bg) for objs in spec.frames])
+
+
+# ----------------------------------------------------------------------------- sampling
+def _combo_ok(shape: str, color: str, holdout: str) -> bool:
+    """holdout: 'exclude' (training), 'only' (generalisation probe) or 'any'."""
+    held = (shape, color) in HOLDOUT_COMBOS
+    return holdout == "any" or (held if holdout == "only" else not held)
+
+
+def _pick_shape_color(rng: np.random.Generator, holdout: str, shapes=SHAPES) -> tuple[str, str]:
+    while True:
+        shape, color = str(rng.choice(shapes)), str(rng.choice(list(COLORS)))
+        if _combo_ok(shape, color, holdout):
+            return shape, color
+
+
+def sample_image_spec(rng: np.random.Generator, holdout: str = "exclude") -> ImageSpec:
+    shape, color = _pick_shape_color(rng, holdout)
+    size = str(rng.choice(list(SIZES)))
+    pos = str(rng.choice(list(GRID)))
+    return _place(shape, color, size, pos, str(rng.choice(list(BACKGROUNDS))), rng)
+
+
+def _place(shape: str, color: str, size: str, pos: str, bg: str, rng: np.random.Generator) -> ImageSpec:
+    half = SIZES[size]
+    col, row = GRID[pos]
+    j = lambda: float(rng.uniform(-3, 3))
+    cx = float(np.clip(CELL_CENTERS[col] + j(), half + 1, RES - half - 1))
+    cy = float(np.clip(CELL_CENTERS[row] + j(), half + 1, RES - half - 1))
+    return ImageSpec(Obj(shape, color, half, cx, cy), bg, pos)
+
+
+def make_image_spec(shape: str, color: str, size: str, pos: str, bg: str, seed: int = 0) -> ImageSpec:
+    return _place(shape, color, size, pos, bg, np.random.default_rng(seed))
+
+
+EDITS = ("color", "move", "size", "background", "shape")
+
+
+def sample_edit(rng: np.random.Generator, src: ImageSpec) -> tuple[str, ImageSpec, str]:
+    """Pick an edit; returns (edit_type, target_spec, canonical instruction)."""
+    o = src.obj
+    while True:
+        kind = str(rng.choice(EDITS))
+        if kind == "color":
+            c = str(rng.choice([c for c in COLORS if c != o.color]))
+            if _combo_ok(o.shape, c, "exclude"):
+                return kind, replace(src, obj=replace(o, color=c)), f"change the color to {c}"
+        elif kind == "shape":
+            s = str(rng.choice([s for s in SHAPES if s != o.shape]))
+            if _combo_ok(s, o.color, "exclude"):
+                return kind, replace(src, obj=replace(o, shape=s)), f"turn it into a {s}"
+        elif kind == "size":
+            z = str(rng.choice([z for z in SIZES if z != o.size_name]))
+            t = _place(o.shape, o.color, z, src.pos, src.bg, rng)
+            return kind, replace(t, obj=replace(t.obj, cx=float(np.clip(o.cx, SIZES[z] + 1, RES - SIZES[z] - 1)),
+                                                cy=float(np.clip(o.cy, SIZES[z] + 1, RES - SIZES[z] - 1)))), f"make it {z}"
+        elif kind == "move":
+            p = str(rng.choice([p for p in GRID if p != src.pos]))
+            return kind, _place(o.shape, o.color, o.size_name, p, src.bg, rng), f"move it to the {p}"
+        else:
+            b = str(rng.choice([b for b in BACKGROUNDS if b != src.bg]))
+            return kind, replace(src, bg=b), f"change the background to {b}"
+
+
+def sample_video_spec(rng: np.random.Generator, n_frames: int = 8, holdout: str = "exclude", motion: str | None = None) -> VideoSpec:
+    motion = motion or str(rng.choice(MOTIONS))
+    bg = str(rng.choice(list(BACKGROUNDS)))
+    T = n_frames
+    if motion == "collide":
+        # exclude: neither ball is a held-out combo; only: at least one is (only orange circles are held out)
+        ok = lambda a, b: holdout == "any" or (("circle", a) in HOLDOUT_COMBOS or ("circle", b) in HOLDOUT_COMBOS) == (holdout == "only")
+        c1, c2 = rng.choice([c for c in COLORS], size=2, replace=False)
+        while not ok(str(c1), str(c2)):
+            c1, c2 = rng.choice([c for c in COLORS], size=2, replace=False)
+        half, v = 6.0, float(rng.uniform(3, 4.5))
+        y = float(rng.uniform(16, 48))
+        x1, x2 = float(rng.uniform(8, 14)), float(rng.uniform(50, 56))
+        frames = []
+        v1, v2 = v, -v
+        for _ in range(T):
+            frames.append([Obj("circle", str(c1), half, x1, y), Obj("circle", str(c2), half, x2, y)])
+            x1, x2 = x1 + v1, x2 + v2
+            if x2 - x1 <= 2 * half:  # elastic collision: swap velocities
+                v1, v2 = v2, v1
+                x1, x2 = x1 + 2 * v1, x2 + 2 * v2
+        return VideoSpec(frames, bg, motion, meta={"colors": [str(c1), str(c2)]})
+
+    shapes = ("circle",) if motion == "bounce" else ("square", "triangle") if motion in ("cw", "ccw") else SHAPES
+    shape, color = _pick_shape_color(rng, holdout, shapes)
+    size = str(rng.choice(["small", "medium"] if motion in ("left", "right", "up", "down", "bounce") else list(SIZES)))
+    half = SIZES[size]
+    lo, hi = half + 1, RES - half - 1
+    frames: list[list[Obj]] = []
+    hdir = ""
+    if motion in DIR_VEC:
+        dx, dy = DIR_VEC[motion]
+        v = float(rng.uniform(3, 5))
+        travel = v * (T - 1)
+        start = lambda d: float(rng.uniform(lo, hi - travel)) if d > 0 else float(rng.uniform(lo + travel, hi)) if d < 0 else float(rng.uniform(lo, hi))
+        x, y = start(dx), start(dy)
+        frames = [[Obj(shape, color, half, x + dx * v * t, y + dy * v * t)] for t in range(T)]
+    elif motion == "bounce":
+        vx = float(rng.uniform(1.5, 3)) * (1 if rng.random() < 0.5 else -1)
+        hdir = "right" if vx > 0 else "left"
+        x = float(rng.uniform(lo, hi - 7 * abs(vx))) if vx > 0 else float(rng.uniform(lo + 7 * abs(vx), hi))
+        y, vy, g = float(rng.uniform(lo, 26)), 0.0, float(rng.uniform(2.5, 4.0))
+        for _ in range(T):
+            frames.append([Obj(shape, color, half, x, y)])
+            for _ in range(4):  # substeps for a clean bounce
+                vy += g / 4
+                y += vy / 4
+                if y > hi:
+                    y, vy = 2 * hi - y, -0.9 * vy
+            x += vx
+    elif motion in ("cw", "ccw"):
+        step = float(rng.uniform(12, 22)) * (1 if motion == "cw" else -1)
+        x, y, a0 = float(rng.uniform(lo, hi)), float(rng.uniform(lo, hi)), float(rng.uniform(0, 90))
+        frames = [[Obj(shape, color, half, x, y, a0 + step * t)] for t in range(T)]
+    elif motion in ("grow", "shrink"):
+        a, b = (4.0, 14.0) if motion == "grow" else (14.0, 4.0)
+        x, y = float(rng.uniform(15, 49)), float(rng.uniform(15, 49))
+        frames = [[Obj(shape, color, a + (b - a) * t / (T - 1), x, y)] for t in range(T)]
+    else:  # still
+        x, y = float(rng.uniform(lo, hi)), float(rng.uniform(lo, hi))
+        frames = [[Obj(shape, color, half, x, y)] for _ in range(T)]
+    return VideoSpec(frames, bg, motion, hdir)
+
+
+def render_clutter(rng: np.random.Generator) -> np.ndarray:
+    """Tokenizer-only augmentation: 1-6 shapes in arbitrary RGB colours on a random gradient.
+
+    Broadens what the visual codec can represent beyond the 8-colour palette (e.g. user uploads).
+    """
+    c0, c1 = rng.integers(0, 256, 3), rng.integers(0, 256, 3)
+    t = np.linspace(0, 1, RES * SS)[:, None, None] if rng.random() < 0.5 else np.linspace(0, 1, RES * SS)[None, :, None]
+    base = (c0 * (1 - t) + c1 * t).astype(np.uint8)
+    im = Image.fromarray(np.broadcast_to(base, (RES * SS, RES * SS, 3)).copy())
+    d = ImageDraw.Draw(im)
+    for _ in range(int(rng.integers(1, 7))):
+        o = Obj(str(rng.choice(SHAPES)), "red", float(rng.uniform(3, 20)), float(rng.uniform(0, RES)), float(rng.uniform(0, RES)), float(rng.uniform(0, 360)))
+        fill = tuple(int(v) for v in rng.integers(0, 256, 3))
+        if o.shape == "circle":
+            d.ellipse([(o.cx - o.half) * SS, (o.cy - o.half) * SS, (o.cx + o.half) * SS, (o.cy + o.half) * SS], fill=fill)
+        else:
+            d.polygon(_polygon(o), fill=fill)
+    return np.asarray(im.resize((RES, RES), Image.BOX))
+
+
+# ----------------------------------------------------------------------------- captions
+def caption_image(spec: ImageSpec, rng: np.random.Generator | None = None) -> str:
+    """Canonical caption when rng is None (used as a target); otherwise a random paraphrase (prompt)."""
+    o = spec.obj
+    if rng is None:
+        return f"a {o.size_name} {o.color} {o.shape} {POS_PHRASE[spec.pos]} on a {spec.bg} background"
+    w = str(rng.choice(SHAPE_WORDS[o.shape]))
+    templates = [
+        f"a {o.size_name} {o.color} {w} {POS_PHRASE[spec.pos]} on a {spec.bg} background",
+        f"a {o.color} {w} {POS_PHRASE[spec.pos]} on a {spec.bg} background, {o.size_name}",
+        f"{spec.bg} background with a {o.size_name} {o.color} {w} {POS_PHRASE[spec.pos]}",
+        f"an image of a {o.size_name} {o.color} {w} {POS_PHRASE[spec.pos]}, {spec.bg} background",
+        f"a {o.size_name} {o.color} {w} {POS_PHRASE[spec.pos]} on a {spec.bg} background",
+    ]
+    return str(rng.choice(templates))
+
+
+def caption_video(spec: VideoSpec, rng: np.random.Generator | None = None) -> str:
+    if spec.motion == "collide":
+        c1, c2 = spec.meta["colors"]
+        return f"a {c1} circle and a {c2} circle colliding on a {spec.bg} background"
+    o = spec.frames[0][0]
+    size = "medium" if spec.motion in ("grow", "shrink") else o.size_name
+    if spec.motion == "bounce":
+        name = "ball" if rng is None else str(rng.choice(["ball", "circle"]))
+        what = f"a {size} {o.color} {name} bouncing to the {spec.hdir}"
+    else:
+        name = o.shape if rng is None else str(rng.choice(SHAPE_WORDS[o.shape]))
+        what = f"a {size} {o.color} {name} {MOTION_PHRASE[spec.motion]}"
+        if spec.motion in ("grow", "shrink"):
+            what = f"a {o.color} {name} {MOTION_PHRASE[spec.motion]}"
+    if rng is None or rng.random() < 0.6:
+        return f"{what} on a {spec.bg} background"
+    return str(rng.choice([f"{what}", f"video of {what} on a {spec.bg} background", f"{spec.bg} background, {what}"]))
+
+
+def parse_caption(text: str) -> dict:
+    """Extract attribute words from (generated) caption text; missing attributes map to None."""
+    t = " " + text.lower().replace(",", " ") + " "
+    find = lambda opts: next((o for o in sorted(opts, key=len, reverse=True) if f" {o} " in t), None)
+    shape = find(["circle", "square", "triangle", "ball", "box", "dot"])
+    shape = {"ball": "circle", "dot": "circle", "box": "square"}.get(shape, shape)
+    pos = next((p for p, ph in sorted(POS_PHRASE.items(), key=lambda kv: -len(kv[1])) if ph in t), None)
+    bg = next((b for b in BACKGROUNDS if f"{b} background" in t), None)
+    motion = next((m for m, ph in MOTION_PHRASE.items() if ph in t), None)
+    if " bouncing" in t:
+        motion = "bounce"
+    if " colliding" in t:
+        motion = "collide"
+    color = next((w for w in t.split() if w in COLORS), None)  # fg and bg colour names never overlap
+    return {"shape": shape, "color": color, "size": find(list(SIZES)), "pos": pos, "bg": bg, "motion": motion}
+
+
+# ----------------------------------------------------------------------------- analysis
+def _nearest(rgb: np.ndarray, table: dict[str, tuple[int, int, int]]) -> str:
+    names = list(table)
+    d = ((np.array([table[n] for n in names], dtype=np.float32) - rgb.astype(np.float32)) ** 2).sum(-1)
+    return names[int(d.argmin())]
+
+
+def analyze_image(img: np.ndarray, min_area: int = 12) -> dict:
+    """Measure background, dominant object colour/shape/size/position of a uint8 [H,W,3] image."""
+    img = np.asarray(img)
+    border = np.concatenate([img[0], img[-1], img[:, 0], img[:, -1]]).astype(np.float32)
+    bg_rgb = np.median(border, axis=0)
+    bg = _nearest(bg_rgb, BACKGROUNDS)
+    dist = np.sqrt(((img.astype(np.float32) - bg_rgb) ** 2).sum(-1))
+    mask = dist > 70
+    area = int(mask.sum())
+    out = {"bg": bg, "area": area, "color": None, "shape": None, "size": None, "pos": None, "cx": None, "cy": None}
+    if area < min_area:
+        return out
+    ys, xs = np.nonzero(mask)
+    w, h = xs.max() - xs.min() + 1, ys.max() - ys.min() + 1
+    cx, cy = (xs.min() + xs.max()) / 2.0, (ys.min() + ys.max()) / 2.0  # bbox centre (triangle mass centroid sits low)
+    fill = area / float(w * h)
+    out.update(
+        color=_nearest(np.median(img[mask], axis=0), COLORS),
+        shape="square" if fill > 0.88 else "circle" if fill > 0.69 else "triangle",
+        size=min(SIZES, key=lambda k: abs(SIZES[k] - max(w, h) / 2)),
+        cx=cx, cy=cy,
+        pos=next(p for p, (c, r) in GRID.items() if c == min(2, int(cx // (RES / 3))) and r == min(2, int(cy // (RES / 3)))),
+    )
+    return out
+
+
+def analyze_video(frames: np.ndarray) -> dict:
+    """Per-frame analysis + trajectory summary for single-object clips."""
+    per = [analyze_image(f) for f in frames]
+    pts = [(p["cx"], p["cy"]) for p in per if p["cx"] is not None]
+    diffs = [float(np.abs(frames[i + 1].astype(np.float32) - frames[i].astype(np.float32)).mean()) for i in range(len(frames) - 1)]
+    out = {"frames": per, "detected": len(pts), "frame_diffs": diffs}
+    if len(pts) >= 2:
+        out["dx"], out["dy"] = pts[-1][0] - pts[0][0], pts[-1][1] - pts[0][1]
+        areas = [p["area"] for p in per if p["cx"] is not None]
+        out["area_ratio"] = areas[-1] / max(1, areas[0])
+    return out
+
+
+def motion_direction(dx: float, dy: float, min_travel: float = 6.0) -> str:
+    """Classify a displacement as left/right/up/down/still."""
+    if max(abs(dx), abs(dy)) < min_travel:
+        return "still"
+    if abs(dx) >= abs(dy):
+        return "right" if dx > 0 else "left"
+    return "down" if dy > 0 else "up"
+
+
+def _self_check() -> None:
+    rng = np.random.default_rng(0)
+    ok = 0
+    for _ in range(200):
+        s = sample_image_spec(rng)
+        a = analyze_image(render([s.obj], s.bg))
+        ok += (a["color"], a["shape"], a["pos"], a["bg"], a["size"]) == (s.obj.color, s.obj.shape, s.pos, s.bg, s.obj.size_name)
+    print(f"analyzer exact-match on ground truth renders: {ok}/200")
+    p = parse_caption(caption_image(s))
+    assert (p["shape"], p["color"], p["pos"], p["bg"], p["size"]) == (s.obj.shape, s.obj.color, s.pos, s.bg, s.obj.size_name), p
+    for m in ("left", "right", "up", "down"):
+        v = sample_video_spec(rng, motion=m)
+        a = analyze_video(render_video(v))
+        assert motion_direction(a["dx"], a["dy"]) == m, (m, a["dx"], a["dy"])
+        assert parse_caption(caption_video(v))["motion"] == m
+
+
+if __name__ == "__main__":
+    _self_check()

quasar/data/text_preparation.py

@@ -0,0 +1,514 @@
+"""Bounded offline preparation of the initial real text/chat/Python corpus.
+
+SQLite spools documents and deduplication keys instead of retaining a corpus in
+RAM. No downloads, archive extraction, GPU calls, or optimizer steps occur here.
+"""
+from __future__ import annotations
+
+from collections import Counter
+from copy import deepcopy
+import hashlib
+import json
+import os
+from pathlib import Path
+import shutil
+import sqlite3
+import time
+from urllib.parse import urlsplit, urlunsplit
+
+import numpy as np
+import pyarrow.parquet as pq
+import yaml
+
+from .adapters import normalize
+from .curated import canonical, digest, verify_record
+from .preprocessing import sha256_file
+from ..tokenizers.text import TextTokenizer
+from ..tokenizers.universal import UniversalVocab
+
+
+def atomic_json(path: Path, value: dict) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_name(path.name + ".tmp")
+    tmp.write_text(json.dumps(value, indent=2, ensure_ascii=False), encoding="utf-8")
+    for attempt in range(8):
+        try:
+            os.replace(tmp, path)
+            return
+        except PermissionError:
+            if attempt == 7:
+                raise
+            time.sleep(0.01 * 2**attempt)
+
+
+def normalized_hash(text: str) -> str:
+    return hashlib.sha256(" ".join(text.split()).encode("utf-8")).hexdigest()
+
+
+def document_group(url: str, text: str) -> str:
+    """Group HTTP/HTTPS aliases and fragments of the same document together."""
+    if not url:
+        return normalized_hash(text)
+    p = urlsplit(url)
+    return urlunsplit(("https", p.netloc.lower(), p.path.rstrip("/"), p.query, ""))
+
+
+def parquet_rows(path: Path):
+    for batch in pq.ParquetFile(path).iter_batches(batch_size=128):
+        yield from batch.to_pylist()
+
+
+def pinned_source_sha(path: Path, root: Path, dataset: str, revision: str, ledger: dict) -> str:
+    """Starter downloads retain pinned LFS hashes in HF's local metadata file."""
+    rel = path.relative_to(root).as_posix()
+    if rel in ledger:
+        expected = ledger[rel]["sha256"]
+    else:
+        relative = path.relative_to(root / dataset)
+        metadata = root / dataset / ".cache/huggingface/download" / (str(relative) + ".metadata")
+        lines = metadata.read_text(encoding="utf-8").splitlines()
+        if len(lines) < 2 or lines[0] != revision:
+            raise ValueError(f"missing or mismatched pinned acquisition metadata: {rel}")
+        expected = lines[1]
+    if len(expected) != 64 or any(c not in "0123456789abcdef" for c in expected):
+        raise ValueError(f"source has no pinned SHA256: {rel}")
+    return expected
+
+
+def text_windows(text: str, tok: TextTokenizer, limit: int):
+    """Segment at Unicode character/whitespace boundaries; never cut byte tokens."""
+    remaining = text.strip()
+    while remaining:
+        encoded = tok.tok.encode(remaining)
+        if len(encoded.ids) <= limit:
+            yield remaining
+            return
+        cut = encoded.offsets[limit - 1][1]
+        # Prefer a nearby complete word while retaining at least ~75% capacity.
+        lower = encoded.offsets[max(0, limit * 3 // 4 - 1)][1]
+        boundary = max(remaining.rfind(" ", lower, cut), remaining.rfind("\n", lower, cut))
+        if boundary > lower:
+            cut = boundary
+        while cut > 0 and len(tok.encode(remaining[:cut])) > limit:
+            cut -= 1
+        if cut <= 0:
+            raise ValueError("cannot segment text within context")
+        yield remaining[:cut].strip()
+        remaining = remaining[cut:].strip()
+
+
+def text_sample(record: dict, tok: TextTokenizer, vocab: UniversalVocab, context: int):
+    """Exact UniversalVocab text packet layout, validated in bulk with NumPy."""
+    s = vocab.special
+    ids = [s("bos"), s("task:" + record["task"]), s("input")]
+    for segment in record["inputs"]:
+        ids.extend([s("text"), *tok.encode(segment["text"]), s("/text")])
+    ids.append(s("output"))
+    target = len(ids)
+    for segment in record["outputs"]:
+        ids.extend([s("text"), *tok.encode(segment["text"]), s("/text")])
+    ids.append(s("eos"))
+    tokens = np.array(ids, dtype="<u2")
+    validate_text_sample(tokens, target, record["task"], vocab, context)
+    return tokens, target
+
+
+def validate_text_sample(tokens: np.ndarray, target: int, task: str, v: UniversalVocab, context: int) -> None:
+    """Strict text-only subset of the universal grammar; reserved IDs are rejected."""
+    s = v.special
+    if task not in ("text", "chat_text", "coding") or not 7 <= len(tokens) <= context + 1:
+        raise ValueError("invalid task or context length")
+    if tokens[:3].tolist() != [s("bos"), s("task:" + task), s("input")] or tokens[-1] != s("eos"):
+        raise ValueError("invalid sequence boundary")
+    if not 4 <= target < len(tokens) or tokens[target - 1] != s("output"):
+        raise ValueError("invalid supervision boundary")
+    for start, end, required in ((3, target - 1, task != "text"), (target, len(tokens) - 1, True)):
+        count = 0
+        while start < end:
+            if tokens[start] != s("text"):
+                raise ValueError("invalid text packet start")
+            closes = np.flatnonzero(tokens[start + 1:end] >= v.text_size)
+            if not len(closes):
+                raise ValueError("unterminated text packet")
+            close = start + 1 + int(closes[0])
+            if tokens[close] != s("/text"):
+                raise ValueError("non-text token in text packet")
+            start, count = close + 1, count + 1
+        if required and count == 0:
+            raise ValueError("empty packet side")
+
+
+class TextPreparation:
+    def __init__(self, config: dict, status_path: Path | None = None):
+        self.cfg, self.status_path = config, status_path
+        self.work, self.out = Path(config["work"]), Path(config["output"])
+        if self.work.exists() or self.out.exists():
+            raise FileExistsError("preserve existing preparation/cache; inspect it before selecting a new work directory")
+        self.guard()
+        self.work.mkdir(parents=True)
+        self.catalog = yaml.safe_load(Path(config["catalog"]).read_text(encoding="utf-8"))
+        self.catalog = deepcopy(self.catalog)
+        self.counts, self.source_files = Counter(), []
+        self.started = time.time()
+        self.db = sqlite3.connect(self.work / "records.sqlite")
+        self.db.executescript("""
+          PRAGMA journal_mode=WAL;
+          PRAGMA cache_size=-32768;
+          CREATE TABLE records (id INTEGER PRIMARY KEY, dataset TEXT, split TEXT, payload TEXT);
+          CREATE TABLE groups (dataset TEXT, name TEXT, split TEXT, PRIMARY KEY(dataset,name));
+          CREATE TABLE contents (hash TEXT PRIMARY KEY, split TEXT);
+        """)
+
+    def guard(self):
+        anchor = self.work if self.work.exists() else self.work.parent
+        if shutil.disk_usage(anchor).free < self.cfg["reserve_gib"] * 2**30:
+            raise RuntimeError("25 GiB free-space reserve reached")
+        if self.work.exists():
+            used = sum(p.stat().st_size for p in self.work.rglob("*") if p.is_file())
+            if self.out.exists():
+                used += sum(p.stat().st_size for p in self.out.rglob("*") if p.is_file())
+            if used > self.cfg["max_work_gib"] * 2**30:
+                raise RuntimeError("bounded preparation disk budget reached")
+
+    def progress(self, phase: str, **fields):
+        self.guard()
+        record = {"phase": phase, "pid": os.getpid(), "updated": time.time(), "elapsed_s": time.time() - self.started,
+                  "training_started": False, "counts": dict(self.counts), **fields}
+        atomic_json(self.work / "progress.json", record)
+        if self.status_path:
+            state = json.loads(self.status_path.read_text(encoding="utf-8")) if self.status_path.exists() else {}
+            state.update(state="preparing_data", updated=time.time(),
+                         message="CPU data preparation: " + phase, preprocessing_elapsed_s=record["elapsed_s"])
+            atomic_json(self.status_path, state)
+        print(json.dumps({k: v for k, v in record.items() if k != "counts"}), flush=True)
+
+    def source(self, path: Path, expected: str | None = None) -> str:
+        sha = sha256_file(path)
+        if expected and sha != expected:
+            raise ValueError(f"source checksum mismatch: {path}")
+        self.source_files.append({"path": str(path), "bytes": path.stat().st_size, "sha256": sha})
+        return sha
+
+    def provenance(self, ds: str, ident: str, split: str, path: Path, sha: str, source_split: str, **extra):
+        d = self.catalog["datasets"][ds]
+        return {"id": ds + ":" + ident, "revision": d.get("revision", d.get("release")), "license": d["license"],
+                "attribution": d.get("attribution", d.get("repo_id", ds)) + "; see pinned dataset card and retained source metadata",
+                "source_url": d["source"], "source_file": str(path), "source_sha256": sha,
+                "source_split": source_split, "split": split, **extra}
+
+    def add(self, record: dict) -> bool:
+        verify_record(record, self.catalog)
+        if any(not segment.get("text", "").strip() for segment in record["inputs"] + record["outputs"]):
+            self.counts["empty_text_excluded"] += 1
+            return False
+        key = normalized_hash("\n".join(x["text"] for x in record["inputs"] + record["outputs"]))
+        old = self.db.execute("SELECT split FROM contents WHERE hash=?", (key,)).fetchone()
+        group = (record["dataset"], str(record["group"]))
+        prior = self.db.execute("SELECT split FROM groups WHERE dataset=? AND name=?", group).fetchone()
+        if old or (prior and prior[0] != record["split"]):
+            self.counts["cross_split_duplicates_excluded" if (old and old[0] != record["split"]) or
+                        (prior and prior[0] != record["split"]) else "duplicates_excluded"] += 1
+            return False
+        self.db.execute("INSERT INTO contents VALUES (?,?)", (key, record["split"]))
+        self.db.execute("INSERT OR IGNORE INTO groups VALUES (?,?,?)", (*group, record["split"]))
+        record["document_content_sha256"] = key
+        self.db.execute("INSERT INTO records(dataset,split,payload) VALUES (?,?,?)",
+                        (record["dataset"], record["split"], canonical(record)))
+        self.counts[f"documents/{record['split']}/{record['dataset']}"] += 1
+        return True
+
+    def acquired_sources(self):
+        root = Path(self.cfg["acquired"])
+        ledger = {}
+        for line in (root / "files.jsonl").open(encoding="utf-8"):
+            r = json.loads(line)
+            ledger[r["path"]] = r
+        for ds in ("mbpp", "smoltalk", "fineweb_edu"):
+            for path in sorted((root / ds).rglob("*.parquet")):
+                expected = pinned_source_sha(path, root, ds, self.catalog["datasets"][ds]["revision"], ledger)
+                self.source(path, expected)
+        hashes = {r["path"]: r["sha256"] for r in self.source_files}
+        # Held-out MBPP precedes train so duplicate solutions cannot contaminate it.
+        for split, source_split in (("test", "test"), ("val", "validation"), ("train", "train")):
+            p = next((root / "mbpp/full").glob(source_split + "-*.parquet"))
+            for row in parquet_rows(p):
+                prov = self.provenance("mbpp", str(row["task_id"]), split, p, hashes[str(p)], source_split)
+                try:
+                    self.add(normalize("mbpp", row, prov)[0])
+                except SyntaxError:
+                    self.counts["mbpp_syntax_rejected/" + split] += 1
+        self.db.commit()
+        for source_split, split, capkey in (("test", "val", "smoltalk_val_per_subset"),
+                                          ("train", "train", "smoltalk_train_per_subset")):
+            for subset in self.catalog["datasets"]["smoltalk"]["subsets"]:
+                accepted = 0
+                for p in sorted((root / "smoltalk/data" / subset).glob(source_split + "-*.parquet")):
+                    for i, row in enumerate(parquet_rows(p)):
+                        prov = self.provenance("smoltalk", f"{subset}/{p.name}/{i}", split, p, hashes[str(p)], source_split,
+                                               component=subset)
+                        try:
+                            accepted += self.add(normalize("smoltalk", row, prov)[0])
+                        except (ValueError, KeyError, TypeError):
+                            self.counts["smoltalk_invalid"] += 1
+                        if accepted >= self.cfg[capkey]:
+                            break
+                    if accepted >= self.cfg[capkey]:
+                        break
+                self.db.commit()
+                self.progress("source_selection", component=subset, split=split, accepted=accepted)
+        # FineWeb source budget is publisher-token based; actual BPE budget is applied later.
+        source_tokens = 0
+        for p in sorted((root / "fineweb_edu").rglob("*.parquet")):
+            for i, row in enumerate(parquet_rows(p)):
+                if row.get("language") != "en" or float(row["score"]) < self.cfg["fineweb_min_score"]:
+                    self.counts["fineweb_quality_rejected"] += 1
+                    continue
+                text = row["text"].strip()
+                if not text or len(text) > 500000:
+                    self.counts["fineweb_size_rejected"] += 1
+                    continue
+                group = document_group(row.get("url", ""), text)
+                split = "val" if int(digest(group)[:8], 16) % 100 < 2 else "train"
+                prov = self.provenance("fineweb_edu", str(row["id"]), split, p, hashes[str(p)], "sample-10BT",
+                                       original_url=row.get("url"), crawl_dump=row.get("dump"), warc_file=row.get("file_path"),
+                                       educational_score=row["score"])
+                record = normalize("fineweb_edu", row, prov)[0]
+                record["group"] = group
+                if self.add(record):
+                    source_tokens += int(row["token_count"])
+                if i % 2000 == 0:
+                    self.db.commit()
+                    self.progress("source_selection", component="fineweb_edu", publisher_tokens=source_tokens)
+                if source_tokens >= self.cfg["fineweb_source_tokens"]:
+                    break
+            if source_tokens >= self.cfg["fineweb_source_tokens"]:
+                break
+        self.db.commit()
+
+    def legacy_sources(self):
+        raw, cache = Path(self.cfg["legacy_raw"]), Path(self.cfg["legacy_cache"])
+        stats = json.loads((cache / "stats.json").read_text())
+        for ds in ("tinystories", "dolly"):
+            p = raw / ds / "provenance.json"
+            src = json.loads(p.read_text())
+            self.source(p)
+            self.catalog["datasets"][ds] = {"role": "core", "repo_id": src["repo_id"], "revision": src["revision"],
+                "license": src["license"], "source": src["url"], "tasks": ["text" if ds == "tinystories" else "chat_text"],
+                "attribution": src["repo_id"], "notes": src["notes"], "retained_builtin": True}
+            if src["license"] not in self.catalog["policy"]["allowed_asset_licenses"]:
+                self.catalog["policy"]["allowed_asset_licenses"].append(src["license"])
+            for filename in src["files"]:
+                self.source(raw / ds / filename, src["sha256"][Path(filename).name])
+        for split, source_split in (("val", "validation"), ("train", "train")):
+            p = next((raw / "tinystories/data").glob(source_split + "-*.parquet"))
+            sha = next(r["sha256"] for r in self.source_files if r["path"] == str(p))
+            for i, row in enumerate(parquet_rows(p)):
+                if i >= stats["text"][split]["documents"]:
+                    break
+                text = row["text"].strip()
+                self.add({**self.provenance("tinystories", f"{source_split}/{i}", split, p, sha, source_split),
+                          "dataset": "tinystories", "task": "text", "group": normalized_hash(text),
+                          "inputs": [], "outputs": [{"type": "text", "text": text}]})
+        # Preserve the actual historical Dolly split and pairs, not a new random split.
+        for split in ("val", "train"):
+            p = cache / f"chat_{split}.json"
+            sha = self.source(p)
+            for i, row in enumerate(json.loads(p.read_text())):
+                self.add({**self.provenance("dolly", f"legacy-{split}/{i}", split, p, sha, "legacy-" + split),
+                          "dataset": "dolly", "task": "chat_text", "group": normalized_hash(row["user"]),
+                          "inputs": [{"type": "text", "text": row["user"]}],
+                          "outputs": [{"type": "text", "text": row["assistant"]}]})
+        self.db.commit()
+
+    def fit_tokenizer(self):
+        fitted = Counter()
+        def corpus():
+            for ds, cap in self.cfg["tokenizer_characters"].items():
+                for (payload,) in self.db.execute("SELECT payload FROM records WHERE dataset=? AND split='train' ORDER BY id", (ds,)):
+                    r = json.loads(payload)
+                    for segment in r["inputs"] + r["outputs"]:
+                        text = segment["text"][:cap - fitted[ds]]
+                        if text:
+                            fitted[ds] += len(text)
+                            yield text
+                    if fitted[ds] >= cap:
+                        break
+        self.progress("fit_training_only_bpe")
+        tok = TextTokenizer.train(corpus(), self.cfg["text_vocab_size"])
+        if tok.vocab_size != self.cfg["text_vocab_size"]:
+            raise ValueError("BPE did not reach the exact shared vocabulary size")
+        tok.save(self.work / "text_tokenizer.json")
+        atomic_json(self.work / "tokenizer_fit.json", {"split": "train", "characters_by_dataset": dict(fitted),
+                                                      "sha256": sha256_file(self.work / "text_tokenizer.json")})
+        return tok
+
+    def compile(self, tok: TextTokenizer):
+        v, context = UniversalVocab(), self.cfg["context"]
+        if tok.vocab_size != v.text_size:
+            raise ValueError("text vocabulary mismatch")
+        self.out.mkdir()
+        (self.out / "INCOMPLETE").write_text("awaiting complete validation")
+        handles, offsets, seen = {}, Counter(), {}
+        manifest_hash = hashlib.sha256()
+        targets, lengths, sources = Counter(), Counter(), Counter()
+        fineweb_train_tokens, total = 0, 0
+        for split in ("train", "val", "test"):
+            handles[split] = ((self.out / f"{split}.bin").open("wb"), (self.out / f"{split}.index.jsonl").open("w", encoding="utf-8"))
+        try:
+            with (self.out / "provenance.jsonl").open("w", encoding="utf-8") as provenance:
+                for row_id, payload in self.db.execute("SELECT id,payload FROM records ORDER BY id"):
+                    manifest_hash.update((payload + "\n").encode("utf-8"))
+                    r = json.loads(payload)
+                    chunks = text_windows(r["outputs"][0]["text"], tok, context - 6) if r["task"] == "text" else [r["outputs"][0]["text"]]
+                    for chunk_id, chunk in enumerate(chunks):
+                        candidate = r | {"outputs": [{"type": "text", "text": chunk}]}
+                        try:
+                            tokens, target = text_sample(candidate, tok, v, context)
+                        except ValueError as exc:
+                            if r["task"] == "text" or "context length" not in str(exc):
+                                raise
+                            self.counts[f"over_context_rejected/{r['split']}/{r['dataset']}"] += 1
+                            continue
+                        if r["dataset"] == "fineweb_edu" and r["split"] == "train":
+                            if fineweb_train_tokens + len(tokens) - target > self.cfg["fineweb_train_tokens"]:
+                                self.counts["fineweb_bpe_budget_excluded"] += 1
+                                continue
+                        content_key = normalized_hash("\n".join(x["text"] for x in candidate["inputs"] + candidate["outputs"]))
+                        if content_key in seen:
+                            self.counts["window_duplicates_excluded"] += 1
+                            continue
+                        seen[content_key] = r["split"]
+                        split, task = r["split"], r["task"]
+                        ident = r["id"] + f"/window-{chunk_id}"
+                        idx = {"offset": offsets[split], "length": len(tokens), "target_start": target,
+                               "task": task, "id": ident, "dataset": r["dataset"], "group": r["group"], "split": split,
+                               "document_content_sha256": r["document_content_sha256"], "content_sha256": content_key,
+                               "tokens_sha256": hashlib.sha256(tokens.tobytes()).hexdigest()}
+                        tokens.tofile(handles[split][0]); handles[split][1].write(canonical(idx) + "\n")
+                        prov = {k: val for k, val in r.items() if k not in ("inputs", "outputs")}
+                        prov.update(id=ident, window=chunk_id, token_offset=offsets[split], token_length=len(tokens),
+                                    tokens_sha256=idx["tokens_sha256"], content_sha256=content_key)
+                        provenance.write(canonical(prov) + "\n")
+                        offsets[split] += len(tokens)
+                        self.counts[f"samples/{split}/{task}"] += 1
+                        targets[f"{split}/{task}"] += len(tokens) - target
+                        sources[f"{split}/{r['dataset']}"] += 1
+                        lengths[str(min(context, max(128, 1 << (len(tokens) - 2).bit_length())))] += 1
+                        if r["dataset"] == "fineweb_edu" and split == "train":
+                            fineweb_train_tokens += len(tokens) - target
+                        total += 1
+                    if row_id % 1000 == 0:
+                        self.progress("encode_cache", documents=row_id, samples=total, fineweb_train_targets=fineweb_train_tokens)
+        finally:
+            for pair in handles.values():
+                for f in pair:
+                    f.close()
+        tok.save(self.out / "text_tokenizer.json")
+        atomic_json(self.out / "vocab_v2.json", v.to_dict())
+        atomic_json(self.out / "catalog.json", self.catalog)
+        report = {"training_started": False, "records": total, "context": context, "counts": dict(self.counts),
+                  "target_tokens_by_task": dict(targets), "samples_by_dataset": dict(sources), "bucket_lengths": dict(lengths),
+                  "input_tokens_by_split": dict(offsets), "token_bytes": sum(offsets.values()) * 2,
+                  "sources": self.catalog["datasets"], "source_files": self.source_files,
+                  "tokenizer_fit": json.loads((self.work / "tokenizer_fit.json").read_text()),
+                  "user_data_included": False, "media_included": False,
+                  "limitations": ["exact whitespace-normalized dedup only; near-duplicate/template overlap is unmeasured",
+                                  "over-context chat/coding examples are rejected, never truncated",
+                                  "FineWeb page rights remain governed by underlying sources; crawl URLs retained"],
+                  "elapsed_preparation_s": time.time() - self.started}
+        report["normalized_manifest_sha256"] = manifest_hash.hexdigest()
+        atomic_json(self.out / "stats.json", report)
+        self.progress("validate_cache", samples=total)
+        validation = validate_cache(self.out)
+        atomic_json(self.out / "validation.json", validation)
+        checksums = {p.name: sha256_file(p) for p in self.out.iterdir() if p.is_file() and p.name != "INCOMPLETE"}
+        atomic_json(self.out / "checksums.json", checksums)
+        (self.out / "INCOMPLETE").unlink()
+        self.progress("complete", samples=total, validation_passed=True)
+        return report
+
+    def run(self):
+        try:
+            self.progress("verify_sources")
+            self.legacy_sources()
+            self.acquired_sources()
+            return self.compile(self.fit_tokenizer())
+        finally:
+            self.db.close()
+
+
+def validate_cache(cache: Path, require_tasks: bool = True, expected_vocab: UniversalVocab | None = None) -> dict:
+    cache = Path(cache)
+    v = UniversalVocab(**json.loads((cache / "vocab_v2.json").read_text()))
+    if v.to_dict() != (expected_vocab or UniversalVocab()).to_dict():
+        raise ValueError("not the exact expected UniversalVocab")
+    tok = TextTokenizer.load(cache / "text_tokenizer.json")
+    if tok.vocab_size != v.text_size:
+        raise ValueError("tokenizer vocabulary mismatch")
+    stats = json.loads((cache / "stats.json").read_text())
+    groups, contents, documents, identities = {}, {}, {}, {}
+    counts, total = Counter(), 0
+    for split in ("train", "val", "test"):
+        p = cache / f"{split}.bin"
+        if p.stat().st_size % 2:
+            raise ValueError("misaligned uint16 stream")
+        tokens = np.memmap(p, dtype="<u2", mode="r") if p.stat().st_size else np.array([], dtype="<u2")
+        offset = 0
+        with (cache / f"{split}.index.jsonl").open(encoding="utf-8") as index:
+            for line in index:
+                r = json.loads(line)
+                if r["offset"] != offset or offset + r["length"] > len(tokens) or r["split"] != split:
+                    raise ValueError("invalid index offset/length/split")
+                sample = tokens[offset:offset + r["length"]]
+                if r["task"] in ("text", "chat_text", "coding"):
+                    validate_text_sample(sample, r["target_start"], r["task"], v, stats["context"])
+                else:
+                    if len(sample)-1 > stats["context"] or not 0 < r["target_start"] < len(sample):
+                        raise ValueError("invalid multimodal context or target boundary")
+                    if int(sample[r["target_start"]-1]) != v.special("output"):
+                        raise ValueError("multimodal targets must start after the output marker")
+                    v.validate(sample.tolist())
+                if hashlib.sha256(sample.tobytes()).hexdigest() != r["tokens_sha256"]:
+                    raise ValueError("sample token checksum mismatch")
+                for key, store in ((r["dataset"] + ":" + r["group"], groups),
+                                   (r["content_sha256"], contents), (r["document_content_sha256"], documents)):
+                    if key in store and store[key] != split:
+                        raise ValueError("train/heldout content or source-group leakage")
+                    store[key] = split
+                if r["id"] in identities:
+                    raise ValueError("duplicate sample identity")
+                identities[r["id"]] = (split, r["tokens_sha256"], r["dataset"], r["group"])
+                offset += r["length"]; counts[split + "/" + r["task"]] += 1; total += 1
+        if offset != len(tokens):
+            raise ValueError("unindexed trailing tokens")
+    catalog = json.loads((cache / "catalog.json").read_text())
+    checked = set()
+    with (cache / "provenance.jsonl").open(encoding="utf-8") as f:
+        for line in f:
+            r = json.loads(line)
+            key = r["id"]
+            if key in checked or identities.get(key) != (r["split"], r["tokens_sha256"], r["dataset"], r["group"]):
+                raise ValueError("provenance/index mismatch")
+            ds = catalog["datasets"][r["dataset"]]
+            if r["revision"] != ds.get("revision", ds.get("release")) or r["license"] != ds["license"]:
+                raise ValueError("provenance identity/license mismatch")
+            if not all(r.get(k) for k in ("attribution", "source_url", "source_file", "source_sha256")):
+                raise ValueError("missing source provenance")
+            if r["dataset"] == "mbpp" and {"train":"train","validation":"val","test":"test"}.get(r["source_split"]) != r["split"]:
+                raise ValueError("MBPP heldout partition changed")
+            checked.add(key)
+    if len(checked) != total or total != stats["records"]:
+        raise ValueError("provenance/statistics count mismatch")
+    if require_tasks:
+        for split in ("train", "val"):
+            for task in ("text", "chat_text", "coding"):
+                if not counts[split + "/" + task]:
+                    raise ValueError("missing required task coverage")
+    checksums = cache / "checksums.json"
+    if checksums.exists():
+        for name, sha in json.loads(checksums.read_text()).items():
+            if Path(name).name != name or sha256_file(cache / name) != sha:
+                raise ValueError("cache file checksum mismatch: " + name)
+    return {"passed": True, "records": total, "counts": dict(counts), "source_groups": len(groups),
+            "content_leakage": 0, "context": stats["context"], "vocab_size": v.size,
+            "validated_at": time.time(), "training_started": False}

quasar/inference/decoding.py

@@ -0,0 +1,163 @@
+"""Constrained decoding: a per-sequence structure state machine + token sampling.
+
+The state machine guarantees every generated media block is well formed: inside <image> only
+visual codes are allowed and exactly G of them; inside a video, frames are
+<frame><t_k> G codes </frame> with consecutive timestamps; text spans can only contain text
+tokens.  Purely structural tokens are *forced* (no sampling); content tokens are sampled from
+the model's distribution restricted to the legal set.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from ..tokenizers.multimodal import MAX_FRAMES, UnifiedVocab
+
+
+@dataclass
+class Sampling:
+    temperature: float = 1.0
+    top_k: int = 0
+    top_p: float = 1.0
+    repetition_penalty: float = 1.0  # applied to previously generated *text* tokens only
+    max_new_tokens: int = 256
+    seed: int = 0
+    cfg_scale: float = 1.0  # classifier-free guidance on visual tokens (1.0 = off)
+
+
+class Masks:
+    """Boolean legal-token sets, precomputed once per vocabulary/device."""
+
+    def __init__(self, v: UnifiedVocab, device):
+        V = v.size
+
+        def mk(ids=(), ranges=()):
+            m = torch.zeros(V, dtype=torch.bool, device=device)
+            for lo, hi in ranges:
+                m[lo:hi] = True
+            if ids:
+                m[list(ids)] = True
+            return m
+
+        self.text = mk([v.eos], [(0, v.text_size)])
+        self.text_media = mk([v.eos, v.img_start, v.vid_start], [(0, v.text_size)])
+        self.visual = mk(ranges=[(v.visual_offset, v.visual_offset + v.visual_size)])
+        self.timestamp = mk(ranges=[(v.timestamp(0), v.timestamp(0) + MAX_FRAMES)])
+        self.frame_or_end = mk([v.frm_start, v.vid_end])
+
+
+class StructureFSM:
+    """Tracks where one sequence is in the media grammar.
+
+    mode:
+      "text"   free text until <eos> (media tags disallowed unless allow_media)
+      "image"  we are right after <image>: G codes, </image>, then <eos> if stop_after_media
+      "video"  we are inside <video> right before the first <frame>
+    n_frames: exact total frame count for videos (None = model decides, bounded by max_timestamp)
+    first_ts: timestamp of the next frame to generate (None = sampled by the model)
+    frames_done: frames already present in the prefix (prefix then ends right after a </frame>)
+    codes_done: image codes already present after <image> (image completion)
+    """
+
+    def __init__(self, v: UnifiedVocab, mode: str = "text", allow_media: bool = False, stop_after_media: bool = True,
+                 n_frames: int | None = None, first_ts: int | None = None, max_timestamp: int = 7, frames_done: int = 0, codes_done: int = 0):
+        self.v, self.allow_media, self.stop_after_media = v, allow_media, stop_after_media
+        self.n_frames, self.next_ts, self.max_ts = n_frames, first_ts, max_timestamp
+        self.done = False
+        self.frames = frames_done
+        self.state = {"text": "text", "image": "img_codes", "video": "after_frame" if frames_done else "frame_start"}[mode]
+        self.count = codes_done if mode == "image" else 0
+
+    def constraint(self, masks: Masks) -> tuple[int | None, torch.Tensor | None]:
+        """(forced_token, None) or (None, legal_mask)."""
+        v, s = self.v, self.state
+        if self.done:
+            return v.pad, None
+        if s == "text":
+            return None, masks.text_media if self.allow_media else masks.text
+        if s in ("img_codes", "frm_codes"):
+            return None, masks.visual
+        if s == "img_end":
+            return v.img_end, None
+        if s == "frm_end":
+            return v.frm_end, None
+        if s == "frame_start":
+            return v.frm_start, None
+        if s == "timestamp":
+            return (v.timestamp(self.next_ts), None) if self.next_ts is not None else (None, masks.timestamp)
+        if s == "after_frame":
+            last = self.next_ts - 1
+            if (self.n_frames is not None and self.frames >= self.n_frames) or last >= self.max_ts:
+                return v.vid_end, None
+            return (v.frm_start, None) if self.n_frames is not None else (None, masks.frame_or_end)
+        if s == "eos":
+            return v.eos, None
+        raise AssertionError(s)
+
+    def advance(self, tok: int) -> None:
+        v, s = self.v, self.state
+        if self.done:
+            return
+        G = v.tokens_per_frame
+        if s == "text":
+            if tok == v.eos:
+                self.done = True
+            elif tok == v.img_start:
+                self.state, self.count = "img_codes", 0
+            elif tok == v.vid_start:
+                self.state = "frame_start"
+        elif s == "img_codes":
+            self.count += 1
+            if self.count == G:
+                self.state = "img_end"
+        elif s == "img_end":
+            self.state = "eos" if self.stop_after_media else "text"
+        elif s == "frame_start":
+            self.state = "timestamp"
+        elif s == "timestamp":
+            self.next_ts = tok - v.timestamp(0) + 1
+            self.state, self.count = "frm_codes", 0
+        elif s == "frm_codes":
+            self.count += 1
+            if self.count == G:
+                self.state = "frm_end"
+        elif s == "frm_end":
+            self.frames += 1
+            self.state = "after_frame"
+        elif s == "after_frame":
+            if tok == v.frm_start:
+                self.state = "timestamp"
+            else:  # </video>
+                self.state = "eos" if self.stop_after_media else "text"
+        elif s == "eos":
+            self.done = True
+
+    @property
+    def in_visual(self) -> bool:
+        return self.state in ("img_codes", "frm_codes")
+
+
+def sample_next(logits: torch.Tensor, legal: torch.Tensor, sp: Sampling, gen: torch.Generator,
+                history: torch.Tensor | None = None, text_size: int = 0) -> torch.Tensor:
+    """logits [B,V] (float32), legal [B,V] bool -> sampled ids [B]."""
+    logits = logits.float().clone()
+    if sp.repetition_penalty != 1.0 and history is not None and history.numel():
+        prev = history.clamp(min=0)
+        is_text = prev < text_size
+        vals = logits.gather(1, prev)
+        pen = torch.where(vals > 0, vals / sp.repetition_penalty, vals * sp.repetition_penalty)
+        logits.scatter_(1, prev, torch.where(is_text, pen, vals))
+    logits = logits.masked_fill(~legal, float("-inf"))
+    if sp.temperature <= 0:
+        return logits.argmax(-1)
+    logits = logits / sp.temperature
+    if sp.top_k > 0:
+        kth = logits.topk(min(sp.top_k, logits.size(-1)), dim=-1).values[:, -1:]
+        logits = logits.masked_fill(logits < kth, float("-inf"))
+    if sp.top_p < 1.0:
+        srt, idx = logits.sort(dim=-1, descending=True)
+        probs = srt.softmax(-1)
+        drop = probs.cumsum(-1) - probs > sp.top_p  # keep the smallest prefix with mass >= top_p
+        logits = logits.masked_fill(drop.scatter(1, idx, drop), float("-inf"))
+    return torch.multinomial(logits.softmax(-1), 1, generator=gen).squeeze(1)

quasar/inference/extended.py

@@ -0,0 +1,148 @@
+"""Audio, geometry and schema-checked calls from the same shared generator.
+
+Recorded training exposure is required. Structural masks constrain packets;
+semantic usefulness must be established separately on trained checkpoints.
+"""
+from pathlib import Path
+import json
+import numpy as np
+import torch
+from .universal import UniversalGenerator
+from .decoding import Sampling, sample_next
+from ..data.curated import validate_calls
+from ..data.text_preparation import sha256_file
+from ..tokenizers.universal import TokenGrammar
+from ..tokenizers.audio_aligned import AlignedAudioTokenizer
+
+
+class ExtendedGenerator(UniversalGenerator):
+    def require_task(self, task):
+        if not self.meta.get('task_target_tokens', {}).get(task):
+            raise RuntimeError('This checkpoint has no recorded training exposure for ' + task)
+
+    def audio(self):
+        if not hasattr(self, '_audio'):
+            identity = self.meta.get('data', {}).get('preparation', {}).get('audio_codec')
+            if not identity or identity.get('implementation') != 'quasar.tokenizers.audio_aligned.AlignedAudioTokenizer':
+                raise RuntimeError('Checkpoint has no compatible frozen audio codec identity')
+            path = Path(identity['path'])
+            if not path.is_absolute():
+                path = (self.path / path).resolve()
+            for name in ('config.json', 'model.safetensors'):
+                if sha256_file(path / name) != identity['sha256'][name]:
+                    raise ValueError('Audio codec differs from training')
+            self._audio = AlignedAudioTokenizer(path, identity['alignment_samples'], identity['codebooks'], device=self.device)
+        return self._audio
+
+    def _candidates(self, grammar, kind, remaining):
+        v, s, state = self.vocab, self.vocab.special, grammar.state
+        if state == 'audio':
+            book = grammar.count % v.audio_books
+            lo = v.ranges['audio'][0] + book * v.audio_size
+            # Keep two boundary guards within the two-second source budget.
+            if grammar.count >= 152 * v.audio_books or (book == 0 and grammar.count >= 64 and remaining <= 5):
+                return [s('/audio')]
+            tokens = list(range(lo, lo + v.audio_size))
+            if grammar.count >= 64 and book == 0:
+                tokens.append(s('/audio'))
+            return tokens
+        if state == 'mesh':
+            return [s('vertices')]
+        if state == 'vertices':
+            if grammar.count >= 9 and grammar.count % 3 == 0 and remaining <= 8:
+                return [s('faces')]
+            return list(range(*v.ranges['coordinate'])) + [s('faces')]
+        if state == 'faces':
+            if grammar.count >= 3 and grammar.count % 3 == 0 and remaining <= 5:
+                return [s('/mesh')]
+            lo = v.ranges['index'][0]
+            return list(range(lo, lo + grammar.vertices)) + [s('/mesh')]
+        if state == 'calls':
+            return [s('/calls')] if remaining <= 2 else list(range(v.text_size)) + [s('/calls')]
+        if state == 'packets':
+            return [s('eos')]
+        raise ValueError('Unsupported generated packet state: ' + state)
+
+    @torch.inference_mode()
+    def packet(self, task, inputs, kind, sp):
+        self.require_task(task)
+        v, s = self.vocab, self.vocab.special
+        prefix = [s('bos'), s('task:' + task), s('input')]
+        for segment in inputs:
+            prefix += v.packet(segment, self.tok.encode)
+        prefix += [s('output'), s(kind)]
+        budget = min(sp.max_new_tokens, self.model.cfg.max_seq_len - len(prefix))
+        minimum = {'audio': 66, 'mesh': 16, 'calls': 3}[kind]
+        if budget < minimum:
+            raise ValueError('Complete packet and conditioning exceed the requested context/output budget')
+        grammar = TokenGrammar(v)
+        for token in prefix:
+            grammar.advance(token)
+        device = next(self.model.parameters()).device
+        amp = device.type == 'cuda'
+        cache = self.model.new_cache(1, dtype=torch.bfloat16 if amp else torch.float32)
+        rng = torch.Generator(device=device).manual_seed(sp.seed)
+        emitted = []
+        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=amp):
+            logits = self.model(torch.tensor([prefix], device=device), cache=cache)[:, -1]
+            while grammar.state != 'done' and len(emitted) < budget:
+                candidates = [t for t in self._candidates(grammar, kind, budget - len(emitted)) if grammar.accepts(t)]
+                if not candidates:
+                    raise RuntimeError('No structurally valid completion fits the packet budget')
+                if len(candidates) == 1:
+                    token = candidates[0]
+                else:
+                    legal = torch.zeros((1, v.size), dtype=torch.bool, device=device)
+                    legal[:, candidates] = True
+                    history = torch.tensor([emitted], dtype=torch.long, device=device)
+                    token = int(sample_next(logits, legal, sp, rng, history, v.text_size)[0])
+                grammar.advance(token); emitted.append(token)
+                if grammar.state != 'done':
+                    logits = self.model(torch.tensor([[token]], device=device), cache=cache, start=cache.next_pos)[:, -1]
+        if grammar.state != 'done':
+            raise RuntimeError('Model did not finish a valid packet within the output budget')
+        v.validate(prefix + emitted)
+        return emitted
+
+    def generate_audio(self, prompt, sp=None):
+        tokens = self.packet('t2a', [{'type': 'text', 'text': prompt}], 'audio',
+                             sp or Sampling(temperature=.8, top_k=64, top_p=.9, max_new_tokens=642))
+        raw = np.asarray(tokens[:-2], dtype=np.int64).reshape(-1, self.vocab.audio_books)
+        offsets = self.vocab.ranges['audio'][0] + np.arange(self.vocab.audio_books) * self.vocab.audio_size
+        packet = {'type': 'audio', 'codes': (raw - offsets).tolist()}
+        codec = self.audio()
+        return codec.decode(packet), codec.sr
+
+    def describe_audio(self, waveform, sample_rate, sp=None):
+        self.require_task('a2t')
+        packet = self.audio().encode(waveform, sample_rate)
+        return self._predict('a2t', [packet], 'text', sp or Sampling(temperature=.5, max_new_tokens=128))[0]
+
+    def generate_mesh(self, prompt, sp=None):
+        t = self.packet('t2m', [{'type': 'text', 'text': prompt}], 'mesh',
+                        sp or Sampling(temperature=.8, top_k=64, top_p=.9, max_new_tokens=780))
+        s = self.vocab.special
+        vertices = np.asarray(t[t.index(s('vertices')) + 1:t.index(s('faces'))]).reshape(-1, 3)
+        faces = np.asarray(t[t.index(s('faces')) + 1:t.index(s('/mesh'))]).reshape(-1, 3)
+        vertices -= self.vocab.ranges['coordinate'][0]; faces -= self.vocab.ranges['index'][0]
+        if len(np.unique(vertices, axis=0)) != len(vertices) or np.any(np.linalg.norm(
+                np.cross(vertices[faces[:, 1]] - vertices[faces[:, 0]], vertices[faces[:, 2]] - vertices[faces[:, 0]]), axis=1) == 0):
+            raise RuntimeError('Generated mesh has collapsed or degenerate geometry; semantic generation is unverified')
+        return {'type': 'mesh', 'vertices': vertices.tolist(), 'faces': faces.tolist()}
+
+    def describe_mesh(self, packet, sp=None):
+        return self._predict('m2t', [packet], 'text', sp or Sampling(temperature=.5, max_new_tokens=128))[0]
+
+    def propose_calls(self, prompt, tools, sp=None):
+        self.require_task('tools')
+        validate_calls(tools, [])
+        tokens = self.packet('tools', [{'type': 'tools', 'value': tools}, {'type': 'text', 'text': 'user: ' + prompt}],
+                             'calls', sp or Sampling(temperature=.5, max_new_tokens=256))
+        text = self.tok.decode([t for t in tokens if t < self.vocab.text_size])
+        def reject_constant(value):
+            raise ValueError('Non-finite JSON constant: ' + value)
+        result = json.loads(text, parse_constant=reject_constant)
+        if not isinstance(result, list):
+            raise ValueError('Generated calls must be a JSON list')
+        validate_calls(tools, result)
+        return result

quasar/inference/generate.py

@@ -0,0 +1,285 @@
+"""Unified generation engine: every capability is one call into the same trained backbone.
+
+    eng = QuasarEngine("checkpoints/quasar-exp/instruct/final")
+    eng.generate_text("Once upon a time")
+    eng.generate_image("a large red circle in the center on a black background")
+    eng.generate_video("a small blue circle moving right on a navy background")
+    eng.describe_image(img); eng.describe_video(frames); eng.continue_video(frames)
+    eng.edit_image(img, "change the color to green"); eng.image_to_video(img, "... moving up ...")
+    eng.chat("Make a video of a blue ball bouncing.")
+
+The visual tokenizer only converts pixels <-> tokens; all content tokens are sampled from the
+Quasar transformer.  Capabilities the checkpoint was never trained on raise CapabilityError
+unless force=True, so untrained noise is never presented as a result.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from ..data.preprocessing import load_video, prepare_image
+from ..tokenizers.text import TextTokenizer
+from ..tokenizers.vision import VisualTokenizer
+from ..training.checkpoint import load_quasar
+from .decoding import Masks, Sampling, StructureFSM, sample_next
+
+CAPABILITY_TASKS = {
+    "generate_text": ("text",), "chat": ("chat_text", "chat_mm"), "generate_image": ("t2i",), "complete_image": ("t2i",),
+    "describe_image": ("i2t",), "edit_image": ("i2i",), "generate_video": ("t2v",), "image_to_video": ("t2v",),
+    "continue_video": ("vcont",), "describe_video": ("v2t",),
+}
+TRAINED_FRAMES = 8  # clips in training have 8 frames (t0..t7)
+
+
+class CapabilityError(RuntimeError):
+    pass
+
+
+# ------------------------------------------------------------------------------- media io
+def save_video(frames: np.ndarray, path: str | Path, fps: int = 4, scale: int = 4) -> Path:
+    """Write uint8 [T,H,W,3] as .mp4 (H.264, browser-playable) or .gif, upscaled `scale`x (nearest)."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    big = np.asarray(frames).repeat(scale, 1).repeat(scale, 2)
+    if path.suffix.lower() == ".gif":
+        ims = [Image.fromarray(f) for f in big]
+        ims[0].save(path, save_all=True, append_images=ims[1:], duration=int(1000 / fps), loop=0)
+        return path
+    import imageio_ffmpeg
+
+    T, H, W, _ = big.shape
+    w = imageio_ffmpeg.write_frames(str(path), (W, H), fps=fps, codec="libx264", pix_fmt_out="yuv420p", macro_block_size=1,
+                                    output_params=["-crf", "18"])
+    w.send(None)
+    for f in big:
+        w.send(np.ascontiguousarray(f))
+    w.close()
+    return path
+
+
+def normalize_prompt(p: str) -> str:
+    """Training captions are lower-case without a trailing period."""
+    return p.strip().rstrip(".").lower()
+
+
+# ------------------------------------------------------------------------------- engine
+class QuasarEngine:
+    def __init__(self, ckpt: str | Path, device: str | None = None):
+        self.ckpt = Path(ckpt)
+        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        self.model, self.vocab, self.meta = load_quasar(self.ckpt, self.device)
+        if getattr(self.vocab, "schema_version", 1) in (2,3):
+            raise CapabilityError("This is a v2 preparation checkpoint. V2 generation/codec reconstruction must be validated and trained before exposing it in the legacy demo.")
+        self.model.eval()
+        self.tok = TextTokenizer.load(self.ckpt / "tokenizers" / "text_tokenizer.json")
+        self.vq = VisualTokenizer.load(self.ckpt / "tokenizers" / "vq", self.device)
+        self.masks = Masks(self.vocab, self.device)
+        self.amp = torch.bfloat16 if self.device.type == "cuda" and torch.cuda.is_bf16_supported() else None
+
+    # -- capability bookkeeping ------------------------------------------------------------
+    def capabilities(self) -> dict[str, dict]:
+        tt: dict[str, int] = {}
+        for k, n in self.meta.get("task_target_tokens", {}).items():  # 'user/t2i' counts toward 't2i'
+            base = k.rpartition("/")[2]
+            tt[base] = tt.get(base, 0) + int(n)
+        return {m: {"trained": all(tt.get(t, 0) > 0 for t in ts), "target_tokens": {t: tt.get(t, 0) for t in ts}}
+                for m, ts in CAPABILITY_TASKS.items()}
+
+    def _require(self, method: str, force: bool) -> None:
+        if force:
+            return
+        cap = self.capabilities()[method]
+        if not cap["trained"]:
+            raise CapabilityError(
+                f"'{method}' has not been trained in this checkpoint (stage '{self.meta.get('stage')}', "
+                f"target tokens seen: {cap['target_tokens']}). Train a stage that includes these tasks, "
+                f"or pass force=True to sample the untrained model anyway.")
+
+    # -- token helpers ---------------------------------------------------------------------
+    def enc(self, s: str) -> list[int]:
+        return self.tok.encode(s)
+
+    def image_codes(self, image) -> np.ndarray:
+        return self.vq.encode(prepare_image(image)[None])[0]
+
+    def video_codes(self, frames) -> np.ndarray:
+        return self.vq.encode(load_video(frames)[None])[0]
+
+    # -- core decode loop ------------------------------------------------------------------
+    @torch.no_grad()
+    def _decode(self, prefix: list[int], fsm_kwargs: dict, n: int, sp: Sampling, uncond_prefix: list[int] | None = None) -> list[list[int]]:
+        m, dev = self.model, self.device
+        if len(prefix) >= m.cfg.max_seq_len:
+            raise ValueError(f"prompt is {len(prefix)} tokens; the model context is {m.cfg.max_seq_len}")
+        max_len = min(m.cfg.max_seq_len, len(prefix) + sp.max_new_tokens)
+        guided = uncond_prefix is not None and sp.cfg_scale != 1.0
+        amp = torch.autocast(dev.type, dtype=self.amp) if self.amp else torch.autocast(dev.type, enabled=False)
+
+        def start(pref):
+            cache = m.new_cache(n, max_len)
+            hist = torch.tensor(pref, device=dev).repeat(n, 1)
+            with amp:
+                return cache, hist, m(hist, cache=cache)[:, -1].float()
+
+        cache, hist, logits = start(prefix)
+        if guided:
+            ucache, uhist, ulogits = start(uncond_prefix)
+        fsms = [StructureFSM(self.vocab, **fsm_kwargs) for _ in range(n)]
+        gen = torch.Generator(device=dev).manual_seed(sp.seed)
+        out = torch.empty(n, 0, dtype=torch.long, device=dev)
+        pos, upos = len(prefix), len(uncond_prefix or [])
+        while pos < max_len and not all(f.done for f in fsms):
+            cons = [f.constraint(self.masks) for f in fsms]
+            forced = torch.tensor([c[0] if c[0] is not None else -1 for c in cons], device=dev)
+            if (forced >= 0).all():
+                nxt = forced
+            else:
+                legal = torch.stack([c[1] if c[1] is not None else self.masks.text for c in cons])
+                lg = logits
+                if guided:
+                    vis = torch.tensor([f.in_visual for f in fsms], device=dev)[:, None]
+                    lg = torch.where(vis, ulogits + sp.cfg_scale * (logits - ulogits), logits)
+                nxt = sample_next(lg, legal, sp, gen, out, self.vocab.text_size)
+                nxt = torch.where(forced >= 0, forced, nxt)
+            for f, t in zip(fsms, nxt.tolist()):
+                f.advance(t)
+            out = torch.cat([out, nxt[:, None]], 1)
+            hist = torch.cat([hist, nxt[:, None]], 1)
+            if all(f.done for f in fsms) or pos + 1 >= max_len:
+                break
+            with amp:
+                logits = m(nxt[:, None], vis_pos=m.visual_positions(hist)[:, -1:], cache=cache, start=pos)[:, -1].float()
+                if guided:
+                    uhist = torch.cat([uhist, nxt[:, None]], 1)
+                    ulogits = m(nxt[:, None], vis_pos=m.visual_positions(uhist)[:, -1:], cache=ucache, start=upos)[:, -1].float()
+            pos += 1
+            upos += 1
+        res = []
+        for row in out.tolist():
+            res.append(row[: row.index(self.vocab.eos)] if self.vocab.eos in row else row)
+        return res
+
+    def _text(self, ids: list[int]) -> str:
+        return self.tok.decode([i for i in ids if i < self.vocab.text_size]).strip()
+
+    def _visual(self, ids: list[int], kind: str) -> np.ndarray:
+        segs = [s for k, s in self.vocab.parse(ids) if k == kind]
+        if not segs:
+            raise RuntimeError(f"generation produced no {kind} block")
+        return self.vq.decode(segs[0][None])[0]
+
+    # -- public API ------------------------------------------------------------------------
+    def generate_text(self, prompt: str, sp: Sampling | None = None, force: bool = False) -> str:
+        """Continue `prompt` with the base language-model head (TinyStories-style text)."""
+        self._require("generate_text", force)
+        sp = sp or Sampling(temperature=0.8, top_k=50, top_p=0.95, repetition_penalty=1.1, max_new_tokens=200)
+        return self._text(self._decode([self.vocab.bos, self.vocab.task("text"), *self.enc(prompt)], {"mode": "text"}, 1, sp)[0])
+
+    def generate_image(self, prompt: str, n: int = 1, sp: Sampling | None = None, force: bool = False) -> list[np.ndarray]:
+        self._require("generate_image", force)
+        sp = sp or Sampling(temperature=1.0, top_k=100, cfg_scale=3.0)
+        v = self.vocab
+        prefix = [v.bos, v.task("t2i"), *self.enc(normalize_prompt(prompt)), v.img_start]
+        rows = self._decode(prefix, {"mode": "image"}, n, sp, uncond_prefix=[v.bos, v.task("t2i"), v.img_start])
+        return [self.vq.decode(np.array(r[: v.tokens_per_frame]).reshape(1, *v.grid) - v.visual_offset)[0] for r in rows]
+
+    def complete_image(self, image, prompt: str = "", keep_rows: int = 4, n: int = 1, sp: Sampling | None = None, force: bool = False) -> list[np.ndarray]:
+        """Keep the top `keep_rows` token rows of `image` and let the model fill in the rest."""
+        self._require("complete_image", force)
+        sp = sp or Sampling(temperature=1.0, top_k=100)
+        v = self.vocab
+        keep = (self.image_codes(image).reshape(-1)[: keep_rows * v.grid[1]] + v.visual_offset).tolist()
+        prefix = [v.bos, v.task("t2i"), *self.enc(normalize_prompt(prompt)), v.img_start, *keep]
+        rows = self._decode(prefix, {"mode": "image", "codes_done": len(keep)}, n, sp)
+        return [self.vq.decode(np.array(keep + r[: v.tokens_per_frame - len(keep)]).reshape(1, *v.grid) - v.visual_offset)[0] for r in rows]
+
+    def describe_image(self, image, sp: Sampling | None = None, force: bool = False) -> str:
+        self._require("describe_image", force)
+        sp = sp or Sampling(temperature=0.0, max_new_tokens=48)
+        v = self.vocab
+        return self._text(self._decode([v.bos, v.task("i2t"), *v.image_tokens(self.image_codes(image))], {"mode": "text"}, 1, sp)[0])
+
+    def edit_image(self, image, instruction: str, n: int = 1, sp: Sampling | None = None, force: bool = False) -> list[np.ndarray]:
+        self._require("edit_image", force)
+        sp = sp or Sampling(temperature=1.0, top_k=100)
+        v = self.vocab
+        prefix = [v.bos, v.task("i2i"), *self.enc(normalize_prompt(instruction)), *v.image_tokens(self.image_codes(image)), v.img_start]
+        rows = self._decode(prefix, {"mode": "image"}, n, sp)
+        return [self.vq.decode(np.array(r[: v.tokens_per_frame]).reshape(1, *v.grid) - v.visual_offset)[0] for r in rows]
+
+    def _check_frames(self, n_frames: int, force: bool) -> None:
+        if n_frames > TRAINED_FRAMES and not force:
+            raise CapabilityError(f"clips were trained with {TRAINED_FRAMES} frames; use continue_video to extend beyond that")
+
+    def generate_video(self, prompt: str, n_frames: int = 8, n: int = 1, sp: Sampling | None = None, force: bool = False) -> list[np.ndarray]:
+        self._require("generate_video", force)
+        self._check_frames(n_frames, force)
+        sp = sp or Sampling(temperature=1.0, top_k=100, cfg_scale=3.0, max_new_tokens=n_frames * (self.vocab.tokens_per_frame + 3) + 4)
+        v = self.vocab
+        prefix = [v.bos, v.task("t2v"), *self.enc(normalize_prompt(prompt)), v.vid_start]
+        rows = self._decode(prefix, {"mode": "video", "n_frames": n_frames, "first_ts": 0}, n, sp, uncond_prefix=[v.bos, v.task("t2v"), v.vid_start])
+        return [self.vq.decode(self.vocab.parse([v.vid_start, *r])[0][1]) for r in rows]
+
+    def image_to_video(self, image, prompt: str, n_frames: int = 8, n: int = 1, sp: Sampling | None = None, force: bool = False) -> list[np.ndarray]:
+        """Use `image` as frame 0 of a text-conditioned clip and generate the remaining frames."""
+        self._require("image_to_video", force)
+        self._check_frames(n_frames, force)
+        sp = sp or Sampling(temperature=1.0, top_k=100, cfg_scale=3.0, max_new_tokens=n_frames * (self.vocab.tokens_per_frame + 3) + 4)
+        v = self.vocab
+        f0 = v.frame_tokens(self.image_codes(image), 0)
+        prefix = [v.bos, v.task("t2v"), *self.enc(normalize_prompt(prompt)), v.vid_start, *f0]
+        rows = self._decode(prefix, {"mode": "video", "n_frames": n_frames, "first_ts": 1, "frames_done": 1}, n, sp,
+                            uncond_prefix=[v.bos, v.task("t2v"), v.vid_start, *f0])
+        return [self.vq.decode(self.vocab.parse([v.vid_start, *f0, *r])[0][1]) for r in rows]
+
+    def continue_video(self, frames, n_new: int = 4, context: int = 4, sp: Sampling | None = None, force: bool = False) -> np.ndarray:
+        """Append `n_new` frames.  Rolls a window of the last `context` frames (re-timestamped from t0),
+        so clips can be extended beyond the 8-frame training length."""
+        self._require("continue_video", force)
+        v = self.vocab
+        codes = self.video_codes(frames)
+        context = min(context, TRAINED_FRAMES - 1)
+        out = list(codes)
+        while len(out) < len(codes) + n_new:
+            ctx = np.stack(out[-context:])
+            k = len(ctx)
+            todo = min(TRAINED_FRAMES - k, len(codes) + n_new - len(out))
+            s = sp or Sampling(temperature=1.0, top_k=100, seed=len(out))
+            s = Sampling(**{**s.__dict__, "max_new_tokens": todo * (v.tokens_per_frame + 3) + 4})
+            prefix = [v.bos, v.task("vcont"), *v.video_tokens(ctx, close=False)]
+            r = self._decode(prefix, {"mode": "video", "n_frames": k + todo, "first_ts": k, "frames_done": k}, 1, s)[0]
+            new = v.parse([v.vid_start, *v.video_tokens(ctx, close=False)[1:], *r])[0][1][k:]
+            out += list(new)
+        return self.vq.decode(np.stack(out))
+
+    def describe_video(self, frames, sp: Sampling | None = None, force: bool = False) -> str:
+        self._require("describe_video", force)
+        sp = sp or Sampling(temperature=0.0, max_new_tokens=48)
+        v = self.vocab
+        return self._text(self._decode([v.bos, v.task("v2t"), *v.video_tokens(self.video_codes(frames))], {"mode": "text"}, 1, sp)[0])
+
+    def chat(self, message: str, image=None, video=None, sp: Sampling | None = None, force: bool = False) -> list[tuple[str, object]]:
+        """Instruction-following turn.  Returns segments: ('text', str) / ('image', array) / ('video', array)."""
+        self._require("chat", force)
+        v = self.vocab
+        sp = sp or Sampling(temperature=0.7, top_k=50, top_p=0.95, repetition_penalty=1.1, max_new_tokens=700)
+        user: list[int] = []
+        if image is not None:
+            user += v.image_tokens(self.image_codes(image))
+        if video is not None:
+            user += v.video_tokens(self.video_codes(video))
+        user += self.enc((" " if user else "") + message.strip())
+        prefix = [v.bos, v.user, *user, v.assistant]
+        ids = self._decode(prefix, {"mode": "text", "allow_media": True, "stop_after_media": False}, 1, sp)[0]
+        segs: list[tuple[str, object]] = []
+        for kind, val in v.parse(ids):
+            if kind == "text":
+                segs.append(("text", self._text(val)))
+            elif kind == "image":
+                segs.append(("image", self.vq.decode(val[None])[0]))
+            else:
+                segs.append(("video", self.vq.decode(val)))
+        return segs

quasar/inference/universal.py

@@ -0,0 +1,103 @@
+"""Generation for the actual v2 packet grammar and a shared Axiom checkpoint."""
+from __future__ import annotations
+from pathlib import Path
+import numpy as np
+import torch
+from ..tokenizers.text import TextTokenizer
+from ..tokenizers.vision import VisualTokenizer
+from ..training.checkpoint import load_quasar
+from .decoding import Sampling,sample_next
+
+
+class UniversalGenerator:
+    def __init__(self,checkpoint: str|Path,device: str='cuda'):
+        self.path=Path(checkpoint)
+        self.model,self.vocab,self.meta=load_quasar(self.path,device)
+        self.model.eval()
+        if not hasattr(self.vocab,'special'):raise ValueError('UniversalGenerator requires a packet-v2 checkpoint')
+        self.tok=TextTokenizer.load(self.path/'tokenizers/text_tokenizer.json')
+        self.codec=None;self.device=device
+
+    def vision(self):
+        if self.codec is None:self.codec=VisualTokenizer.load(self.path/'tokenizers/vq',self.device)
+        return self.codec
+
+    @torch.inference_mode()
+    def _predict(self,task,inputs,kind,sp: Sampling,frames=8,resolution=64):
+        if not self.meta.get('task_target_tokens',{}).get(task):
+            raise RuntimeError(f'This checkpoint has no recorded training exposure for {task}')
+        v=self.vocab;s=v.special
+        prefix=[s('bos'),s('task:'+task),s('input')]
+        for segment in inputs:prefix+=v.packet(segment,self.tok.encode)
+        prefix+=[s('output'),s(kind)]
+        # This is structural forcing only. All content codes/tokens are predicted by Quasar.
+        if kind=='image':
+            if resolution not in (64,128) or (resolution==128 and v.schema_version!=3):
+                raise ValueError('128px image generation requires a schema-v3 checkpoint')
+            if resolution==128:prefix += [s('resolution:128')]
+            plan=[('vision',(resolution//8)**2),('/image',1),('eos',1)]
+        elif kind=='video':
+            if not 1<=frames<=v.max_frames:raise ValueError('frame count exceeds vocabulary')
+            plan=[]
+            for i in range(frames):plan += [('frame',1),('time:'+str(i),1),('vision',64),('/frame',1)]
+            plan += [('/video',1),('eos',1)]
+        else:plan=[('free_text',sp.max_new_tokens)]
+        needed=sum(n for _,n in plan)
+        if len(prefix)+needed>self.model.cfg.max_seq_len:raise ValueError('requested input/output exceeds trained context')
+        device=next(self.model.parameters()).device;use_amp=device.type=='cuda'
+        cache=self.model.new_cache(1,dtype=torch.bfloat16 if use_amp else torch.float32)
+        gen=torch.Generator(device=device).manual_seed(sp.seed)
+        emitted=[];content=[];stop=False
+        with torch.autocast(device.type,dtype=torch.bfloat16,enabled=use_amp):
+            logits=self.model(torch.tensor([prefix],device=device),cache=cache)[:,-1]
+            for state,n in plan:
+                for _ in range(n):
+                    if state in ('vision','free_text'):
+                        legal=torch.zeros((1,v.size),dtype=torch.bool,device=device)
+                        lo,hi=v.ranges['vision' if state=='vision' else 'text'];legal[:,lo:hi]=True
+                        if state=='free_text':legal[:,s('/text')]=True
+                        history=torch.tensor([content],dtype=torch.long,device=device)
+                        token=int(sample_next(logits,legal,sp,gen,history,v.text_size)[0])
+                        if state=='free_text' and token==s('/text'):
+                            emitted += [token,s('eos')];stop=True;break
+                        content.append(token)
+                    else:token=s(state)
+                    emitted.append(token)
+                    # A one-token cache append cannot infer its location in the
+                    # current visual run. Preserve the full-prefix spatial slot,
+                    # resetting at image/frame boundaries exactly as training does.
+                    history_tokens=torch.tensor([prefix+emitted],device=device)
+                    position=self.model.visual_positions(history_tokens)[:,-1:]
+                    logits=self.model(history_tokens[:,-1:],vis_pos=position,
+                                      cache=cache,start=cache.next_pos)[:,-1]
+                if stop:break
+        if kind=='text':return self.tok.decode(content),{'learned_text_stop':stop,'token_limit':sp.max_new_tokens}
+        v.validate(prefix+emitted)
+        codes=np.asarray(content,dtype=np.int64)-v.visual_offset
+        return codes.reshape((resolution//8,resolution//8) if kind=='image' else (frames,8,8)),{'structurally_valid':True}
+
+    def generate_text(self,prompt,sp=None):
+        return self._predict('chat_text',[{'type':'text','text':prompt}],'text',sp or Sampling(temperature=.5,top_p=.9))[0]
+
+    def generate_image(self,prompt,sp=None,resolution=64):
+        if resolution==128 and not self.meta.get('image_128_trained',False):
+            raise RuntimeError('128px image generation has no verified training exposure in this checkpoint')
+        codes,_=self._predict('t2i',[{'type':'text','text':prompt}],'image',sp or Sampling(),resolution=resolution)
+        return self.vision().decode(codes[None])[0]
+
+    def generate_video(self,prompt,frames=8,sp=None):
+        codes,_=self._predict('t2v',[{'type':'text','text':prompt}],'video',sp or Sampling(),frames)
+        return self.vision().decode(codes)
+
+    def describe_image(self,pixels,sp=None):
+        codes=self.vision().encode(np.asarray(pixels)[None])[0]
+        return self._predict('i2t',[{'type':'image','codes':codes}],'text',sp or Sampling(temperature=.5))[0]
+
+    def describe_video(self,frames,sp=None):
+        codes=self.vision().encode(np.asarray(frames))
+        return self._predict('v2t',[{'type':'video','codes':codes}],'text',sp or Sampling(temperature=.5))[0]
+
+    def continue_video(self,frames,new_frames=4,sp=None):
+        codes=self.vision().encode(np.asarray(frames))
+        generated,_=self._predict('vcont',[{'type':'video','codes':codes}],'video',sp or Sampling(),new_frames)
+        return self.vision().decode(generated)

quasar/model/attention.py

@@ -0,0 +1,58 @@
+"""Causal (grouped-query) self-attention on top of torch's scaled_dot_product_attention.
+
+SDPA picks the best available backend at runtime (cuDNN / memory-efficient / flash / math);
+nothing here requires a specific kernel, so CPU and any CUDA build work.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import ModelConfig
+from .layers import apply_rope
+
+
+class KVCache:
+    """Pre-allocated per-layer key/value buffers stored with n_kv_heads (GQA saves cache memory)."""
+
+    def __init__(self, cfg: ModelConfig, batch: int, max_len: int, device, dtype):
+        shape = (batch, cfg.n_kv_heads, max_len, cfg.head_dim)
+        self.k = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(cfg.n_layers)]
+        self.v = [torch.zeros(shape, device=device, dtype=dtype) for _ in range(cfg.n_layers)]
+        self.max_len = max_len
+
+
+class Attention(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.n_heads, self.n_kv, self.hd = cfg.n_heads, cfg.n_kv_heads, cfg.head_dim
+        self.wq = nn.Linear(cfg.d_model, cfg.n_heads * cfg.head_dim, bias=False)
+        self.wk = nn.Linear(cfg.d_model, cfg.n_kv_heads * cfg.head_dim, bias=False)
+        self.wv = nn.Linear(cfg.d_model, cfg.n_kv_heads * cfg.head_dim, bias=False)
+        self.wo = nn.Linear(cfg.n_heads * cfg.head_dim, cfg.d_model, bias=False)
+
+    def forward(self, x, cos, sin, cache: KVCache | None = None, layer: int = 0, start: int = 0) -> torch.Tensor:
+        B, L, _ = x.shape
+        q = self.wq(x).view(B, L, self.n_heads, self.hd).transpose(1, 2)
+        k = self.wk(x).view(B, L, self.n_kv, self.hd).transpose(1, 2)
+        v = self.wv(x).view(B, L, self.n_kv, self.hd).transpose(1, 2)
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+
+        if cache is not None:
+            cache.k[layer][:, :, start : start + L] = k
+            cache.v[layer][:, :, start : start + L] = v
+            k, v = cache.k[layer][:, :, : start + L], cache.v[layer][:, :, : start + L]
+
+        if self.n_kv != self.n_heads:  # expand KV groups; works with every SDPA backend
+            k = k.repeat_interleave(self.n_heads // self.n_kv, dim=1)
+            v = v.repeat_interleave(self.n_heads // self.n_kv, dim=1)
+
+        if L == 1:  # single decode step attends to the whole cache
+            o = F.scaled_dot_product_attention(q, k, v)
+        elif start == 0:
+            o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:  # chunk appended after a cached prefix: causal w.r.t. absolute positions
+            mask = torch.ones(L, start + L, dtype=torch.bool, device=x.device).tril(diagonal=start)
+            o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        return self.wo(o.transpose(1, 2).reshape(B, L, -1))

quasar/model/axiom.py

@@ -0,0 +1,124 @@
+"""Axiom-1.1 token mixers and dense LawFFN, adapted to Quasar's interface.
+
+Architecture reference: user-owned Axiom/axiom/lm.py and lm_variants.py.
+No fission, dormant slots, or Mamba dependency. FP32 recurrence state; ordinary
+PyTorch SDPA attention. Cache appends support both one token and causal chunks.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .config import ModelConfig
+from .layers import RMSNorm
+
+
+def scan(log_a: torch.Tensor, u: torch.Tensor, chunk: int = 64, initial: torch.Tensor | None = None):
+    """Parallel-within-chunk stable diagonal recurrence; inputs [B,T,H], [B,T,H,P]."""
+    B, T, H, P = u.shape
+    pad = (-T) % chunk
+    la = F.pad(log_a.float(), (0, 0, 0, pad)).reshape(B, -1, chunk, H)
+    uu = F.pad(u.float(), (0, 0, 0, 0, 0, pad)).reshape(B, -1, chunk, H, P)
+    A = la.cumsum(2)
+    causal = torch.ones(chunk, chunk, dtype=torch.bool, device=u.device).tril()
+    diff = (A[:, :, :, None] - A[:, :, None]).masked_fill(~causal[None, None, :, :, None], -math.inf)
+    weights = diff.exp() * -torch.expm1(la)[:, :, None]
+    local = torch.einsum("bnijh,bnjhp->bnihp", weights, uu)
+    state = u.new_zeros(B, H, P, dtype=torch.float32) if initial is None else initial.float()
+    chunks = []
+    for k in range(la.shape[1]):
+        out = local[:, k] + A[:, k].exp()[..., None] * state[:, None]
+        state = out[:, -1]
+        chunks.append(out)
+    result = torch.cat(chunks, 1)[:, :T]
+    return result, result[:, -1]
+
+
+class HybridCache:
+    """Only attention layers store growing KV; recurrent layers store one FP32 vector."""
+    def __init__(self, cfg: ModelConfig, batch: int, max_len: int, device, dtype):
+        self.max_len, self.next_pos = max_len, 0
+        shape = (batch, cfg.n_heads, max_len, cfg.head_dim)
+        self.k = {i: torch.zeros(shape, device=device, dtype=dtype) for i in cfg.attention_layers}
+        self.v = {i: torch.zeros(shape, device=device, dtype=dtype) for i in cfg.attention_layers}
+        self.states = {i: torch.zeros(batch, cfg.n_heads, cfg.head_dim, device=device, dtype=torch.float32)
+                       for i in range(cfg.n_layers) if i not in cfg.attention_layers}
+
+
+class Recurrence(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.h, self.chunk = cfg.n_heads, cfg.scan_chunk
+        self.inp = nn.Linear(cfg.d_model, 2 * cfg.d_model + cfg.n_heads, bias=False)
+        self.out = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+        self.decay_bias = nn.Parameter(torch.linspace(math.log(math.expm1(1.0)), math.log(math.expm1(1e-3)), self.h))
+
+    def forward(self, x, cache=None, layer=0):
+        B, T, D = x.shape
+        u, g, z = self.inp(x).split([D, D, self.h], -1)
+        with torch.autocast(x.device.type, enabled=False):
+            la = -F.softplus(z.float() + self.decay_bias)
+            s, final = scan(la, u.reshape(B, T, self.h, D // self.h), self.chunk,
+                            None if cache is None else cache.states[layer])
+        if cache is not None:
+            cache.states[layer] = final.detach()
+        return self.out(s.reshape(B, T, D).to(g.dtype) * F.silu(g))
+
+
+class AxiomAttention(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.h = cfg.n_heads
+        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False)
+        self.out = nn.Linear(cfg.d_model, cfg.d_model, bias=False)
+
+    def forward(self, x, cos, sin, cache=None, layer=0, start=0):
+        B, T, D = x.shape
+        q, k, v = self.qkv(x).view(B, T, 3, self.h, D // self.h).permute(2, 0, 3, 1, 4)
+        def rotate(t):  # Axiom uses interleaved pairs, rather than Quasar-v1's half pairs
+            a, b = t.float()[..., 0::2], t.float()[..., 1::2]
+            return torch.stack((a * cos - b * sin, a * sin + b * cos), -1).flatten(-2).to(t.dtype)
+        q, k = rotate(q), rotate(k)
+        if cache is not None:
+            cache.k[layer][:, :, start:start + T] = k
+            cache.v[layer][:, :, start:start + T] = v
+            k, v = cache.k[layer][:, :, :start + T], cache.v[layer][:, :, :start + T]
+        mask = None
+        if start and T > 1:
+            mask = torch.ones(T, start + T, dtype=torch.bool, device=x.device).tril(start)
+        o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=start == 0 and T > 1)
+        return self.out(o.transpose(1, 2).reshape(B, T, D))
+
+
+class LawFFN(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.router = nn.Linear(cfg.d_model, cfg.n_cells)
+        self.cells = nn.ModuleList(nn.Sequential(nn.Linear(cfg.d_model, cfg.d_cell), nn.GELU(),
+                                                nn.Linear(cfg.d_cell, cfg.d_model)) for _ in range(cfg.n_cells))
+
+    def forward(self, x):
+        weights = self.router(x).float().softmax(-1)
+        y = sum(cell(x) * weights[..., i:i + 1].to(x.dtype) for i, cell in enumerate(self.cells))
+        # Switch-style fraction is 1/K for top-K == K. This is identically 1;
+        # retain source semantics without claiming it regularizes expert weights.
+        balance = weights.mean(dim=(0, 1)).sum()
+        return y, balance
+
+
+class HybridBlock(nn.Module):
+    def __init__(self, cfg: ModelConfig, index: int):
+        super().__init__()
+        self.is_attention = index in cfg.attention_layers
+        self.n1, self.n2 = RMSNorm(cfg.d_model, cfg.norm_eps), RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.mix = AxiomAttention(cfg) if self.is_attention else Recurrence(cfg)
+        self.ffn = LawFFN(cfg)
+
+    def forward(self, x, cos, sin, cache=None, layer=0, start=0):
+        y = self.mix(self.n1(x), cos, sin, cache, layer, start) if self.is_attention else self.mix(self.n1(x), cache, layer)
+        x = x + y
+        f, balance = self.ffn(self.n2(x))
+        return x + f, balance

quasar/model/config.py

@@ -0,0 +1,73 @@
+"""Model configuration and YAML profile loading."""
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, fields
+from pathlib import Path
+
+import yaml
+
+
+@dataclass
+class ModelConfig:
+    architecture: str = "transformer"  # old checkpoints default to the original transformer
+    vocab_size: int = 9280
+    d_model: int = 512
+    n_layers: int = 8
+    n_heads: int = 8
+    n_kv_heads: int = 8  # < n_heads => grouped-query attention
+    ffn_mult: float = 8 / 3  # SwiGLU hidden = ffn_mult * d_model, rounded up to a multiple of 128
+    max_seq_len: int = 1024
+    rope_theta: float = 10000.0
+    norm_eps: float = 1e-5
+    vis_positions: int = 64  # spatial slots per frame (8x8 grid); 0 disables visual position embedding
+    visual_offset: int = 0  # first visual-code token id (from UnifiedVocab)
+    visual_size: int = 0  # number of visual-code ids
+    tie_embeddings: bool = True
+    init_std: float = 0.02
+    grad_checkpoint: bool = False
+    attention_layers: tuple[int, ...] = (3, 7, 11)  # zero-based Axiom-1.1 indices
+    d_cell: int = 1280
+    n_cells: int = 2
+    top_k: int = 2
+    scan_chunk: int = 64
+    balance_weight: float = 0.01
+
+    def __post_init__(self):
+        assert self.d_model % self.n_heads == 0, "d_model must divide by n_heads"
+        assert self.n_heads % self.n_kv_heads == 0, "n_heads must divide by n_kv_heads"
+        if self.architecture not in ("transformer", "axiom1_1"):
+            raise ValueError(f"unknown architecture: {self.architecture}")
+        if self.architecture == "axiom1_1":
+            self.attention_layers = tuple(self.attention_layers)
+            if self.n_cells != 2 or self.top_k != 2:
+                raise ValueError("Axiom-1.1 uses two cells, both active; sparse routing is not implemented")
+            if self.n_kv_heads != self.n_heads or self.head_dim % 2:
+                raise ValueError("Axiom-1.1 requires MHA and an even head dimension")
+            if len(set(self.attention_layers)) != len(self.attention_layers) or any(i < 0 or i >= self.n_layers for i in self.attention_layers):
+                raise ValueError("attention_layers must contain distinct zero-based layer indices")
+            if self.scan_chunk < 1 or self.d_cell < 1:
+                raise ValueError("scan_chunk and d_cell must be positive")
+
+    @property
+    def head_dim(self) -> int:
+        return self.d_model // self.n_heads
+
+    @property
+    def ffn_hidden(self) -> int:
+        return -(-int(self.ffn_mult * self.d_model) // 128) * 128
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "ModelConfig":
+        names = {f.name for f in fields(cls)}
+        unknown = set(d) - names
+        if unknown:
+            raise ValueError(f"unknown model config keys: {sorted(unknown)}")
+        return cls(**d)
+
+
+def load_yaml(path: str | Path) -> dict:
+    with open(path) as f:
+        return yaml.safe_load(f) or {}

quasar/model/layers.py

@@ -0,0 +1,42 @@
+"""Transformer building blocks: RMSNorm, rotary position embedding, SwiGLU."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        xf = x.float()
+        return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x) * self.weight
+
+
+def rope_tables(head_dim: int, max_len: int, theta: float, device=None) -> tuple[torch.Tensor, torch.Tensor]:
+    """cos/sin tables [max_len, head_dim/2] (float32)."""
+    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
+    ang = torch.outer(torch.arange(max_len, device=device, dtype=torch.float32), inv)
+    return ang.cos(), ang.sin()
+
+
+def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """x: [B, H, L, D]; cos/sin: [L, D/2]. Rotates interleaved-as-halves pairs (GPT-NeoX style)."""
+    x1, x2 = x.float().chunk(2, dim=-1)
+    out = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+    return out.type_as(x)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim: int, hidden: int):
+        super().__init__()
+        self.w_gate = nn.Linear(dim, hidden, bias=False)
+        self.w_up = nn.Linear(dim, hidden, bias=False)
+        self.w_down = nn.Linear(hidden, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))

quasar/model/transformer.py

@@ -0,0 +1,119 @@
+"""Quasar: one decoder-only transformer over the unified text/image/video token space.
+
+Pre-norm blocks (RMSNorm -> GQA attention with RoPE -> residual -> RMSNorm -> SwiGLU -> residual),
+tied input/output embeddings, and one extra signal for visual tokens: a learned embedding of the
+token's spatial slot inside its frame (0 for non-visual tokens).  The slot is a pure function of
+the token sequence (index within the current contiguous run of visual codes), so it carries no
+information the tokens do not already imply; it just saves a small model from learning to count.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+from .attention import Attention, KVCache
+from .axiom import HybridBlock, HybridCache
+from .config import ModelConfig
+from .layers import RMSNorm, SwiGLU, rope_tables
+
+
+class Block(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.attn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.attn = Attention(cfg)
+        self.ffn_norm = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.ffn = SwiGLU(cfg.d_model, cfg.ffn_hidden)
+
+    def forward(self, x, cos, sin, cache=None, layer=0, start=0):
+        x = x + self.attn(self.attn_norm(x), cos, sin, cache, layer, start)
+        return x + self.ffn(self.ffn_norm(x))
+
+
+class Quasar(nn.Module):
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.embed = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.vis_embed = nn.Embedding(cfg.vis_positions + 1, cfg.d_model, padding_idx=0) if cfg.vis_positions else None
+        self.hybrid = cfg.architecture == "axiom1_1"
+        self.blocks = nn.ModuleList(HybridBlock(cfg, i) if self.hybrid else Block(cfg) for i in range(cfg.n_layers))
+        self.norm = RMSNorm(cfg.d_model, cfg.norm_eps)
+        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        if cfg.tie_embeddings:
+            self.head.weight = self.embed.weight
+        cos, sin = rope_tables(cfg.head_dim, cfg.max_seq_len, cfg.rope_theta)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+        self.apply(self._init)
+        for b in self.blocks:  # scale residual-branch outputs (GPT-2 style)
+            residuals = [b.mix.out.weight, *(c[2].weight for c in b.ffn.cells)] if self.hybrid else [b.attn.wo.weight, b.ffn.w_down.weight]
+            for w in residuals:
+                nn.init.normal_(w, std=cfg.init_std / math.sqrt(2 * cfg.n_layers))
+
+    def _init(self, m: nn.Module) -> None:
+        if isinstance(m, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(m.weight, std=self.cfg.init_std)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.zeros_(m.bias)
+            if isinstance(m, nn.Embedding) and m.padding_idx is not None:
+                with torch.no_grad():
+                    m.weight[m.padding_idx].zero_()
+
+    def num_params(self, include_embedding: bool = True) -> int:
+        n = sum(p.numel() for p in self.parameters())  # tied weights counted once
+        return n if include_embedding else n - self.embed.weight.numel()
+
+    def visual_positions(self, tokens: torch.Tensor) -> torch.Tensor:
+        """1-based index of each visual token inside its contiguous run (0 for other tokens)."""
+        lo, hi = self.cfg.visual_offset, self.cfg.visual_offset + self.cfg.visual_size
+        is_vis = (tokens >= lo) & (tokens < hi)
+        idx = torch.arange(tokens.shape[-1], device=tokens.device).expand_as(tokens)
+        run_start = torch.where(is_vis, torch.zeros_like(idx), idx + 1).cummax(dim=-1).values
+        return torch.where(is_vis, idx - run_start + 1, torch.zeros_like(idx)).clamp(max=self.cfg.vis_positions)
+
+    def forward(self, tokens: torch.Tensor, vis_pos: torch.Tensor | None = None, cache: KVCache | HybridCache | None = None, start: int = 0,
+                return_aux: bool = False):
+        """tokens [B, L] -> logits [B, L, V].  With a cache, `start` is the absolute position of tokens[:, 0]."""
+        B, L = tokens.shape
+        if start + L > self.cfg.max_seq_len:
+            raise ValueError(f"sequence length {start + L} exceeds max_seq_len {self.cfg.max_seq_len}")
+        if cache is not None and self.hybrid:
+            if not isinstance(cache, HybridCache) or start != cache.next_pos or start + L > cache.max_len:
+                raise ValueError("hybrid cache requires consecutive appends within its capacity; use a fresh cache for a new sequence")
+        x = self.embed(tokens)
+        if self.vis_embed is not None:
+            x = x + self.vis_embed(self.visual_positions(tokens) if vis_pos is None else vis_pos)
+        cos, sin = self.rope_cos[start : start + L], self.rope_sin[start : start + L]
+        balances = []
+        for i, blk in enumerate(self.blocks):
+            if self.cfg.grad_checkpoint and self.training and cache is None:
+                result = checkpoint(blk, x, cos, sin, use_reentrant=False)
+            else:
+                result = blk(x, cos, sin, cache, i, start)
+            if self.hybrid:
+                x, balance = result
+                balances.append(balance)
+            else:
+                x = result
+        if cache is not None and self.hybrid:
+            cache.next_pos = start + L
+        logits = self.head(self.norm(x))
+        aux = torch.stack(balances).mean() if balances else logits.new_zeros(())
+        return (logits, aux) if return_aux else logits
+
+    def loss(self, tokens: torch.Tensor, labels: torch.Tensor, include_aux: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
+        """Mean target-token CE, optionally plus Axiom balance loss; (objective, n_targets)."""
+        logits, aux = self(tokens, return_aux=True)
+        n = (labels != -100).sum()
+        loss = F.cross_entropy(logits.float().reshape(-1, logits.size(-1)), labels.reshape(-1), ignore_index=-100)
+        return loss + (self.cfg.balance_weight * aux if include_aux else 0), n
+
+    def new_cache(self, batch: int, max_len: int | None = None, dtype=None) -> KVCache:
+        p = next(self.parameters())
+        cache_type = HybridCache if self.hybrid else KVCache
+        return cache_type(self.cfg, batch, max_len or self.cfg.max_seq_len, p.device, dtype or p.dtype)

quasar/tokenizers/audio.py

@@ -0,0 +1,48 @@
+"""Frozen DAC interface. Loads ONLY explicitly supplied local safetensors weights.
+
+No implicit downloads. The shared Quasar model predicts the codes; this module
+only reconstructs waveforms. Four-of-32 codebook quality is unmeasured here.
+"""
+from pathlib import Path
+import math
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+class AudioTokenizer:
+    def __init__(self, local_path: str | Path, codebooks: int = 4, max_seconds: float = 2.0, device: str = "cpu"):
+        p = Path(local_path)
+        if not (p / "config.json").exists() or not (p / "model.safetensors").exists():
+            raise FileNotFoundError(f"DAC is not installed at {p}; provision the pinned codec explicitly before audio preprocessing")
+        from transformers import DacModel
+        self.model = DacModel.from_pretrained(str(p), local_files_only=True, use_safetensors=True).to(device).eval()
+        self.model.requires_grad_(False)
+        self.sr, self.books, self.max_seconds = self.model.config.sampling_rate, codebooks, max_seconds
+        self.hop = math.prod(self.model.config.downsampling_ratios)
+        if not 1 <= codebooks <= self.model.config.n_codebooks:
+            raise ValueError("requested codebooks exceed the codec")
+        self.device = device
+
+    @torch.no_grad()
+    def encode(self, waveform, sample_rate: int) -> dict:
+        x = torch.as_tensor(waveform, dtype=torch.float32, device=self.device)
+        if sample_rate != self.sr:
+            raise ValueError(f"resample to {self.sr} Hz before encoding; received {sample_rate}")
+        if x.ndim != 1 or not x.numel() or not torch.isfinite(x).all() or x.abs().max() > 1.001:
+            raise ValueError("expected finite mono waveform in [-1,1]")
+        if x.numel() > self.sr * self.max_seconds:
+            raise ValueError("audio exceeds duration budget; use transcript-aligned clips")
+        n = x.numel()
+        x = F.pad(x, (0, (-n) % self.hop))[None, None]
+        codes = self.model.encode(x, n_quantizers=self.books).audio_codes[0].T
+        return {"type": "audio", "codes": codes.cpu().numpy().astype(np.int32).tolist(), "samples": n, "sample_rate": self.sr}
+
+    @torch.no_grad()
+    def decode(self, packet: dict) -> np.ndarray:
+        a = np.asarray(packet["codes"])
+        if a.ndim != 2 or a.shape[1] != self.books or a.dtype.kind not in "iu" or not a.size or a.min() < 0 or a.max() >= self.model.config.codebook_size:
+            raise ValueError("invalid DAC codes")
+        codes = torch.tensor(a.T[None], dtype=torch.long, device=self.device)
+        y = self.model.decode(audio_codes=codes).audio_values[0]
+        return y[:packet.get("samples", len(y))].cpu().numpy()

quasar/tokenizers/audio_aligned.py

@@ -0,0 +1,35 @@
+"""Frozen DAC interface with explicit, training-calibrated boundary alignment.
+
+Source utterances remain whole. Silent boundary guards are codec padding, and
+the fixed decoder offset is recorded separately from the generative weights.
+"""
+import numpy as np
+from .audio import AudioTokenizer
+
+
+class AlignedAudioTokenizer(AudioTokenizer):
+    def __init__(self, local_path, alignment_samples: int, codebooks=4, max_seconds=2.0, device='cpu'):
+        super().__init__(local_path, codebooks, max_seconds, device)
+        if self.sr != 24000 or self.hop != 320 or abs(alignment_samples) > 32:
+            raise ValueError('alignment is only supported for the pinned 24 kHz/320-hop codec')
+        self.source_max_seconds = max_seconds
+        self.max_seconds = max_seconds + 2 * self.hop / self.sr
+        self.alignment_samples = alignment_samples
+
+    def encode(self, waveform, sample_rate):
+        x = np.asarray(waveform)
+        if x.ndim != 1 or not 0 < len(x) <= self.source_max_seconds * self.sr:
+            raise ValueError('complete source utterance exceeds duration budget')
+        packet = super().encode(np.pad(x, (self.hop, self.hop)), sample_rate)
+        packet['samples'] = len(x)
+        return packet
+
+    def decode(self, packet):
+        # Ask the base decoder for all samples, before applying boundary guards.
+        expanded = dict(packet, samples=(len(packet['codes']) + 2) * self.hop)
+        raw = super().decode(expanded)
+        start = self.hop + self.alignment_samples
+        n = packet.get('samples', min(int(self.source_max_seconds * self.sr), len(raw) - 2 * self.hop))
+        if not 0 < n <= self.source_max_seconds * self.sr or start + n > len(raw):
+            raise ValueError('decoded audio cannot preserve the complete source duration')
+        return raw[start:start + n]

quasar/tokenizers/audio_complete.py

@@ -0,0 +1,23 @@
+"""DAC boundary padding for complete source utterances; no source cropping."""
+import numpy as np
+from .audio import AudioTokenizer
+
+
+class CompleteAudioTokenizer(AudioTokenizer):
+    """Guard the decoder's eight-sample edge shortfall with a full silent hop.
+
+The original source must fit the two-second budget. Only codec boundary padding
+extends beyond it; the packet retains the original sample count for decoding.
+"""
+    def __init__(self, local_path, codebooks=4, max_seconds=2.0, device='cpu'):
+        super().__init__(local_path, codebooks, max_seconds, device)
+        self.source_max_seconds = max_seconds
+        self.max_seconds = max_seconds + self.hop / self.sr
+
+    def encode(self, waveform, sample_rate):
+        x = np.asarray(waveform)
+        if x.ndim != 1 or not 0 < len(x) <= self.source_max_seconds * self.sr:
+            raise ValueError('complete source utterance exceeds duration budget')
+        packet = super().encode(np.pad(x, (0, self.hop)), sample_rate)
+        packet['samples'] = len(x)
+        return packet

quasar/tokenizers/mesh.py

@@ -0,0 +1,35 @@
+"""Deterministic quantized triangle-mesh codec; no pretrained generator or mesh simplification."""
+from pathlib import Path
+import numpy as np
+
+
+def quantize_mesh(vertices, faces, bins: int = 512, max_vertices: int = 128, max_faces: int = 128):
+    v, f = np.asarray(vertices, dtype=np.float64), np.asarray(faces)
+    if v.ndim != 2 or v.shape[1] != 3 or not 3 <= len(v) <= max_vertices or not np.isfinite(v).all():
+        raise ValueError("expected finite mesh vertices; simplify externally if over the vertex budget")
+    if f.ndim != 2 or f.shape[1] != 3 or not 1 <= len(f) <= max_faces or f.dtype.kind not in "iu" or f.min() < 0 or f.max() >= len(v):
+        raise ValueError("invalid triangle faces or face budget exceeded")
+    center = (v.min(0) + v.max(0)) / 2
+    scale = float(np.ptp(v, axis=0).max())
+    if scale <= 1e-12:
+        raise ValueError("zero-extent mesh")
+    q = np.rint(((v - center) / scale + 0.5) * (bins - 1)).clip(0, bins - 1).astype(np.int32)
+    # Quantization can merge vertices or collapse skinny triangles: reject these samples.
+    if len(np.unique(q, axis=0)) != len(q) or np.any(np.linalg.norm(np.cross(q[f[:, 1]] - q[f[:, 0]], q[f[:, 2]] - q[f[:, 0]]), axis=1) == 0):
+        raise ValueError("degenerate geometry after quantization")
+    return {"type": "mesh", "vertices": q.tolist(), "faces": f.tolist()}, {"center": center.tolist(), "scale": scale, "bins": bins}
+
+
+def decode_mesh(packet: dict, transform: dict | None = None, bins: int = 512):
+    v = np.asarray(packet["vertices"], dtype=np.float64) / (bins - 1) - 0.5
+    if transform is not None:
+        v = v * transform["scale"] + transform["center"]
+    return v, np.asarray(packet["faces"], dtype=np.int32)
+
+
+def write_obj(packet: dict, path: str | Path, transform: dict | None = None):
+    """Export geometry as a real OBJ; materials/textures are deliberately outside this codec."""
+    v, f = decode_mesh(packet, transform)
+    text = "".join(f"v {x:.7f} {y:.7f} {z:.7f}\n" for x, y, z in v)
+    text += "".join(f"f {a+1} {b+1} {c+1}\n" for a, b, c in f)
+    Path(path).write_text(text, encoding="utf-8")

quasar/tokenizers/multimodal.py

@@ -0,0 +1,197 @@
+"""Unified token space shared by text, images and video.
+
+ID layout (non-overlapping ranges):
+
+    [0, T)                  text BPE tokens
+    [T, T + 64)             structural / special tokens (64 reserved slots)
+    [T + 64, T + 64 + V)    visual codes from the frozen FSQ tokenizer
+    padded up to a multiple of 64 (unused tail ids are never produced or targeted)
+
+Sequence formats (one per task; loss is applied from `target_start` onward):
+
+    text   <bos> text... <eos>
+    t2i    <bos><task:t2i> caption <image> v*G </image> <eos>
+    i2t    <bos><task:i2t> <image> v*G </image> caption <eos>
+    i2i    <bos><task:i2i> instruction <image> src </image> <image> tgt </image> <eos>
+    t2v    <bos><task:t2v> caption <video> (<frame><t_k> v*G </frame>)*F </video> <eos>
+    v2t    <bos><task:v2t> <video> ... </video> caption <eos>
+    vcont  <bos><task:vcont> <video> ... </video> <eos>
+    chat   <bos><user> text [media] <assistant> text [media] <eos>
+"""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+
+import numpy as np
+
+N_SPECIAL_SLOTS = 64
+MAX_FRAMES = 32
+SPECIALS = [
+    "<pad>", "<bos>", "<eos>",
+    "<task:text>", "<task:t2i>", "<task:i2t>", "<task:i2i>", "<task:t2v>", "<task:v2t>", "<task:vcont>",
+    "<user>", "<assistant>",
+    "<image>", "</image>", "<video>", "</video>", "<frame>", "</frame>",
+] + [f"<t{i}>" for i in range(MAX_FRAMES)]  # frame timestamps (index; seconds = k / fps)
+assert len(SPECIALS) <= N_SPECIAL_SLOTS
+
+
+@dataclass
+class Sample:
+    tokens: np.ndarray  # int32 [L]
+    target_start: int  # loss on predicting tokens[target_start:]
+
+
+class UnifiedVocab:
+    def __init__(self, text_size: int, visual_size: int, grid: tuple[int, int] = (8, 8), fps: int = 4):
+        self.text_size = text_size
+        self.visual_size = visual_size
+        self.grid = tuple(grid)
+        self.fps = fps
+        self.special_offset = text_size
+        self.visual_offset = text_size + N_SPECIAL_SLOTS
+        self.size = -(-(self.visual_offset + visual_size) // 64) * 64
+        self.ids = {name: self.special_offset + i for i, name in enumerate(SPECIALS)}
+        self.names = {v: k for k, v in self.ids.items()}
+        self.pad, self.bos, self.eos = self.ids["<pad>"], self.ids["<bos>"], self.ids["<eos>"]
+        self.img_start, self.img_end = self.ids["<image>"], self.ids["</image>"]
+        self.vid_start, self.vid_end = self.ids["<video>"], self.ids["</video>"]
+        self.frm_start, self.frm_end = self.ids["<frame>"], self.ids["</frame>"]
+        self.user, self.assistant = self.ids["<user>"], self.ids["<assistant>"]
+
+    @property
+    def tokens_per_frame(self) -> int:
+        return self.grid[0] * self.grid[1]
+
+    def task(self, name: str) -> int:
+        return self.ids[f"<task:{name}>"]
+
+    def timestamp(self, k: int) -> int:
+        return self.ids[f"<t{k}>"]
+
+    def is_visual(self, t) -> np.ndarray:
+        t = np.asarray(t)
+        return (t >= self.visual_offset) & (t < self.visual_offset + self.visual_size)
+
+    def is_text(self, t) -> np.ndarray:
+        t = np.asarray(t)
+        return (t >= 0) & (t < self.text_size)
+
+    # ------------------------------------------------------------------ media -> tokens
+    def image_tokens(self, codes: np.ndarray) -> list[int]:
+        codes = np.asarray(codes).reshape(-1)
+        assert codes.size == self.tokens_per_frame and codes.min() >= 0 and codes.max() < self.visual_size
+        return [self.img_start, *(codes + self.visual_offset).tolist(), self.img_end]
+
+    def frame_tokens(self, codes: np.ndarray, k: int) -> list[int]:
+        codes = np.asarray(codes).reshape(-1)
+        return [self.frm_start, self.timestamp(k), *(codes + self.visual_offset).tolist(), self.frm_end]
+
+    def video_tokens(self, codes: np.ndarray, first_index: int = 0, close: bool = True) -> list[int]:
+        out = [self.vid_start]
+        for k, f in enumerate(np.asarray(codes)):
+            out += self.frame_tokens(f, first_index + k)
+        return out + ([self.vid_end] if close else [])
+
+    # ------------------------------------------------------------------ task builders
+    def build(self, task: str, text: list[int] | None = None, image=None, image2=None, video=None) -> Sample:
+        b = [self.bos]
+        if task == "t2i":
+            p = b + [self.task("t2i"), *text]
+            return self._s(p, self.image_tokens(image) + [self.eos])
+        if task == "i2t":
+            return self._s(b + [self.task("i2t"), *self.image_tokens(image)], [*text, self.eos])
+        if task == "i2i":
+            return self._s(b + [self.task("i2i"), *text, *self.image_tokens(image)], self.image_tokens(image2) + [self.eos])
+        if task == "t2v":
+            return self._s(b + [self.task("t2v"), *text], self.video_tokens(video) + [self.eos])
+        if task == "v2t":
+            return self._s(b + [self.task("v2t"), *self.video_tokens(video)], [*text, self.eos])
+        if task == "vcont":
+            return self._s(b + [self.task("vcont")], self.video_tokens(video) + [self.eos])
+        raise ValueError(f"unknown task {task}")
+
+    def chat(self, user: list[int], assistant: list[int]) -> Sample:
+        """user/assistant are already-tokenized turn contents (text and/or media tokens)."""
+        return self._s([self.bos, self.user, *user, self.assistant], [*assistant, self.eos])
+
+    @staticmethod
+    def _s(prompt: list[int], target: list[int]) -> Sample:
+        return Sample(np.array(prompt + target, dtype=np.int32), len(prompt))
+
+    # ------------------------------------------------------------------ structure
+    def check_structure(self, tokens) -> None:
+        """Validate media blocks: right number of codes, balanced tags, ordered timestamps."""
+        t = np.asarray(tokens).tolist()
+        G = self.tokens_per_frame
+        i, n = 0, len(t)
+        while i < n:
+            tok = t[i]
+            if tok == self.img_start:
+                body = t[i + 1 : i + 1 + G]
+                if len(body) != G or not self.is_visual(body).all() or i + 1 + G >= n or t[i + 1 + G] != self.img_end:
+                    raise ValueError(f"malformed image block at {i}")
+                i += G + 2
+            elif tok == self.vid_start:
+                i += 1
+                k = None
+                while i < n and t[i] == self.frm_start:
+                    ts = t[i + 1] - self.timestamp(0) if i + 1 < n else -1
+                    if not (0 <= ts < MAX_FRAMES) or (k is not None and ts != k + 1):
+                        raise ValueError(f"bad frame timestamp at {i + 1}")
+                    k = ts
+                    body = t[i + 2 : i + 2 + G]
+                    if len(body) != G or not self.is_visual(body).all() or i + 2 + G >= n or t[i + 2 + G] != self.frm_end:
+                        raise ValueError(f"malformed frame at {i}")
+                    i += G + 3
+                if i >= n or t[i] != self.vid_end or k is None:
+                    raise ValueError("video block not closed or empty")
+                i += 1
+            elif self.is_visual(tok) or tok in (self.img_end, self.vid_end, self.frm_start, self.frm_end):
+                raise ValueError(f"visual/structural token {tok} outside a media block at {i}")
+            else:
+                i += 1
+
+    def parse(self, tokens) -> list[tuple[str, object]]:
+        """Split a generated sequence into ('text', ids) / ('image', codes[h,w]) / ('video', codes[T,h,w])."""
+        t = list(np.asarray(tokens).tolist())
+        G, (h, w) = self.tokens_per_frame, self.grid
+        out: list[tuple[str, object]] = []
+        txt: list[int] = []
+        i = 0
+        while i < len(t):
+            if t[i] == self.img_start and i + G < len(t):
+                out += [("text", txt)] if txt else []
+                txt = []
+                out.append(("image", (np.array(t[i + 1 : i + 1 + G]) - self.visual_offset).reshape(h, w)))
+                i += G + 2
+            elif t[i] == self.vid_start:
+                out += [("text", txt)] if txt else []
+                txt = []
+                frames = []
+                i += 1
+                while i + 2 + G <= len(t) and t[i] == self.frm_start:
+                    frames.append((np.array(t[i + 2 : i + 2 + G]) - self.visual_offset).reshape(h, w))
+                    i += G + 3
+                if frames:
+                    out.append(("video", np.stack(frames)))
+                i += 1 if i < len(t) and t[i] == self.vid_end else 0
+            else:
+                if self.is_text(t[i]):
+                    txt.append(t[i])
+                i += 1
+        return out + ([("text", txt)] if txt else [])
+
+    def to_dict(self) -> dict:
+        return {"text_size": self.text_size, "visual_size": self.visual_size, "grid": list(self.grid), "fps": self.fps,
+                "special_offset": self.special_offset, "visual_offset": self.visual_offset, "size": self.size, "specials": SPECIALS}
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "UnifiedVocab":
+        v = cls(d["text_size"], d["visual_size"], tuple(d["grid"]), d.get("fps", 4))
+        assert v.to_dict()["specials"] == d.get("specials", SPECIALS), "special-token table changed since checkpoint"
+        return v
+
+    def save(self, path) -> None:
+        with open(path, "w") as f:
+            json.dump(self.to_dict(), f, indent=2)

quasar/tokenizers/multires_emu.py

@@ -0,0 +1,49 @@
+"""Isolated candidate codec adapter; not registered with the live v2 trainer.
+
+Frozen full Emu3 weights support square 64/128 images. This changes shapes only;
+packet grammar, spatial embedding migration and training admission are separate.
+"""
+import numpy as np
+import torch
+from .restricted_emu import RestrictedEmuCodec
+
+
+class MultiResolutionEmuCodec(RestrictedEmuCodec):
+    def __init__(self, path, device='cpu'):
+        super().__init__(path, device)
+        if self.config.get('codec') != 'emu3_full_32768' or self.codebook_size != 32768:
+            raise ValueError('Only the pinned full visual codebook is supported')
+
+    @torch.inference_mode()
+    def encode(self, images, batch=1):
+        images = np.asarray(images)
+        if images.dtype != np.uint8 or images.ndim < 4 or images.shape[-3:] not in ((64, 64, 3), (128, 128, 3)):
+            raise ValueError('Expected uint8 square 64px or 128px RGB batches')
+        if batch < 1 or images.size == 0:
+            raise ValueError('Empty image batch or invalid microbatch')
+        lead, resolution = images.shape[:-3], images.shape[-2]
+        flat, out = images.reshape(-1, resolution, resolution, 3), []
+        for i in range(0, len(flat), batch):
+            x = torch.tensor(flat[i:i + batch], device=self.selected.device).permute(0, 3, 1, 2).float() / 127.5 - 1
+            codes = self.full_to_local[self.model.encode(x)].cpu().numpy()
+            expected = (len(x), resolution // 8, resolution // 8)
+            if codes.shape != expected or codes.min() < 0 or codes.max() >= self.codebook_size:
+                raise ValueError('Codec produced unexpected shape or token range')
+            out.append(codes)
+        return np.concatenate(out).reshape(*lead, resolution // 8, resolution // 8)
+
+    @torch.inference_mode()
+    def decode(self, codes, batch=1):
+        codes = np.asarray(codes)
+        if (codes.ndim < 3 or codes.shape[-2:] not in ((8, 8), (16, 16)) or codes.dtype.kind not in 'iu'
+                or codes.size == 0 or batch < 1 or codes.min() < 0 or codes.max() >= self.codebook_size):
+            raise ValueError('Invalid visual codes or token grid')
+        lead, grid = codes.shape[:-2], codes.shape[-1]
+        flat, out = codes.reshape(-1, grid, grid), []
+        for i in range(0, len(flat), batch):
+            full = self.selected[torch.as_tensor(flat[i:i + batch], device=self.selected.device, dtype=torch.long)]
+            x = self.model.decode(full)
+            if x.shape != (len(full), 3, grid * 8, grid * 8):
+                raise ValueError('Codec decoder returned unexpected spatial shape')
+            out.append(((x.float().clamp(-1, 1) + 1) * 127.5).round().byte().permute(0, 2, 3, 1).cpu().numpy())
+        return np.concatenate(out).reshape(*lead, grid * 8, grid * 8, 3)

quasar/tokenizers/restricted_emu.py

@@ -0,0 +1,52 @@
+"""Frozen pinned Emu3 pixel codec with a training-only 1,000-entry subset.
+
+Quasar still predicts every image code. The frozen codec maps those codes back
+to pixels. No pretrained generative backbone is loaded. Per-frame mode retains
+the existing 64-code grammar and avoids changing the trained text vocabulary.
+"""
+from pathlib import Path
+import json
+import numpy as np
+import torch
+from torch import nn
+from safetensors.torch import load_file
+
+
+class RestrictedEmuCodec(nn.Module):
+    codebook_size = 1000
+    downsample = 8
+
+    def __init__(self, path, device='cpu'):
+        super().__init__()
+        from transformers import AutoModel
+        path = Path(path)
+        self.config = json.loads((path/'vq_config.json').read_text())
+        self.codebook_size = self.config['codebook_size']
+        base = Path(self.config['base_path'])
+        if not base.is_absolute():
+            base = (path / base).resolve()
+        self.model = AutoModel.from_pretrained(str(base), trust_remote_code=True,
+                                              local_files_only=True).to(device).eval()
+        self.model.requires_grad_(False)
+        mapping = load_file(str(path/'vq.safetensors'))
+        self.register_buffer('selected', mapping['selected'].to(device))
+        self.register_buffer('full_to_local', mapping['full_to_local'].to(device))
+
+    @torch.inference_mode()
+    def encode(self, images, batch=8):
+        images=np.asarray(images);lead=images.shape[:-3];flat=images.reshape(-1,*images.shape[-3:]);out=[]
+        for i in range(0,len(flat),batch):
+            x=torch.tensor(flat[i:i+batch],device=self.selected.device).permute(0,3,1,2).float()/127.5-1
+            full=self.model.encode(x)
+            out.append(self.full_to_local[full].cpu().numpy())
+        return np.concatenate(out).reshape(*lead,8,8)
+
+    @torch.inference_mode()
+    def decode(self, codes, batch=8):
+        codes=np.asarray(codes);lead=codes.shape[:-2];flat=codes.reshape(-1,8,8);out=[]
+        if codes.dtype.kind not in 'iu' or codes.min()<0 or codes.max()>=self.codebook_size:raise ValueError('invalid codec codes')
+        for i in range(0,len(flat),batch):
+            full=self.selected[torch.tensor(flat[i:i+batch],device=self.selected.device,dtype=torch.long)]
+            x=self.model.decode(full)
+            out.append(((x.float().clamp(-1,1)+1)*127.5).round().byte().permute(0,2,3,1).cpu().numpy())
+        return np.concatenate(out).reshape(*lead,64,64,3)

quasar/tokenizers/text.py

@@ -0,0 +1,47 @@
+"""Byte-level BPE text tokenizer (trained locally with HF `tokenizers`).
+
+Byte-level means every UTF-8 string is encodable (no UNK).  The tokenizer has no special
+tokens of its own: structural tokens live in the unified vocabulary (multimodal.py), above the
+text range, so text ids are exactly [0, vocab_size).
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable
+
+from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
+
+
+class TextTokenizer:
+    def __init__(self, tok: Tokenizer):
+        self.tok = tok
+
+    @classmethod
+    def train(cls, texts: Iterable[str], vocab_size: int = 8192) -> "TextTokenizer":
+        tok = Tokenizer(models.BPE())
+        tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+        tok.decoder = decoders.ByteLevel()
+        trainer = trainers.BpeTrainer(vocab_size=vocab_size, initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), show_progress=False)
+        tok.train_from_iterator(texts, trainer)
+        return cls(tok)
+
+    @property
+    def vocab_size(self) -> int:
+        return self.tok.get_vocab_size()
+
+    def encode(self, text: str) -> list[int]:
+        return self.tok.encode(text).ids
+
+    def encode_batch(self, texts: list[str]) -> list[list[int]]:
+        return [e.ids for e in self.tok.encode_batch(texts)]
+
+    def decode(self, ids: Iterable[int]) -> str:
+        return self.tok.decode([i for i in ids if 0 <= i < self.vocab_size])
+
+    def save(self, path: str | Path) -> None:
+        Path(path).parent.mkdir(parents=True, exist_ok=True)
+        self.tok.save(str(path))
+
+    @classmethod
+    def load(cls, path: str | Path) -> "TextTokenizer":
+        return cls(Tokenizer.from_file(str(path)))