ImportSafeTensors

Import safetensors binary file

DefinitionResourceCreateNotebook

(* A hack, should not be needed when full support of Real16 is there.
Also there is a problem with scalar tensors that causes working with NumericArray produced by this to crash kernel.
And casting every array type like this would work faster.
Would be nice to have ByteArrayToNumericArray or something like ImportByteArray[ByteArray[...], LibraryDataType[NumericArray, "Real16"]] that doesn't do ByteArrayToString under the hood like now. *)

$Lib := $Lib = (
	Needs["CCompilerDriver`"];
	CCompilerDriver`CreateLibrary["
		#include \"WolframLibrary.h\"
		#include \"WolframCompileLibrary.h\"
		#include \"WolframNumericArrayLibrary.h\"
		#include <string.h>
		#include <stdlib.h>
		DLLEXPORT int toReal16(WolframLibraryData libData, mint Argc, MArgument *Args, MArgument Res) {
		  MNumericArray array = MArgument_getMNumericArray(Args[0]);
		  MTensor dims = MArgument_getMNumericArray(Args[1]);
		  array->data_type = MNumericArray_Type_Real16;
          array->tensor_property_type = 3;
		  if (dims->nelems == 0) {
            array->rank = 1;
            array->dims[0] = 1;
            array->nelems = 1;
          }
          else {
            array->rank = dims->nelems;
            array->dims = malloc(sizeof(mint) * dims->nelems);
            memcpy(array->dims, dims->data, sizeof(mint) * dims->nelems);
            array->nelems = array->nelems / 2;
          }
          
          *((unsigned long long*)&array->prec) = 0xfff0000000000000; // -∞
          MArgument_setMNumericArray(Res, array);
		  return LIBRARY_NO_ERROR;
		}",
		"real16"
	]
)

toReal16 := toReal16 = Enclose @ LibraryFunctionLoad[Confirm @ $Lib, "toReal16", {{LibraryDataType[ByteArray], "Shared"}, {Integer, 1}}, LibraryDataType[NumericArray, "Real16"]]

dtypeToNumericArrayType[dtype_] := Replace[dtype, {
	"F64" -> "Real64", "F32" -> "Real32", "F16" -> "Real16",
	"I64" -> "Integer64", "I32" -> "Integer32", "I16" -> "Integer16", "I8" -> "Integer8",
	"U64" -> "UnsignedInteger64", "U32" -> "UnsignedInteger32", "U16" -> "UnsignedInteger16", "U8" -> "UnsignedInteger8",
	"BOOL" -> "Boolean",
	_ -> $Failed
}]

ImportSafeTensors::missing = "Following tensors are missing: ``";

reportMissing[header_List] := With[{missing = Select[header, MissingQ][[All, 2]]},
	If[ missing =!= {},
		Message[ImportSafeTensors::missing, StringRiffle[missing, ", "]]
	];
	header
]

ImportSafeTensors[file : _String | _File, return : _String | {___String} | All | "Header" : All] := Enclose @ Block[{
	stream, headerSize, header, result, progressBar
},
	WithCleanup[
		stream = Confirm @ OpenRead[file, BinaryFormat -> True],

		headerSize = ConfirmBy[BinaryRead[stream, "Integer64"], IntegerQ];
		header = ConfirmMatch[ImportByteArray[ReadByteArray[stream, headerSize], "RawJSON"], KeyValuePattern[_ -> KeyValuePattern[{"dtype" -> _String, "shape" -> {___Integer}, "data_offsets" -> {_Integer, _Integer}}]]];
		header = KeyMap[Replace[{"dtype" -> "NumericArrayType", "shape" -> "Dimensions", "data_offsets" -> "Range"}]] /@ header;
		header = MapAt[Confirm @* dtypeToNumericArrayType, header, {All, Key["NumericArrayType"]}];
		If[return === "Header", Return[header]];
		header = If[return === All, header, reportMissing @ Lookup[header, Developer`ToList[return]]];
		DynamicModule[{progress = 0},
			
			progressBar = PrintTemporary[ProgressIndicator[Dynamic[progress], {0, Length[header]}]];
				
			result = MapIndexed[(
				progress = #2[[1]];
				If[ MissingQ[#],
					#,
					Enclose @ Block[{
						range = #["Range"],
						type = #["NumericArrayType"],
						dims = #["Dimensions"],
						start, length, bytes, tensor
					},
						start = range[[1]] + 8 + headerSize;
						length = range[[2]] - range[[1]];
						ConfirmQuiet @ SetStreamPosition[stream, start];
						bytes = ConfirmBy[ReadByteArray[stream, length], ByteArrayQ];
						tensor = If[type === "Real16",
							ConfirmAssert[Length[bytes] / 2 === Times @@ dims];
							(* First sometimes crashes for some reason *)
							(*If[dims === {}, First, Identity] @ *)toReal16[bytes, dims],
							ArrayReshape[ImportByteArray[bytes, Replace[type, "Boolean" -> "Integer8"]], dims]
						];
						Which[
							type === "Boolean",
							ResourceFunction["InverseBoole"][tensor],
							dims === {} || type === "Real16",
							tensor,
							True,
							NumericArray[tensor, type]
						]
					]
				]) &,
				header
			];
			If[ MatchQ[result, {_}] && MatchQ[return, _String],
				result = First[result]
			];

		];
		result,
		Close[stream]
	]	
]

