mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

can't allocate this at once, need parts back

Browse source
This commit is contained in:
Ahmed Harmouche 2024-12-13 13:51:18 +01:00
commit d4bce5be8c
2 changed files with 15 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

15
examples/webgpu/stable_diffusion/index.html
View file

@ -363,7 +363,7 @@
return res.arrayBuffer();
};
const getAndDecompressF16Safetensors = async (device, progress, f16safeTensor) => {
const decompressf16Safetensor = async (device, progress, f16safeTensor) => {
let totalLoaded = 0;
let totalSize = 0;
let partSize = {};
@ -419,18 +419,15 @@
}
const device = await getDevice();
f16decomp = await f16tof32().load(device, ''),
safetensorParts = await getAndDecompressF16Safetensors(device, progress);
modelDlTitle.innerHTML = "Compiling model"
let models = ["textModel", "diffusor", "decoder"];
let netText = await textModel.load(device, "./net_textModel.safetensors");
let netDiffusor = await diffusor.load(device, "./net_diffusor_f16.safetensors");
let netDecoder = await decoder.load(device, "./net_decoder.safetensors");
let funcF16Decomp = await f16tof32.load(device);
nets = await timer(() => Promise.all([
textModel.load(device, safetensorParts),
diffusor().setup(device, safetensorParts),
decoder().setup(device, safetensorParts)
]).then((loadedModels) => loadedModels.reduce((acc, model, index) => { acc[models[index]] = model; return acc; }, {})), "(compilation)")
decompressf16Safetensor(device, progress, diffusor.getWeights());
progress(1, 1);

12
extra/export_model.py
View file

@ -94,6 +94,7 @@ def export_model_webgpu(functions, statements, bufs, weight_names, input_names,
output_return = '[{}]'.format(",".join([f'resultBuffer{i}' for i in range(len(output_names))]))
return f"""
const {exported_name} = (() => {{
let weights = null;
const getTensorBuffer = (safetensorBuffer, tensorMetadata) => {{
return safetensorBuffer.subarray(...tensorMetadata.data_offsets);
}};
@ -146,7 +147,8 @@ const addComputePass = (device, commandEncoder, pipeline, layout, infinityUnifor
{kernel_code}
const setupNet = async (device, safetensor) => {{
const metadata = getTensorMetadata(safetensor);
weights = safetensor;
const metadata = safetensor ? getTensorMetadata(safetensor) : null;
const infinityBuf = createInfinityUniformBuf(device);
{layouts}
@ -184,8 +186,12 @@ const setupNet = async (device, safetensor) => {{
return {output_return};
}}
}}
const load = async (device, weight_path) => {{ return await fetch(weight_path).then(x => x.arrayBuffer()).then(x => setupNet(device, new Uint8Array(x))); }}
return {{ load }};
const load = async (device, weight_path) =>
{{
const buffer = weight_path ? await fetch(weight_path).then(x => x.arrayBuffer()) : null;
return setupNet(device, buffer ? new Uint8Array(buffer) : null);
}}
return {{ load, getWeights: () => weights }};
}})();
export default {exported_name};
"""