Load sage sheet and use it on training with 50% dropout so it generates sheet music and handles without it.

toolkit/models/registry.py

@@ -55,7 +55,7 @@ ARCH_REGISTRY: Dict[str, dict] = {
     "ltx2.3": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.3/ltx-2.3-22b-dev.safetensors", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}},
     "ltx2.5": {"modality": "video", "model": {"name_or_path": "Lightricks/LTX-2.5", "quantize": True, "quantize_te": True}, "sample": {**VID, "num_inference_steps": 25, "guidance_scale": 3.0, "num_frames": 49, "fps": 24}},
     "ace_step_15": {"modality": "audio", "model": {"name_or_path": "ostris/ace_step_1.5_ComfyUI_files/ace_step_1.5_base_aio.safetensors", "quantize": True, "quantize_te": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 20, "guidance_scale": 4.0}},
-    "yue2": {"modality": "audio", "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_bf16.safetensors", "quantize": True}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}},
+    "yue2": {"modality": "audio", "model": {"name_or_path": "Comfy-Org/YuE2/checkpoints/yue2_3b_int8_convrot.safetensors", "quantize": True, "qtype": "convrot8"}, "sample": {"width": 512, "height": 512, "num_inference_steps": 32, "guidance_scale": 1.0, "duration": 120}},
 }
 
 OUTPUT_EXT = {"image": "png", "video": "mp4", "audio": "mp3"}

toolkit/util/comfy_quant_import.py

@@ -247,7 +247,31 @@ def import_comfy_quantized_layers(
                 )
             parent_path, _, attr = module_path.rpartition(".")
             parent = root.get_submodule(parent_path) if parent_path else root
-            setattr(parent, attr, Int8Embedding(weight, weight_scale, orig_dtype))
+            if conf.get("convrot"):
+                # rows are stored rotated like a linear's weight: attach to a stand-in linear
+                # (same shape) so the convrot backend undoes the rotation, then keep the table
+                # as a plain embedding (per-row lookups cannot apply the group rotation)
+                stand_in = torch.nn.Linear(
+                    module.embedding_dim, module.num_embeddings, bias=False, device="meta"
+                )
+                quantizer = get_ostris_quantizer("convrot8")
+                stand_in.register_buffer("cr8_qdata", weight.contiguous(), persistent=False)
+                stand_in.register_buffer(
+                    "cr8_scales",
+                    weight_scale.detach().float().reshape(-1).contiguous().view(torch.uint8),
+                    persistent=False,
+                )
+                stand_in.cr8_rot_size = int(conf.get("convrot_groupsize", 256))
+                _to_ostris(stand_in, quantizer, orig_dtype)
+                table = quantizer.dequantize_folded(stand_in).to(orig_dtype)
+                new = torch.nn.Embedding(
+                    module.num_embeddings, module.embedding_dim, padding_idx=module.padding_idx,
+                    device="meta", dtype=orig_dtype,
+                )
+                new.weight = torch.nn.Parameter(table, requires_grad=False)
+                setattr(parent, attr, new)
+            else:
+                setattr(parent, attr, Int8Embedding(weight, weight_scale, orig_dtype))
             converted += 1
             continue
 

toolkit/util/streamed_safetensors.py

@@ -0,0 +1,62 @@
+import json
+import struct
+from typing import Dict, Optional
+
+import torch
+
+# safetensors dtype tags
+_DTYPE_TAGS = {
+    torch.float64: "F64",
+    torch.float32: "F32",
+    torch.float16: "F16",
+    torch.bfloat16: "BF16",
+    torch.int64: "I64",
+    torch.int32: "I32",
+    torch.int16: "I16",
+    torch.int8: "I8",
+    torch.uint8: "U8",
+    torch.bool: "BOOL",
+}
+
+_CHUNK_BYTES = 256 * 1024 * 1024
+
+
+def save_file_streamed(
+    tensors: Dict[str, torch.Tensor],
+    path: str,
+    metadata: Optional[Dict[str, str]] = None,
+):
+    """Write a safetensors file with sequential writes instead of ftruncate+mmap.
+    Some filesystems (kernel ntfs3) reject one large ftruncate; incremental writes succeed."""
+    header = {}
+    offset = 0
+    flat = {}
+    for name, t in tensors.items():
+        if t.dtype not in _DTYPE_TAGS:
+            raise ValueError(f"Unsupported dtype {t.dtype} for tensor {name}")
+        t = t.detach().to("cpu").contiguous()
+        nbytes = t.numel() * t.element_size()
+        header[name] = {
+            "dtype": _DTYPE_TAGS[t.dtype],
+            "shape": list(t.shape),
+            "data_offsets": [offset, offset + nbytes],
+        }
+        flat[name] = t
+        offset += nbytes
+    if metadata is not None:
+        header["__metadata__"] = {str(k): str(v) for k, v in metadata.items()}
+
+    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    # pad header to 8-byte alignment like the reference writer
+    pad = (8 - len(header_bytes) % 8) % 8
+    header_bytes += b" " * pad
+
+    with open(path, "wb") as f:
+        f.write(struct.pack("<Q", len(header_bytes)))
+        f.write(header_bytes)
+        for name, t in flat.items():
+            if t.numel() == 0:
+                continue
+            buf = memoryview(t.reshape(-1).view(torch.uint8).numpy())
+            for start in range(0, len(buf), _CHUNK_BYTES):
+                f.write(buf[start:start + _CHUNK_BYTES])