Wolfram Cloud Resource Object

Function Resource

Function Repository Resource:

 ImportSafeTensors

Source Notebook

Import safetensors binary file

Contributed by: Nikolay Murzin

ResourceFunction["ImportSafeTensors"][file]

return association with tensors from a safetensors binary file.

ResourceFunction["ImportSafeTensors"][file,name]

return a specific tensor with a given name.

ResourceFunction["ImportSafeTensors"][file,{name1,name2,…}]

return multiple files.

ResourceFunction["ImportSafeTensors"][file,"Header"]

return a header containing each tensor's type, dimensions and file offset.

Details and Options

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/.

Examples

Basic Examples (3) 

Import tensors from the wild:

In[1]:=
ResourceFunction[ CloudObject[ "https://www.wolframcloud.com/obj/nikm/DeployedResources/Function/ImportSafeTensors"]][ "https://civitai.com/api/download/models/107366"] // Short
Out[1]=

Import a single tensor:

In[2]:=
ResourceFunction[ CloudObject[ "https://www.wolframcloud.com/obj/nikm/DeployedResources/Function/ImportSafeTensors"]]["https://civitai.com/api/download/models/107366", "lora_unet_up_blocks_3_attentions_2_transformer_blocks_0_ff_net_2.lora_up.weight"]
Out[2]=

Only import a header:

In[3]:=
ResourceFunction[ CloudObject[ "https://www.wolframcloud.com/obj/nikm/DeployedResources/Function/ImportSafeTensors"]]["https://civitai.com/api/download/models/107366", "Header"] // Short
Out[3]=

Neat Examples (4) 

Import DreamShaper weights:

In[4]:=
tensors = ResourceFunction[ CloudObject[ "https://www.wolframcloud.com/obj/nikm/DeployedResources/Function/ImportSafeTensors"]]["~/Downloads/dreamshaper_7.safetensors"];

Modify Stable Diffusion V1 parts:

In[5]:=
textEncoder = NetModel[{"CLIP Multi-domain Feature Extractor", "InputDomain" -> "Text", "Architecture" -> "ViT-L/14"}];
In[6]:=
unet = NetModel["Stable Diffusion V1"];
In[7]:=
decoder = NetModel[{"Stable Diffusion V1", "Part" -> "Decoder"}];
In[8]:=
modifyNet[net_, mappings_, tensors_] := NetReplacePart[net, KeyValueMap[#1 -> ArrayReshape[If[ListQ[#2], Normal, Identity]@Lookup[tensors, #2], Dimensions[NetExtract[net, #1]]] &, DeleteMissing@mappings]]
In[9]:=

Compare vanilla and modified weights:

In[10]:=
ResourceFunction["StableDiffusionSynthesize"]["spaceship", 3]
Out[10]=
In[11]:=
ResourceFunction["StableDiffusionSynthesize"]["spaceship", 3, "TextEncoder" -> modifiedTextEncoder, "UNet" -> modifiedUNet, "Decoder" -> modifiedDecoder]
Out[11]=

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]
In[12]:=
LoRATensors = ResourceFunction[ CloudObject[ "https://www.wolframcloud.com/obj/nikm/DeployedResources/Function/ImportSafeTensors"]]["https://civitai.com/api/download/models/87153"];
In[13]:=
textEncoderLoRA = DeleteMissing[ AssociationMap[computeLoRA[modifiedTextEncoder, #, LoRATensors] &, Keys@Information[modifiedTextEncoder, "Arrays"]]];
In[14]:=
unetLoRA = DeleteMissing[ AssociationMap[computeLoRA[modifiedUNet, #, LoRATensors] &, Keys@Information[modifiedUNet, "Arrays"]]];
In[15]:=
modifiedTextEncoder2 = NetReplacePart[modifiedTextEncoder, textEncoderLoRA]; modifiedUNet2 = NetReplacePart[modifiedUNet, unetLoRA];
In[16]:=
ResourceFunction["StableDiffusionSynthesize"]["spaceship", 3, "TextEncoder" -> modifiedTextEncoder2, "UNet" -> modifiedUNet2, "Decoder" -> modifiedDecoder]
Out[16]=