├── .gitignore
├── LICENSE.md
├── README.md
├── clip_tokenizer.js
└── index.html


/.gitignore:
--------------------------------------------------------------------------------
1 | .DS_Store


--------------------------------------------------------------------------------
/LICENSE.md:
--------------------------------------------------------------------------------
1 | MIT License
2 | 
3 | Copyright (c) 2023 Softwired Technologies Ltd.
4 | 
5 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
6 | 
7 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
8 | 
9 | THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Stable Diffusion on [tinygrad](https://github.com/tinygrad/tinygrad) WebGPU
 2 | 
 3 | NOTE: This is an outdated model export, please try these instead:
 4 | - new f32 version (~10% faster): https://github.com/wpmed92/stable-diffusion-wgpu-tinygrad
 5 | - f16 version (~40% faster): https://github.com/wpmed92/stable-diffusion-tinygrad-f16
 6 | 
 7 | This is a WebGPU port of Stable Diffusion in [tinygrad](https://github.com/tinygrad/tinygrad).<br>
 8 | Try it out [here](https://softwiredtech.github.io/stable-diffusion-webgpu/)!<br>
 9 | The python code I wrote to compile and export the model can be found [here](https://github.com/tinygrad/tinygrad/blob/master/examples/webgpu/stable_diffusion/compile.py)
10 | 
11 | ## How it works
12 | 
13 | The Stable Diffusion model is exported in three parts:
14 | - textModel
15 | - diffusor
16 | - decoder
17 | 
18 | If you open [net.js](./net.js) you can see all the WebGPU kernels that are involved in the inference.<br>
19 | When you open the page for the first time, the model will be downloaded from [huggingface](https://huggingface.co/wpmed/tinygrad-sd-f16/resolve/main) in tinygrad's safetensor format. The model is in f16 to optimize download speed.<br>
20 | Since the computation is done in `f32` in the model, and since `shader-f16` is not yet supported in production Chrome, the model is decompressed to `f32` using [f16-to-f32-gpu](https://github.com/wpmed92/f16-to-f32-gpu) WebGPU-based `f16` decompression library.<br>
21 | When you open the site a second time, the model will be loaded from the `IndexedDB` cache into which it was saved on first visit. If you have a full cache hit, the model will be decompressed, compiled and ready to use. If you have a cache miss (usually due to `QuotaExceededError`), the model will be redownloaded.
22 | 
23 | ## License
24 | 
25 | MIT
26 | 


--------------------------------------------------------------------------------
/clip_tokenizer.js:
--------------------------------------------------------------------------------
  1 | import { Html5Entities as htmlEntities } from "https://deno.land/x/html_entities@v1.0/mod.js";
  2 | import bpeVocabData from "./bpe_simple_vocab_16e6.mjs";
  3 | // import ftfy from "https://deno.land/x/ftfy_pyodide@v0.1.1/mod.js";
  4 | 
  5 | 
  6 | function ord(c) {
  7 |   return c.charCodeAt(0);
  8 | }
  9 | function range(start, stop, step=1) {
 10 |   if(stop === undefined) {
 11 |     stop = start;
 12 |     start = 0;
 13 |   }
 14 | 
 15 |   if((step > 0 && start >= stop) || (step < 0 && start <= stop)) {
 16 |     return [];
 17 |   }
 18 | 
 19 |   const result = [];
 20 |   for(let i = start; step > 0 ? i < stop : i > stop; i += step) {
 21 |     result.push(i);
 22 |   }
 23 | 
 24 |   return result;
 25 | }
 26 | 
 27 | 
 28 | 
 29 | function bytesToUnicode() {
 30 |   let bs = [
 31 |     ...range(ord("!"), ord("~") + 1),
 32 |     ...range(ord("¡"), ord("¬") + 1),
 33 |     ...range(ord("®"), ord("ÿ") + 1),
 34 |   ];
 35 |   let cs = bs.slice(0);
 36 |   let n = 0;
 37 |   for(let b of range(2**8)) {
 38 |     if(!bs.includes(b)) {
 39 |       bs.push(b);
 40 |       cs.push(2**8 + n);
 41 |       n += 1;
 42 |     }
 43 |   }
 44 |   cs = cs.map(n => String.fromCharCode(n));
 45 |   return Object.fromEntries(bs.map((v, i) => [v, cs[i]]));
 46 | }
 47 | 
 48 | function getPairs(word) {
 49 |   let pairs = [];
 50 |   let prevChar = word[0];
 51 |   for(let char of word.slice(1)) {
 52 |     pairs.push([prevChar, char]);
 53 |     prevChar = char;
 54 |   }
 55 |   return pairs;
 56 | }
 57 | 
 58 | function basicClean(text) {
 59 |   // text = ftfy.fix_text(text);
 60 |   text = htmlEntities.decode(htmlEntities.decode(text));
 61 |   return text.trim();
 62 | }
 63 | 
 64 | function whitespaceClean(text) {
 65 |   return text.replace(/\s+/g, " ").trim();
 66 | }
 67 | 
 68 | 
 69 | export default class ClipTokenizer {
 70 |   constructor() {
 71 |     this.byteEncoder = bytesToUnicode();
 72 |     this.byteDecoder = Object.fromEntries(Object.entries(this.byteEncoder).map(([k,v]) => [v,k]));
 73 |     let merges = bpeVocabData.text.split("\n");
 74 |     merges = merges.slice(1, 49152-256-2+1);
 75 |     merges = merges.map(merge => merge.split(" "));
 76 |     // There was a bug related to the ordering of Python's .values() output. I'm lazy do I've just copy-pasted the Python output:
 77 |     let vocab = ['!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', ':', ';', '<', '=', '>', '?', '@', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '[', '\\', ']', '^', '_', '`', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', '{', '|', '}', '~', '¡', '¢', '£', '¤', '¥', '¦', '§', '¨', '©', 'ª', '«', '¬', '®', '¯', '°', '±', '²', '³', '´', 'µ', '¶', '·', '¸', '¹', 'º', '»', '¼', '½', '¾', '¿', 'À', 'Á', 'Â', 'Ã', 'Ä', 'Å', 'Æ', 'Ç', 'È', 'É', 'Ê', 'Ë', 'Ì', 'Í', 'Î', 'Ï', 'Ð', 'Ñ', 'Ò', 'Ó', 'Ô', 'Õ', 'Ö', '×', 'Ø', 'Ù', 'Ú', 'Û', 'Ü', 'Ý', 'Þ', 'ß', 'à', 'á', 'â', 'ã', 'ä', 'å', 'æ', 'ç', 'è', 'é', 'ê', 'ë', 'ì', 'í', 'î', 'ï', 'ð', 'ñ', 'ò', 'ó', 'ô', 'õ', 'ö', '÷', 'ø', 'ù', 'ú', 'û', 'ü', 'ý', 'þ', 'ÿ', 'Ā', 'ā', 'Ă', 'ă', 'Ą', 'ą', 'Ć', 'ć', 'Ĉ', 'ĉ', 'Ċ', 'ċ', 'Č', 'č', 'Ď', 'ď', 'Đ', 'đ', 'Ē', 'ē', 'Ĕ', 'ĕ', 'Ė', 'ė', 'Ę', 'ę', 'Ě', 'ě', 'Ĝ', 'ĝ', 'Ğ', 'ğ', 'Ġ', 'ġ', 'Ģ', 'ģ', 'Ĥ', 'ĥ', 'Ħ', 'ħ', 'Ĩ', 'ĩ', 'Ī', 'ī', 'Ĭ', 'ĭ', 'Į', 'į', 'İ', 'ı', 'IJ', 'ij', 'Ĵ', 'ĵ', 'Ķ', 'ķ', 'ĸ', 'Ĺ', 'ĺ', 'Ļ', 'ļ', 'Ľ', 'ľ', 'Ŀ', 'ŀ', 'Ł', 'ł', 'Ń'];
 78 |     vocab = [...vocab, ...vocab.map(v => v+'</w>')];
 79 |     for(let merge of merges) {
 80 |       vocab.push(merge.join(""));
 81 |     }
 82 |     vocab.push('<|startoftext|>', '<|endoftext|>');
 83 |     this.encoder = Object.fromEntries(vocab.map((v,i) => [v,i]));
 84 |     this.decoder = Object.fromEntries(Object.entries(this.encoder).map(([k,v]) => [v,k]));
 85 |     this.bpeRanks = Object.fromEntries(merges.map((v,i) => [v.join("·😎·"),i])); // ·😎· because js doesn't yet have tuples
 86 |     this.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'};
 87 |     this.pat = /<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+/gui;
 88 |   }
 89 | 
 90 |   bpe(token) {
 91 |     if(this.cache[token] !== undefined) {
 92 |       return this.cache[token];
 93 |     }
 94 | 
 95 |     let word = [...token.slice(0, -1), token.slice(-1)+'</w>'];
 96 |     let pairs = getPairs(word);
 97 | 
 98 |     if(pairs.length === 0) {
 99 |       return token+'</w>';
100 |     }
101 | 
102 |     while(1) {
103 | 
104 |       let bigram = null;
105 |       let minRank = Infinity;
106 |       for(let p of pairs) {
107 |         let r = this.bpeRanks[p.join("·😎·")];
108 |         if(r === undefined) continue;
109 |         if(r < minRank) {
110 |           minRank = r;
111 |           bigram = p;
112 |         }
113 |       }
114 | 
115 |       if(bigram === null) {
116 |         break;
117 |       }
118 | 
119 |       let [first, second] = bigram;
120 |       let newWord = [];
121 |       let i = 0;
122 |       while(i < word.length) {
123 | 
124 |         let j = word.indexOf(first, i);
125 | 
126 |         if(j === -1) {
127 |           newWord.push(...word.slice(i));
128 |           break;
129 |         }
130 | 
131 |         newWord.push(...word.slice(i, j));
132 |         i = j;
133 | 
134 |         if(word[i] === first && i < word.length-1 && word[i+1] === second) {
135 |           newWord.push(first+second);
136 |           i += 2;
137 |         } else {
138 |           newWord.push(word[i]);
139 |           i += 1;
140 |         }
141 |       }
142 |       word = newWord;
143 |       if(word.length === 1) {
144 |         break;
145 |       } else {
146 |         pairs = getPairs(word);
147 |       }
148 |     }
149 |     word = word.join(" ");
150 |     this.cache[token] = word;
151 |     return word;
152 |   }
153 | 
154 |   encode(text) {
155 |     let bpeTokens = []
156 |     text = whitespaceClean(text).toLowerCase();
157 |     for(let token of [...text.matchAll(this.pat)].map(m => m[0])) {
158 |       token = [...token].map(b => this.byteEncoder[b.charCodeAt(0)]).join("");
159 |       bpeTokens.push(...this.bpe(token).split(' ').map(bpe_token => this.encoder[bpe_token]));
160 |     }
161 |     return bpeTokens;
162 |   }
163 |   
164 |   // adds start and end token, and adds padding 0's and ensures it's 77 tokens long
165 |   encodeForCLIP(text) {
166 |     let tokens = this.encode(text);
167 |     tokens.unshift(49406); // start token
168 |     tokens = tokens.slice(0, 76);
169 |     tokens.push(49407); // end token
170 |     while(tokens.length < 77) tokens.push(49407);
171 |     return tokens;
172 |   }
173 | 
174 |   decode(tokens) {
175 |     let text = tokens.map(token => this.decoder[token]).join("");
176 |     text = [...text].map(c => this.byteDecoder[c]).map(v => String.fromCharCode(v)).join("").replaceAll('</w>', ' ');
177 |     return text;
178 |   }
179 | }


--------------------------------------------------------------------------------
/index.html:
--------------------------------------------------------------------------------
  1 | <!DOCTYPE html>
  2 | <html lang="en">
  3 | 
  4 | <head>
  5 |     <meta charset="UTF-8">
  6 |     <meta name="viewport" content="width=device-width, initial-scale=1.0">
  7 |     <title>tinygrad has WebGPU</title>
  8 | 
  9 |     <style>
 10 |         body {
 11 |             font-family: 'Arial', sans-serif;
 12 |             text-align: center;
 13 |             padding: 30px;
 14 |         }
 15 | 
 16 |         a {
 17 |             text-decoration: none;
 18 |             color: #4A90E2;
 19 |         }
 20 | 
 21 |         h1 {
 22 |             font-size: 36px;
 23 |             font-weight: normal;
 24 |             margin-bottom: 20px;
 25 |         }
 26 | 
 27 |         #mybox {
 28 |             display: flex;
 29 |             flex-direction: column;
 30 |             align-items: center;
 31 |             gap: 20px;
 32 |             width: 50%;
 33 |             margin: 0 auto;
 34 |         }
 35 | 
 36 |         #promptText, #stepRange, #btnRunNet, #guidanceRange {
 37 |             font-size: 18px;
 38 |             width: 100%;
 39 |         }
 40 | 
 41 |         #result {
 42 |             font-size: 48px;
 43 |         }
 44 | 
 45 |         #time {
 46 |             font-size: 16px;
 47 |             color: grey;
 48 |         }
 49 | 
 50 |         canvas {
 51 |             margin-top: 20px;
 52 |             border: 1px solid #000; 
 53 |         }
 54 | 
 55 |         label {
 56 |             display: flex;
 57 |             align-items: center;
 58 |             gap: 10px;
 59 |             width: 100%;
 60 |         }
 61 |         
 62 |         #sliderValue {
 63 |             margin-right: 10px;/
 64 |         }
 65 |     </style>
 66 | 
 67 |     <script type="module">
 68 |         import ClipTokenizer from './clip_tokenizer.js';
 69 |         window.clipTokenizer = new ClipTokenizer();
 70 |     </script>
 71 |     <script type="module">
 72 |         import { f16tof32GPU } from 'https://unpkg.com/f16-to-f32-gpu@0.1.0/src/index.js';
 73 |         window.f16tof32GPU = f16tof32GPU;
 74 |     </script>
 75 |     <script src="./net.js"></script>
 76 | </head>
 77 | 
 78 | <body>
 79 | <h1 id="wgpuError" style="display: none; color: red;">WebGPU is not supported in this browser</h1>
 80 | <h1 id="sdTitle">StableDiffusion by <a href="https://github.com/tinygrad/tinygrad" target="_blank">tinygrad</a> WebGPU</h1>
 81 | <div id="mybox">
 82 |     <input id="promptText" type="text" placeholder="Enter your prompt here" value="a horse sized cat eating a bagel">
 83 |     
 84 |     <label>
 85 |         Steps: <span id="stepValue">8</span>
 86 |         <input id="stepRange" type="range" min="5" max="20" value="8" step="1">
 87 |     </label>
 88 | 
 89 |     <label>
 90 |         Guidance: <span id="guidanceValue">7.5</span>
 91 |         <input id="guidanceRange" type="range" min="3" max="15" value="7.5" step="0.1">
 92 |     </label>
 93 | 
 94 |     <label>
 95 |         <input id="showStepCheckBox" type="checkbox">
 96 |         Show image at each step
 97 |     </label>
 98 |     
 99 |     <input id="btnRunNet" type="button" value="Run" disabled>
100 | 
101 |     <div id="divModelDl" style="display: flex; align-items: center; width: 100%; gap: 10px;">
102 |         <span id="modelDlTitle">Downloading model</span>
103 |         <progress id="modelDlProgressBar" value="0" max="100" style="flex-grow: 1;"></progress>
104 |         <span id="modelDlProgressValue"></span>
105 |     </div>
106 | 
107 |     <div id="divStepProgress" style="display: none; align-items: center; width: 100%; gap: 10px;">
108 |         <progress id="progressBar" value="0" max="100" style="flex-grow: 1;"></progress>
109 |         <span id="progressFraction"></span>
110 |     </div>
111 | </div>
112 | <canvas id="canvas" width="512" height="512"></canvas>
113 | 
114 | <script>
115 |     function initDb() {
116 |         return new Promise((resolve, reject) => {
117 |             let db;
118 |             const request = indexedDB.open('tinydb', 1);
119 |             request.onerror = (event) => {
120 |                 console.error('Database error:', event.target.error);
121 |                 resolve(null);
122 |             };
123 | 
124 |             request.onsuccess = (event) => {
125 |                 db = event.target.result;
126 |                 console.log("Db initialized.");
127 |                 resolve(db);
128 |             };
129 | 
130 |             request.onupgradeneeded = (event) => {
131 |                 db = event.target.result;
132 |                 if (!db.objectStoreNames.contains('tensors')) {
133 |                     db.createObjectStore('tensors', { keyPath: 'id' });
134 |                 }
135 |             };
136 |         });
137 |     }
138 | 
139 |     function saveTensorToDb(db, id, tensor) {
140 |         return new Promise((resolve, reject) => {
141 |             if (db == null) {
142 |                 resolve(null);
143 |             }
144 | 
145 |             const transaction = db.transaction(['tensors'], 'readwrite');
146 |             const store = transaction.objectStore('tensors');
147 |             const request = store.put({ id: id, content: tensor });
148 | 
149 |             transaction.onabort = (event) => {
150 |                 console.log("Transaction error while saving tensor: " + event.target.error);
151 |                 resolve(null);
152 |             };
153 | 
154 |             request.onsuccess = () => {
155 |                 console.log('Tensor saved successfully.');
156 |                 resolve();
157 |             };
158 | 
159 |             request.onerror = (event) => {
160 |                 console.error('Tensor save failed:', event.target.error);
161 |                 resolve(null);
162 |             };
163 |         });
164 |     }
165 | 
166 |     function readTensorFromDb(db, id) {
167 |         return new Promise((resolve, reject) => {
168 |             if (db == null) {
169 |                 resolve(null);
170 |             }
171 |             
172 |             const transaction = db.transaction(['tensors'], 'readonly');
173 |             const store = transaction.objectStore('tensors');
174 |             const request = store.get(id);
175 | 
176 |             transaction.onabort = (event) => {
177 |                 console.log("Transaction error while reading tensor: " + event.target.error);
178 |                 resolve(null);
179 |             };
180 | 
181 |             request.onsuccess = (event) => {
182 |                 const result = event.target.result;
183 |                 if (result) {
184 |                     console.log("Cache hit: " + id);
185 |                     resolve(result);
186 |                 } else {
187 |                     console.log("Cache miss: " + id);
188 |                     resolve(null);
189 |                 }
190 |             };
191 | 
192 |             request.onerror = (event) => {
193 |                 console.error('Tensor retrieve failed: ', event.target.error);
194 |                 resolve(null);
195 |             };
196 |         });
197 |     }
198 | 
199 |     window.addEventListener('load', async function() {
200 |         if (!navigator.gpu) {
201 |             document.getElementById("wgpuError").style.display = "";
202 |             document.getElementById("sdTitle").style.display = "none";
203 |             return;
204 |         }
205 | 
206 |         let db = await initDb();
207 | 
208 |         const ctx = document.getElementById("canvas").getContext("2d", { willReadFrequently: true });
209 |         let labels, nets, safetensorParts;
210 | 
211 |         const getDevice = async () => {
212 |             const adapter = await navigator.gpu.requestAdapter();
213 |             const requiredLimits = {};
214 |             const maxBufferSizeInSDModel = 1073741824;
215 |             requiredLimits.maxStorageBufferBindingSize = maxBufferSizeInSDModel;
216 |             requiredLimits.maxBufferSize = maxBufferSizeInSDModel;
217 |             
218 |             return await adapter.requestDevice({
219 |                 requiredLimits
220 |             });
221 |         };
222 | 
223 |         const timer = async (func, label = "") => {
224 |             const start = performance.now();
225 |             const out = await func();
226 |             const delta = (performance.now() - start).toFixed(1)
227 |             console.log(`${delta} ms ${label}`);
228 |             return out;
229 |         }
230 | 
231 |         const getProgressDlForPart = async (part, progressCallback) => {
232 |             const response = await fetch(part);
233 |             const contentLength = response.headers.get('content-length');
234 |             const total = parseInt(contentLength, 10);
235 | 
236 |             const res = new Response(new ReadableStream({
237 |                 async start(controller) {
238 |                     const reader = response.body.getReader();
239 |                     for (;;) {
240 |                         const { done, value } = await reader.read();
241 |                         if (done) break;
242 |                         progressCallback(part, value.byteLength, total);
243 |                         controller.enqueue(value);
244 |                     }
245 |                     
246 |                     controller.close();
247 |                 },
248 |             }));
249 |         
250 |             return res.arrayBuffer();
251 |         };
252 |   
253 |         const getAndDecompressF16Safetensors = async (device, progress) => {
254 |             let totalLoaded = 0;
255 |             let totalSize = 0;
256 |             let partSize = {};
257 | 
258 |             const progressCallback = (part, loaded, total) => {
259 |                 totalLoaded += loaded;
260 | 
261 |                 if (!partSize[part]) {
262 |                     totalSize += total;
263 |                     partSize[part] = true;
264 |                 }
265 |                 
266 |                 progress(totalLoaded, totalSize);
267 |             };
268 | 
269 |             let combinedBuffer = await readTensorFromDb(db, "net.f16");
270 |             let textModelU8 = await readTensorFromDb(db, "net.text");
271 |             let textModelFetched = false;
272 | 
273 |             if (combinedBuffer == null) {
274 |                 let dlParts = [
275 |                     getProgressDlForPart(window.MODEL_BASE_URL + '/net_part0.safetensors', progressCallback),
276 |                     getProgressDlForPart(window.MODEL_BASE_URL + '/net_part1.safetensors', progressCallback),
277 |                     getProgressDlForPart(window.MODEL_BASE_URL + '/net_part2.safetensors', progressCallback),
278 |                     getProgressDlForPart(window.MODEL_BASE_URL + '/net_part3.safetensors', progressCallback)
279 |                 ];
280 | 
281 |                 if (textModelU8 == null) {
282 |                     dlParts.push(getProgressDlForPart(window.MODEL_BASE_URL + '/net_textmodel.safetensors', progressCallback));
283 |                 }
284 | 
285 |                 let buffers = await Promise.all(dlParts);
286 | 
287 |                 // Combine everything except for text model, since that's alreafy f32
288 |                 const totalLength = buffers.reduce((acc, buffer, index, array) => {
289 |                     if (index < 4) {
290 |                         return acc + buffer.byteLength;
291 |                     } else {
292 |                         return acc;
293 |                     }
294 |                     }, 0
295 |                 );
296 | 
297 |                 combinedBuffer = new Uint8Array(totalLength);
298 | 
299 |                 let offset = 0;
300 |                 buffers.forEach((buffer, index) => {
301 |                     if (index < 4) {
302 |                         combinedBuffer.set(new Uint8Array(buffer), offset);
303 |                         offset += buffer.byteLength;
304 |                         buffer = null;
305 |                     }
306 |                 });
307 | 
308 |                 // Disabling caches as this cause OOM in Chrome currently
309 |                 //await saveTensorToDb(db, "net.f16", combinedBuffer);
310 | 
311 |                 if (textModelU8 == null) {
312 |                     textModelFetched = true;
313 |                     textModelU8 = new Uint8Array(buffers[4]);
314 |                     //await saveTensorToDb(db, "net.text", textModelU8);
315 |                 }
316 |             } else {
317 |                 combinedBuffer = combinedBuffer.content;
318 |             }
319 | 
320 |             if (textModelU8 == null) {
321 |                 textModelU8 = new Uint8Array(await getProgressDlForPart(window.MODEL_BASE_URL + '/net_textmodel.safetensors', progressCallback));
322 |                 //await saveTensorToDb(db, "net.text", textModelU8);
323 |             } else if (!textModelFetched) {
324 |                 textModelU8 = textModelU8.content;
325 |             }
326 | 
327 |             document.getElementById("modelDlTitle").innerHTML = "Decompressing model";
328 | 
329 |             const textModelOffset = 3772703308;
330 |             const metadataLength = Number(new DataView(combinedBuffer.buffer).getBigUint64(0, true));
331 |             const metadata = JSON.parse(new TextDecoder("utf8").decode(combinedBuffer.subarray(8, 8 + metadataLength)));
332 | 
333 |             const allToDecomp = combinedBuffer.byteLength - (8 + metadataLength);
334 |             const decodeChunkSize = 67107840;
335 |             const numChunks = Math.ceil(allToDecomp/decodeChunkSize);
336 | 
337 |             console.log(allToDecomp + " bytes to decompress");
338 |             console.log("Will be decompressed in " + numChunks+ " chunks");
339 | 
340 |             let partOffsets = [{start: 0, end: 1131408336}, {start: 1131408336, end: 2227518416}, {start: 2227518416, end: 3308987856}, {start: 3308987856, end: 4265298864}];
341 |             let parts = [];
342 | 
343 |             for (let offsets of partOffsets) {
344 |                 parts.push(new Uint8Array(offsets.end-offsets.start));
345 |             }
346 |             parts[0].set(new Uint8Array(new BigUint64Array([BigInt(metadataLength)]).buffer), 0);
347 |             parts[0].set(combinedBuffer.subarray(8, 8 + metadataLength), 8);
348 |             parts[3].set(textModelU8, textModelOffset+8+metadataLength - partOffsets[3].start);
349 | 
350 |             let start = Date.now();
351 |             let cursor = 0;
352 | 
353 |             for (let i = 0; i < numChunks; i++) {
354 |                 progress(i, numChunks);
355 |                 let chunkStartF16 = 8 + metadataLength + (decodeChunkSize * i);
356 |                 let chunkEndF16 = chunkStartF16 + decodeChunkSize;
357 |                 let chunk = combinedBuffer.subarray(chunkStartF16, chunkEndF16);
358 |                 let result = await f16tof32GPU(chunk);
359 |                 let resultUint8 = new Uint8Array(result.buffer);
360 |                 let chunkStartF32 = 8 + metadataLength + (decodeChunkSize * i * 2);
361 |                 let chunkEndF32 = chunkStartF32 + resultUint8.byteLength;
362 |                 let offsetInPart = chunkStartF32 - partOffsets[cursor].start;
363 | 
364 |                 if (chunkEndF32 < partOffsets[cursor].end || cursor === parts.length - 1) {
365 |                     parts[cursor].set(resultUint8, offsetInPart);
366 |                 } else {
367 |                     let spaceLeftInCurrentPart = partOffsets[cursor].end - chunkStartF32;
368 |                     parts[cursor].set(resultUint8.subarray(0, spaceLeftInCurrentPart), offsetInPart);
369 | 
370 |                     cursor++;
371 | 
372 |                     if (cursor < parts.length) {
373 |                         let nextPartOffset = spaceLeftInCurrentPart;
374 |                         let nextPartLength = resultUint8.length - nextPartOffset;
375 |                         parts[cursor].set(resultUint8.subarray(nextPartOffset, nextPartOffset + nextPartLength), 0);
376 |                     }
377 |                 }
378 | 
379 |                 resultUint8 = null;
380 |                 result = null;
381 |             }
382 | 
383 |             combinedBuffer = null;
384 | 
385 |             let end = Date.now();
386 |             console.log("Decoding took: " + ((end - start) / 1000) + " s");
387 |             console.log("Avarage " + ((end - start) / numChunks) + " ms per chunk");
388 | 
389 |             return parts;
390 |         };
391 | 
392 |         const loadNet = async () => {
393 |             const modelDlTitle = document.getElementById("modelDlTitle");
394 | 
395 |             const progress = (loaded, total) => {
396 |                 document.getElementById("modelDlProgressBar").value = (loaded/total) * 100
397 |                 document.getElementById("modelDlProgressValue").innerHTML = Math.trunc((loaded/total) * 100) + "%"
398 |             }
399 | 
400 |             const device = await getDevice();
401 |             safetensorParts = await getAndDecompressF16Safetensors(device, progress);
402 | 
403 |             modelDlTitle.innerHTML = "Compiling model"
404 | 
405 |             let models = ["textModel", "diffusor", "decoder"];
406 | 
407 |             nets = await timer(() => Promise.all([
408 |                 textModel().setup(device, safetensorParts),
409 |                 diffusor().setup(device, safetensorParts),
410 |                 decoder().setup(device, safetensorParts)
411 |             ]).then((loadedModels) => loadedModels.reduce((acc, model, index) => { acc[models[index]] = model; return acc; }, {})), "(compilation)")
412 | 
413 |             progress(1, 1);
414 | 
415 |             modelDlTitle.innerHTML = "Model ready"
416 |             setTimeout(() => {
417 |                 document.getElementById("modelDlProgressBar").style.display = "none";
418 |                 document.getElementById("modelDlProgressValue").style.display = "none";
419 |                 document.getElementById("divStepProgress").style.display = "flex";
420 |             }, 1000);
421 |             document.getElementById("btnRunNet").disabled = false;
422 |         }
423 | 
424 |         function runStableDiffusion(prompt, steps, guidance, showStep) {
425 |             return new Promise(async (resolve, reject) => {
426 |                 let context = await timer(() => nets["textModel"](clipTokenizer.encodeForCLIP(prompt)));
427 |                 let unconditional_context = await timer(() => nets["textModel"](clipTokenizer.encodeForCLIP("")));
428 | 
429 |                 let timesteps = [];
430 | 
431 |                 for (let i = 1; i < 1000; i += (1000/steps)) {
432 |                     timesteps.push(i);
433 |                 }
434 | 
435 |                 console.log("Timesteps: " + timesteps);
436 | 
437 |                 let alphasCumprod = getWeight(safetensorParts,"alphas_cumprod");
438 |                 let alphas = [];
439 | 
440 |                 for (t of timesteps) {
441 |                     alphas.push(alphasCumprod[Math.floor(t)]);
442 |                 }
443 | 
444 |                 alphas_prev = [1.0];
445 | 
446 |                 for (let i = 0; i < alphas.length-1; i++) {
447 |                     alphas_prev.push(alphas[i]);
448 |                 }
449 | 
450 |                 let inpSize = 4*64*64;
451 |                 latent = new Float32Array(inpSize);
452 | 
453 |                 for (let i = 0; i < inpSize; i++) {
454 |                     latent[i] = Math.sqrt(-2.0 * Math.log(Math.random())) * Math.cos(2.0 * Math.PI * Math.random());
455 |                 }
456 |                 
457 |                 for (let i = timesteps.length - 1; i >= 0; i--) {
458 |                     let timestep = new Float32Array([timesteps[i]]);
459 |                     let x_prev = await timer(() => nets["diffusor"](unconditional_context, context, latent, timestep, new Float32Array([alphas[i]]), new Float32Array([alphas_prev[i]]), new Float32Array([guidance])));
460 |                     latent = x_prev;
461 | 
462 |                     if (showStep != null) {
463 |                         showStep(await nets["decoder"](latent));
464 |                     }
465 | 
466 |                     document.getElementById("progressBar").value = ((steps - i) / steps) * 100
467 |                     document.getElementById("progressFraction").innerHTML = (steps - i) + "/" + steps
468 |                 }
469 |                 
470 |                 resolve(await timer(() => nets["decoder"](latent)));
471 |             });
472 |         }
473 | 
474 |         function renderImage(e, image) {
475 |             let pixels = []
476 |             let pixelCounter = 0
477 | 
478 |             for (let j = 0; j < 512; j++) {
479 |                 for (let k = 0; k < 512; k++) {
480 |                     pixels.push(image[pixelCounter])
481 |                     pixels.push(image[pixelCounter+1])
482 |                     pixels.push(image[pixelCounter+2])
483 |                     pixels.push(255)
484 |                     pixelCounter += 3
485 |                 }
486 |             }
487 | 
488 |             ctx.putImageData(new ImageData(new Uint8ClampedArray(pixels), 512, 512), 0, 0);
489 |             console.log(image);
490 |             console.log("Success");
491 |             e.target.disabled = false;
492 |         }
493 | 
494 |         document.getElementById("btnRunNet").addEventListener("click", function(e) {
495 |             e.target.disabled = true;
496 |             const canvas = document.getElementById("canvas");
497 |             ctx.clearRect(0, 0, canvas.width, canvas.height);
498 | 
499 |             runStableDiffusion(
500 |                 document.getElementById("promptText").value,
501 |                 document.getElementById("stepRange").value,
502 |                 document.getElementById("guidanceRange").value,
503 |                 // Decode at each step
504 |                 (document.getElementById("showStepCheckBox").checked) ? (image) => renderImage(e, image) : null
505 |             ).then((image) => {
506 |                 renderImage(e, image);
507 |             });
508 |         }, false);
509 | 
510 |         const stepSlider = document.getElementById('stepRange');
511 |         const stepValue = document.getElementById('stepValue');
512 | 
513 |         stepSlider.addEventListener('input', function() {
514 |             stepValue.textContent = stepSlider.value;
515 |         });
516 | 
517 |         const guidanceSlider = document.getElementById('guidanceRange');
518 |         const guidanceValue = document.getElementById('guidanceValue');
519 | 
520 |         guidanceSlider.addEventListener('input', function() {
521 |             guidanceValue.textContent = guidanceSlider.value;
522 |         });
523 | 
524 |         loadNet();
525 |     });
526 | </script>
527 | </body>
528 | </html>
529 | 


--------------------------------------------------------------------------------