Documentation

UsageResourceCreateNotebook

Details & OptionsResourceCreateNotebook

Safetensors is the format by HuggingFace for a simple, safe way to store and distribute tensors.

Safetensors become a very popular format to distribute weigths for Stable Diffusion neural network on https://civitai.com/.

ExamplesResourceCreateNotebook

Basic Examples

Import tensors from the wild:

ImportSafeTensors["https://civitai.com/api/download/models/107366"]//Short

Import a single tensor:

ImportSafeTensors["https://civitai.com/api/download/models/107366","lora_unet_up_blocks_3_attentions_2_transformer_blocks_0_ff_net_2.lora_up.weight"]

Only import a header:

ImportSafeTensors["https://civitai.com/api/download/models/107366","Header"]//Short

Scope

Options

Applications

Properties and Relations

Possible Issues

Neat Examples

Import DreamShaper weights:

tensors=ImportSafeTensors["~/Downloads/dreamshaper_7.safetensors"];

Modify Stable Diffusion V1 parts:

textEncoder=NetModel[{"CLIP Multi-domain Feature Extractor", "InputDomain" -> "Text", "Architecture" -> "ViT-L/14"}];

unet=NetModel["Stable Diffusion V1"];

decoder=NetModel[{"Stable Diffusion V1", "Part" -> "Decoder"}];

