mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-11-04 09:57:44 -06:00
Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Xuan-Son NguyenandPascal
|
5851a8e505 | ||
|
Aleksander Grygier
|
efbf6e80a5 | ||
|
Masashi Yoshimura
|
b393d73bb0 | ||
|
Neo ZhangandSigbjørn Skjæret
|
00df7c22d9 | ||
|
Aaron Thomas
|
034e2535a9 | ||
|
Adrien Gallouët
|
52e53a76df | ||
|
Georgi Gerganov
|
cd959bae99 |
@@ -64,7 +64,7 @@ RUN mkdir -p /app/full \
|
||||
&& cp requirements.txt /app/full \
|
||||
&& cp .devops/tools.sh /app/full/tools.sh
|
||||
|
||||
FROM docker.io/intel/oneapi-toolkit:$ONEAPI_VERSION AS base
|
||||
FROM docker.io/intel/deep-learning-essentials:$ONEAPI_VERSION AS base
|
||||
|
||||
ARG BUILD_DATE=N/A
|
||||
ARG APP_VERSION=N/A
|
||||
@@ -79,6 +79,8 @@ LABEL org.opencontainers.image.created=$BUILD_DATE \
|
||||
org.opencontainers.image.url=$IMAGE_URL \
|
||||
org.opencontainers.image.source=$IMAGE_SOURCE
|
||||
|
||||
COPY --from=build /opt/intel/oneapi/dnnl /opt/intel/oneapi/dnnl
|
||||
|
||||
#Following versions are for multiple GPUs, since 26.x has known issue:
|
||||
# https://github.com/ggml-org/llama.cpp/issues/21747,
|
||||
# https://github.com/intel/compute-runtime/issues/921.
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
# date: Wed Oct 28 12:38:26 EET 2026
|
||||
# date: Wed Nov 4 12:03:31 EET 2026
|
||||
# this file is auto-generated by scripts/gen-authors.sh
|
||||
|
||||
Нияз Гарифзянов <112617865+garrnizon@users.noreply.github.com>
|
||||
@@ -48,6 +48,7 @@ Abheek Gulati <abheekg@hotmail.com>
|
||||
abhijain1204fujitsu <139222713+abhijain1204fujitsu@users.noreply.github.com>
|
||||
Abhijit Ramesh <abhijitramesh2k@gmail.com>
|
||||
abhijitb11 <113058133+abhijitb11@users.noreply.github.com>
|
||||
Abhijith Jajam <129675815+jajamabhijith@users.noreply.github.com>
|
||||
Abhilash Majumder <30946547+abhilash1910@users.noreply.github.com>
|
||||
Abhinay Krishna <abhinaykrishna60@gmail.com>
|
||||
Abhiram <78226909+geckguy@users.noreply.github.com>
|
||||
@@ -203,14 +204,17 @@ Ankit Khandelwal <ankk98@gmail.com>
|
||||
Ankur Verma <31362771+ankurvdev@users.noreply.github.com>
|
||||
anon998 <131767832+anon998@users.noreply.github.com>
|
||||
Anri Lombard <anri.m.lombard@gmail.com>
|
||||
Anthony Maltsev <77708730+anthonymaltsev@users.noreply.github.com>
|
||||
Anthony Umfer <aumfer@gmail.com>
|
||||
Anthony Van de Gejuchte <anthonyvdgent@gmail.com>
|
||||
Anthony Wang <anthonywang903@gmail.com>
|
||||
antichristHater <142441588+antichristHater@users.noreply.github.com>
|
||||
Antoine Viallon <antoine@lesviallon.fr>
|
||||
Anton I <inflyer@hotmail.com>
|
||||
Anton Mitkov <anton_b_mitkov@abv.bg>
|
||||
Anton Mitkov <anton.mitkov@codeplay.com>
|
||||
Antonis Makropoulos <benuix@gmail.com>
|
||||
anubhavfujitsu <Anubhav.Srivastava@fujitsu.com>
|
||||
Anudit Nagar <nagaranudit@gmail.com>
|
||||
Anuj Attri <anujattri01@gmail.com>
|
||||
anujj <ajalota@nvidia.com>
|
||||
@@ -293,6 +297,7 @@ Bernat Vadell <hounter.caza@gmail.com>
|
||||
Bernhard M. Wiedemann <githubbmwprimary@lsmod.de>
|
||||
Bert Wagner <github@bertwagner.com>
|
||||
Bertay Eren <39909689+bertaye@users.noreply.github.com>
|
||||
Bhargav Krish <bkrishnamurt@nvidia.com>
|
||||
Bhavik Sharda <10757940+BLSharda@users.noreply.github.com>
|
||||
bhubbb <79117352+bhubbb@users.noreply.github.com>
|
||||
Bill Sideris <bill88t@feline.gr>
|
||||
@@ -533,6 +538,7 @@ Don Mahurin <dmahurin@users.noreply.github.com>
|
||||
Dong Won Kim <63934649+ddwkim@users.noreply.github.com>
|
||||
Donghyeon Jeong <54725479+djeong20@users.noreply.github.com>
|
||||
Dongliang Wei <121270393+wdl339@users.noreply.github.com>
|
||||
Dongsheng Luo <luods1995@outlook.com>
|
||||
Doomsdayrs <38189170+Doomsdayrs@users.noreply.github.com>
|
||||
DooWoong Lee (David) <manics99@naver.com>
|
||||
DorianRudolph <dorianrudo97@googlemail.com>
|
||||
@@ -612,6 +618,7 @@ Faez Shakil <faez.shakil@gmail.com>
|
||||
fairydreaming <166155368+fairydreaming@users.noreply.github.com>
|
||||
Faisal Zaghloul <faisal.zaghloul@gmail.com>
|
||||
Faisal Zaghloul <quic_fzaghlou@quicinc.com>
|
||||
FamousMonsterr <famousmonster@ya.ru>
|
||||
Fan Shupei <dymarkfan@outlook.com>
|
||||
FantasyGmm <16450052+FantasyGmm@users.noreply.github.com>
|
||||
fanyang <fanyang89@outlook.com>
|
||||
@@ -654,6 +661,7 @@ Frederik Vogel <Schaltfehler@users.noreply.github.com>
|
||||
Fredrik Hultin <noname@nurd.se>
|
||||
fredzillman <fzillman@gmail.com>
|
||||
frob <rick+github@frob.com.au>
|
||||
frogscooper <frogscooper@gmail.com>
|
||||
Frost-54 <91168148+Frost-54@users.noreply.github.com>
|
||||
Frosty40 <newjordan@gmail.com>
|
||||
Funtowicz Morgan <mfuntowicz@users.noreply.github.com>
|
||||
@@ -694,6 +702,7 @@ Gilad S. <7817232+giladgd@users.noreply.github.com>
|
||||
Giovanni Rivera <rivera.giovanni271@gmail.com>
|
||||
github-actions[bot] <github-actions[bot]@users.noreply.github.com>
|
||||
GittyBurstein <g0534163997@gmail.com>
|
||||
Giulianno <lernexbusiness@gmail.com>
|
||||
Giuseppe Scrivano <giuseppe@scrivano.org>
|
||||
Giuseppe Scrivano <gscrivan@redhat.com>
|
||||
GiviMAD <GiviMAD@users.noreply.github.com>
|
||||
@@ -762,6 +771,7 @@ Henry147147 <44851451+Henry147147@users.noreply.github.com>
|
||||
Herman Semenoff <GermanAizek@yandex.ru>
|
||||
Herman Semenov <GermanAizek@yandex.ru>
|
||||
Hesen Peng <hesen.peng@gmail.com>
|
||||
hesiru <siruhe666@gmail.com>
|
||||
hey-gm <heitorgm@outlook.com>
|
||||
HighDoping <highdoping@gmail.com>
|
||||
HimariO <dsfhe49854@gmail.com>
|
||||
@@ -919,6 +929,7 @@ Jesse Ikonen <jesse.ikonen@gmail.com>
|
||||
Jesse Jojo Johnson <williamsaintgeorge@gmail.com>
|
||||
Jesse LaRose <jesse@taey.ai>
|
||||
Jesse Posner <jesse.posner@gmail.com>
|
||||
Jesse Symons <155509533+Torchit1@users.noreply.github.com>
|
||||
Jesus Gulfo <jesdga95@gmail.com>
|
||||
Jesus Talavera <145992175+jesus-talavera-ibm@users.noreply.github.com>
|
||||
Jetson Tan <tanzongyouyi@outlook.com>
|
||||
@@ -1071,6 +1082,7 @@ Kevin Liu <4396kevinliu@gmail.com>
|
||||
Kevin Lo <kevlo@kevlo.org>
|
||||
Kevin Pouget <kpouget@redhat.com>
|
||||
Kevin Wang <kevmo314@gmail.com>
|
||||
Keyu Zhu <inuebisu@outlook.com>
|
||||
Khashayar Ghafouri <43180261+khashayarghafouri@users.noreply.github.com>
|
||||
khimaros <me@khimaros.com>
|
||||
Kilian Hu <90606809+kilian-hu@users.noreply.github.com>
|
||||
@@ -1078,6 +1090,7 @@ Kilian Krampf <kilian@krampf.de>
|
||||
kiltyj <kiltyj@gmail.com>
|
||||
Kim S. <polydecay@users.noreply.github.com>
|
||||
kimminsu <80271594+kimminsu38oo@users.noreply.github.com>
|
||||
Kirill A. Korinsky <kirill@korins.ky>
|
||||
KITAITI Makoto <KitaitiMakoto@gmail.com>
|
||||
kiwi <122582483+kiwi142857@users.noreply.github.com>
|
||||
klosax <131523366+klosax@users.noreply.github.com>
|
||||
@@ -1136,6 +1149,7 @@ leejet <leejet714@gmail.com>
|
||||
Leng Yue <lengyue@lengyue.me>
|
||||
Lennart Austenfeld <53152202+l-austenfeld@users.noreply.github.com>
|
||||
leo-pony <nengjunma@outlook.com>
|
||||
Leon <118652405+LeON-Nie-code@users.noreply.github.com>
|
||||
Leon Knauer <git@leonknauer.com>
|
||||
Leonard Mosescu <tlemo@users.noreply.github.com>
|
||||
leonardHONG <2695316095@qq.com>
|
||||
@@ -1146,8 +1160,10 @@ lexasub <lexakopp2212@gmail.com>
|
||||
lgai-exaone <exaonemodels@lgresearch.ai>
|
||||
lhez <lih@qti.qualcomm.com>
|
||||
lhez <quic_lih@quicinc.com>
|
||||
Li Chuyang <1418587903@qq.com>
|
||||
Li Pengzhan <151381994+Lpzhan931@users.noreply.github.com>
|
||||
Li Tan <tanliboy@gmail.com>
|
||||
li-lizhe <82501526+li-lizhe@users.noreply.github.com>
|
||||
liangliangchang <lichang@amd.com>
|
||||
liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
|
||||
limitedAtonement <limitedAtonement@users.noreply.github.com>
|
||||
@@ -1166,6 +1182,7 @@ LoganDark <github@logandark.mozmail.com>
|
||||
Loïc Carrère <loic.carrere@gmail.com>
|
||||
lon <114724657+longregen@users.noreply.github.com>
|
||||
loonerin <132926317+loonerin@users.noreply.github.com>
|
||||
Lorenzo Bona <lorenzobona@tuta.com>
|
||||
losewayy <1134285612@qq.com>
|
||||
LostRuins <39025047+LostRuins@users.noreply.github.com>
|
||||
LostRuins Concedo <39025047+LostRuins@users.noreply.github.com>
|
||||
@@ -1253,6 +1270,7 @@ Masaya, Kato <62578291+msy-kato@users.noreply.github.com>
|
||||
mashdragon <122402293+mashdragon@users.noreply.github.com>
|
||||
Mason Milburn <masonmilby@gmail.com>
|
||||
MasterYi1024 <39848311+MasterYi1024@users.noreply.github.com>
|
||||
Maternion <98753158+maternion@users.noreply.github.com>
|
||||
Mateusz Charytoniuk <mateusz.charytoniuk@protonmail.com>
|
||||
Matheus C. França <matheus-catarino@hotmail.com>
|
||||
Matheus Gabriel Alves Silva <matheusgasource@gmail.com>
|
||||
@@ -1350,6 +1368,7 @@ Miwa / Ensan <63481257+ensan-hcl@users.noreply.github.com>
|
||||
miyan <1138989048@qq.com>
|
||||
mj-shifu <77107165+mj-shifu@users.noreply.github.com>
|
||||
mkoker <132301062+mkoker@users.noreply.github.com>
|
||||
mmalpe-nv <mmalpe@nvidia.com>
|
||||
mmyjona <jonathan.gonse@gmail.com>
|
||||
mnehete32 <33429707+mnehete32@users.noreply.github.com>
|
||||
Mohamed Elashri <git@elashri.com>
|
||||
@@ -1368,6 +1387,7 @@ mtmcp <141645996+mtmcp@users.noreply.github.com>
|
||||
muggle-stack <promuggle@qq.com>
|
||||
Muhammad Saad <saadydev@gmail.com>
|
||||
Muhammad Salem <salem.ebo@gmail.com>
|
||||
MuoDoo <36472129+MuoDoo@users.noreply.github.com>
|
||||
Murilo Santana <mvrilo@gmail.com>
|
||||
Musab Gultekin <musabgultekin@users.noreply.github.com>
|
||||
musoles <135031143+musoles@users.noreply.github.com>
|
||||
@@ -1416,10 +1436,12 @@ Nicolas Budyn <16596345+nbud@users.noreply.github.com>
|
||||
Nicolas Mowen <nickmowen213@gmail.com>
|
||||
Nicolás Pérez <nicolas_perez@brown.edu>
|
||||
Nicolò Scipione <nicolo.scipione@codeplay.com>
|
||||
Nidhi <nidhisharma16635@gmail.com>
|
||||
Nigel Bosch <pnigelb@gmail.com>
|
||||
Nigel Jones <nigel.l.jones+git@gmail.com>
|
||||
Nik Bogatyrev <50734723+notforu777@users.noreply.github.com>
|
||||
Nikhil Jain <nikhil.jain0987@gmail.com>
|
||||
Nikhil Sharma <150424968+nikhilfujitsu@users.noreply.github.com>
|
||||
Nikita Gordeev <iamfobey@gmail.com>
|
||||
Nikita Sarychev <42014488+sARY77@users.noreply.github.com>
|
||||
Niklas Korz <niklas@niklaskorz.de>
|
||||
@@ -1592,6 +1614,7 @@ Raman Shinde <raman.shinde15@gmail.com>
|
||||
Rand Xie <randxiexyy29@gmail.com>
|
||||
Randall Fitzgerald <randall@dasaku.net>
|
||||
Random Fly <renfei8@live.cn>
|
||||
randomized <286843573@qq.com>
|
||||
rankaiyx <rankaiyx@foxmail.com>
|
||||
rankaiyx <rankaiyx@rankaiyx.com>
|
||||
RapidMark <32768622+RapidMark@users.noreply.github.com>
|
||||
@@ -2009,6 +2032,7 @@ Tungsten842 <quantmint@protonmail.com>
|
||||
Tushar <ditsuke@protonmail.com>
|
||||
tv1wnd <55383215+tv1wnd@users.noreply.github.com>
|
||||
tyronecai <tyronecai@163.com>
|
||||
Uber <20715993+YagUber@users.noreply.github.com>
|
||||
ubergarm <leimgrub@gmail.com>
|
||||
ubik2 <ubik2@users.noreply.github.com>
|
||||
UEXTM.com <84163508+uextm@users.noreply.github.com>
|
||||
@@ -2242,4 +2266,5 @@ Zijun Yu <zijun.yu@intel.com>
|
||||
zql <37731799+zqlcode@users.noreply.github.com>
|
||||
zrm <trustiosity.zrm@gmail.com>
|
||||
Zsapi <martin1.zsapka@gmail.com>
|
||||
zyang-dev <267119621+zyang-dev@users.noreply.github.com>
|
||||
zzzzwc <tiddar@foxmail.com>
|
||||
|
||||
+1
-1
@@ -4,7 +4,7 @@ include(CheckIncludeFileCXX)
|
||||
|
||||
### llama.cpp version
|
||||
set(LLAMA_VERSION_MAJOR 0)
|
||||
set(LLAMA_VERSION_MINOR 8)
|
||||
set(LLAMA_VERSION_MINOR 9)
|
||||
set(LLAMA_VERSION_PATCH 0)
|
||||
set(LLAMA_VERSION_BASE "${LLAMA_VERSION_MAJOR}.${LLAMA_VERSION_MINOR}.${LLAMA_VERSION_PATCH}")
|
||||
|
||||
|
||||
+1
-1
@@ -710,7 +710,7 @@ static void common_params_apply_system_config(common_params & params, llama_exam
|
||||
std::vector<std::filesystem::path> paths;
|
||||
|
||||
#if defined(_WIN32)
|
||||
const std::filesystem::path program_data = common_get_path_from_env("PROGRAMDATA");
|
||||
const std::filesystem::path program_data = common_get_env("PROGRAMDATA");
|
||||
if (!program_data.empty()) {
|
||||
paths.push_back(program_data / "llama.cpp" / "config.ini");
|
||||
}
|
||||
|
||||
+6
-17
@@ -983,20 +983,9 @@ void common_set_env(const std::string & name, const std::string & value) {
|
||||
#endif
|
||||
}
|
||||
|
||||
std::filesystem::path common_get_path_from_env(const std::string & name) {
|
||||
#if defined(_WIN32)
|
||||
const std::wstring wname = utf8_to_wstring(name);
|
||||
const wchar_t * wvalue = _wgetenv(wname.c_str());
|
||||
return wvalue ? std::filesystem::path(wvalue) : std::filesystem::path();
|
||||
#else
|
||||
const char * value = std::getenv(name.c_str());
|
||||
return value ? std::filesystem::path(value) : std::filesystem::path();
|
||||
#endif
|
||||
}
|
||||
|
||||
#if !defined(_WIN32)
|
||||
static std::filesystem::path get_home_directory() {
|
||||
std::filesystem::path home = common_get_path_from_env("HOME");
|
||||
std::filesystem::path home = common_get_env("HOME");
|
||||
if (!home.empty()) {
|
||||
return home;
|
||||
}
|
||||
@@ -1009,19 +998,19 @@ static std::filesystem::path get_home_directory() {
|
||||
#endif
|
||||
|
||||
std::filesystem::path fs_get_cache_directory() {
|
||||
std::filesystem::path cache_directory = common_get_path_from_env("LLAMA_CACHE");
|
||||
std::filesystem::path cache_directory = common_get_env("LLAMA_CACHE");
|
||||
if (!cache_directory.empty()) {
|
||||
return cache_directory;
|
||||
}
|
||||
#if defined(_WIN32)
|
||||
cache_directory = common_get_path_from_env("LOCALAPPDATA");
|
||||
cache_directory = common_get_env("LOCALAPPDATA");
|
||||
if (cache_directory.empty()) {
|
||||
throw std::runtime_error("Failed to find %LOCALAPPDATA% directory");
|
||||
}
|
||||
#elif defined(__APPLE__)
|
||||
cache_directory = get_home_directory() / "Library/Caches";
|
||||
#else
|
||||
cache_directory = common_get_path_from_env("XDG_CACHE_HOME");
|
||||
cache_directory = common_get_env("XDG_CACHE_HOME");
|
||||
if (cache_directory.empty()) {
|
||||
cache_directory = get_home_directory() / ".cache";
|
||||
}
|
||||
@@ -1032,12 +1021,12 @@ std::filesystem::path fs_get_cache_directory() {
|
||||
std::filesystem::path fs_get_config_directory() {
|
||||
std::filesystem::path config_directory;
|
||||
#if defined(_WIN32)
|
||||
config_directory = common_get_path_from_env("APPDATA");
|
||||
config_directory = common_get_env("APPDATA");
|
||||
if (config_directory.empty()) {
|
||||
throw std::runtime_error("Failed to find %APPDATA% directory");
|
||||
}
|
||||
#else
|
||||
config_directory = common_get_path_from_env("XDG_CONFIG_HOME");
|
||||
config_directory = common_get_env("XDG_CONFIG_HOME");
|
||||
if (config_directory.empty()) {
|
||||
config_directory = get_home_directory() / ".config";
|
||||
}
|
||||
|
||||
@@ -919,9 +919,6 @@ std::string wstring_to_utf8(const std::wstring & str);
|
||||
std::string common_get_env(const std::string & name);
|
||||
void common_set_env(const std::string & name, const std::string & value);
|
||||
|
||||
// reads a path from the environment, an unset variable gives an empty path
|
||||
std::filesystem::path common_get_path_from_env(const std::string & name);
|
||||
|
||||
//
|
||||
// Filesystem utils
|
||||
//
|
||||
|
||||
+1
-1
@@ -45,7 +45,7 @@ static fs::path get_cache_directory() {
|
||||
{HOME_DIR, fs::path(".cache") / "huggingface" / "hub"}
|
||||
};
|
||||
for (const auto & entry : entries) {
|
||||
if (fs::path base = common_get_path_from_env(entry.var); !base.empty()) {
|
||||
if (fs::path base = common_get_env(entry.var); !base.empty()) {
|
||||
return entry.path.empty() ? base : base / entry.path;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -95,6 +95,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
|
||||
"FalconForCausalLM": "falcon",
|
||||
"FalconH1ForCausalLM": "falcon_h1",
|
||||
"FalconMambaForCausalLM": "mamba",
|
||||
"FishOmniModel": "drama3",
|
||||
"GPT2LMHeadModel": "gpt2",
|
||||
"GPTBigCodeForCausalLM": "starcoder",
|
||||
"GPTNeoXForCausalLM": "gptneox",
|
||||
@@ -346,6 +347,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
|
||||
"DotsOCRForCausalLM": "dotsocr",
|
||||
"EmbeddingGemma2Model": "gemma",
|
||||
"Exaone4_5_ForConditionalGeneration": "exaone",
|
||||
"FishOmniModel": "drama3",
|
||||
"Gemma3ForConditionalGeneration": "gemma",
|
||||
"Gemma3nForConditionalGeneration": "gemma",
|
||||
"Gemma4ForConditionalGeneration": "gemma",
|
||||
|
||||
+1
-1
@@ -494,7 +494,7 @@ class ModelBase:
|
||||
s = self.model_tensors[name]
|
||||
self.model_tensors[weight_name] = lambda w=w, s=s: dequant_bitnet(w(), s())
|
||||
tensors_to_remove.append(name)
|
||||
elif quant_method == "fp8":
|
||||
elif quant_method in ("fp8", "bsgemm-moe"): # bsgemm-moe: drama-3, same block layout
|
||||
block_size = quant_config.get("weight_block_size")
|
||||
for name in self.model_tensors.keys():
|
||||
if name.endswith("_scale_inv"):
|
||||
|
||||
@@ -0,0 +1,251 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Callable, Iterable, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch import Tensor
|
||||
|
||||
from .base import ModelBase, MmprojModel, gguf
|
||||
from .qwen import Qwen3_5MoeTextModel
|
||||
|
||||
# Fish Audio Drama 3 (FishOmniModel)
|
||||
# text model: Qwen3.5-MoE backbone, samples the main token of each frame
|
||||
# mmproj: codebook tower and Fish DAC decoder (codec.pth)
|
||||
|
||||
_TOWER_SHARED = "audio_decoder_mtp_shared.layers."
|
||||
_TOWER_RESIDUAL = "audio_decoder_mtp.layers."
|
||||
|
||||
|
||||
@ModelBase.register("FishOmniModel")
|
||||
@ModelBase.example("fishaudio/drama-3")
|
||||
class Drama3TextModel(Qwen3_5MoeTextModel):
|
||||
model_arch = gguf.MODEL_ARCH.QWEN35MOE
|
||||
no_mtp = True # the checkpoint has no MTP head
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
if item[0].startswith("audio_decoder"):
|
||||
return None
|
||||
return super().filter_tensors(item)
|
||||
|
||||
|
||||
@ModelBase.register("FishOmniModel")
|
||||
@ModelBase.example("fishaudio/drama-3")
|
||||
class Drama3MmprojModel(MmprojModel):
|
||||
has_vision_encoder = False
|
||||
has_audio_encoder = False
|
||||
|
||||
_LAYER_MAP = {
|
||||
"enorm.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_ENORM, True),
|
||||
"hnorm.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_HNORM, True),
|
||||
"eh_proj.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_EH_PROJ, False),
|
||||
"semantic_embedding.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_BLK_EMBD, False),
|
||||
"codebook_embedding.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_BLK_EMBD, False),
|
||||
"final_layernorm.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_BLK_OUT_NORM, True),
|
||||
"transformer_layer.attention_norm.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_NORM, True),
|
||||
"transformer_layer.ffn_norm.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_FFN_NORM, True),
|
||||
"transformer_layer.attention.q_norm.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM, True),
|
||||
"transformer_layer.attention.k_norm.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM, True),
|
||||
"transformer_layer.attention.wo.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_OUT, False),
|
||||
"transformer_layer.feed_forward.w1.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_FFN_GATE, False),
|
||||
"transformer_layer.feed_forward.w3.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_FFN_UP, False),
|
||||
"transformer_layer.feed_forward.w2.weight": (gguf.MODEL_TENSOR.A_GEN_CODE_FFN_DOWN, False),
|
||||
}
|
||||
|
||||
_heads: dict[int, Tensor]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._heads = {}
|
||||
|
||||
def get_audio_config(self) -> dict[str, Any] | None:
|
||||
cfg = dict(self.global_config["audio_decoder_config"])
|
||||
cfg["num_hidden_layers"] = cfg["shared_enrichment_layers"] + cfg["num_codebooks"] - 1
|
||||
return cfg
|
||||
|
||||
def set_gguf_parameters(self):
|
||||
cfg = self.hparams
|
||||
self.gguf_writer.add_file_type(self.ftype)
|
||||
self.gguf_writer.add_clip_has_gen_audio_encoder(True)
|
||||
self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.DRAMA3_GEN)
|
||||
self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)
|
||||
self.gguf_writer.add_gen_audio_embedding_length(cfg["dim"])
|
||||
self.gguf_writer.add_gen_audio_feed_forward_length(cfg["intermediate_size"])
|
||||
self.gguf_writer.add_gen_audio_block_count(cfg["num_hidden_layers"])
|
||||
self.gguf_writer.add_gen_audio_head_count(cfg["n_head"])
|
||||
self.gguf_writer.add_gen_audio_head_count_kv(cfg["n_local_heads"])
|
||||
self.gguf_writer.add_gen_audio_attention_layernorm_eps(cfg["norm_eps"])
|
||||
# other tower and codec hparams are hardcoded in clip.cpp
|
||||
|
||||
@classmethod
|
||||
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
|
||||
name = item[0]
|
||||
if name.startswith("audio_decoder") or name == "model.language_model.embed_tokens.weight":
|
||||
return item
|
||||
return None
|
||||
|
||||
def tensor_force_quant(self, name, new_name, bid, n_dims):
|
||||
# conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path
|
||||
if new_name.endswith(".weight") and (
|
||||
new_name in ("a.gen.wav.dac.entry.weight", "a.gen.wav.dac.post_conv.weight")
|
||||
or (".up.blk." in new_name and new_name.endswith(".dwconv.weight"))
|
||||
or (".dac.blk." in new_name and (new_name.endswith(".conv1.weight") or new_name.endswith(".conv2.weight")))
|
||||
):
|
||||
return gguf.GGMLQuantizationType.F16
|
||||
# keep ConvTranspose1d kernels and folded codebook tables in F32
|
||||
if new_name.endswith(".conv.weight") and (".up.blk." in new_name or ".dac.blk." in new_name):
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
if new_name.startswith("a.gen.wav.quant."):
|
||||
return gguf.GGMLQuantizationType.F32
|
||||
return super().tensor_force_quant(name, new_name, bid, n_dims)
|
||||
|
||||
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
|
||||
T = gguf.MODEL_TENSOR
|
||||
cfg = self.hparams
|
||||
|
||||
if name.startswith("a.gen.wav."):
|
||||
yield (name, data_torch)
|
||||
return
|
||||
|
||||
if name == "model.language_model.embed_tokens.weight":
|
||||
# semantic band of the backbone embeddings
|
||||
begin = self.global_config["semantic_start_token_id"]
|
||||
yield (self.format_tensor_name(T.A_GEN_CODE_OUT_EMBD), data_torch[begin:begin + cfg["vocab_size"]])
|
||||
return
|
||||
|
||||
if name == "audio_decoder.codebook_embeddings.weight":
|
||||
n_cb = cfg["num_codebooks"]
|
||||
yield (self.format_tensor_name(T.A_GEN_CODE_EMBD), data_torch.reshape(n_cb, -1, data_torch.shape[-1]))
|
||||
return
|
||||
|
||||
if name.startswith("audio_decoder_mtp.heads."):
|
||||
n_heads = cfg["num_codebooks"] - 1
|
||||
self._heads[int(name.split(".")[2])] = data_torch
|
||||
if len(self._heads) < n_heads:
|
||||
return
|
||||
yield (self.format_tensor_name(T.A_GEN_CODE_HEAD), torch.stack([self._heads.pop(i) for i in range(n_heads)]))
|
||||
return
|
||||
|
||||
for prefix, offset in ((_TOWER_SHARED, 0), (_TOWER_RESIDUAL, cfg["shared_enrichment_layers"])):
|
||||
if not name.startswith(prefix):
|
||||
continue
|
||||
idx, key = name[len(prefix):].split(".", 1)
|
||||
il = int(idx) + offset
|
||||
if key == "transformer_layer.attention.wqkv.weight":
|
||||
# rows are [query | gate | key | value]
|
||||
n_q = cfg["n_head"] * cfg["head_dim"]
|
||||
n_kv = cfg["n_local_heads"] * cfg["head_dim"]
|
||||
q, gate, k, v = data_torch.split([n_q, n_q, n_kv, n_kv], dim=0)
|
||||
yield (self.format_tensor_name(T.A_GEN_CODE_ATTN_Q, il), q)
|
||||
yield (self.format_tensor_name(T.A_GEN_CODE_ATTN_GATE, il), gate)
|
||||
yield (self.format_tensor_name(T.A_GEN_CODE_ATTN_K, il), k)
|
||||
yield (self.format_tensor_name(T.A_GEN_CODE_ATTN_V, il), v)
|
||||
return
|
||||
tensor, zero_centered = self._LAYER_MAP[key]
|
||||
if zero_centered:
|
||||
data_torch = data_torch.float() + 1
|
||||
yield (self.format_tensor_name(tensor, il), data_torch)
|
||||
return
|
||||
|
||||
raise ValueError(f"unexpected tensor {name}")
|
||||
|
||||
def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
|
||||
# only the DAC decoder, the encoder is for voice cloning (not supported yet)
|
||||
sd = torch.load(self.dir_model / "codec.pth", map_location="cpu", mmap=True, weights_only=True)
|
||||
sd = sd.get("state_dict", sd)
|
||||
T = gguf.MODEL_TENSOR
|
||||
|
||||
def get(name: str) -> Tensor:
|
||||
return sd[name].float()
|
||||
|
||||
def conv(prefix: str) -> Tensor:
|
||||
# fold weight norm
|
||||
if prefix + ".conv.weight" in sd:
|
||||
return get(prefix + ".conv.weight")
|
||||
g = get(prefix + ".conv.parametrizations.weight.original0")
|
||||
v = get(prefix + ".conv.parametrizations.weight.original1")
|
||||
return g * v / v.norm(dim=(1, 2), keepdim=True)
|
||||
|
||||
def quant_table(prefix: str) -> Tensor:
|
||||
# fold out_proj into the codebook, so decoding is a row lookup
|
||||
g = get(prefix + ".out_proj.weight_g")
|
||||
v = get(prefix + ".out_proj.weight_v")
|
||||
w = (g * v / v.norm(dim=(1, 2), keepdim=True)).squeeze(-1)
|
||||
return get(prefix + ".codebook.weight") @ w.T + get(prefix + ".out_proj.bias")
|
||||
|
||||
def snake(alpha: Tensor) -> tuple[Tensor, Tensor]:
|
||||
alpha = alpha.flatten()
|
||||
return alpha, 1.0 / (alpha + 1e-9)
|
||||
|
||||
q = "quantizer."
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_CB), quant_table(q + "semantic_quantizer.quantizers.0"))
|
||||
n_rest = sum(1 for k in sd if k.startswith(q + "quantizer.quantizers.") and k.endswith(".codebook.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_REST_CB),
|
||||
torch.stack([quant_table(f"{q}quantizer.quantizers.{i}") for i in range(n_rest)]))
|
||||
|
||||
# post_module uses interleaved RoPE, permute q/k rows so NEOX RoPE gives the same result
|
||||
tfm = q + "post_module."
|
||||
n_head, d_head = 16, 64
|
||||
order = torch.cat([torch.arange(0, d_head, 2), torch.arange(1, d_head, 2)])
|
||||
|
||||
def neox_rows(w: Tensor) -> Tensor:
|
||||
return w.reshape(n_head, d_head, -1)[:, order].reshape(w.shape)
|
||||
|
||||
n_tfm = sum(1 for k in sd if k.startswith(tfm + "layers.") and k.endswith(".ffn_norm.weight"))
|
||||
for il in range(n_tfm):
|
||||
p = f"{tfm}layers.{il}."
|
||||
wq, wk, wv = get(p + "attention.wqkv.weight").chunk(3, dim=0)
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_Q, il), neox_rows(wq))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_K, il), neox_rows(wk))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_V, il), wv)
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_OUT, il), get(p + "attention.wo.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_NORM, il), get(p + "attention_norm.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_ATTN_SCALE, il), get(p + "attention_layer_scale.gamma"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_FFN_NORM, il), get(p + "ffn_norm.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_FFN_GATE, il), get(p + "feed_forward.w1.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_FFN_UP, il), get(p + "feed_forward.w3.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_FFN_DOWN, il), get(p + "feed_forward.w2.weight"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_FFN_SCALE, il), get(p + "ffn_layer_scale.gamma"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUTPUT_NORM), get(tfm + "norm.weight"))
|
||||
|
||||
# upsample: 2x (ConvTranspose1d + ConvNeXt)
|
||||
for il in range(2):
|
||||
p = f"{q}upsample.{il}."
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_UP_CONV, il, suffix=".weight"), conv(p + "0"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_UP_CONV, il, suffix=".bias"), get(p + "0.conv.bias"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_UP_DWCONV, il, suffix=".weight"), conv(p + "1.dwconv"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_UP_DWCONV, il, suffix=".bias"), get(p + "1.dwconv.conv.bias"))
|
||||
for key, tensor in (("norm", T.A_GEN_WAV_UP_NORM), ("pwconv1", T.A_GEN_WAV_UP_PW1), ("pwconv2", T.A_GEN_WAV_UP_PW2)):
|
||||
yield (self.format_tensor_name(tensor, il, suffix=".weight"), get(f"{p}1.{key}.weight"))
|
||||
yield (self.format_tensor_name(tensor, il, suffix=".bias"), get(f"{p}1.{key}.bias"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_UP_GAMMA, il, suffix=""), get(p + "1.gamma"))
|
||||
|
||||
# DAC: entry conv, 4x (Snake, ConvTranspose1d, 3 residual units), Snake, post conv
|
||||
d = "decoder.model."
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_ENTRY, suffix=".weight"), conv(d + "0"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_ENTRY, suffix=".bias"), get(d + "0.conv.bias"))
|
||||
for bid in range(4):
|
||||
p = f"{d}{bid + 1}.block."
|
||||
a, b = snake(get(p + "0.alpha"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_SNAKE, bid, suffix=".alpha"), a)
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_SNAKE, bid, suffix=".beta"), b)
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_CONV, bid, suffix=".weight"), conv(p + "1"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_CONV, bid, suffix=".bias"), get(p + "1.conv.bias"))
|
||||
for xid in range(3):
|
||||
r = f"{p}{xid + 2}.block."
|
||||
for act, tensor in (("0", T.A_GEN_WAV_DAC_RES_ACT1), ("2", T.A_GEN_WAV_DAC_RES_ACT2)):
|
||||
a, b = snake(get(r + act + ".alpha"))
|
||||
name = gguf.TENSOR_NAMES[tensor].format(bid=bid, xid=xid)
|
||||
yield (name + ".alpha", a)
|
||||
yield (name + ".beta", b)
|
||||
for cv, tensor in (("1", T.A_GEN_WAV_DAC_RES_CONV1), ("3", T.A_GEN_WAV_DAC_RES_CONV2)):
|
||||
name = gguf.TENSOR_NAMES[tensor].format(bid=bid, xid=xid)
|
||||
yield (name + ".weight", conv(r + cv))
|
||||
yield (name + ".bias", get(r + cv + ".conv.bias"))
|
||||
a, b = snake(get(d + "5.alpha"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_SNAKE, suffix=".alpha"), a)
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_SNAKE, suffix=".beta"), b)
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_CONV, suffix=".weight"), conv(d + "6"))
|
||||
yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_CONV, suffix=".bias"), get(d + "6.conv.bias"))
|
||||
@@ -182,7 +182,6 @@ struct webgpu_capabilities {
|
||||
uint32_t subgroup_size = 0;
|
||||
uint32_t min_subgroup_size = 0;
|
||||
uint32_t max_subgroup_size = 0;
|
||||
size_t memset_bytes_per_thread;
|
||||
};
|
||||
|
||||
// Stores global webgpu members
|
||||
@@ -256,8 +255,6 @@ struct webgpu_context_struct {
|
||||
wgpu::ComputePassEncoder active_compute_pass;
|
||||
bool batch_compute_passes = true;
|
||||
|
||||
size_t memset_bytes_per_thread;
|
||||
|
||||
#ifdef GGML_WEBGPU_GPU_PROFILE
|
||||
// Profiling: per-shader GPU time in ms
|
||||
std::unordered_map<std::string, double> shader_gpu_time_ms;
|
||||
@@ -672,38 +669,54 @@ static void ggml_backend_webgpu_buffer_memset(webgpu_global_context & ctx,
|
||||
size_t size) {
|
||||
std::lock_guard<std::recursive_mutex> lock(ctx->mutex);
|
||||
|
||||
std::vector<uint32_t> params = { (uint32_t) offset, (uint32_t) size, value };
|
||||
std::vector<wgpu::BindGroupEntry> entries = { ggml_webgpu_make_bind_group_entry(0, buf, 0, buf.GetSize()) };
|
||||
size_t bytes_per_wg =
|
||||
ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup * ctx->capabilities.memset_bytes_per_thread;
|
||||
uint32_t wg_x = CEIL_DIV(size + 3, bytes_per_wg);
|
||||
const size_t max_bind_size = ctx->capabilities.limits.maxStorageBufferBindingSize;
|
||||
const size_t max_wg_size = ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
|
||||
|
||||
ctx->queue.WriteBuffer(ctx->memset_params_buf, 0, params.data(), params.size() * sizeof(uint32_t));
|
||||
while (size > 0) {
|
||||
const size_t align_offset = offset & ~(ctx->capabilities.limits.minStorageBufferOffsetAlignment - 1);
|
||||
const size_t lead = offset - align_offset;
|
||||
const size_t chunk_size = std::min(size, max_bind_size - lead);
|
||||
const size_t bind_size = ROUNDUP_POW2(chunk_size + lead, WEBGPU_STORAGE_BUF_BINDING_MULT);
|
||||
|
||||
wgpu::BindGroupEntry params_entry = {};
|
||||
params_entry.binding = 1;
|
||||
params_entry.buffer = ctx->memset_params_buf;
|
||||
params_entry.offset = 0;
|
||||
params_entry.size = WEBGPU_PARAMS_BUF_SIZE_BYTES;
|
||||
entries.push_back(params_entry);
|
||||
std::vector<uint32_t> params = { (uint32_t) lead, (uint32_t) chunk_size, value };
|
||||
std::vector<wgpu::BindGroupEntry> entries = { ggml_webgpu_make_bind_group_entry(0, buf, align_offset,
|
||||
bind_size) };
|
||||
uint32_t wg_x;
|
||||
uint32_t wg_y;
|
||||
uint32_t n_values = CEIL_DIV(chunk_size + lead, sizeof(uint32_t));
|
||||
uint32_t total_wg = CEIL_DIV(n_values, max_wg_size);
|
||||
compute_2d_workgroups(total_wg, ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
|
||||
|
||||
wgpu::BindGroupDescriptor bind_group_desc;
|
||||
bind_group_desc.layout = ctx->memset_pipeline.pipeline.GetBindGroupLayout(0);
|
||||
bind_group_desc.entryCount = entries.size();
|
||||
bind_group_desc.entries = entries.data();
|
||||
bind_group_desc.label = ctx->memset_pipeline.name.c_str();
|
||||
wgpu::BindGroup bind_group = ctx->device.CreateBindGroup(&bind_group_desc);
|
||||
ctx->queue.WriteBuffer(ctx->memset_params_buf, 0, params.data(), params.size() * sizeof(uint32_t));
|
||||
|
||||
wgpu::CommandEncoder encoder = ctx->device.CreateCommandEncoder();
|
||||
wgpu::ComputePassEncoder pass = encoder.BeginComputePass();
|
||||
pass.SetPipeline(ctx->memset_pipeline.pipeline);
|
||||
pass.SetBindGroup(0, bind_group);
|
||||
pass.DispatchWorkgroups(wg_x, 1, 1);
|
||||
pass.End();
|
||||
wgpu::BindGroupEntry params_entry = {};
|
||||
params_entry.binding = 1;
|
||||
params_entry.buffer = ctx->memset_params_buf;
|
||||
params_entry.offset = 0;
|
||||
params_entry.size = WEBGPU_PARAMS_BUF_SIZE_BYTES;
|
||||
entries.push_back(params_entry);
|
||||
|
||||
wgpu::CommandBuffer command = encoder.Finish();
|
||||
std::vector<wgpu::CommandBuffer> commands = { command };
|
||||
ctx->queue.Submit(commands.size(), commands.data());
|
||||
wgpu::BindGroupDescriptor bind_group_desc;
|
||||
bind_group_desc.layout = ctx->memset_pipeline.pipeline.GetBindGroupLayout(0);
|
||||
bind_group_desc.entryCount = entries.size();
|
||||
bind_group_desc.entries = entries.data();
|
||||
bind_group_desc.label = ctx->memset_pipeline.name.c_str();
|
||||
wgpu::BindGroup bind_group = ctx->device.CreateBindGroup(&bind_group_desc);
|
||||
|
||||
wgpu::CommandEncoder encoder = ctx->device.CreateCommandEncoder();
|
||||
wgpu::ComputePassEncoder pass = encoder.BeginComputePass();
|
||||
pass.SetPipeline(ctx->memset_pipeline.pipeline);
|
||||
pass.SetBindGroup(0, bind_group);
|
||||
pass.DispatchWorkgroups(wg_x, wg_y, 1);
|
||||
pass.End();
|
||||
|
||||
wgpu::CommandBuffer command = encoder.Finish();
|
||||
std::vector<wgpu::CommandBuffer> commands = { command };
|
||||
ctx->queue.Submit(commands.size(), commands.data());
|
||||
|
||||
offset += chunk_size;
|
||||
size -= chunk_size;
|
||||
}
|
||||
}
|
||||
|
||||
/** End WebGPU Actions */
|
||||
@@ -4006,12 +4019,10 @@ static size_t ggml_backend_webgpu_buffer_type_get_alignment(ggml_backend_buffer_
|
||||
return dev_ctx->webgpu_global_ctx->capabilities.limits.minStorageBufferOffsetAlignment;
|
||||
}
|
||||
|
||||
// maxBufferSize might be larger, but you can't bind more than
|
||||
// maxStorageBufferBindingSize to a single binding.
|
||||
static size_t ggml_backend_webgpu_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_webgpu_device_context * dev_ctx =
|
||||
static_cast<ggml_backend_webgpu_device_context *>(buft->device->context);
|
||||
return dev_ctx->webgpu_global_ctx->capabilities.limits.maxStorageBufferBindingSize;
|
||||
return dev_ctx->webgpu_global_ctx->capabilities.limits.maxBufferSize;
|
||||
}
|
||||
|
||||
static size_t ggml_backend_webgpu_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft,
|
||||
@@ -4176,17 +4187,9 @@ static ggml_guid_t ggml_backend_webgpu_guid(void) {
|
||||
}
|
||||
|
||||
static void ggml_webgpu_init_memset_pipeline(webgpu_global_context & ctx) {
|
||||
// we use the maximum workgroup size for the memset pipeline
|
||||
size_t max_threads = ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup *
|
||||
ctx->capabilities.limits.maxComputeWorkgroupsPerDimension;
|
||||
// Size the bytes_per_thread so that the largest buffer size can be handled
|
||||
ctx->capabilities.memset_bytes_per_thread =
|
||||
CEIL_DIV(ctx->capabilities.limits.maxStorageBufferBindingSize, max_threads);
|
||||
std::vector<wgpu::ConstantEntry> constants(2);
|
||||
std::vector<wgpu::ConstantEntry> constants(1);
|
||||
constants[0].key = "wg_size";
|
||||
constants[0].value = ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
|
||||
constants[1].key = "bytes_per_thread";
|
||||
constants[1].value = ctx->capabilities.memset_bytes_per_thread;
|
||||
ctx->memset_pipeline = ggml_webgpu_create_pipeline(ctx->device, wgsl_memset, "memset", constants);
|
||||
}
|
||||
|
||||
|
||||
@@ -11,30 +11,33 @@ struct Params {
|
||||
var<uniform> params: Params;
|
||||
|
||||
override wg_size: u32;
|
||||
override bytes_per_thread: u32;
|
||||
|
||||
@compute @workgroup_size(wg_size)
|
||||
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||
let i = gid.x * bytes_per_thread;
|
||||
fn main(@builtin(workgroup_id) wid: vec3<u32>,
|
||||
@builtin(num_workgroups) num_wg: vec3<u32>,
|
||||
@builtin(local_invocation_id) lid: vec3<u32>) {
|
||||
|
||||
let i = (wid.x + wid.y * num_wg.x) * wg_size + lid.x;
|
||||
let start = params.offset;
|
||||
let end = params.offset + params.size;
|
||||
let id = (start >> 2u) + i;
|
||||
let id_u32 = id * 4;
|
||||
|
||||
for (var j: u32 = 0u; j < bytes_per_thread; j += 4) {
|
||||
let byte_index = start + i + j;
|
||||
if (byte_index + 4 <= end) {
|
||||
output_buffer[byte_index >> 2] = params.value;
|
||||
} else {
|
||||
// Handle tail (unaligned)
|
||||
for (var k: u32 = 0; k < 4; k++) {
|
||||
let idx = byte_index + k;
|
||||
if (idx < end) {
|
||||
let word_idx = idx >> 2;
|
||||
let bit_offset = (idx & 3) * 8u;
|
||||
let mask = ~(0xffu << bit_offset);
|
||||
let existing = output_buffer[word_idx];
|
||||
output_buffer[word_idx] = (existing & mask) | (params.value & (0xffu << bit_offset));
|
||||
}
|
||||
}
|
||||
}
|
||||
if (id_u32 >= end) {
|
||||
return;
|
||||
}
|
||||
|
||||
let lo = max(start, id_u32) - id_u32;
|
||||
let hi = min(end, id_u32 + 4u) - id_u32;
|
||||
var mask = 0xFFFFFFFFu << (8u * lo);
|
||||
if (hi < 4u) {
|
||||
mask &= (1u << (8u * hi)) - 1u;
|
||||
}
|
||||
|
||||
if (mask == 0xFFFFFFFFu) {
|
||||
output_buffer[id] = params.value;
|
||||
} else {
|
||||
let existing = output_buffer[id];
|
||||
output_buffer[id] = (existing & ~(mask)) | (params.value & mask);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1295,6 +1295,11 @@ class MODEL_TENSOR(IntEnum):
|
||||
A_GEN_CODE_ATTN_GATE = auto() # lfm3a: per-head sigmoid gate on the attention output
|
||||
A_GEN_CODE_HEAD_NORM = auto() # lfm3a: per-codebook norm before the output head, merged 2D [n_codebooks, dim]
|
||||
A_GEN_CODE_MASK_EMBD = auto() # lfm3a: MaskGIT mask embedding
|
||||
A_GEN_CODE_ENORM = auto() # drama3: per-layer norm of the code embedding
|
||||
A_GEN_CODE_HNORM = auto() # drama3: per-layer norm of the input hidden state
|
||||
A_GEN_CODE_EH_PROJ = auto() # drama3: per-layer [code embd; hidden] -> hidden projection
|
||||
A_GEN_CODE_BLK_EMBD = auto() # drama3: per-layer code embedding table
|
||||
A_GEN_CODE_BLK_OUT_NORM = auto() # drama3: per-layer output norm
|
||||
# qwen3tts code2wav: RVQ codes -> raw PCM
|
||||
A_GEN_WAV_QUANT_FIRST_IN = auto() # semantic RVQ, in_proj (1x1 conv, loaded as 2D)
|
||||
A_GEN_WAV_QUANT_FIRST_OUT = auto() # semantic RVQ, out_proj
|
||||
@@ -2198,6 +2203,11 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_GATE: "a.gen.code.blk.{bid}.attn_gate",
|
||||
MODEL_TENSOR.A_GEN_CODE_HEAD_NORM: "a.gen.code.head_norm",
|
||||
MODEL_TENSOR.A_GEN_CODE_MASK_EMBD: "a.gen.code.mask_embd",
|
||||
MODEL_TENSOR.A_GEN_CODE_ENORM: "a.gen.code.blk.{bid}.enorm",
|
||||
MODEL_TENSOR.A_GEN_CODE_HNORM: "a.gen.code.blk.{bid}.hnorm",
|
||||
MODEL_TENSOR.A_GEN_CODE_EH_PROJ: "a.gen.code.blk.{bid}.eh_proj",
|
||||
MODEL_TENSOR.A_GEN_CODE_BLK_EMBD: "a.gen.code.blk.{bid}.embd",
|
||||
MODEL_TENSOR.A_GEN_CODE_BLK_OUT_NORM: "a.gen.code.blk.{bid}.out_norm",
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_IN: "a.gen.wav.quant.first.in_proj",
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_OUT: "a.gen.wav.quant.first.out_proj",
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_CB: "a.gen.wav.quant.first.codebook",
|
||||
@@ -2615,6 +2625,11 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.A_GEN_CODE_ATTN_GATE,
|
||||
MODEL_TENSOR.A_GEN_CODE_HEAD_NORM,
|
||||
MODEL_TENSOR.A_GEN_CODE_MASK_EMBD,
|
||||
MODEL_TENSOR.A_GEN_CODE_ENORM,
|
||||
MODEL_TENSOR.A_GEN_CODE_HNORM,
|
||||
MODEL_TENSOR.A_GEN_CODE_EH_PROJ,
|
||||
MODEL_TENSOR.A_GEN_CODE_BLK_EMBD,
|
||||
MODEL_TENSOR.A_GEN_CODE_BLK_OUT_NORM,
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_IN,
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_OUT,
|
||||
MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_CB,
|
||||
@@ -6919,6 +6934,7 @@ class VisionProjectorType:
|
||||
CHATTERBOX_SPKENC = "chatterbox_spkenc" # audio: CAMPPlus speaker encoder
|
||||
CHATTERBOX = "chatterbox" # audio generation: s3gen flow matching + HiFT vocoder
|
||||
SOPRANO = "soprano" # audio generation: Vocos decoder
|
||||
DRAMA3_GEN = "drama3_gen" # audio generation: codebook tower + Fish DAC decoder
|
||||
HUNYUANVL = "hunyuanvl"
|
||||
PARAKEET = "parakeet" # audio
|
||||
MINIMAXM3 = "minimax_m3"
|
||||
|
||||
@@ -221,10 +221,12 @@ class SpecialVocab:
|
||||
if isinstance(merges, list) and merges:
|
||||
if isinstance(merges[0], str):
|
||||
self.merges = merges
|
||||
elif isinstance(merges[0], list) and len(merges[0]) == 2 and isinstance(merges[0][0], str):
|
||||
elif isinstance(merges[0], list):
|
||||
# New format since transformers 4.45 to support spaces in merges
|
||||
# ref: https://github.com/ggml-org/llama.cpp/issues/9692
|
||||
# TODO: internally store as the new format instead of converting to old
|
||||
if not all(isinstance(pair, list) and len(pair) == 2 and all(isinstance(s, str) for s in pair) for pair in merges):
|
||||
raise ValueError("Unknown tokenizer merges format")
|
||||
if any(' ' in s for pair in merges for s in pair):
|
||||
logger.warning(f'Spaces in merges detected, encoding as {chr(ord(" ") + 256)!r}')
|
||||
self.merges = [
|
||||
|
||||
@@ -70,6 +70,7 @@ add_library(mtmd
|
||||
models/mimo-audio.cpp
|
||||
models/qwen3tts-spkenc.cpp
|
||||
models/qwen3tts-gen.cpp
|
||||
models/drama3-gen.cpp
|
||||
models/pockettts-seanet.cpp
|
||||
models/pockettts-spkenc.cpp
|
||||
models/pockettts-gen.cpp
|
||||
|
||||
@@ -261,6 +261,13 @@
|
||||
#define TN_A_GEN_CODE_MASK_EMBD "a.gen.code.mask_embd"
|
||||
#define TN_ATTN_GATE "%s.blk.%d.attn_gate.%s"
|
||||
|
||||
// drama3 tower, the rest of the block uses the generic names with prefix "a.gen.code"
|
||||
#define TN_A_GEN_CODE_ENORM "a.gen.code.blk.%d.enorm.%s"
|
||||
#define TN_A_GEN_CODE_HNORM "a.gen.code.blk.%d.hnorm.%s"
|
||||
#define TN_A_GEN_CODE_EH_PROJ "a.gen.code.blk.%d.eh_proj.%s"
|
||||
#define TN_A_GEN_CODE_BLK_EMBD "a.gen.code.blk.%d.embd.%s"
|
||||
#define TN_A_GEN_CODE_BLK_OUT_NORM "a.gen.code.blk.%d.out_norm.%s"
|
||||
|
||||
// qwen3tts code2wav (RVQ codes -> raw PCM)
|
||||
// pre_transformer layers use the generic TN_ATTN_*/TN_FFN_*/TN_LN_*/TN_LS_* macros, prefix "a.gen.wav.tfm"
|
||||
#define TN_A_GEN_WAV_QUANT_FIRST_IN "a.gen.wav.quant.first.in_proj.%s"
|
||||
@@ -595,6 +602,7 @@ enum projector_type {
|
||||
PROJECTOR_TYPE_POCKETTTS_GEN,
|
||||
PROJECTOR_TYPE_SOPRANO,
|
||||
PROJECTOR_TYPE_LFM3A_GEN,
|
||||
PROJECTOR_TYPE_DRAMA3_GEN,
|
||||
PROJECTOR_TYPE_MUSE_GLIMMER,
|
||||
PROJECTOR_TYPE_SENSENOVA_U1,
|
||||
PROJECTOR_TYPE_COHERE2V,
|
||||
@@ -679,6 +687,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
|
||||
{ PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"},
|
||||
{ PROJECTOR_TYPE_SOPRANO, "soprano"},
|
||||
{ PROJECTOR_TYPE_LFM3A_GEN, "lfm3a_gen"},
|
||||
{ PROJECTOR_TYPE_DRAMA3_GEN, "drama3_gen"},
|
||||
{ PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"},
|
||||
{ PROJECTOR_TYPE_SENSENOVA_U1, "sensenova_u1"},
|
||||
{ PROJECTOR_TYPE_COHERE2V, "cohere2v"},
|
||||
|
||||
@@ -841,6 +841,15 @@ struct clip_chatterbox {
|
||||
ggml_tensor * spk_affine_b = nullptr;
|
||||
};
|
||||
|
||||
// drama3 tower: extra tensors around the generic block
|
||||
struct clip_drama3_layer {
|
||||
ggml_tensor * enorm_w = nullptr;
|
||||
ggml_tensor * hnorm_w = nullptr;
|
||||
ggml_tensor * eh_proj_w = nullptr; // [code embd; hidden] -> hidden
|
||||
ggml_tensor * embd_w = nullptr;
|
||||
ggml_tensor * out_norm_w = nullptr;
|
||||
};
|
||||
|
||||
struct clip_vocos {
|
||||
ggml_tensor * input_w = nullptr;
|
||||
ggml_tensor * input_b = nullptr;
|
||||
@@ -1107,6 +1116,7 @@ struct clip_model {
|
||||
|
||||
// qwen3tts code2wav: RVQ codes -> raw PCM
|
||||
clip_code2wav c2w;
|
||||
std::vector<clip_drama3_layer> drama3_layers;
|
||||
|
||||
// pocket-tts: SEANet stack, shared by the encoder (speaker path) and the decoder (gen path)
|
||||
clip_seanet seanet;
|
||||
|
||||
+209
-91
@@ -1237,6 +1237,15 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
|
||||
}
|
||||
builder = std::make_unique<clip_graph_lfm3a_gen>(ctx, img, gen_process, all_known, n_frames);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DRAMA3_GEN:
|
||||
{
|
||||
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
|
||||
const int64_t n_cb = ctx->model.gen_code_embd_w->ne[2];
|
||||
const int n_frames = params && params->codes ? (int) (params->codes->size() / n_cb) : 1;
|
||||
builder = std::make_unique<clip_graph_drama3_gen>(ctx, img, gen_process, params ? params->gen_state : nullptr, n_frames,
|
||||
params ? params->top_k : 30, params ? params->top_p : 0.9f,
|
||||
params ? params->temp : 1.0f);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
{
|
||||
const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE;
|
||||
@@ -2257,6 +2266,22 @@ struct clip_model_loader {
|
||||
// matches the reference decoder's sliding_window (speech_tokenizer/config.json)
|
||||
hparams.wav_tfm_swa = 72;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DRAMA3_GEN:
|
||||
{
|
||||
// Fish DAC decoder, codec.pth has no config
|
||||
hparams.audio_sample_rate = 44100;
|
||||
hparams.wav_tfm_n_layer = 8;
|
||||
hparams.wav_tfm_n_embd = 1024;
|
||||
hparams.wav_tfm_n_ff = 3072;
|
||||
hparams.wav_tfm_n_head = 16;
|
||||
hparams.wav_tfm_n_head_kv = 16;
|
||||
hparams.wav_tfm_eps = 1e-5f;
|
||||
hparams.wav_tfm_rope_theta = 10000.0f;
|
||||
hparams.wav_tfm_swa = 128;
|
||||
hparams.wav_upsample_n_block = 2;
|
||||
hparams.wav_dac_n_block = 4;
|
||||
hparams.wav_dac_n_res = 3;
|
||||
} break;
|
||||
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN:
|
||||
{
|
||||
@@ -2785,6 +2810,92 @@ struct clip_model_loader {
|
||||
}
|
||||
}
|
||||
|
||||
// code2wav: RVQ codes -> raw PCM, drama3 has no quant/pre_conv/tfm projections
|
||||
auto load_code2wav = [&](bool with_proj) {
|
||||
auto & c2w = model.c2w;
|
||||
|
||||
c2w.quant_first_in_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_IN, "weight"), with_proj);
|
||||
c2w.quant_first_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_OUT, "weight"), with_proj);
|
||||
c2w.quant_first_cb_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_CB, "weight"));
|
||||
c2w.quant_rest_in_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_REST_IN, "weight"), with_proj);
|
||||
c2w.quant_rest_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_REST_OUT, "weight"), with_proj);
|
||||
c2w.quant_rest_cb_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_REST_CB, "weight"));
|
||||
|
||||
c2w.pre_conv_w = get_tensor(string_format(TN_A_GEN_WAV_PRE_CONV, "weight"), with_proj);
|
||||
c2w.pre_conv_b = get_tensor(string_format(TN_A_GEN_WAV_PRE_CONV, "bias"), with_proj);
|
||||
|
||||
c2w.tfm_in_proj_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_IN_PROJ, "weight"), with_proj);
|
||||
c2w.tfm_in_proj_b = get_tensor(string_format(TN_A_GEN_WAV_TFM_IN_PROJ, "bias"), with_proj);
|
||||
c2w.tfm_out_proj_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_PROJ, "weight"), with_proj);
|
||||
c2w.tfm_out_proj_b = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_PROJ, "bias"), with_proj);
|
||||
c2w.tfm_output_norm_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_NORM, "weight"));
|
||||
|
||||
// loaded manually, the generic model.layers loop is taken by code_predictor
|
||||
c2w.tfm_layers.resize(hparams.wav_tfm_n_layer);
|
||||
for (int il = 0; il < hparams.wav_tfm_n_layer; il++) {
|
||||
auto & layer = c2w.tfm_layers[il];
|
||||
const char * p = "a.gen.wav.tfm";
|
||||
layer.q_w = get_tensor(string_format(TN_ATTN_Q, p, il, "weight"));
|
||||
layer.k_w = get_tensor(string_format(TN_ATTN_K, p, il, "weight"));
|
||||
layer.v_w = get_tensor(string_format(TN_ATTN_V, p, il, "weight"));
|
||||
layer.o_w = get_tensor(string_format(TN_ATTN_OUTPUT, p, il, "weight"));
|
||||
layer.ln_1_w = get_tensor(string_format(TN_LN_1, p, il, "weight"));
|
||||
layer.ln_2_w = get_tensor(string_format(TN_LN_2, p, il, "weight"));
|
||||
layer.ls_1_w = get_tensor(string_format(TN_LS_1, p, il, "weight"));
|
||||
layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight"));
|
||||
layer.ff_gate_w = get_tensor(string_format(TN_FFN_GATE, p, il, "weight"));
|
||||
layer.ff_up_w = get_tensor(string_format(TN_FFN_UP, p, il, "weight"));
|
||||
layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, p, il, "weight"));
|
||||
}
|
||||
|
||||
// upsample: 2x (causal ConvTranspose1d + ConvNeXt block)
|
||||
c2w.upsample.resize(hparams.wav_upsample_n_block);
|
||||
for (int il = 0; il < hparams.wav_upsample_n_block; il++) {
|
||||
auto & up = c2w.upsample[il];
|
||||
up.conv_w = get_tensor(string_format(TN_A_GEN_WAV_UP_CONV, il, "weight"));
|
||||
up.conv_b = get_tensor(string_format(TN_A_GEN_WAV_UP_CONV, il, "bias"));
|
||||
up.dwconv_w = get_tensor(string_format(TN_A_GEN_WAV_UP_DWCONV, il, "weight"));
|
||||
up.dwconv_b = get_tensor(string_format(TN_A_GEN_WAV_UP_DWCONV, il, "bias"));
|
||||
up.norm_w = get_tensor(string_format(TN_A_GEN_WAV_UP_NORM, il, "weight"));
|
||||
up.norm_b = get_tensor(string_format(TN_A_GEN_WAV_UP_NORM, il, "bias"));
|
||||
up.pw1_w = get_tensor(string_format(TN_A_GEN_WAV_UP_PW1, il, "weight"));
|
||||
up.pw1_b = get_tensor(string_format(TN_A_GEN_WAV_UP_PW1, il, "bias"));
|
||||
up.pw2_w = get_tensor(string_format(TN_A_GEN_WAV_UP_PW2, il, "weight"));
|
||||
up.pw2_b = get_tensor(string_format(TN_A_GEN_WAV_UP_PW2, il, "bias"));
|
||||
up.gamma = get_tensor(string_format(TN_A_GEN_WAV_UP_GAMMA, il));
|
||||
}
|
||||
|
||||
// DAC decoder: conv_pre + n upsample blocks (each with n_res residual units) + conv_post
|
||||
c2w.dac_entry_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_ENTRY, "weight"));
|
||||
c2w.dac_entry_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_ENTRY, "bias"));
|
||||
|
||||
c2w.dac.resize(hparams.wav_dac_n_block);
|
||||
for (int il = 0; il < hparams.wav_dac_n_block; il++) {
|
||||
auto & blk = c2w.dac[il];
|
||||
blk.snake_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_SNAKE, il, "alpha"));
|
||||
blk.snake_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_SNAKE, il, "beta"));
|
||||
blk.conv_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_CONV, il, "weight"));
|
||||
blk.conv_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_CONV, il, "bias"));
|
||||
|
||||
blk.res.resize(hparams.wav_dac_n_res);
|
||||
for (int ir = 0; ir < hparams.wav_dac_n_res; ir++) {
|
||||
auto & res = blk.res[ir];
|
||||
res.act1_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT1, il, ir, "alpha"));
|
||||
res.act1_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT1, il, ir, "beta"));
|
||||
res.conv1_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV1, il, ir, "weight"));
|
||||
res.conv1_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV1, il, ir, "bias"));
|
||||
res.act2_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT2, il, ir, "alpha"));
|
||||
res.act2_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT2, il, ir, "beta"));
|
||||
res.conv2_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV2, il, ir, "weight"));
|
||||
res.conv2_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV2, il, ir, "bias"));
|
||||
}
|
||||
}
|
||||
|
||||
c2w.dac_post_snake_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_SNAKE, "alpha"));
|
||||
c2w.dac_post_snake_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_SNAKE, "beta"));
|
||||
c2w.dac_post_conv_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_CONV, "weight"));
|
||||
c2w.dac_post_conv_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_CONV, "bias"));
|
||||
};
|
||||
|
||||
switch (model.proj_type) {
|
||||
case PROJECTOR_TYPE_MLP:
|
||||
@@ -3656,91 +3767,25 @@ struct clip_model_loader {
|
||||
model.gen_code_norm_w = get_tensor(string_format(TN_A_GEN_CODE_NORM, "weight"));
|
||||
|
||||
// code2wav: RVQ codes -> raw PCM, lives in the same ctx as code_predictor
|
||||
{
|
||||
auto & c2w = model.c2w;
|
||||
|
||||
c2w.quant_first_in_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_IN, "weight"));
|
||||
c2w.quant_first_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_OUT, "weight"));
|
||||
c2w.quant_first_cb_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_CB, "weight"));
|
||||
c2w.quant_rest_in_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_REST_IN, "weight"));
|
||||
c2w.quant_rest_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_REST_OUT, "weight"));
|
||||
c2w.quant_rest_cb_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_REST_CB, "weight"));
|
||||
|
||||
c2w.pre_conv_w = get_tensor(string_format(TN_A_GEN_WAV_PRE_CONV, "weight"));
|
||||
c2w.pre_conv_b = get_tensor(string_format(TN_A_GEN_WAV_PRE_CONV, "bias"));
|
||||
|
||||
c2w.tfm_in_proj_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_IN_PROJ, "weight"));
|
||||
c2w.tfm_in_proj_b = get_tensor(string_format(TN_A_GEN_WAV_TFM_IN_PROJ, "bias"));
|
||||
c2w.tfm_out_proj_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_PROJ, "weight"));
|
||||
c2w.tfm_out_proj_b = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_PROJ, "bias"));
|
||||
c2w.tfm_output_norm_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_NORM, "weight"));
|
||||
|
||||
// loaded manually, the generic model.layers loop is taken by code_predictor
|
||||
c2w.tfm_layers.resize(hparams.wav_tfm_n_layer);
|
||||
for (int il = 0; il < hparams.wav_tfm_n_layer; il++) {
|
||||
auto & layer = c2w.tfm_layers[il];
|
||||
const char * p = "a.gen.wav.tfm";
|
||||
layer.q_w = get_tensor(string_format(TN_ATTN_Q, p, il, "weight"));
|
||||
layer.k_w = get_tensor(string_format(TN_ATTN_K, p, il, "weight"));
|
||||
layer.v_w = get_tensor(string_format(TN_ATTN_V, p, il, "weight"));
|
||||
layer.o_w = get_tensor(string_format(TN_ATTN_OUTPUT, p, il, "weight"));
|
||||
layer.ln_1_w = get_tensor(string_format(TN_LN_1, p, il, "weight"));
|
||||
layer.ln_2_w = get_tensor(string_format(TN_LN_2, p, il, "weight"));
|
||||
layer.ls_1_w = get_tensor(string_format(TN_LS_1, p, il, "weight"));
|
||||
layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight"));
|
||||
layer.ff_gate_w = get_tensor(string_format(TN_FFN_GATE, p, il, "weight"));
|
||||
layer.ff_up_w = get_tensor(string_format(TN_FFN_UP, p, il, "weight"));
|
||||
layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, p, il, "weight"));
|
||||
}
|
||||
|
||||
// upsample: 2x (causal ConvTranspose1d + ConvNeXt block)
|
||||
c2w.upsample.resize(hparams.wav_upsample_n_block);
|
||||
for (int il = 0; il < hparams.wav_upsample_n_block; il++) {
|
||||
auto & up = c2w.upsample[il];
|
||||
up.conv_w = get_tensor(string_format(TN_A_GEN_WAV_UP_CONV, il, "weight"));
|
||||
up.conv_b = get_tensor(string_format(TN_A_GEN_WAV_UP_CONV, il, "bias"));
|
||||
up.dwconv_w = get_tensor(string_format(TN_A_GEN_WAV_UP_DWCONV, il, "weight"));
|
||||
up.dwconv_b = get_tensor(string_format(TN_A_GEN_WAV_UP_DWCONV, il, "bias"));
|
||||
up.norm_w = get_tensor(string_format(TN_A_GEN_WAV_UP_NORM, il, "weight"));
|
||||
up.norm_b = get_tensor(string_format(TN_A_GEN_WAV_UP_NORM, il, "bias"));
|
||||
up.pw1_w = get_tensor(string_format(TN_A_GEN_WAV_UP_PW1, il, "weight"));
|
||||
up.pw1_b = get_tensor(string_format(TN_A_GEN_WAV_UP_PW1, il, "bias"));
|
||||
up.pw2_w = get_tensor(string_format(TN_A_GEN_WAV_UP_PW2, il, "weight"));
|
||||
up.pw2_b = get_tensor(string_format(TN_A_GEN_WAV_UP_PW2, il, "bias"));
|
||||
up.gamma = get_tensor(string_format(TN_A_GEN_WAV_UP_GAMMA, il));
|
||||
}
|
||||
|
||||
// DAC decoder: conv_pre + n upsample blocks (each with n_res residual units) + conv_post
|
||||
c2w.dac_entry_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_ENTRY, "weight"));
|
||||
c2w.dac_entry_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_ENTRY, "bias"));
|
||||
|
||||
c2w.dac.resize(hparams.wav_dac_n_block);
|
||||
for (int il = 0; il < hparams.wav_dac_n_block; il++) {
|
||||
auto & blk = c2w.dac[il];
|
||||
blk.snake_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_SNAKE, il, "alpha"));
|
||||
blk.snake_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_SNAKE, il, "beta"));
|
||||
blk.conv_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_CONV, il, "weight"));
|
||||
blk.conv_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_CONV, il, "bias"));
|
||||
|
||||
blk.res.resize(hparams.wav_dac_n_res);
|
||||
for (int ir = 0; ir < hparams.wav_dac_n_res; ir++) {
|
||||
auto & res = blk.res[ir];
|
||||
res.act1_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT1, il, ir, "alpha"));
|
||||
res.act1_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT1, il, ir, "beta"));
|
||||
res.conv1_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV1, il, ir, "weight"));
|
||||
res.conv1_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV1, il, ir, "bias"));
|
||||
res.act2_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT2, il, ir, "alpha"));
|
||||
res.act2_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_ACT2, il, ir, "beta"));
|
||||
res.conv2_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV2, il, ir, "weight"));
|
||||
res.conv2_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_RES_CONV2, il, ir, "bias"));
|
||||
}
|
||||
}
|
||||
|
||||
c2w.dac_post_snake_alpha = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_SNAKE, "alpha"));
|
||||
c2w.dac_post_snake_beta = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_SNAKE, "beta"));
|
||||
c2w.dac_post_conv_w = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_CONV, "weight"));
|
||||
c2w.dac_post_conv_b = get_tensor(string_format(TN_A_GEN_WAV_DAC_POST_CONV, "bias"));
|
||||
load_code2wav(true);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DRAMA3_GEN:
|
||||
{
|
||||
// generic a.gen.code.blk.* tensors are loaded by the loop above
|
||||
model.drama3_layers.resize(hparams.n_layer);
|
||||
for (int il = 0; il < hparams.n_layer; ++il) {
|
||||
auto & tl = model.drama3_layers[il];
|
||||
model.layers[il].attn_gate_w = get_tensor(string_format(TN_ATTN_GATE, prefix, il, "weight"));
|
||||
tl.enorm_w = get_tensor(string_format(TN_A_GEN_CODE_ENORM, il, "weight"));
|
||||
tl.hnorm_w = get_tensor(string_format(TN_A_GEN_CODE_HNORM, il, "weight"));
|
||||
tl.eh_proj_w = get_tensor(string_format(TN_A_GEN_CODE_EH_PROJ, il, "weight"));
|
||||
tl.embd_w = get_tensor(string_format(TN_A_GEN_CODE_BLK_EMBD, il, "weight"));
|
||||
tl.out_norm_w = get_tensor(string_format(TN_A_GEN_CODE_BLK_OUT_NORM, il, "weight"));
|
||||
}
|
||||
model.gen_code_embd_w = get_tensor(string_format(TN_A_GEN_CODE_EMBD, "weight")); // merged 3D, one table per codebook
|
||||
model.gen_code_head_w = get_tensor(string_format(TN_A_GEN_CODE_HEAD, "weight")); // merged 3D, codebooks 1..
|
||||
model.gen_code_out_embd_w = get_tensor(string_format(TN_A_GEN_CODE_OUT_EMBD, "weight")); // backbone rows of the semantic band
|
||||
load_code2wav(false);
|
||||
} break;
|
||||
case PROJECTOR_TYPE_CHATTERBOX:
|
||||
{
|
||||
@@ -5520,6 +5565,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
case PROJECTOR_TYPE_LFM3A_GEN:
|
||||
case PROJECTOR_TYPE_DRAMA3_GEN:
|
||||
{
|
||||
// one hidden-state vector fed back to the talker per call
|
||||
n_patches = 1;
|
||||
@@ -5587,6 +5633,7 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
|
||||
static std::vector<c2w_state_slot> list_gen_state_slots(const clip_hparams & hparams, const clip_model & model) {
|
||||
switch (model.proj_type) {
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN: return list_c2w_state_slots(hparams, model);
|
||||
case PROJECTOR_TYPE_DRAMA3_GEN: return list_drama3_state_slots(hparams, model);
|
||||
case PROJECTOR_TYPE_POCKETTTS_GEN: return list_pockettts_state_slots(hparams, model);
|
||||
case PROJECTOR_TYPE_LFM3A_GEN: return list_lfm3a_gen_state_slots(hparams, model);
|
||||
default: return {};
|
||||
@@ -5772,6 +5819,14 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
LOG_ERR("%s: audio gen requires a gen state\n", __func__);
|
||||
return false;
|
||||
}
|
||||
if (ctx->proj_type() == PROJECTOR_TYPE_DRAMA3_GEN && params->gen_process == CLIP_GEN_PROCESS_GEN_CODE) {
|
||||
float pos = 0.0f;
|
||||
ggml_backend_tensor_get(params->gen_state->slots[0], &pos, 0, sizeof(float));
|
||||
if ((int) pos >= clip_graph_drama3_gen::n_pos_max) {
|
||||
LOG_ERR("%s: the codebook tower holds at most %d rows\n", __func__, clip_graph_drama3_gen::n_pos_max);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// build the inference graph
|
||||
ggml_backend_sched_reset(ctx->sched.get());
|
||||
@@ -5845,6 +5900,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
size_t i = 0;
|
||||
for (const auto & slot : list_gen_state_slots(hparams, model)) {
|
||||
GGML_ASSERT(i < slots.size());
|
||||
if (slot.in_place) {
|
||||
i++;
|
||||
continue;
|
||||
}
|
||||
ggml_backend_tensor_copy(slots[i++], get_inp_tensor(("state_in_" + slot.name).c_str()));
|
||||
}
|
||||
};
|
||||
@@ -6900,6 +6959,58 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
}
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_DRAMA3_GEN:
|
||||
{
|
||||
const int64_t n_cb = model.gen_code_embd_w->ne[2];
|
||||
const int64_t n_vocab = model.gen_code_embd_w->ne[1];
|
||||
if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) {
|
||||
// frame-major to group-major, clamp: codec residual codebooks are smaller than the tower vocab
|
||||
GGML_ASSERT(params->codes != nullptr && params->codes->size() % n_cb == 0);
|
||||
const int64_t n_frames = (int64_t) params->codes->size() / n_cb;
|
||||
const int64_t n_first = model.c2w.quant_first_cb_w->ne[1];
|
||||
const int64_t n_rest = model.c2w.quant_rest_cb_w->ne[1];
|
||||
std::vector<int32_t> codes(params->codes->size());
|
||||
for (int64_t f = 0; f < n_frames; f++) {
|
||||
for (int64_t g = 0; g < n_cb; g++) {
|
||||
const int64_t n = g == 0 ? n_first : n_rest;
|
||||
codes[g * n_frames + f] = std::clamp<int32_t>((*params->codes)[f * n_cb + g], 0, (int32_t) n - 1);
|
||||
}
|
||||
}
|
||||
set_input_i32("inp_codes", codes);
|
||||
set_gen_state_in();
|
||||
} else {
|
||||
// code0 < 0: next row is text (prompt), the tower only fills its KV
|
||||
// codes: input frame codes for a speech row, empty for a text row
|
||||
const bool next_speech = params->code0 >= 0;
|
||||
const bool cur_speech = params->codes && (int64_t) params->codes->size() == n_cb;
|
||||
if (params->code0 >= n_vocab) {
|
||||
LOG_ERR("%s: code0 out of range (%d, vocab %lld)\n", __func__, params->code0, (long long) n_vocab);
|
||||
return false;
|
||||
}
|
||||
std::vector<int32_t> code0 = { std::max(params->code0, 0) };
|
||||
set_input_i32("inp_code0", code0);
|
||||
|
||||
std::vector<int32_t> extra(n_cb - 2);
|
||||
for (int64_t g = 2; g < n_cb; g++) {
|
||||
const int32_t c = cur_speech ? std::clamp<int32_t>((*params->codes)[g], 0, (int32_t) n_vocab - 1) : 0;
|
||||
extra[g - 2] = c + (int32_t) (g * n_vocab);
|
||||
}
|
||||
set_input_i32("inp_extra", extra);
|
||||
set_input_f32("inp_scale", { next_speech ? 1.0f : 0.0f, cur_speech ? 1.0f : 0.0f });
|
||||
|
||||
float pos = 0.0f;
|
||||
ggml_backend_tensor_get(params->gen_state->slots[0], &pos, 0, sizeof(float));
|
||||
std::vector<int32_t> inp_pos = { (int32_t) pos };
|
||||
set_input_i32("inp_pos", inp_pos);
|
||||
|
||||
std::uniform_real_distribution<float> dist(0.0f, 1.0f);
|
||||
std::vector<float> r(model.gen_code_head_w->ne[2]);
|
||||
for (auto & v : r) {
|
||||
v = dist(params->gen_state->rng);
|
||||
}
|
||||
set_input_f32("inp_rand", r);
|
||||
}
|
||||
} break;
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
{
|
||||
if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) {
|
||||
@@ -7700,13 +7811,15 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
clip_graph_soprano::decode_spectrum(model, spectrum, out_audio);
|
||||
}
|
||||
|
||||
// drop the tail audio that comes from the code-0 rear padding
|
||||
const int64_t n_codes = params->codes ? model.gen_code_head_w->ne[2] + 1 : 0;
|
||||
const int64_t n_frames_w = hparams.wav_tfm_swa;
|
||||
const int64_t n_frames = params->codes ? (int64_t) params->codes->size() / n_codes : n_frames_w;
|
||||
if (n_frames < n_frames_w) {
|
||||
const size_t hop = out_audio.size() / n_frames_w;
|
||||
out_audio.resize((size_t) n_frames * hop);
|
||||
// drop the tail audio that comes from the code-0 rear padding, drama3 has no padding
|
||||
if (model.proj_type != PROJECTOR_TYPE_DRAMA3_GEN) {
|
||||
const int64_t n_codes = params->codes ? model.gen_code_head_w->ne[2] + 1 : 0;
|
||||
const int64_t n_frames_w = hparams.wav_tfm_swa;
|
||||
const int64_t n_frames = params->codes ? (int64_t) params->codes->size() / n_codes : n_frames_w;
|
||||
if (n_frames < n_frames_w) {
|
||||
const size_t hop = out_audio.size() / n_frames_w;
|
||||
out_audio.resize((size_t) n_frames * hop);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (params->enc_state != nullptr && model.proj_type == PROJECTOR_TYPE_LFM3A) {
|
||||
@@ -7744,6 +7857,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
|
||||
const auto & slots = params->gen_state->slots;
|
||||
size_t i = 0;
|
||||
for (const auto & slot : list_gen_state_slots(hparams, model)) {
|
||||
if (slot.in_place) {
|
||||
i++;
|
||||
continue;
|
||||
}
|
||||
ggml_tensor * t = ggml_graph_get_tensor(gf, ("state_out_" + slot.name).c_str());
|
||||
if (t == nullptr) {
|
||||
GGML_ABORT("graph has no \"state_out_%s\" tensor", slot.name.c_str());
|
||||
@@ -7924,6 +8041,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
|
||||
return ctx->model.mm_fc_w->ne[2];
|
||||
case PROJECTOR_TYPE_QWEN3TTS_GEN:
|
||||
case PROJECTOR_TYPE_LFM3A_GEN:
|
||||
case PROJECTOR_TYPE_DRAMA3_GEN:
|
||||
return ctx->model.gen_code_out_embd_w->ne[0];
|
||||
case PROJECTOR_TYPE_POCKETTTS_SPKENC:
|
||||
return ctx->model.spk_proj_w->ne[1];
|
||||
|
||||
@@ -0,0 +1,256 @@
|
||||
#include "models.h"
|
||||
|
||||
#include <cmath>
|
||||
#include <string>
|
||||
|
||||
// tower RoPE: interleaved, first 64 dims of each head
|
||||
static constexpr int DRAMA3_N_ROT = 64;
|
||||
static constexpr float DRAMA3_ROPE_THETA = 1e7f;
|
||||
|
||||
std::vector<c2w_state_slot> list_drama3_state_slots(const clip_hparams & hparams, const clip_model & model) {
|
||||
std::vector<c2w_state_slot> slots;
|
||||
slots.push_back({"tower_pos", 1, 1, true});
|
||||
for (size_t il = 0; il < model.layers.size(); il++) {
|
||||
const int64_t n_kv = model.layers[il].k_w->ne[1];
|
||||
slots.push_back({"tower_k_" + std::to_string(il), n_kv, clip_graph_drama3_gen::n_pos_max, true});
|
||||
slots.push_back({"tower_v_" + std::to_string(il), n_kv, clip_graph_drama3_gen::n_pos_max, true});
|
||||
}
|
||||
for (auto & slot : list_c2w_state_slots(hparams, model)) {
|
||||
slots.push_back(slot);
|
||||
}
|
||||
return slots;
|
||||
}
|
||||
|
||||
// h_out = out_norm(Block(eh_proj([enorm(mask * E[code] + extra); hnorm(h_in)])))
|
||||
ggml_tensor * clip_graph_drama3_gen::tower_layer(ggml_tensor * x, int il, ggml_tensor * code, ggml_tensor * mask, ggml_tensor * extra,
|
||||
ggml_tensor * inp_pos, ggml_tensor * k_slot, ggml_tensor * v_slot, int pos) {
|
||||
const auto & layer = model.layers[il];
|
||||
const auto & tl = model.drama3_layers[il];
|
||||
const int n_head = hparams.n_head;
|
||||
const int n_head_kv = hparams.n_head_kv;
|
||||
const int64_t d_head = layer.q_w->ne[1] / n_head;
|
||||
const int64_t n_kv = d_head * n_head_kv;
|
||||
const float eps = hparams.eps;
|
||||
|
||||
ggml_tensor * e = ggml_mul(ctx0, ggml_get_rows(ctx0, tl.embd_w, code), mask);
|
||||
if (extra) {
|
||||
e = ggml_add(ctx0, e, extra);
|
||||
}
|
||||
e = ggml_mul(ctx0, ggml_rms_norm(ctx0, e, eps), tl.enorm_w);
|
||||
ggml_tensor * h = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, eps), tl.hnorm_w);
|
||||
ggml_tensor * mixed = ggml_mul_mat(ctx0, tl.eh_proj_w, ggml_concat(ctx0, e, h, 0));
|
||||
|
||||
ggml_tensor * cur = ggml_mul(ctx0, ggml_rms_norm(ctx0, mixed, eps), layer.ln_1_w);
|
||||
ggml_tensor * q = ggml_reshape_3d(ctx0, ggml_mul_mat(ctx0, layer.q_w, cur), d_head, n_head, 1);
|
||||
ggml_tensor * k = ggml_reshape_3d(ctx0, ggml_mul_mat(ctx0, layer.k_w, cur), d_head, n_head_kv, 1);
|
||||
ggml_tensor * v = ggml_mul_mat(ctx0, layer.v_w, cur);
|
||||
ggml_tensor * gate = ggml_mul_mat(ctx0, layer.attn_gate_w, cur);
|
||||
|
||||
q = ggml_mul(ctx0, ggml_rms_norm(ctx0, q, eps), layer.q_norm);
|
||||
k = ggml_mul(ctx0, ggml_rms_norm(ctx0, k, eps), layer.k_norm);
|
||||
q = ggml_rope_ext(ctx0, q, inp_pos, nullptr, DRAMA3_N_ROT, GGML_ROPE_TYPE_NORMAL, 0, DRAMA3_ROPE_THETA, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
k = ggml_rope_ext(ctx0, k, inp_pos, nullptr, DRAMA3_N_ROT, GGML_ROPE_TYPE_NORMAL, 0, DRAMA3_ROPE_THETA, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
|
||||
|
||||
// write K/V of this row to the slot, attention uses the cached rows + this row
|
||||
ggml_tensor * k_new = ggml_reshape_2d(ctx0, k, n_kv, 1);
|
||||
ggml_build_forward_expand(gf, ggml_cpy(ctx0, k_new, ggml_view_2d(ctx0, k_slot, n_kv, 1, k_slot->nb[1], (size_t) pos * k_slot->nb[1])));
|
||||
ggml_build_forward_expand(gf, ggml_cpy(ctx0, v, ggml_view_2d(ctx0, v_slot, n_kv, 1, v_slot->nb[1], (size_t) pos * v_slot->nb[1])));
|
||||
ggml_tensor * k_all = k_new;
|
||||
ggml_tensor * v_all = v;
|
||||
if (pos > 0) {
|
||||
k_all = ggml_concat(ctx0, ggml_view_2d(ctx0, k_slot, n_kv, pos, k_slot->nb[1], 0), k_new, 1);
|
||||
v_all = ggml_concat(ctx0, ggml_view_2d(ctx0, v_slot, n_kv, pos, v_slot->nb[1], 0), v, 1);
|
||||
}
|
||||
|
||||
ggml_tensor * attn = build_attn(nullptr, nullptr,
|
||||
ggml_reshape_4d(ctx0, q, d_head, n_head, 1, 1),
|
||||
ggml_reshape_4d(ctx0, k_all, d_head, n_head_kv, pos + 1, 1),
|
||||
ggml_reshape_4d(ctx0, v_all, d_head, n_head_kv, pos + 1, 1),
|
||||
nullptr, 1.0f / sqrtf((float) d_head), il);
|
||||
attn = ggml_mul(ctx0, attn, ggml_sigmoid(ctx0, gate));
|
||||
cur = ggml_add(ctx0, mixed, ggml_mul_mat(ctx0, layer.o_w, attn));
|
||||
|
||||
ggml_tensor * ff = ggml_mul(ctx0, ggml_rms_norm(ctx0, cur, eps), layer.ln_2_w);
|
||||
ff = ggml_swiglu_split(ctx0, ggml_mul_mat(ctx0, layer.ff_gate_w, ff), ggml_mul_mat(ctx0, layer.ff_up_w, ff));
|
||||
cur = ggml_add(ctx0, cur, ggml_mul_mat(ctx0, layer.ff_down_w, ff));
|
||||
|
||||
cur = ggml_mul(ctx0, ggml_rms_norm(ctx0, cur, eps), tl.out_norm_w);
|
||||
cb(cur, "tower_out", il);
|
||||
return cur;
|
||||
}
|
||||
|
||||
// one row: backbone hidden + next main code -> other codes of that frame + next backbone input
|
||||
ggml_cgraph * clip_graph_drama3_gen::build_tower() {
|
||||
const int n_layer = (int) model.layers.size();
|
||||
const int n_heads = (int) model.gen_code_head_w->ne[2]; // one per residual codebook
|
||||
const int n_shared = n_layer - n_heads;
|
||||
const int64_t n_cb = model.gen_code_embd_w->ne[2];
|
||||
const int64_t n_embd_cb = model.gen_code_embd_w->ne[0];
|
||||
|
||||
int pos = 0;
|
||||
std::vector<ggml_tensor *> slots;
|
||||
const auto slot_list = list_drama3_state_slots(hparams, model);
|
||||
if (gen_state) {
|
||||
float pos_f = 0.0f;
|
||||
ggml_backend_tensor_get(gen_state->slots[0], &pos_f, 0, sizeof(float));
|
||||
pos = (int) pos_f;
|
||||
slots = gen_state->slots;
|
||||
} else {
|
||||
// buffer reserve: no gen state, use stand-in slots
|
||||
for (const auto & slot : slot_list) {
|
||||
ggml_tensor * t = nullptr;
|
||||
if (slot.in_place) {
|
||||
t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1);
|
||||
ggml_set_input(t);
|
||||
}
|
||||
slots.push_back(t);
|
||||
}
|
||||
}
|
||||
GGML_ASSERT(pos < n_pos_max);
|
||||
|
||||
ggml_tensor * h = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_mmproj_embd);
|
||||
ggml_set_name(h, "inp_raw");
|
||||
ggml_set_input(h);
|
||||
|
||||
ggml_tensor * inp_code0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 1);
|
||||
ggml_set_name(inp_code0, "inp_code0");
|
||||
ggml_set_input(inp_code0);
|
||||
|
||||
// codes 2.. of the input frame, as row ids of the merged codebook table
|
||||
ggml_tensor * inp_extra = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_cb - 2);
|
||||
ggml_set_name(inp_extra, "inp_extra");
|
||||
ggml_set_input(inp_extra);
|
||||
|
||||
// [next frame is speech, input frame is speech]
|
||||
ggml_tensor * inp_scale = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 2);
|
||||
ggml_set_name(inp_scale, "inp_scale");
|
||||
ggml_set_input(inp_scale);
|
||||
|
||||
ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 1);
|
||||
ggml_set_name(inp_pos, "inp_pos");
|
||||
ggml_set_input(inp_pos);
|
||||
|
||||
ggml_tensor * inp_rand = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_heads);
|
||||
ggml_set_name(inp_rand, "inp_rand");
|
||||
ggml_set_input(inp_rand);
|
||||
|
||||
ggml_tensor * mask = ggml_view_1d(ctx0, inp_scale, 1, 0);
|
||||
ggml_tensor * extra_on = ggml_view_1d(ctx0, inp_scale, 1, sizeof(float));
|
||||
|
||||
ggml_tensor * cb_all = ggml_reshape_2d(ctx0, model.gen_code_embd_w, n_embd_cb, model.gen_code_embd_w->ne[1] * n_cb);
|
||||
ggml_tensor * extra = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_get_rows(ctx0, cb_all, inp_extra)));
|
||||
extra = ggml_mul(ctx0, ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, extra), n_embd_cb, 1), extra_on);
|
||||
|
||||
clip_graph_qwen3tts_gen::code_gen cg(*this, top_k, top_p);
|
||||
ggml_tensor * out_codes = cg.cache_set(ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_cb), 0, inp_code0);
|
||||
|
||||
ggml_tensor * code = inp_code0;
|
||||
ggml_tensor * x = ggml_reshape_2d(ctx0, h, n_mmproj_embd, 1);
|
||||
for (int il = 0; il < n_layer; il++) {
|
||||
x = tower_layer(x, il, code, mask, il == 0 ? extra : nullptr, inp_pos, slots[1 + 2 * il], slots[2 + 2 * il], pos);
|
||||
if (il < n_shared) {
|
||||
continue;
|
||||
}
|
||||
const int ih = il - n_shared;
|
||||
ggml_tensor * head = ggml_view_2d(ctx0, model.gen_code_head_w, model.gen_code_head_w->ne[0], model.gen_code_head_w->ne[1],
|
||||
model.gen_code_head_w->nb[1], (size_t) ih * model.gen_code_head_w->nb[2]);
|
||||
ggml_tensor * logits = ggml_mul_mat(ctx0, head, x);
|
||||
if (temp > 0.0f && temp != 1.0f) {
|
||||
logits = ggml_scale(ctx0, logits, 1.0f / temp);
|
||||
}
|
||||
code = cg.do_sampling(logits, ggml_view_1d(ctx0, inp_rand, 1, (size_t) ih * sizeof(float)));
|
||||
out_codes = cg.cache_set(out_codes, ih + 1, code);
|
||||
}
|
||||
|
||||
ggml_tensor * pos_slot = slots[0];
|
||||
ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_scale_bias(ctx0, pos_slot, 1.0f, 1.0f), pos_slot));
|
||||
|
||||
out_codes = ggml_cont(ctx0, out_codes);
|
||||
ggml_set_name(out_codes, "out_codes");
|
||||
ggml_set_output(out_codes);
|
||||
ggml_build_forward_expand(gf, out_codes);
|
||||
|
||||
// backbone input of the next row: (E_lm[main] + E_cb1[code 1]) / sqrt(2)
|
||||
ggml_tensor * code1 = ggml_view_1d(ctx0, out_codes, 1, out_codes->nb[1]);
|
||||
ggml_tensor * cb1 = ggml_view_2d(ctx0, model.gen_code_embd_w, n_embd_cb, model.gen_code_embd_w->ne[1],
|
||||
model.gen_code_embd_w->nb[1], model.gen_code_embd_w->nb[2]);
|
||||
ggml_tensor * out_embd = ggml_add(ctx0, ggml_get_rows(ctx0, model.gen_code_out_embd_w, inp_code0), ggml_get_rows(ctx0, cb1, code1));
|
||||
out_embd = ggml_scale(ctx0, out_embd, 1.0f / sqrtf(2.0f));
|
||||
cb(out_embd, "gen_audio_out", -1);
|
||||
ggml_build_forward_expand(gf, out_embd);
|
||||
|
||||
return gf;
|
||||
}
|
||||
|
||||
// frames -> PCM, qwen3tts code2wav blocks without pre_conv and tfm projections
|
||||
ggml_cgraph * clip_graph_drama3_gen::build_codec() {
|
||||
const auto & c2w = model.c2w;
|
||||
const int64_t n_cb = model.gen_code_embd_w->ne[2];
|
||||
|
||||
ggml_tensor * inp_codes = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_frames, n_cb);
|
||||
ggml_set_name(inp_codes, "inp_codes");
|
||||
ggml_set_input(inp_codes);
|
||||
|
||||
clip_graph_qwen3tts_gen::code2wav dec(*this);
|
||||
dec.gelu_erf = true;
|
||||
for (const auto & slot : list_c2w_state_slots(hparams, model)) {
|
||||
ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1);
|
||||
ggml_set_name(t, ("state_in_" + slot.name).c_str());
|
||||
ggml_set_input(t);
|
||||
dec.state_in[slot.name] = t;
|
||||
}
|
||||
|
||||
// out_proj is folded into the codebook tables at conversion
|
||||
auto group_ids = [&](int g) {
|
||||
return ggml_view_1d(ctx0, inp_codes, n_frames, (size_t) g * inp_codes->nb[1]);
|
||||
};
|
||||
ggml_tensor * cur = ggml_get_rows(ctx0, c2w.quant_first_cb_w, group_ids(0)); // [dim, T]
|
||||
for (int g = 1; g < n_cb; g++) {
|
||||
ggml_tensor * table = ggml_view_2d(ctx0, c2w.quant_rest_cb_w, c2w.quant_rest_cb_w->ne[0], c2w.quant_rest_cb_w->ne[1],
|
||||
c2w.quant_rest_cb_w->nb[1], (size_t) (g - 1) * c2w.quant_rest_cb_w->nb[2]);
|
||||
cur = ggml_add(ctx0, cur, ggml_get_rows(ctx0, table, group_ids(g)));
|
||||
}
|
||||
|
||||
for (int il = 0; il < hparams.wav_tfm_n_layer; il++) {
|
||||
cur = dec.tfm_layer_forward(cur, c2w.tfm_layers[il], il);
|
||||
}
|
||||
cur = ggml_mul(ctx0, ggml_rms_norm(ctx0, cur, hparams.wav_tfm_eps), c2w.tfm_output_norm_w);
|
||||
cb(cur, "wav_tfm_out", -1);
|
||||
|
||||
ggml_tensor * x = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [T, dim]
|
||||
for (size_t il = 0; il < c2w.upsample.size(); il++) {
|
||||
const auto & up = c2w.upsample[il];
|
||||
x = dec.causal_conv_transpose1d(x, up.conv_w, up.conv_b, 2, "");
|
||||
x = dec.convnext_block(x, up, "up" + std::to_string(il));
|
||||
}
|
||||
|
||||
static constexpr int DAC_DILATIONS[3] = { 1, 3, 9 };
|
||||
x = dec.causal_conv1d(x, c2w.dac_entry_w, c2w.dac_entry_b, 1, "dac_entry");
|
||||
for (size_t il = 0; il < c2w.dac.size(); il++) {
|
||||
const auto & blk = c2w.dac[il];
|
||||
const int stride = (int) (blk.conv_w->ne[0] / 2);
|
||||
const std::string blk_name = "dac" + std::to_string(il);
|
||||
x = dec.snake(x, blk.snake_alpha, blk.snake_beta);
|
||||
x = dec.causal_conv_transpose1d(x, blk.conv_w, blk.conv_b, stride, blk_name + "_tail");
|
||||
for (size_t ir = 0; ir < blk.res.size(); ir++) {
|
||||
x = dec.dac_res_unit(x, blk.res[ir], DAC_DILATIONS[ir], blk_name + "_res" + std::to_string(ir));
|
||||
}
|
||||
}
|
||||
x = dec.snake(x, c2w.dac_post_snake_alpha, c2w.dac_post_snake_beta);
|
||||
x = dec.causal_conv1d(x, c2w.dac_post_conv_w, c2w.dac_post_conv_b, 1, "dac_post_conv");
|
||||
x = ggml_tanh(ctx0, ggml_reshape_1d(ctx0, x, x->ne[0]));
|
||||
ggml_set_name(x, "out_audio");
|
||||
ggml_set_output(x);
|
||||
ggml_build_forward_expand(gf, x);
|
||||
|
||||
for (auto & slot : dec.state_out) {
|
||||
ggml_set_name(slot.second, ("state_out_" + slot.first).c_str());
|
||||
ggml_set_output(slot.second);
|
||||
ggml_build_forward_expand(gf, slot.second);
|
||||
}
|
||||
return gf;
|
||||
}
|
||||
|
||||
ggml_cgraph * clip_graph_drama3_gen::build() {
|
||||
return gen_process == CLIP_GEN_PROCESS_GEN_WAV ? build_codec() : build_tower();
|
||||
}
|
||||
@@ -424,6 +424,9 @@ struct clip_graph_qwen3tts_gen : clip_graph {
|
||||
ggml_tensor * convnext_block(ggml_tensor * x, const clip_code2wav::upsample_block & blk, const std::string & state_prefix) const;
|
||||
ggml_tensor * dac_res_unit(ggml_tensor * x, const clip_code2wav::dac_res & res, int dilation, const std::string & state_name) const;
|
||||
|
||||
// drama3: erf GELU in ConvNeXt, not the tanh approximation
|
||||
bool gelu_erf = false;
|
||||
|
||||
// inp_codes [1, n_codes] I32 -> this frame's audio samples [n_samples] F32, clamped to [-1, 1]
|
||||
ggml_tensor * decode(ggml_tensor * inp_codes) const;
|
||||
};
|
||||
@@ -487,12 +490,43 @@ struct c2w_state_slot {
|
||||
std::string name;
|
||||
int64_t ne0;
|
||||
int64_t ne1;
|
||||
bool in_place = false; // graph writes the slot directly, no state_in/state_out copy
|
||||
};
|
||||
std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model);
|
||||
|
||||
// drama3: tower KV slots, then list_c2w_state_slots()
|
||||
std::vector<c2w_state_slot> list_drama3_state_slots(const clip_hparams & hparams, const clip_model & model);
|
||||
|
||||
// same, for the streaming mimi decoder (pocket-tts GEN_WAV)
|
||||
std::vector<c2w_state_slot> list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model);
|
||||
|
||||
//
|
||||
// drama3 generation:
|
||||
// GEN_CODE = one codebook tower row, each layer attends over the whole sequence (KV in gen state)
|
||||
// GEN_WAV = Fish DAC decoder, frames -> PCM
|
||||
//
|
||||
struct clip_graph_drama3_gen : clip_graph {
|
||||
clip_graph_drama3_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, clip_gen_state * gen_state,
|
||||
int n_frames, int top_k, float top_p, float temp)
|
||||
: clip_graph(ctx, img), gen_process(gen_process), gen_state(gen_state), n_frames(n_frames), top_k(top_k), top_p(top_p), temp(temp) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
// tower KV rows, shared by the prompt and the generated frames
|
||||
static constexpr int n_pos_max = 4096;
|
||||
|
||||
clip_gen_process_type gen_process;
|
||||
clip_gen_state * gen_state; // nullptr on buffer reserve
|
||||
int n_frames; // GEN_WAV only
|
||||
int top_k;
|
||||
float top_p;
|
||||
float temp;
|
||||
|
||||
ggml_cgraph * build_tower();
|
||||
ggml_cgraph * build_codec();
|
||||
ggml_tensor * tower_layer(ggml_tensor * x, int il, ggml_tensor * code, ggml_tensor * mask, ggml_tensor * extra,
|
||||
ggml_tensor * inp_pos, ggml_tensor * k_slot, ggml_tensor * v_slot, int pos);
|
||||
};
|
||||
|
||||
struct clip_graph_soprano : clip_graph {
|
||||
clip_graph_soprano(clip_ctx * ctx, const clip_image_f32 & img, int n_frames) : clip_graph(ctx, img), n_frames(n_frames) {}
|
||||
ggml_cgraph * build() override;
|
||||
|
||||
@@ -527,7 +527,7 @@ ggml_tensor * clip_graph_qwen3tts_gen::code2wav::convnext_block(ggml_tensor * x,
|
||||
|
||||
ggml_tensor * g = ggml_mul_mat(ctx0, blk.pw1_w, hc);
|
||||
g = ggml_add(ctx0, g, blk.pw1_b);
|
||||
g = ggml_gelu(ctx0, g);
|
||||
g = gelu_erf ? ggml_gelu_erf(ctx0, g) : ggml_gelu(ctx0, g);
|
||||
g = ggml_mul_mat(ctx0, blk.pw2_w, g);
|
||||
g = ggml_add(ctx0, g, blk.pw2_b);
|
||||
g = ggml_mul(ctx0, g, blk.gamma);
|
||||
@@ -617,7 +617,7 @@ std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, c
|
||||
const auto & c2w = model.c2w;
|
||||
std::vector<c2w_state_slot> slots;
|
||||
|
||||
if (c2w.pre_conv_w == nullptr) {
|
||||
if (c2w.dac_entry_w == nullptr) {
|
||||
return slots; // not a code2wav model, it keeps no state between calls
|
||||
}
|
||||
|
||||
@@ -632,7 +632,9 @@ std::vector<c2w_state_slot> list_c2w_state_slots(const clip_hparams & hparams, c
|
||||
slots.push_back({"tfm_v_" + std::to_string(il), kv_ch, prefix});
|
||||
}
|
||||
|
||||
slots.push_back({"pre_conv", c2w.pre_conv_w->ne[0] - 1, c2w.pre_conv_w->ne[1]});
|
||||
if (c2w.pre_conv_w) {
|
||||
slots.push_back({"pre_conv", c2w.pre_conv_w->ne[0] - 1, c2w.pre_conv_w->ne[1]});
|
||||
}
|
||||
|
||||
for (size_t il = 0; il < c2w.upsample.size(); il++) {
|
||||
const auto & up = c2w.upsample[il];
|
||||
|
||||
@@ -2077,6 +2077,345 @@ private:
|
||||
bool wav_header_sent = false;
|
||||
};
|
||||
|
||||
// Drama 3: backbone samples the main token of each frame, codebook tower samples the other codes, Fish DAC makes the PCM
|
||||
// the tower attends over the whole sequence, so it also runs on every prompt row
|
||||
class drama3_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
|
||||
public:
|
||||
using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline;
|
||||
|
||||
void reset() override {
|
||||
seq_id = 0;
|
||||
pos = 0;
|
||||
prompt_tokens.clear();
|
||||
prompt_pos = 0;
|
||||
frame.clear();
|
||||
hist.clear();
|
||||
codes_buf.clear();
|
||||
n_decoded = 0;
|
||||
gen_state.reset();
|
||||
pcm_sent = 0;
|
||||
wav_header_sent = false;
|
||||
audio_pcm.clear();
|
||||
h_state_buf.clear();
|
||||
out_buf.clear();
|
||||
}
|
||||
|
||||
int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
|
||||
reset();
|
||||
seq_id = inp->seq_id;
|
||||
out_type = inp->out_type;
|
||||
stream = inp->stream;
|
||||
if (inp->speaker_ref) {
|
||||
LOG_WRN("mtmd_helper_gen_audio: voice cloning is not supported yet, the speaker file is ignored\n");
|
||||
}
|
||||
|
||||
sem_begin = find_special_token(vocab, "<|semantic:0|>");
|
||||
eos = find_special_token(vocab, "<|im_end|>");
|
||||
if (sem_begin == LLAMA_TOKEN_NULL || eos == LLAMA_TOKEN_NULL) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: missing Drama 3 tokens in vocab\n");
|
||||
return 1;
|
||||
}
|
||||
|
||||
// reference-bracket-v1 prompt without references, tokenize each segment alone as the reference does
|
||||
std::string text(inp->prompt, inp->prompt_len);
|
||||
if (text.rfind("<|speaker:", 0) != 0) {
|
||||
text = "<|speaker:0|> " + text;
|
||||
}
|
||||
const char * segments[] = {
|
||||
"<|im_start|>system\n", "convert the provided text to speech", "<|im_end|>\n",
|
||||
"<|im_start|>user\n", text.c_str(), "<|im_end|>\n", "<|im_start|>assistant\n<|voice|>",
|
||||
};
|
||||
for (const char * seg : segments) {
|
||||
const int32_t len = (int32_t) strlen(seg);
|
||||
std::vector<llama_token> ids(len + 16);
|
||||
const int32_t n = llama_tokenize(vocab, seg, len, ids.data(), (int32_t) ids.size(), false, true);
|
||||
if (n < 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
|
||||
return 1;
|
||||
}
|
||||
prompt_tokens.insert(prompt_tokens.end(), ids.begin(), ids.begin() + n);
|
||||
}
|
||||
if (prompt_tokens.size() >= llama_n_ctx_seq(lctx) || prompt_tokens.size() > n_tower_rows) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: prompt does not fit in the context\n");
|
||||
return 1;
|
||||
}
|
||||
if (!llama_memory_seq_rm(llama_get_memory(lctx), seq_id, 0, -1)) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: cannot reset the sequence\n");
|
||||
return 1;
|
||||
}
|
||||
rng.seed(inp->seed == UINT32_MAX ? std::random_device{}() : inp->seed);
|
||||
return init_gen_state(inp->seed) ? 0 : 1;
|
||||
}
|
||||
|
||||
// run the tower on every prompt row to fill its KV
|
||||
int32_t step_prompt(int32_t n_batch) override {
|
||||
GGML_ASSERT(n_batch > 0);
|
||||
const int n_prompt = (int) prompt_tokens.size();
|
||||
if (prompt_pos >= n_prompt) {
|
||||
return 0;
|
||||
}
|
||||
// small batches: every row also outputs logits
|
||||
const int n = std::min({ n_batch, n_prompt - prompt_pos, 64 });
|
||||
llama_batch batch = llama_batch_init(n, 0, 1);
|
||||
for (int i = 0; i < n; i++) {
|
||||
batch.token [i] = prompt_tokens[(size_t) (prompt_pos + i)];
|
||||
batch.pos [i] = pos + i;
|
||||
batch.n_seq_id[i] = 1;
|
||||
batch.seq_id [i][0] = seq_id;
|
||||
batch.logits [i] = 1;
|
||||
}
|
||||
batch.n_tokens = n;
|
||||
const int ret = llama_decode(lctx, batch);
|
||||
llama_batch_free(batch);
|
||||
if (ret != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n");
|
||||
return -1;
|
||||
}
|
||||
for (int i = 0; i < n; i++) {
|
||||
const float * h = llama_get_embeddings_ith(lctx, i);
|
||||
if (prompt_pos + i == n_prompt - 1) {
|
||||
h_state_buf.assign(h, h + n_embd); // first frame comes from the last prompt row
|
||||
break;
|
||||
}
|
||||
if (!run_tower(-1, h)) {
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
pos += n;
|
||||
prompt_pos += n;
|
||||
return n_prompt - prompt_pos;
|
||||
}
|
||||
|
||||
// sample the main token here, ignore the token of the caller
|
||||
int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
|
||||
GGML_UNUSED(sampled);
|
||||
GGML_UNUSED(h_state_in);
|
||||
*h_state_out = nullptr;
|
||||
*out_stop = false;
|
||||
|
||||
const llama_token main = sample_main(llama_get_logits_ith(lctx, -1));
|
||||
const size_t n_rows = prompt_tokens.size() - 1 + codes_buf.size() / n_cb;
|
||||
if (main == eos || (uint32_t) pos >= llama_n_ctx_seq(lctx) || n_rows >= n_tower_rows) {
|
||||
*out_stop = true;
|
||||
return 0;
|
||||
}
|
||||
hist.push_back(main);
|
||||
if (hist.size() > ras_window) {
|
||||
hist.erase(hist.begin());
|
||||
}
|
||||
|
||||
const float * next_embd = run_tower(main - sem_begin, h_state_buf.data());
|
||||
if (!next_embd) {
|
||||
return 1;
|
||||
}
|
||||
codes_buf.insert(codes_buf.end(), frame.begin(), frame.end());
|
||||
|
||||
// text/speech stream: all mrope sections use the same position
|
||||
const bool mrope = llama_model_rope_type(model) == LLAMA_ROPE_TYPE_MROPE ||
|
||||
llama_model_rope_type(model) == LLAMA_ROPE_TYPE_IMROPE;
|
||||
decode_embd_batch batch(next_embd, 1, mrope ? 4 : 1, n_embd);
|
||||
if (mrope) {
|
||||
batch.set_position_mrope_1d(pos, seq_id);
|
||||
} else {
|
||||
batch.set_position_normal(pos, seq_id);
|
||||
}
|
||||
batch.logits[0] = 1;
|
||||
if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.render(lctx, 0, 1)) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: decode failed\n");
|
||||
return 1;
|
||||
}
|
||||
pos++;
|
||||
const float * h = llama_get_embeddings_ith(lctx, -1);
|
||||
h_state_buf.assign(h, h + n_embd);
|
||||
*h_state_out = h_state_buf.data();
|
||||
|
||||
if (stream && codes_buf.size() / n_cb - n_decoded >= stream_frames) {
|
||||
return decode_wav() ? 0 : 1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
|
||||
*out_sample_rate = info.sample_rate;
|
||||
|
||||
if (!stream) {
|
||||
if (flush() != 0) {
|
||||
return 1;
|
||||
}
|
||||
if (out_n_samples) {
|
||||
*out_n_samples = (int64_t) audio_pcm.size();
|
||||
}
|
||||
if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
|
||||
*out_data = (const char *) audio_pcm.data();
|
||||
*out_data_len = audio_pcm.size() * sizeof(float);
|
||||
return 0;
|
||||
}
|
||||
out_buf.clear();
|
||||
if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
|
||||
return 1;
|
||||
}
|
||||
*out_data = out_buf.data();
|
||||
*out_data_len = out_buf.size();
|
||||
return 0;
|
||||
}
|
||||
|
||||
// streaming: only the audio since the previous call
|
||||
const size_t n_new = audio_pcm.size() - pcm_sent;
|
||||
if (out_n_samples) {
|
||||
*out_n_samples = (int64_t) n_new;
|
||||
}
|
||||
if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
|
||||
*out_data = (const char *) (audio_pcm.data() + pcm_sent);
|
||||
*out_data_len = n_new * sizeof(float);
|
||||
pcm_sent = audio_pcm.size();
|
||||
return 0;
|
||||
}
|
||||
out_buf.clear();
|
||||
if (!wav_header_sent) {
|
||||
write_wav16_header(out_buf, UINT32_MAX, info.sample_rate);
|
||||
wav_header_sent = true;
|
||||
}
|
||||
append_wav16_pcm(out_buf, audio_pcm.data() + pcm_sent, n_new);
|
||||
pcm_sent = audio_pcm.size();
|
||||
*out_data = out_buf.data();
|
||||
*out_data_len = out_buf.size();
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t flush() override {
|
||||
return decode_wav() ? 0 : 1;
|
||||
}
|
||||
|
||||
private:
|
||||
// repetition-aware sampling: if a semantic draw repeats one of the last ras_window main tokens, draw again with fixed settings
|
||||
static constexpr size_t ras_window = 16;
|
||||
static constexpr size_t n_cb = 10;
|
||||
// codebook tower KV rows (clip_graph_drama3_gen::n_pos_max): one per prompt row but the last, one per frame
|
||||
static constexpr size_t n_tower_rows = 4096;
|
||||
// frames per codec call when streaming (0.33 s)
|
||||
static constexpr size_t stream_frames = 7;
|
||||
|
||||
// top-p keeps the token that crosses the threshold, as the reference does
|
||||
llama_token draw(const std::vector<float> & cand_logits, float temp, int32_t top_k, float top_p) {
|
||||
std::vector<std::pair<float, int>> cand(cand_logits.size());
|
||||
for (size_t i = 0; i < cand.size(); i++) {
|
||||
cand[i] = { cand_logits[i] / temp, (int) i };
|
||||
}
|
||||
std::sort(cand.begin(), cand.end(), [](const auto & a, const auto & b) { return a.first > b.first; });
|
||||
const float lmax = cand[0].first;
|
||||
double sum = 0.0;
|
||||
std::vector<double> prob(cand.size());
|
||||
for (size_t i = 0; i < cand.size(); i++) {
|
||||
prob[i] = std::exp((double) (cand[i].first - lmax));
|
||||
sum += prob[i];
|
||||
}
|
||||
size_t n_keep = cand.size();
|
||||
if (top_k > 0) {
|
||||
n_keep = std::min(n_keep, (size_t) top_k);
|
||||
}
|
||||
double cum = 0.0;
|
||||
for (size_t i = 0; i < n_keep; i++) {
|
||||
if (cum > top_p * sum) {
|
||||
n_keep = i;
|
||||
break;
|
||||
}
|
||||
cum += prob[i];
|
||||
}
|
||||
double kept = 0.0;
|
||||
for (size_t i = 0; i < n_keep; i++) {
|
||||
kept += prob[i];
|
||||
}
|
||||
double r = std::uniform_real_distribution<double>(0.0, 1.0)(rng) * kept;
|
||||
for (size_t i = 0; i < n_keep; i++) {
|
||||
r -= prob[i];
|
||||
if (r <= 0.0) {
|
||||
return cand[i].second;
|
||||
}
|
||||
}
|
||||
return cand[n_keep - 1].second;
|
||||
}
|
||||
|
||||
// candidates: eos, then semantic codes 1.. (the reference never samples code 0)
|
||||
llama_token sample_main(const float * logits) {
|
||||
std::vector<float> cand(n_cb_vocab);
|
||||
cand[0] = logits[eos];
|
||||
for (int i = 1; i < n_cb_vocab; i++) {
|
||||
cand[i] = logits[sem_begin + i];
|
||||
}
|
||||
auto to_token = [&](int i) { return i == 0 ? eos : sem_begin + i; };
|
||||
const llama_token first = to_token(draw(cand, 1.0f, 0, 0.9f));
|
||||
if (first != eos && std::find(hist.begin(), hist.end(), first) != hist.end()) {
|
||||
return to_token(draw(cand, 1.0f, 30, 0.95f));
|
||||
}
|
||||
return first;
|
||||
}
|
||||
|
||||
// code0 < 0 for a prompt row, else the main code of the next frame
|
||||
// returns the next backbone input, the codes of the next frame go to frame
|
||||
const float * run_tower(int32_t code0, const float * h) {
|
||||
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE;
|
||||
inp.code0 = code0;
|
||||
inp.embd = const_cast<float *>(h);
|
||||
inp.codes = frame.empty() ? nullptr : frame.data();
|
||||
inp.n_codes = frame.size();
|
||||
mtmd_gen_out out{};
|
||||
if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0 || (code0 >= 0 && out.n_codes != n_cb)) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: codebook tower failed\n");
|
||||
return nullptr;
|
||||
}
|
||||
if (code0 >= 0) {
|
||||
frame.assign(out.codes, out.codes + out.n_codes);
|
||||
}
|
||||
return out.embd;
|
||||
}
|
||||
|
||||
// the codec is stateful, so frames can be decoded in pieces
|
||||
bool decode_wav() {
|
||||
const size_t n_frames = codes_buf.size() / n_cb;
|
||||
while (n_decoded < n_frames) {
|
||||
const size_t n = std::min(n_frames - n_decoded, (size_t) 64);
|
||||
mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
|
||||
inp.codes = codes_buf.data() + n_decoded * n_cb;
|
||||
inp.n_codes = n * n_cb;
|
||||
mtmd_gen_out out{};
|
||||
if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
|
||||
LOG_ERR("mtmd_helper_gen_audio: codec decode failed\n");
|
||||
return false;
|
||||
}
|
||||
audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples);
|
||||
n_decoded += n;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static constexpr int n_cb_vocab = 4096; // <|semantic:0..4095|>
|
||||
|
||||
llama_token sem_begin = LLAMA_TOKEN_NULL;
|
||||
llama_token eos = LLAMA_TOKEN_NULL;
|
||||
|
||||
llama_seq_id seq_id = 0;
|
||||
llama_pos pos = 0;
|
||||
std::vector<llama_token> prompt_tokens;
|
||||
int prompt_pos = 0;
|
||||
|
||||
std::mt19937 rng;
|
||||
std::vector<int32_t> frame; // codes of the last frame, input of the next tower row
|
||||
std::vector<llama_token> hist; // last main tokens, for repetition-aware sampling
|
||||
std::vector<int32_t> codes_buf; // frame-major
|
||||
size_t n_decoded = 0; // frames already through the codec
|
||||
|
||||
std::vector<float> audio_pcm;
|
||||
std::vector<float> h_state_buf;
|
||||
mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
|
||||
std::vector<char> out_buf;
|
||||
bool stream = false;
|
||||
size_t pcm_sent = 0; // samples already returned by get_output()
|
||||
bool wav_header_sent = false;
|
||||
};
|
||||
|
||||
static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lctx, mtmd_context * mctx) {
|
||||
switch (mtmd_gen_audio_get_info(mctx).type) {
|
||||
case MTMD_GEN_AUDIO_TYPE_QWEN3TTS:
|
||||
@@ -2089,6 +2428,8 @@ static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lc
|
||||
return std::unique_ptr<mtmd_gen_audio_pipeline>(new lfm3a_gen_audio_pipeline(lctx, mctx));
|
||||
case MTMD_GEN_AUDIO_TYPE_CHATTERBOX:
|
||||
return std::unique_ptr<mtmd_gen_audio_pipeline>(new chatterbox_gen_audio_pipeline(lctx, mctx));
|
||||
case MTMD_GEN_AUDIO_TYPE_DRAMA3:
|
||||
return std::unique_ptr<mtmd_gen_audio_pipeline>(new drama3_gen_audio_pipeline(lctx, mctx));
|
||||
default:
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
@@ -2373,6 +2373,10 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_CHATTERBOX;
|
||||
info.sample_rate = 24000;
|
||||
break;
|
||||
case PROJECTOR_TYPE_DRAMA3_GEN:
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_DRAMA3;
|
||||
info.sample_rate = 44100;
|
||||
break;
|
||||
default:
|
||||
info.type = MTMD_GEN_AUDIO_TYPE_NONE;
|
||||
break;
|
||||
@@ -2776,6 +2780,12 @@ mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) {
|
||||
case PROJECTOR_TYPE_SOPRANO:
|
||||
inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
|
||||
break;
|
||||
case PROJECTOR_TYPE_DRAMA3_GEN:
|
||||
// codebook heads, same as the reference (sglang-omni fishaudio_drama3)
|
||||
inp.top_k = 30;
|
||||
inp.top_p = 0.9f;
|
||||
inp.temp = 1.0f;
|
||||
break;
|
||||
case PROJECTOR_TYPE_LFM3A_GEN:
|
||||
{
|
||||
// MaskGIT sampling runs in the helper, with the settings of the checkpoint
|
||||
|
||||
@@ -467,6 +467,7 @@ enum mtmd_gen_audio_type {
|
||||
MTMD_GEN_AUDIO_TYPE_SOPRANO,
|
||||
MTMD_GEN_AUDIO_TYPE_LFM3A,
|
||||
MTMD_GEN_AUDIO_TYPE_CHATTERBOX,
|
||||
MTMD_GEN_AUDIO_TYPE_DRAMA3,
|
||||
};
|
||||
|
||||
struct mtmd_gen_audio_info {
|
||||
|
||||
@@ -128,13 +128,10 @@ static std::string path_to_utf8(const fs::path & p) {
|
||||
static const std::string & home_dir() {
|
||||
static const std::string home = [] {
|
||||
#ifdef _WIN32
|
||||
// the narrow getenv would return the profile path in the active code page
|
||||
const wchar_t * w = _wgetenv(L"HOME");
|
||||
if (w == nullptr) w = _wgetenv(L"USERPROFILE");
|
||||
return w ? path_to_utf8(fs::path(w)) : std::string();
|
||||
const std::string h = common_get_env("HOME");
|
||||
return h.empty() ? common_get_env("USERPROFILE") : h;
|
||||
#else
|
||||
const char * h = getenv("HOME");
|
||||
return h ? std::string(h) : std::string();
|
||||
return common_get_env("HOME");
|
||||
#endif
|
||||
}();
|
||||
return home;
|
||||
|
||||
@@ -154,3 +154,34 @@ The upstream repositories ship raw checkpoints without `config.json`. Download t
|
||||
python convert_hf_to_gguf.py path/to/chatterbox --outfile chatterbox.gguf
|
||||
python convert_hf_to_gguf.py path/to/chatterbox --mmproj --outfile mmproj-chatterbox.gguf
|
||||
```
|
||||
|
||||
## Drama 3
|
||||
|
||||
[Fish Audio Drama 3](https://huggingface.co/fishaudio/drama-3): a Qwen3.5-MoE backbone (35B, 3B active) samples each frame's main token, a codebook tower samples the 9 other codes, and the Fish DAC decodes them to 44.1 kHz audio.
|
||||
|
||||
Available params:
|
||||
- Speaker tags `<|speaker:0|>` to `<|speaker:15|>` in the text give each character its own voice, a text without a tag starts with `<|speaker:0|>`
|
||||
- The language is the one of the text, `--tts-lang` is not used
|
||||
- Note: `--tts-speaker-file` is not supported yet (voice cloning)
|
||||
- Note: `--top-k` and `--top-p` are not used, frames are sampled with the settings of the reference (main token: temperature 1.0, top-p 0.9, repetition-aware redraw over the last 16 frames; other codes: temperature 1.0, top-k 30, top-p 0.9)
|
||||
- Each frame is 2048 samples (about 46 ms), generation stops at the end-of-speech token
|
||||
- The codebook tower keeps the whole sequence in its KV cache, the prompt and the generated frames are limited to 4096 positions
|
||||
|
||||
Example usage:
|
||||
|
||||
```sh
|
||||
llama-tts -m drama3.gguf \
|
||||
-mm mmproj-drama3.gguf \
|
||||
-p "<|speaker:0|> Did you remember to buy the bread? <|speaker:1|> Of course I did, it is on the kitchen table." \
|
||||
-n 1000 \
|
||||
--output out.wav
|
||||
```
|
||||
|
||||
**Note for GGUF conversion:**
|
||||
|
||||
The routed experts of the checkpoint are FP8, they are dequantized at conversion. The mmproj holds the codebook tower and the decoder half of `codec.pth`, it is best kept in BF16:
|
||||
|
||||
```sh
|
||||
python convert_hf_to_gguf.py path/to/drama-3 --outtype q8_0 --outfile drama3.gguf
|
||||
python convert_hf_to_gguf.py path/to/drama-3 --mmproj --outtype bf16 --outfile mmproj-drama3.gguf
|
||||
```
|
||||
|
||||
+11
-7
@@ -106,18 +106,22 @@
|
||||
{readonly}
|
||||
width={imageWidth}
|
||||
/>
|
||||
{:else if isPdfFile(item.attachment, item.uploadedFile)}
|
||||
<ChatAttachmentsListItemThumbnailFile
|
||||
attachment={item.attachment}
|
||||
class="flex-shrink-0 cursor-pointer {className} {scrollClasses}"
|
||||
{:else if isPdfFile(item.attachment, item.uploadedFile) && item.preview}
|
||||
{@const pageCount =
|
||||
item.attachment?.type === AttachmentType.PDF ? (item.attachment.images?.length ?? 0) : 0}
|
||||
|
||||
<ChatAttachmentsListItemThumbnailImage
|
||||
class="{layoutClasses} cursor-pointer {className} {scrollClasses}"
|
||||
height={imageHeight}
|
||||
id={item.id}
|
||||
{imageClass}
|
||||
label={pageCount > 0 ? `PDF, ${pageCount} ${pageCount === 1 ? 'page' : 'pages'}` : 'PDF'}
|
||||
name={item.name}
|
||||
onRemove={onFileRemove}
|
||||
onclick={() => onPreview?.(item)}
|
||||
preview={item.preview}
|
||||
{readonly}
|
||||
size={item.size}
|
||||
textContent={item.textContent}
|
||||
uploadedFile={item.uploadedFile}
|
||||
width={imageWidth}
|
||||
/>
|
||||
{:else}
|
||||
<ChatAttachmentsListItemThumbnailFile
|
||||
|
||||
+10
@@ -7,6 +7,7 @@
|
||||
height?: string;
|
||||
id: string;
|
||||
imageClass?: string;
|
||||
label?: string;
|
||||
onclick?: (event?: MouseEvent) => void;
|
||||
onRemove?: (id: string) => void;
|
||||
name: string;
|
||||
@@ -20,6 +21,7 @@
|
||||
height = 'h-16',
|
||||
id,
|
||||
imageClass = '',
|
||||
label,
|
||||
name,
|
||||
onclick,
|
||||
onRemove,
|
||||
@@ -49,6 +51,14 @@
|
||||
{@render image()}
|
||||
{/if}
|
||||
|
||||
{#if label}
|
||||
<span
|
||||
class="pointer-events-none absolute bottom-1 left-1 rounded bg-black/60 px-1.5 py-0.5 text-[10px] font-medium text-white"
|
||||
>
|
||||
{label}
|
||||
</span>
|
||||
{/if}
|
||||
|
||||
{#if !readonly}
|
||||
<div
|
||||
class="absolute top-1 right-1 flex items-center justify-center opacity-0 transition-opacity group-focus-within:opacity-100 group-hover:opacity-100"
|
||||
|
||||
+2
@@ -32,6 +32,7 @@
|
||||
isImage: boolean;
|
||||
isAudio: boolean;
|
||||
isVideo: boolean;
|
||||
isPdf: boolean;
|
||||
}
|
||||
|
||||
interface Props {
|
||||
@@ -58,6 +59,7 @@
|
||||
...item,
|
||||
isAudio: isAudioFile(item.attachment, item.uploadedFile),
|
||||
isImage: isImageFile(item.attachment, item.uploadedFile),
|
||||
isPdf: isPdfFile(item.attachment, item.uploadedFile),
|
||||
isVideo: isVideoFile(item.attachment, item.uploadedFile)
|
||||
})
|
||||
)
|
||||
|
||||
+130
-73
@@ -3,11 +3,11 @@
|
||||
import { SyntaxHighlightedCode } from '$lib/components/app';
|
||||
import * as Alert from '$lib/components/ui/alert';
|
||||
import { Button } from '$lib/components/ui/button';
|
||||
import { ICON_CLASS_DEFAULT } from '$lib/constants';
|
||||
import { PdfViewMode } from '$lib/enums';
|
||||
import { ICON_CLASS_DEFAULT, PDF_AS_IMAGES_LEGACY_CONTENT_REGEX } from '$lib/constants';
|
||||
import { AttachmentType, MimeTypeApplication, PdfViewMode } from '$lib/enums';
|
||||
import type { ChatAttachmentDisplayItem } from '$lib/types';
|
||||
import { getLanguageFromFilename } from '$lib/utils';
|
||||
import { convertPDFToImage } from '$lib/utils/browser-only';
|
||||
import { convertPDFToImage, convertPDFToText } from '$lib/utils/browser-only';
|
||||
|
||||
interface Props {
|
||||
currentItem: ChatAttachmentDisplayItem | null;
|
||||
@@ -23,54 +23,70 @@
|
||||
let pdfViewMode = $state<PdfViewMode>(PdfViewMode.PAGES);
|
||||
let pdfImages = $state<string[]>([]);
|
||||
let pdfImagesLoading = $state(false);
|
||||
// a failure stays on screen: the effect does not retry while an error is set
|
||||
let pdfImagesError = $state<string | null>(null);
|
||||
let pdfText = $state<string | null>(null);
|
||||
let pdfTextLoading = $state(false);
|
||||
let pdfTextError = $state<string | null>(null);
|
||||
|
||||
let language = $derived(getLanguageFromFilename(displayName));
|
||||
|
||||
// older PDFs sent as images stored a placeholder in `content`, so their text is extracted again
|
||||
let hasStoredText = $derived(
|
||||
currentItem?.attachment?.type === AttachmentType.PDF
|
||||
? !(
|
||||
currentItem.attachment.processedAsImages &&
|
||||
PDF_AS_IMAGES_LEGACY_CONTENT_REGEX.test(currentItem.attachment.content)
|
||||
)
|
||||
: !!displayTextContent
|
||||
);
|
||||
|
||||
let isLoading = $derived(pdfImagesLoading || pdfTextLoading);
|
||||
|
||||
function getPdfFile(): File | null {
|
||||
if (currentItem?.uploadedFile?.file) {
|
||||
return currentItem.uploadedFile.file;
|
||||
}
|
||||
|
||||
if (
|
||||
currentItem?.attachment &&
|
||||
'base64Data' in currentItem.attachment &&
|
||||
currentItem.attachment.base64Data
|
||||
) {
|
||||
const byteArray = Uint8Array.from(atob(currentItem.attachment.base64Data), (char) =>
|
||||
char.charCodeAt(0)
|
||||
);
|
||||
|
||||
return new File([byteArray], displayName, { type: MimeTypeApplication.PDF });
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
async function loadPdfImages() {
|
||||
if (pdfImages.length > 0 || pdfImagesLoading || !currentItem) return;
|
||||
|
||||
if (
|
||||
currentItem.attachment?.type === AttachmentType.PDF &&
|
||||
currentItem.attachment.images?.length
|
||||
) {
|
||||
pdfImages = currentItem.attachment.images;
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (pdfImagesError) return;
|
||||
|
||||
pdfImagesLoading = true;
|
||||
pdfImagesError = null;
|
||||
|
||||
try {
|
||||
let file: File | null = null;
|
||||
const file = getPdfFile();
|
||||
|
||||
if (currentItem.uploadedFile?.file) {
|
||||
file = currentItem.uploadedFile.file;
|
||||
} else if (currentItem.attachment) {
|
||||
// Check if we have pre-processed images
|
||||
if (
|
||||
'images' in currentItem.attachment &&
|
||||
currentItem.attachment.images &&
|
||||
Array.isArray(currentItem.attachment.images) &&
|
||||
currentItem.attachment.images.length > 0
|
||||
) {
|
||||
pdfImages = currentItem.attachment.images;
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
// Convert base64 back to File for processing
|
||||
if ('base64Data' in currentItem.attachment && currentItem.attachment.base64Data) {
|
||||
const base64Data = currentItem.attachment.base64Data;
|
||||
const byteCharacters = atob(base64Data);
|
||||
const byteNumbers = new Array(byteCharacters.length);
|
||||
|
||||
for (let i = 0; i < byteCharacters.length; i++) {
|
||||
byteNumbers[i] = byteCharacters.charCodeAt(i);
|
||||
}
|
||||
const byteArray = new Uint8Array(byteNumbers);
|
||||
|
||||
file = new File([byteArray], displayName, { type: 'application/pdf' });
|
||||
}
|
||||
}
|
||||
|
||||
if (file) {
|
||||
pdfImages = await convertPDFToImage(file);
|
||||
} else {
|
||||
if (!file) {
|
||||
throw new Error('No PDF file available for conversion');
|
||||
}
|
||||
|
||||
pdfImages = await convertPDFToImage(file);
|
||||
} catch (error) {
|
||||
pdfImagesError = error instanceof Error ? error.message : 'Failed to load PDF images';
|
||||
} finally {
|
||||
@@ -78,26 +94,62 @@
|
||||
}
|
||||
}
|
||||
|
||||
async function loadPdfText() {
|
||||
if (pdfText !== null || pdfTextLoading || !currentItem) return;
|
||||
|
||||
if (hasStoredText) {
|
||||
pdfText = displayTextContent ?? '';
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (pdfTextError) return;
|
||||
|
||||
pdfTextLoading = true;
|
||||
|
||||
try {
|
||||
const file = getPdfFile();
|
||||
|
||||
if (!file) {
|
||||
throw new Error('No PDF file available for text extraction');
|
||||
}
|
||||
|
||||
pdfText = await convertPDFToText(file);
|
||||
} catch (error) {
|
||||
pdfTextError = error instanceof Error ? error.message : 'Failed to extract PDF text';
|
||||
} finally {
|
||||
pdfTextLoading = false;
|
||||
}
|
||||
}
|
||||
|
||||
$effect(() => {
|
||||
if (pdfViewMode === PdfViewMode.PAGES) {
|
||||
loadPdfImages();
|
||||
} else {
|
||||
loadPdfText();
|
||||
}
|
||||
});
|
||||
</script>
|
||||
|
||||
<div class="mb-4 flex items-center justify-end gap-2">
|
||||
<Button
|
||||
disabled={pdfImagesLoading}
|
||||
disabled={isLoading}
|
||||
onclick={() => (pdfViewMode = PdfViewMode.TEXT)}
|
||||
size="sm"
|
||||
variant={pdfViewMode === PdfViewMode.TEXT ? 'default' : 'outline'}
|
||||
>
|
||||
<FileText class="mr-1 {ICON_CLASS_DEFAULT}" />
|
||||
{#if pdfTextLoading}
|
||||
<div
|
||||
class="mr-1 {ICON_CLASS_DEFAULT} animate-spin rounded-full border-2 border-current border-t-transparent"
|
||||
></div>
|
||||
{:else}
|
||||
<FileText class="mr-1 {ICON_CLASS_DEFAULT}" />
|
||||
{/if}
|
||||
Text
|
||||
</Button>
|
||||
|
||||
<Button
|
||||
disabled={pdfImagesLoading}
|
||||
disabled={isLoading}
|
||||
onclick={() => (pdfViewMode = PdfViewMode.PAGES)}
|
||||
size="sm"
|
||||
variant={pdfViewMode === PdfViewMode.PAGES ? 'default' : 'outline'}
|
||||
@@ -136,28 +188,24 @@
|
||||
</Alert.Root>
|
||||
{/if}
|
||||
|
||||
{#if pdfImagesLoading}
|
||||
<div class="flex flex-1 items-center justify-center p-8">
|
||||
<div class="text-center">
|
||||
<div
|
||||
class="mx-auto mb-4 h-8 w-8 animate-spin rounded-full border-4 border-white border-t-transparent"
|
||||
></div>
|
||||
|
||||
<p class="text-white/70">Converting PDF to images...</p>
|
||||
{#if pdfViewMode === PdfViewMode.TEXT}
|
||||
{#if pdfTextLoading}
|
||||
{@render loadingStatus('Extracting PDF text...')}
|
||||
{:else if pdfTextError}
|
||||
{@render status('Failed to extract PDF text', pdfTextError)}
|
||||
{:else if pdfText}
|
||||
<div class="px-4 pb-4">
|
||||
<SyntaxHighlightedCode class="max-w-4xl" code={pdfText} {language} maxHeight="none" />
|
||||
</div>
|
||||
</div>
|
||||
{:else}
|
||||
{@render status('No text found in this PDF')}
|
||||
{/if}
|
||||
{:else if pdfImagesLoading}
|
||||
{@render loadingStatus('Converting PDF to images...')}
|
||||
{:else if pdfImagesError}
|
||||
<div class="flex flex-1 items-center justify-center p-8">
|
||||
<div class="text-center">
|
||||
<FileText class="mx-auto mb-4 h-16 w-16 text-white/50" />
|
||||
|
||||
<p class="mb-4 text-white/70">Failed to load PDF images</p>
|
||||
|
||||
<p class="text-sm text-white/50">{pdfImagesError}</p>
|
||||
</div>
|
||||
</div>
|
||||
{@render status('Failed to load PDF images', pdfImagesError)}
|
||||
{:else if pdfImages.length > 0}
|
||||
{#each pdfImages as image, index (image)}
|
||||
{#each pdfImages as image, index (index)}
|
||||
<p class="mb-2 text-sm text-white/50">Page {index + 1}</p>
|
||||
|
||||
<img alt="PDF Page {index + 1}" class="mx-auto max-w-[85vw] rounded-lg shadow-lg" src={image} />
|
||||
@@ -165,22 +213,31 @@
|
||||
<div class="h-4"></div>
|
||||
{/each}
|
||||
{:else}
|
||||
{@render status('No PDF pages available')}
|
||||
{/if}
|
||||
|
||||
{#snippet loadingStatus(message: string)}
|
||||
<div class="flex flex-1 items-center justify-center p-8">
|
||||
<div class="text-center">
|
||||
<div
|
||||
class="mx-auto mb-4 h-8 w-8 animate-spin rounded-full border-4 border-white border-t-transparent"
|
||||
></div>
|
||||
|
||||
<p class="text-white/70">{message}</p>
|
||||
</div>
|
||||
</div>
|
||||
{/snippet}
|
||||
|
||||
{#snippet status(message: string, detail?: string)}
|
||||
<div class="flex flex-1 items-center justify-center p-8">
|
||||
<div class="text-center">
|
||||
<FileText class="mx-auto mb-4 h-16 w-16 text-white/50" />
|
||||
|
||||
<p class="text-white/70">No PDF pages available</p>
|
||||
<p class={['text-white/70', detail && 'mb-4']}>{message}</p>
|
||||
|
||||
{#if detail}
|
||||
<p class="text-sm text-white/50">{detail}</p>
|
||||
{/if}
|
||||
</div>
|
||||
</div>
|
||||
{/if}
|
||||
|
||||
{#if pdfViewMode === PdfViewMode.TEXT && displayTextContent}
|
||||
<div class="px-4 pb-4">
|
||||
<SyntaxHighlightedCode
|
||||
class="max-w-4xl"
|
||||
code={displayTextContent}
|
||||
{language}
|
||||
maxHeight="none"
|
||||
/>
|
||||
</div>
|
||||
{/if}
|
||||
{/snippet}
|
||||
|
||||
+2
-1
@@ -10,6 +10,7 @@
|
||||
isImage: boolean;
|
||||
isAudio: boolean;
|
||||
isVideo: boolean;
|
||||
isPdf: boolean;
|
||||
preview?: string;
|
||||
}
|
||||
|
||||
@@ -46,7 +47,7 @@
|
||||
]}
|
||||
onclick={() => onNavigate(index)}
|
||||
>
|
||||
{#if item.isImage && item.preview}
|
||||
{#if (item.isImage || item.isPdf) && item.preview}
|
||||
<img alt={item.name} class="h-12 w-12 object-cover" src={item.preview} />
|
||||
{:else}
|
||||
<div
|
||||
|
||||
@@ -74,6 +74,9 @@ export const IMAGE_FILE_TYPES = {
|
||||
}
|
||||
} as const;
|
||||
|
||||
/** `content` that messages stored before the text was kept carry for a PDF sent as images. */
|
||||
export const PDF_AS_IMAGES_LEGACY_CONTENT_REGEX = /^PDF file with \d+ pages$/;
|
||||
|
||||
export const PDF_FILE_TYPES = {
|
||||
[FileTypePdf.PDF]: {
|
||||
extensions: [FileExtensionPdf.PDF],
|
||||
|
||||
@@ -71,13 +71,22 @@ export function getAttachmentDisplayItems(
|
||||
for (const [index, attachment] of attachments.entries()) {
|
||||
const isImage = isImageFile(attachment);
|
||||
|
||||
let preview: string | undefined;
|
||||
|
||||
if (isImage && 'base64Url' in attachment) {
|
||||
preview = attachment.base64Url;
|
||||
} else if (attachment.type === AttachmentType.PDF && attachment.processedAsImages) {
|
||||
// first rendered page is the thumbnail of a PDF sent as images
|
||||
preview = attachment.images?.[0];
|
||||
}
|
||||
|
||||
items.push({
|
||||
attachment,
|
||||
attachmentIndex: index,
|
||||
id: `attachment-${index}`,
|
||||
isImage,
|
||||
name: attachment.name,
|
||||
preview: isImage && 'base64Url' in attachment ? attachment.base64Url : undefined,
|
||||
preview,
|
||||
size: 'size' in attachment ? attachment.size : undefined,
|
||||
textContent: 'content' in attachment ? attachment.content : undefined
|
||||
});
|
||||
|
||||
@@ -140,6 +140,8 @@ export async function parseFilesToMessageExtras(
|
||||
// Process PDF as images (only for vision models)
|
||||
try {
|
||||
const images = await convertPDFToImage(file.file);
|
||||
// kept for the preview's Text view; the model only receives the images
|
||||
const content = file.textContent ?? (await convertPDFToText(file.file).catch(() => ''));
|
||||
|
||||
// Show success toast for PDF image processing
|
||||
toast.success(
|
||||
@@ -151,7 +153,7 @@ export async function parseFilesToMessageExtras(
|
||||
|
||||
extras.push({
|
||||
base64Data: base64Data,
|
||||
content: `PDF file with ${images.length} pages`,
|
||||
content,
|
||||
images: images,
|
||||
name: file.name,
|
||||
processedAsImages: true,
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
<script lang="ts" module>
|
||||
import { defineMeta } from '@storybook/addon-svelte-csf';
|
||||
import ChatAttachmentsPreview from '$lib/components/app/chat/ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreview.svelte';
|
||||
import { AttachmentType } from '$lib/enums';
|
||||
import { expect, waitFor } from 'storybook/test';
|
||||
|
||||
const { Story } = defineMeta({
|
||||
component: ChatAttachmentsPreview,
|
||||
parameters: {
|
||||
layout: 'fullscreen'
|
||||
},
|
||||
title: 'Components/ChatScreen/ChatAttachments/ChatAttachmentsPreview'
|
||||
});
|
||||
|
||||
// same backdrop as DialogChatAttachmentsPreview
|
||||
const frameClass = 'flex h-screen flex-col bg-neutral-950';
|
||||
|
||||
// valid base64 that is not a PDF: pdf.js throws InvalidPDFException
|
||||
const brokenPdfAttachment: DatabaseMessageExtraPdfFile = {
|
||||
base64Data: btoa('this is not a pdf'),
|
||||
content: 'Extracted text of the broken PDF.',
|
||||
name: 'broken.pdf',
|
||||
processedAsImages: false,
|
||||
size: 17,
|
||||
type: AttachmentType.PDF
|
||||
};
|
||||
</script>
|
||||
|
||||
<!-- conversion fails once: shows the error and does not retry -->
|
||||
<Story
|
||||
asChild
|
||||
name="PdfConversionError"
|
||||
play={async ({ canvas }) => {
|
||||
await waitFor(() => expect(canvas.getByText('Failed to load PDF images')).toBeInTheDocument(), {
|
||||
timeout: 5000
|
||||
});
|
||||
|
||||
// a retry loop would flip back to the loading state
|
||||
await new Promise((resolve) => setTimeout(resolve, 500));
|
||||
await expect(canvas.queryByText('Converting PDF to images...')).not.toBeInTheDocument();
|
||||
await expect(canvas.getByText('Failed to load PDF images')).toBeInTheDocument();
|
||||
}}
|
||||
>
|
||||
<div class={frameClass}>
|
||||
<ChatAttachmentsPreview attachments={[brokenPdfAttachment]} class="min-h-0 flex-1" />
|
||||
</div>
|
||||
</Story>
|
||||
Reference in New Issue
Block a user