Skip to content

Commit 5352d4d

Browse files
committed
Downloader: gguf support
1 parent 4a6c9e5 commit 5352d4d

File tree

2 files changed

+32
-20
lines changed

2 files changed

+32
-20
lines changed

src/WebAPI/ModelsAPI.cs

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -572,6 +572,16 @@ public static async Task<JObject> DoModelDownloadWS(Session session, WebSocket w
572572
await ws.SendJson(new JObject() { ["error"] = "Invalid type." }, API.WebsocketTimeout);
573573
return null;
574574
}
575+
string extension = "safetensors";
576+
string folder = handler.DownloadFolderPath;
577+
if (url.EndsWith(".gguf"))
578+
{
579+
extension = "gguf";
580+
if (type == "Stable-Diffusion")
581+
{
582+
folder += "/../diffusion_models"; // Hacky but oughtta do, gguf in diffusion_models is a silly special case
583+
}
584+
}
575585
string originalUrl = url;
576586
url = url.Before('#');
577587
Dictionary<string, string> headers = [];
@@ -598,18 +608,19 @@ public static async Task<JObject> DoModelDownloadWS(Session session, WebSocket w
598608
}
599609
try
600610
{
601-
string outPath = $"{handler.DownloadFolderPath}/{name}.safetensors";
611+
string outPath = $"{folder}/{name}.{extension}";
602612
if (File.Exists(outPath))
603613
{
604614
await ws.SendJson(new JObject() { ["error"] = "Model at that save path already exists." }, API.WebsocketTimeout);
605615
return null;
606616
}
607-
string tempPath = $"{handler.DownloadFolderPath}/{name}.download.tmp";
617+
string tempPath = $"{folder}/{name}.download.tmp";
608618
if (File.Exists(tempPath))
609619
{
610620
File.Delete(tempPath);
611621
}
612622
Directory.CreateDirectory(Path.GetDirectoryName(outPath));
623+
Logs.Debug($"Will download model from '{url}' to '{Path.GetFullPath(outPath)}'");
613624
using CancellationTokenSource canceller = new();
614625
Task downloading = Utilities.DownloadFile(url, tempPath, (progress, total, perSec) =>
615626
{
@@ -652,13 +663,13 @@ public static async Task<JObject> DoModelDownloadWS(Session session, WebSocket w
652663
File.Move(tempPath, outPath);
653664
if (!string.IsNullOrWhiteSpace(metadata))
654665
{
655-
File.WriteAllText($"{handler.DownloadFolderPath}/{name}.swarm.json", metadata);
666+
File.WriteAllText($"{folder}/{name}.swarm.json", metadata);
656667
}
657668
using (ManyReadOneWriteLock.WriteClaim claim = Program.RefreshLock.LockWrite())
658669
{
659670
handler.Refresh();
660671
}
661-
if (Program.ServerSettings.Paths.DownloaderAlwaysResave)
672+
if (Program.ServerSettings.Paths.DownloaderAlwaysResave && extension == "safetensors")
662673
{
663674
if (handler.Models.TryGetValue($"{name}.safetensors", out T2IModel model))
664675
{

src/wwwroot/js/genpage/utiltab.js

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -132,14 +132,11 @@ class LoraExtractorUtil {
132132
if (outName.startsWith('/')) {
133133
outName = outName.substring(1);
134134
}
135-
if (outName.endsWith('.safetensors')) {
136-
outName = outName.substring(0, outName.length - '.safetensors'.length);
137-
}
138-
if (outName.endsWith('.sft')) {
139-
outName = outName.substring(0, outName.length - '.sft'.length);
140-
}
141-
if (outName.endsWith('.ckpt')) {
142-
outName = outName.substring(0, outName.length - '.ckpt'.length);
135+
for (let extension of ['.safetensors', '.sft', '.gguf', '.ckpt']) {
136+
if (outName.endsWith(extension)) {
137+
outName = outName.substring(0, outName.length - extension.length);
138+
break;
139+
}
143140
}
144141
if (!baseModel || !otherModel || !outName) {
145142
this.textArea.innerText = "Missing required values, cannot extract.";
@@ -276,17 +273,17 @@ class ModelDownloaderUtil {
276273
baseLoop:
277274
for (let vers of rawData.modelVersions) {
278275
for (let vFile of vers.files) {
279-
if (vFile.name.endsWith(`.safetensors`) || vFile.name.endsWith(`.sft`)) {
276+
if (vFile.name.endsWith(`.safetensors`) || vFile.name.endsWith(`.sft`) || vFile.name.endsWith(`.gguf`)) {
280277
rawVersion = vers;
281278
file = vFile;
282279
break baseLoop;
283280
}
284281
}
285282
}
286283
}
287-
if (validateSafe && !file.name.endsWith('.safetensors') && !file.name.endsWith('.sft')) {
284+
if (validateSafe && !file.name.endsWith('.safetensors') && !file.name.endsWith('.sft') && !file.name.endsWith('.gguf')) {
288285
console.log(`refuse civitai url because download url is ${file.downloadUrl} / ${file.name} / ${identifier}`);
289-
doError(`Cannot download model from that URL because it is not a safetensors file. Filename is '${file.name}'`);
286+
doError(`Cannot download model from that URL because it is not a safetensors or GGUF file. Filename is '${file.name}'`);
290287
return;
291288
}
292289
if (rawData.type == 'Checkpoint') { modelType = 'Stable-Diffusion'; }
@@ -295,6 +292,10 @@ class ModelDownloaderUtil {
295292
if (rawData.type == 'ControlNet') { modelType = 'ControlNet'; }
296293
if (rawData.type == 'VAE') { modelType = 'VAE'; }
297294
let imgs = rawVersion.images ? rawVersion.images.filter(img => img.type == 'image') : [];
295+
let downloadUrl = file.downloadUrl;
296+
if (file.name.endsWith('.gguf')) {
297+
downloadUrl += `#.gguf`;
298+
}
298299
let applyMetadata = (img) => {
299300
let url = versId ? `${this.civitPrefix}models/${id}?modelVersionId=${versId}` : `${this.civitPrefix}models/${id}`;
300301
metadata = {
@@ -317,7 +318,7 @@ class ModelDownloaderUtil {
317318
if (['Illustrious', 'Pony'].includes(rawVersion.baseModel)) {
318319
metadata['modelspec.usage_hint'] = rawVersion.baseModel;
319320
}
320-
callback(rawData, rawVersion, metadata, modelType, file.downloadUrl, img, imgs.map(x => x.url), null);
321+
callback(rawData, rawVersion, metadata, modelType, downloadUrl, img, imgs.map(x => x.url), null);
321322
}
322323
if (imgs.length > 0) {
323324
imageToData(imgs[0].url, img => applyMetadata(img), true);
@@ -398,8 +399,8 @@ class ModelDownloaderUtil {
398399
parts[4] = parts[4].substring(0, parts[4].length - '?download=true'.length);
399400
this.url.value = `${this.hfPrefix}${parts.join('/')}`;
400401
}
401-
if (!parts[4].endsWith('.safetensors') && !parts[4].endsWith('.sft')) {
402-
this.urlStatusArea.innerText = "URL appears to be a huggingface link, but not a safetensors file. Only safetensors can be auto-downloaded.";
402+
if (!parts[4].endsWith('.safetensors') && !parts[4].endsWith('.sft') && !parts[4].endsWith('.gguf')) {
403+
this.urlStatusArea.innerText = "URL appears to be a huggingface link, but not a safetensors file. Only safetensors and GGUF can be auto-downloaded.";
403404
this.button.disabled = true;
404405
return;
405406
}
@@ -408,14 +409,14 @@ class ModelDownloaderUtil {
408409
this.url.value = `${this.hfPrefix}${parts.join('/')}`;
409410
this.urlStatusArea.innerText = "URL appears to be a huggingface link, and has been autocorrected to a download link.";
410411
this.button.disabled = false;
411-
this.name.value = parts.slice(4).join('/').replaceAll('.safetensors', '').replaceAll('.sft', '');
412+
this.name.value = parts.slice(4).join('/').replaceAll('.safetensors', '').replaceAll('.sft', '').replaceAll('.gguf', '');
412413
this.nameInput();
413414
return;
414415
}
415416
if (parts[2] == 'resolve') {
416417
this.urlStatusArea.innerText = "URL appears to be a valid HuggingFace download link.";
417418
this.button.disabled = false;
418-
this.name.value = parts.slice(4).join('/').replaceAll('.safetensors', '').replaceAll('.sft', '');
419+
this.name.value = parts.slice(4).join('/').replaceAll('.safetensors', '').replaceAll('.sft', '').replaceAll('.gguf', '');
419420
this.nameInput();
420421
return;
421422
}

0 commit comments

Comments
 (0)