modifyNet[net_,mappings_,tensors_]:=NetReplacePart[net,KeyValueMap[#1->ArrayReshape[If[ListQ[#2],Normal,Identity]@Lookup[tensors,#2],Dimensions[NetExtract[net,#1]]]&,DeleteMissing@mappings]]

modifiedTextEncoder=modifyNet[textEncoder,,tensors];
modifiedUNet=modifyNet[unet,,tensors];
modifiedDecoder=modifyNet[decoder,,tensors];

Compare vanilla and modified weights:

ResourceFunction["StableDiffusionSynthesize"]["spaceship",3]

ResourceFunction["StableDiffusionSynthesize"]["spaceship",3,"TextEncoder"->modifiedTextEncoder,"UNet"->modifiedUNet,"Decoder"->modifiedDecoder]

Add more details with LoRA:

toMatrix[tensor_] := If[ArrayDepth[tensor] > 2, ReshapeLayer[Dimensions[tensor][[;; 2]]][tensor], tensor]

LoRA[weight_, {alpha_, down_, up_}] := With[{scale = Dimensions[up][[2]]}, FunctionLayer[#weight + #alpha / scale #up . #down &][<|"weight" -> weight, "alpha" -> alpha, "down" -> toMatrix[down], "up" -> toMatrix[up]|>]]

computeLoRA[net_, name : {"transformer", id_, "self-attention", "input_project", "Net", "Weights"}, tensors_] := With[{array = ReshapeLayer[{3, Automatic, 768}] @ NetExtract[net, name]},
	CatenateLayer[] @ MapIndexed[
		LoRA[PartLayer[#2[[1]]] @ array, #1] &,
		Table[
			tensors[StringTemplate["lora_te_text_model_encoder_layers_``_self_attn_``_proj.``"][id - 1, layer, lora]],
			{layer, {"q", "k", "v"}}, {lora, {"alpha", "lora_down.weight", "lora_up.weight"}}
		]
	]
]

computeLoRA[net_, name : {"transformer", id_, "self-attention", "output_project", "Net", "Weights"}, tensors_] :=
	computeLoRA[net, name, StringTemplate @ StringTemplate["lora_te_text_model_encoder_layers_``_self_attn_out_proj.``"][id - 1, "``"], tensors]

computeLoRA[net_, name : {"transformer", id_, "mlp", linear : "linear1" | "linear2", "Net", "Weights"}, tensors_] :=
	computeLoRA[net, name, StringTemplate @ StringTemplate["lora_te_text_model_encoder_layers_``_mlp_``.``"][id - 1, StringReplace[linear, "linear" -> "fc"], "``"], tensors]

computeLoRA[net_, name : {block : "up" | "down" | "cross_mid", bid_Integer : -1, transformerId_String, rest : PatternSequence[___, "Weights"]}, tensors_] := With[{
	suffix = Replace[{rest}, {
		{"proj_in", "Weights"} -> "proj_in",
		{"proj_out", "Weights"} -> "proj_out",
		{"transformerBlock1", "ff", "Net", linear : "proj" | "linear", "Weights"} :>
			"transformer_blocks_0_ff_" <> Replace[linear, {"proj" -> "net_0_proj", "linear" -> "net_2"}],
		{"transformerBlock1", attn : "self-attention" | "cross-attention", layer : "query" | "value" | "key" | "output", "Net", "Weights"} :>
			StringTemplate["transformer_blocks_0_``_to_``"][
				Replace[attn, {"self-attention" -> "attn1", "cross-attention" -> "attn2"}],
				Replace[layer, {"query" -> "q", "value" -> "v", "key" -> "k", "output" -> "out_0"}]
			],
		_ :> Missing[name]
	}]
},
	If[MissingQ[suffix], Return[suffix]];
	computeLoRA[
		net, name,
		StringTemplate @ StringTemplate["lora_unet_````_attentions_``_" <> suffix <> ".``"][
			Replace[block, {"cross_mid" -> "mid_block", upOrDown_ :> upOrDown <> "_blocks_"}],
			If[bid > 0, bid - 1, ""],
			Interpreter["Integer"][StringTake[transformerId, -1]] - 1,
			"``"
		],
		tensors
	]
]

computeLoRA[net_, name_, template_TemplateObject, tensors_] :=
	With[{array = NetExtract[net, name]}, LoRA[array, tensors[template[#]] & /@ {"alpha", "lora_down.weight", "lora_up.weight"}]]

computeLoRA[_, name_, _] := Missing[name]

LoRATensors=ImportSafeTensors["https://civitai.com/api/download/models/87153"];

textEncoderLoRA=DeleteMissing[AssociationMap[computeLoRA[modifiedTextEncoder,#,LoRATensors]&,Keys@Information[modifiedTextEncoder,"Arrays"]]];

unetLoRA=DeleteMissing[AssociationMap[computeLoRA[modifiedUNet,#,LoRATensors]&,Keys@Information[modifiedUNet,"Arrays"]]];

modifiedTextEncoder2=NetReplacePart[modifiedTextEncoder,textEncoderLoRA];
modifiedUNet2=NetReplacePart[modifiedUNet,unetLoRA];

ResourceFunction["StableDiffusionSynthesize"]["spaceship",3,"TextEncoder"->modifiedTextEncoder2,"UNet"->modifiedUNet2,"Decoder"->modifiedDecoder]

Source & Additional Information

Contributed ByResourceCreateNotebook

Nikolay Murzin

KeywordsResourceCreateNotebook

safetensors

safe

tensors

import

binary

Categories

Related SymbolsResourceCreateNotebook

Import

NumericArray

Related Resource ObjectsResourceCreateNotebook

StableDiffusionSynthesize

Source/Reference CitationResourceCreateNotebook

Source, reference or citation information

LinksResourceCreateNotebook

GitHub - huggingface/safetensors: Simple, safe way to store and distribute tensors

Civitai | Stable Diffusion models, embeddings, LoRAs and more

TestsResourceCreateNotebook

ResourceFunction["PythonPackageInstall"]["safetensors"]

import torch
from safetensors.torch import save_file

tensors = {
   "weight1": 2.543 * torch.ones((3,5)),
   "weight2": 1.23 + torch.zeros((4,3), dtype=torch.float16),
   "bools": torch.tensor([[False, False, ], [False, True], [True, False], [True, True]])
}
save_file(tensors, "model.safetensors")

VerificationTest[
Normal/@ImportSafeTensors["model.safetensors"],\[LeftAssociation]"weight1"->{{2.5429999828338623`,2.5429999828338623`,2.5429999828338623`,2.5429999828338623`,2.5429999828338623`},{2.5429999828338623`,2.5429999828338623`,2.5429999828338623`,2.5429999828338623`,2.5429999828338623`},{2.5429999828338623`,2.5429999828338623`,2.5429999828338623`,2.5429999828338623`,2.5429999828338623`}},"weight2"->{{1.23046875`,1.23046875`,1.23046875`},{1.23046875`,1.23046875`,1.23046875`},{1.23046875`,1.23046875`,1.23046875`},{1.23046875`,1.23046875`,1.23046875`}},"bools"->{{False,False},{False,True},{True,False},{True,True}}\[RightAssociation]
]

CompatibilityResourceCreateNotebook

Wolfram Language VersionResourceCreateNotebook

13.0+

Operating SystemResourceCreateNotebook

Required FeaturesResourceCreateNotebook

EnvironmentsResourceCreateNotebook

Cloud SupportResourceCreateNotebook

Author NotesResourceCreateNotebook

Additional information about limitations, issues, etc.

Submission NotesResourceCreateNotebook

Additional information for the reviewer.