diff --git a/examples/webgpu/stable_diffusion/index.html b/examples/webgpu/stable_diffusion/index.html
index e009663f5a..455613b0b1 100644
--- a/examples/webgpu/stable_diffusion/index.html
+++ b/examples/webgpu/stable_diffusion/index.html
@@ -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);
 
diff --git a/extra/export_model.py b/extra/export_model.py
index 51b48c6f28..78d77bcaa0 100644
--- a/extra/export_model.py
+++ b/extra/export_model.py
@@ -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};
 """