mirror of
https://github.com/MoreTore/openpilot.git
synced 2026-09-30 19:33:49 +08:00
FrogPilot 0.9.7
This commit is contained in:
James
@@ -1,12 +1,13 @@
|
||||
import os
|
||||
from extra.export_model import compile_net, jit_model
|
||||
from extra.export_model import compile_net, jit_model, dtype_to_js_type
|
||||
from extra.f16_decompress import u32_to_f16
|
||||
from examples.stable_diffusion import StableDiffusion
|
||||
from tinygrad.nn.state import get_state_dict, safe_save, safe_load_metadata, torch_load, load_state_dict
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.ops import Device
|
||||
from extra.utils import download_file
|
||||
from tinygrad import Device, dtypes
|
||||
from tinygrad.helpers import fetch
|
||||
from typing import NamedTuple, Any, List
|
||||
from pathlib import Path
|
||||
import requests
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
@@ -28,10 +29,8 @@ def convert_f32_to_f16(input_file, output_file):
|
||||
front_float16_values.tofile(f)
|
||||
rest_float32_values.tofile(f)
|
||||
|
||||
FILENAME = Path(__file__).parent.parent.parent.parent / "weights/sd-v1-4.ckpt"
|
||||
|
||||
def split_safetensor(fn):
|
||||
_, json_len, metadata = safe_load_metadata(fn)
|
||||
_, data_start, metadata = safe_load_metadata(fn)
|
||||
text_model_offset = 3772703308
|
||||
chunk_size = 536870912
|
||||
|
||||
@@ -40,7 +39,7 @@ def split_safetensor(fn):
|
||||
if (metadata[k]["data_offsets"][0] < text_model_offset):
|
||||
metadata[k]["data_offsets"][0] = int(metadata[k]["data_offsets"][0]/2)
|
||||
metadata[k]["data_offsets"][1] = int(metadata[k]["data_offsets"][1]/2)
|
||||
|
||||
|
||||
last_offset = 0
|
||||
part_end_offsets = []
|
||||
|
||||
@@ -51,38 +50,42 @@ def split_safetensor(fn):
|
||||
break
|
||||
|
||||
part_offset = offset - last_offset
|
||||
|
||||
|
||||
if (part_offset >= chunk_size):
|
||||
part_end_offsets.append(8+json_len+offset)
|
||||
part_end_offsets.append(data_start+offset)
|
||||
last_offset = offset
|
||||
|
||||
text_model_start = int(text_model_offset/2)
|
||||
net_bytes = bytes(open(fn, 'rb').read())
|
||||
part_end_offsets.append(text_model_start+8+json_len)
|
||||
part_end_offsets.append(text_model_start+data_start)
|
||||
cur_pos = 0
|
||||
|
||||
|
||||
for i, end_pos in enumerate(part_end_offsets):
|
||||
with open(f'./net_part{i}.safetensors', "wb+") as f:
|
||||
with open(os.path.join(os.path.dirname(__file__), f'./net_part{i}.safetensors'), "wb+") as f:
|
||||
f.write(net_bytes[cur_pos:end_pos])
|
||||
cur_pos = end_pos
|
||||
|
||||
with open(f'./net_textmodel.safetensors', "wb+") as f:
|
||||
f.write(net_bytes[text_model_start+8+json_len:])
|
||||
|
||||
with open(os.path.join(os.path.dirname(__file__), f'./net_textmodel.safetensors'), "wb+") as f:
|
||||
f.write(net_bytes[text_model_start+data_start:])
|
||||
|
||||
return part_end_offsets
|
||||
|
||||
def fetch_dep(file, url):
|
||||
with open(file, "w", encoding="utf-8") as f:
|
||||
f.write(requests.get(url).text.replace("https://huggingface.co/wpmed/tinygrad-sd-f16/raw/main/bpe_simple_vocab_16e6.mjs", "./bpe_simple_vocab_16e6.mjs"))
|
||||
|
||||
if __name__ == "__main__":
|
||||
fetch_dep(os.path.join(os.path.dirname(__file__), "clip_tokenizer.js"), "https://huggingface.co/wpmed/tinygrad-sd-f16/raw/main/clip_tokenizer.js")
|
||||
fetch_dep(os.path.join(os.path.dirname(__file__), "bpe_simple_vocab_16e6.mjs"), "https://huggingface.co/wpmed/tinygrad-sd-f16/raw/main/bpe_simple_vocab_16e6.mjs")
|
||||
parser = argparse.ArgumentParser(description='Run Stable Diffusion', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--remoteweights', action='store_true', help="Use safetensors from Huggingface, or from local")
|
||||
args = parser.parse_args()
|
||||
Device.DEFAULT = "WEBGPU"
|
||||
|
||||
Tensor.no_grad = True
|
||||
model = StableDiffusion()
|
||||
|
||||
# load in weights
|
||||
download_file('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', FILENAME)
|
||||
load_state_dict(model, torch_load(FILENAME)['state_dict'], strict=False)
|
||||
load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], strict=False)
|
||||
|
||||
class Step(NamedTuple):
|
||||
name: str = ""
|
||||
@@ -90,34 +93,48 @@ if __name__ == "__main__":
|
||||
forward: Any = None
|
||||
|
||||
sub_steps = [
|
||||
Step(name = "textModel", input = [Tensor.randn(1, 77)], forward = model.cond_stage_model.transformer.text_model),
|
||||
Step(name = "textModel", input = [Tensor.randn(1, 77)], forward = model.cond_stage_model.transformer.text_model),
|
||||
Step(name = "diffusor", input = [Tensor.randn(1, 77, 768), Tensor.randn(1, 77, 768), Tensor.randn(1,4,64,64), Tensor.rand(1), Tensor.randn(1), Tensor.randn(1), Tensor.randn(1)], forward = model),
|
||||
Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode)
|
||||
Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode),
|
||||
Step(name = "f16tof32", input = [Tensor.randn(2097120, dtype=dtypes.uint32)], forward = u32_to_f16)
|
||||
]
|
||||
|
||||
|
||||
prg = ""
|
||||
|
||||
def fixup_code(code, key):
|
||||
code = code.replace(key, 'main')\
|
||||
.replace("var<uniform> INFINITY : f32;\n", "fn inf(a: f32) -> f32 { return a/0.0; }\n")\
|
||||
.replace("@group(0) @binding(0)", "")\
|
||||
.replace("INFINITY", "inf(1.0)")
|
||||
|
||||
for i in range(1,9): code = code.replace(f"binding({i})", f"binding({i-1})")
|
||||
return code
|
||||
|
||||
def compile_step(model, step: Step):
|
||||
run, special_names = jit_model(step, *step.input)
|
||||
functions, statements, bufs, _ = compile_net(run, special_names)
|
||||
state = get_state_dict(model)
|
||||
weights = {id(x.lazydata.realized): name for name, x in state.items()}
|
||||
kernel_code = '\n\n'.join([f"const {key} = `{code.replace(key, 'main')}`;" for key, code in functions.items()])
|
||||
weights = {id(x.uop.base.realized): name for name, x in state.items()}
|
||||
kernel_code = '\n\n'.join([f"const {key} = `{fixup_code(code, key)}`;" for key, code in functions.items()])
|
||||
kernel_names = ', '.join([name for (name, _, _, _) in statements])
|
||||
input_names = [name for _,name in special_names.items() if "input" in name]
|
||||
output_names = [name for _,name in special_names.items() if "output" in name]
|
||||
input_buf_types = [dtype_to_js_type(bufs[inp_name][1]) for inp_name in input_names]
|
||||
output_buf_types = [dtype_to_js_type(bufs[out_name][1]) for out_name in output_names]
|
||||
kernel_calls = '\n '.join([f"addComputePass(device, commandEncoder, piplines[{i}], [{', '.join(args)}], {global_size});" for i, (_name, args, global_size, _local_size) in enumerate(statements) ])
|
||||
bufs = '\n '.join([f"const {name} = " + (f"createEmptyBuf(device, {size});" if _key not in weights else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weights[_key]}'], '{weights[_key]}'))") + ";" for name,(size,dtype,_key) in bufs.items()])
|
||||
exported_bufs = '\n '.join([f"const {name} = " + (f"createEmptyBuf(device, {size});" if _key not in weights else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weights[_key]}'], '{weights[_key]}'))") + ";" for name,(size,dtype,_key) in bufs.items()])
|
||||
gpu_write_bufs = '\n '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:input{i}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i,(_,value) in enumerate(special_names.items()) if "output" not in value])
|
||||
input_writer = '\n '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n new Float32Array(gpuWriteBuffer{i}.getMappedRange()).set(" + f'data{i});' + f"\n gpuWriteBuffer{i}.unmap();\ncommandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, input{i}, 0, gpuWriteBuffer{i}.size);" for i,(_,value) in enumerate(special_names.items()) if value != "output0"])
|
||||
input_writer = '\n '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n new {input_buf_types[i]}(gpuWriteBuffer{i}.getMappedRange()).set(" + f'data{i});' + f"\n gpuWriteBuffer{i}.unmap();\ncommandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, input{i}, 0, gpuWriteBuffer{i}.size);" for i,_ in enumerate(input_names)])
|
||||
return f"""\n var {step.name} = function() {{
|
||||
|
||||
|
||||
{kernel_code}
|
||||
|
||||
return {{
|
||||
"setup": async (device, safetensor) => {{
|
||||
const metadata = getTensorMetadata(safetensor[0]);
|
||||
const metadata = safetensor ? getTensorMetadata(safetensor[0]) : null;
|
||||
|
||||
{exported_bufs}
|
||||
|
||||
{bufs}
|
||||
|
||||
{gpu_write_bufs}
|
||||
const gpuReadBuffer = device.createBuffer({{ size: output0.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ }});
|
||||
|
||||
@@ -135,12 +152,12 @@ if __name__ == "__main__":
|
||||
device.queue.submit([gpuCommands]);
|
||||
|
||||
await gpuReadBuffer.mapAsync(GPUMapMode.READ);
|
||||
const resultBuffer = new Float32Array(gpuReadBuffer.size/4);
|
||||
resultBuffer.set(new Float32Array(gpuReadBuffer.getMappedRange()));
|
||||
const resultBuffer = new {output_buf_types[0]}(gpuReadBuffer.size/{bufs[output_names[0]][1].itemsize});
|
||||
resultBuffer.set(new {output_buf_types[0]}(gpuReadBuffer.getMappedRange()));
|
||||
gpuReadBuffer.unmap();
|
||||
return resultBuffer;
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
"""
|
||||
@@ -151,14 +168,14 @@ if __name__ == "__main__":
|
||||
|
||||
if step.name == "diffusor":
|
||||
if args.remoteweights:
|
||||
base_url = "https://huggingface.co/wpmed/tinygrad-sd-f16/resolve/main"
|
||||
base_url = "https://huggingface.co/wpmed/stable-diffusion-f16-new/resolve/main"
|
||||
else:
|
||||
state = get_state_dict(model)
|
||||
safe_save(state, os.path.join(os.path.dirname(__file__), "net.safetensors"))
|
||||
convert_f32_to_f16("./net.safetensors", "./net_conv.safetensors")
|
||||
split_safetensor("./net_conv.safetensors")
|
||||
os.remove("net.safetensors")
|
||||
os.remove("net_conv.safetensors")
|
||||
convert_f32_to_f16(os.path.join(os.path.dirname(__file__), "./net.safetensors"), os.path.join(os.path.dirname(__file__), "./net_conv.safetensors"))
|
||||
split_safetensor(os.path.join(os.path.dirname(__file__), "./net_conv.safetensors"))
|
||||
os.remove(os.path.join(os.path.dirname(__file__), "net.safetensors"))
|
||||
os.remove(os.path.join(os.path.dirname(__file__), "net_conv.safetensors"))
|
||||
base_url = "."
|
||||
|
||||
prekernel = f"""
|
||||
@@ -188,20 +205,6 @@ if __name__ == "__main__":
|
||||
counter++;
|
||||
}}
|
||||
|
||||
let allZero = true;
|
||||
let out = safetensorParts[selectedPart].subarray(...correctedOffsets);
|
||||
|
||||
for (let i = 0; i < out.length; i++) {{
|
||||
if (out[i] !== 0) {{
|
||||
allZero = false;
|
||||
break;
|
||||
}}
|
||||
}}
|
||||
|
||||
if (allZero) {{
|
||||
console.log("Error: weight '" + key + "' is all zero.");
|
||||
}}
|
||||
|
||||
return safetensorParts[selectedPart].subarray(...correctedOffsets);
|
||||
}}
|
||||
|
||||
|
||||
@@ -1,64 +0,0 @@
|
||||
const f16tof32 = `
|
||||
fn u16_to_f16(x: u32) -> f32 {
|
||||
let sign = f32((x >> 15) & 0x1);
|
||||
let exponent = f32((x >> 10) & 0x1F);
|
||||
let fraction = f32(x & 0x3FF);
|
||||
|
||||
let sign_multiplier = select(1.0, -1.0, sign == 1.0);
|
||||
if (exponent == 0.0) {
|
||||
return sign_multiplier * 6.103515625e-5 * (fraction / 1024.0);
|
||||
} else {
|
||||
return sign_multiplier * exp2(exponent - 15.0) * (1.0 + fraction / 1024.0);
|
||||
}
|
||||
}
|
||||
|
||||
@group(0) @binding(0) var<storage,read_write> data0: array<u32>;
|
||||
@group(0) @binding(1) var<storage,read_write> data1: array<f32>;
|
||||
@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||
let gidx = gid.x;
|
||||
let outgidx = gidx*2;
|
||||
|
||||
if (gidx >= arrayLength(&data0)) {
|
||||
return;
|
||||
}
|
||||
|
||||
let oo = data0[gidx];
|
||||
let oo1 = (oo >> 16);
|
||||
let oo2 = oo & 0xFFFFu;
|
||||
|
||||
let f1 = u16_to_f16(oo2);
|
||||
let f2 = u16_to_f16(oo1);
|
||||
|
||||
data1[outgidx] = f1;
|
||||
data1[outgidx + 1] = f2;
|
||||
}`;
|
||||
|
||||
window.f16tof32GPU = async(device, inf16) => {
|
||||
const input = device.createBuffer({size: inf16.length, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });
|
||||
const output = device.createBuffer({size: inf16.length*2, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });
|
||||
|
||||
const gpuWriteBuffer = device.createBuffer({size: input.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE });
|
||||
const gpuReadBuffer = device.createBuffer({ size: output.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });
|
||||
const commandEncoder = device.createCommandEncoder();
|
||||
await gpuWriteBuffer.mapAsync(GPUMapMode.WRITE);
|
||||
|
||||
const alignedUint32View = new Uint32Array(inf16.buffer, inf16.byteOffset, inf16.length / 4);
|
||||
new Uint32Array(gpuWriteBuffer.getMappedRange()).set(alignedUint32View);
|
||||
|
||||
gpuWriteBuffer.unmap();
|
||||
commandEncoder.copyBufferToBuffer(gpuWriteBuffer, 0, input, 0, gpuWriteBuffer.size);
|
||||
const pipeline = await device.createComputePipelineAsync({layout: "auto", compute: { module: device.createShaderModule({ code: f16tof32 }), entryPoint: "main" }});
|
||||
|
||||
addComputePass(device, commandEncoder, pipeline, [input, output], [Math.ceil(inf16.length/(4*256)), 1, 1]);
|
||||
|
||||
commandEncoder.copyBufferToBuffer(output, 0, gpuReadBuffer, 0, output.size);
|
||||
const gpuCommands = commandEncoder.finish();
|
||||
device.queue.submit([gpuCommands]);
|
||||
|
||||
await gpuReadBuffer.mapAsync(GPUMapMode.READ);
|
||||
const resultBuffer = new Float32Array(gpuReadBuffer.size/4);
|
||||
resultBuffer.set(new Float32Array(gpuReadBuffer.getMappedRange()));
|
||||
gpuReadBuffer.unmap();
|
||||
|
||||
return resultBuffer;
|
||||
}
|
||||
@@ -5,105 +5,215 @@
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>tinygrad has WebGPU</title>
|
||||
|
||||
<style>
|
||||
body {
|
||||
font-family: 'Arial', sans-serif;
|
||||
text-align: center;
|
||||
padding: 30px;
|
||||
/* General Reset */
|
||||
* {
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
box-sizing: border-box;
|
||||
}
|
||||
|
||||
a {
|
||||
text-decoration: none;
|
||||
color: #4A90E2;
|
||||
body {
|
||||
font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
|
||||
background: #f4f7fb;
|
||||
color: #333;
|
||||
display: flex;
|
||||
justify-content: center;
|
||||
align-items: center;
|
||||
min-height: 100vh;
|
||||
padding: 20px;
|
||||
flex-direction: column;
|
||||
text-align: center;
|
||||
}
|
||||
|
||||
h1 {
|
||||
font-size: 36px;
|
||||
font-weight: normal;
|
||||
font-size: 2.2rem;
|
||||
margin-bottom: 20px;
|
||||
color: #4A90E2;
|
||||
}
|
||||
|
||||
#wgpuError {
|
||||
color: red;
|
||||
font-size: 1.2rem;
|
||||
margin-top: 20px;
|
||||
display: none;
|
||||
}
|
||||
|
||||
#sdTitle {
|
||||
font-size: 1.5rem;
|
||||
margin-bottom: 30px;
|
||||
}
|
||||
|
||||
#mybox {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
align-items: center;
|
||||
gap: 20px;
|
||||
width: 50%;
|
||||
margin: 0 auto;
|
||||
background: #ffffff;
|
||||
padding: 15px;
|
||||
border-radius: 10px;
|
||||
box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
|
||||
width: 120%;
|
||||
max-width: 550px;
|
||||
margin-bottom: 10px;
|
||||
}
|
||||
|
||||
#promptText, #stepRange, #btnRunNet, #guidanceRange {
|
||||
font-size: 18px;
|
||||
input[type="text"] {
|
||||
width: 100%;
|
||||
padding: 12px;
|
||||
margin-bottom: 20px;
|
||||
font-size: 1rem;
|
||||
border: 1px solid #ccc;
|
||||
border-radius: 8px;
|
||||
outline: none;
|
||||
transition: all 0.3s ease;
|
||||
}
|
||||
|
||||
#result {
|
||||
font-size: 48px;
|
||||
}
|
||||
|
||||
#time {
|
||||
font-size: 16px;
|
||||
color: grey;
|
||||
}
|
||||
|
||||
canvas {
|
||||
margin-top: 20px;
|
||||
border: 1px solid #000;
|
||||
input[type="text"]:focus {
|
||||
border-color: #4A90E2;
|
||||
}
|
||||
|
||||
label {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
font-size: 1rem;
|
||||
margin-bottom: 15px;
|
||||
align-items: center;
|
||||
gap: 10px;
|
||||
width: 100%;
|
||||
}
|
||||
|
||||
#sliderValue {
|
||||
margin-right: 10px;/
|
||||
|
||||
input[type="range"] {
|
||||
width: 100%;
|
||||
margin-left: 10px;
|
||||
-webkit-appearance: none;
|
||||
appearance: none;
|
||||
height: 8px;
|
||||
border-radius: 4px;
|
||||
background: #ddd;
|
||||
outline: none;
|
||||
transition: background 0.3s ease;
|
||||
}
|
||||
|
||||
input[type="range"]:focus {
|
||||
background: #4A90E2;
|
||||
}
|
||||
|
||||
#stepRange,
|
||||
#guidanceRange {
|
||||
width: 80%;
|
||||
}
|
||||
|
||||
span {
|
||||
font-size: 1.1rem;
|
||||
font-weight: 600;
|
||||
color: #333;
|
||||
}
|
||||
|
||||
input[type="button"] {
|
||||
padding: 12px 25px;
|
||||
background-color: #4A90E2;
|
||||
color: #fff;
|
||||
font-size: 1.2rem;
|
||||
border: none;
|
||||
border-radius: 8px;
|
||||
cursor: pointer;
|
||||
transition: background-color 0.3s ease;
|
||||
width: 100%;
|
||||
margin-bottom: 20px;
|
||||
}
|
||||
|
||||
input[type="button"]:disabled {
|
||||
background-color: #ccc;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
|
||||
input[type="button"]:hover {
|
||||
background-color: #357ABD;
|
||||
}
|
||||
|
||||
#divModelDl, #divStepProgress {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
align-items: center;
|
||||
margin-bottom: 20px;
|
||||
}
|
||||
|
||||
#modelDlProgressBar,
|
||||
#progressBar {
|
||||
width: 80%;
|
||||
height: 12px;
|
||||
border-radius: 6px;
|
||||
background-color: #e0e0e0;
|
||||
}
|
||||
|
||||
#modelDlProgressBar::-webkit-progress-bar,
|
||||
#progressBar::-webkit-progress-bar {
|
||||
border-radius: 6px;
|
||||
}
|
||||
|
||||
#modelDlProgressValue, #progressFraction {
|
||||
font-size: 1rem;
|
||||
font-weight: 600;
|
||||
color: #333;
|
||||
}
|
||||
|
||||
canvas {
|
||||
max-width: 100%;
|
||||
max-height: 450px;
|
||||
margin-top: 10px;
|
||||
border-radius: 8px;
|
||||
border: 1px solid #ddd;
|
||||
}
|
||||
</style>
|
||||
|
||||
<script type="module">
|
||||
import ClipTokenizer from 'https://softwired.nyc3.cdn.digitaloceanspaces.com/sd/clip_tokenizer.js';
|
||||
import ClipTokenizer from './clip_tokenizer.js';
|
||||
window.clipTokenizer = new ClipTokenizer();
|
||||
</script>
|
||||
<script src="./f16_to_f32.js"></script>
|
||||
<script src="./net.js"></script>
|
||||
</head>
|
||||
|
||||
<body>
|
||||
<h1 id="wgpuError" style="display: none; color: red;">WebGPU is not supported in this browser</h1>
|
||||
<h1 id="sdTitle">StableDiffusion by <a href="https://github.com/tinygrad/tinygrad" target="_blank">tinygrad</a> WebGPU</h1>
|
||||
<div id="mybox">
|
||||
<input id="promptText" type="text" placeholder="Enter your prompt here" value="a horse sized cat eating a bagel">
|
||||
<h1 id="wgpuError" style="display: none;">WebGPU is not supported in this browser</h1>
|
||||
<h1 id="sdTitle">StableDiffusion powered by <a href="https://github.com/tinygrad/tinygrad" target="_blank" style="color: #4A90E2;">tinygrad</a></h1>
|
||||
<a href="https://github.com/tinygrad/tinygrad" target="_blank" style="position: absolute; top: 20px; right: 20px;">
|
||||
<img src="https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg"
|
||||
alt="GitHub Logo"
|
||||
style="width: 32px; height: 32px;">
|
||||
</a>
|
||||
|
||||
<label>
|
||||
Steps: <span id="stepValue">8</span>
|
||||
<input id="stepRange" type="range" min="5" max="20" value="8" step="1">
|
||||
</label>
|
||||
<div id="mybox">
|
||||
<form id="promptForm">
|
||||
<input id="promptText" type="text" placeholder="Enter your prompt here" value="a human standing on the surface of mars">
|
||||
|
||||
<label>
|
||||
Guidance: <span id="guidanceValue">7.5</span>
|
||||
<input id="guidanceRange" type="range" min="3" max="15" value="7.5" step="0.1">
|
||||
</label>
|
||||
|
||||
<input id="btnRunNet" type="button" value="Run" disabled>
|
||||
<label>
|
||||
Steps: <span id="stepValue">9</span>
|
||||
<input id="stepRange" type="range" min="5" max="20" value="9" step="1">
|
||||
</label>
|
||||
|
||||
<div id="divModelDl" style="display: flex; align-items: center; width: 100%; gap: 10px;">
|
||||
<span id="modelDlTitle">Downloading model</span>
|
||||
<progress id="modelDlProgressBar" value="0" max="100" style="flex-grow: 1;"></progress>
|
||||
<span id="modelDlProgressValue"></span>
|
||||
<label>
|
||||
Guidance: <span id="guidanceValue">8.0</span>
|
||||
<input id="guidanceRange" type="range" min="3" max="15" value="8.0" step="0.1">
|
||||
</label>
|
||||
|
||||
<input id="btnRunNet" type="button" value="Run" disabled>
|
||||
</form>
|
||||
|
||||
<div id="divModelDl" style="display: flex; align-items: center; width: 100%; gap: 10px;">
|
||||
<span id="modelDlTitle">Downloading model</span>
|
||||
<progress id="modelDlProgressBar" value="0" max="100" style="flex-grow: 1;"></progress>
|
||||
<span id="modelDlProgressValue"></span>
|
||||
</div>
|
||||
|
||||
<div id="divStepProgress" style="display: none; align-items: center; width: 100%; gap: 10px;">
|
||||
<progress id="progressBar" value="0" max="100" style="flex-grow: 1;"></progress>
|
||||
<span id="progressFraction"></span>
|
||||
</div>
|
||||
|
||||
<div id="divStepTime" style="display: none; align-items: center; width: 100%; gap: 10px;">
|
||||
<span id="stepTimeValue">0 ms</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div id="divStepProgress" style="display: none; align-items: center; width: 100%; gap: 10px;">
|
||||
<progress id="progressBar" value="0" max="100" style="flex-grow: 1;"></progress>
|
||||
<span id="progressFraction"></span>
|
||||
</div>
|
||||
</div>
|
||||
<canvas id="canvas" width="512" height="512"></canvas>
|
||||
<canvas id="canvas" width="512" height="512"></canvas>
|
||||
|
||||
<script>
|
||||
let f16decomp = null;
|
||||
|
||||
function initDb() {
|
||||
return new Promise((resolve, reject) => {
|
||||
let db;
|
||||
@@ -129,30 +239,36 @@
|
||||
}
|
||||
|
||||
function saveTensorToDb(db, id, tensor) {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (db == null) {
|
||||
resolve(null);
|
||||
return readTensorFromDb(db, id).then((result) => {
|
||||
if (!result) {
|
||||
new Promise((resolve, reject) => {
|
||||
if (db == null) {
|
||||
resolve(null);
|
||||
}
|
||||
|
||||
const transaction = db.transaction(['tensors'], 'readwrite');
|
||||
const store = transaction.objectStore('tensors');
|
||||
const request = store.put({ id: id, content: tensor });
|
||||
|
||||
transaction.onabort = (event) => {
|
||||
console.log("Transaction error while saving tensor: " + event.target.error);
|
||||
resolve(null);
|
||||
};
|
||||
|
||||
request.onsuccess = () => {
|
||||
console.log('Tensor saved successfully.');
|
||||
resolve();
|
||||
};
|
||||
|
||||
request.onerror = (event) => {
|
||||
console.error('Tensor save failed:', event.target.error);
|
||||
resolve(null);
|
||||
};
|
||||
});
|
||||
} else {
|
||||
return null;
|
||||
}
|
||||
|
||||
const transaction = db.transaction(['tensors'], 'readwrite');
|
||||
const store = transaction.objectStore('tensors');
|
||||
const request = store.put({ id: id, content: tensor });
|
||||
|
||||
transaction.onabort = (event) => {
|
||||
console.log("Transaction error while saving tensor: " + event.target.error);
|
||||
resolve(null);
|
||||
};
|
||||
|
||||
request.onsuccess = () => {
|
||||
console.log('Tensor saved successfully.');
|
||||
resolve();
|
||||
};
|
||||
|
||||
request.onerror = (event) => {
|
||||
console.error('Tensor save failed:', event.target.error);
|
||||
resolve(null);
|
||||
};
|
||||
});
|
||||
}).catch(()=> null);
|
||||
}
|
||||
|
||||
function readTensorFromDb(db, id) {
|
||||
@@ -173,10 +289,8 @@
|
||||
request.onsuccess = (event) => {
|
||||
const result = event.target.result;
|
||||
if (result) {
|
||||
console.log("Cache hit: " + id);
|
||||
resolve(result);
|
||||
} else {
|
||||
console.log("Cache miss: " + id);
|
||||
resolve(null);
|
||||
}
|
||||
};
|
||||
@@ -190,7 +304,7 @@
|
||||
|
||||
window.addEventListener('load', async function() {
|
||||
if (!navigator.gpu) {
|
||||
document.getElementById("wgpuError").style.display = "";
|
||||
document.getElementById("wgpuError").style.display = "block";
|
||||
document.getElementById("sdTitle").style.display = "none";
|
||||
return;
|
||||
}
|
||||
@@ -208,7 +322,9 @@
|
||||
requiredLimits.maxBufferSize = maxBufferSizeInSDModel;
|
||||
|
||||
return await adapter.requestDevice({
|
||||
requiredLimits
|
||||
requiredLimits,
|
||||
requiredFeatures: ["shader-f16"],
|
||||
powerPreference: "high-performance"
|
||||
});
|
||||
};
|
||||
|
||||
@@ -247,6 +363,18 @@
|
||||
let totalSize = 0;
|
||||
let partSize = {};
|
||||
|
||||
const getPart = async(key) => {
|
||||
let part = await readTensorFromDb(db, key);
|
||||
|
||||
if (part) {
|
||||
console.log(`Cache hit: ${key}`);
|
||||
return Promise.resolve(part.content);
|
||||
} else {
|
||||
console.log(`Cache miss: ${key}`);
|
||||
return getProgressDlForPart(`${window.MODEL_BASE_URL}/${key}.safetensors`, progressCallback);
|
||||
}
|
||||
}
|
||||
|
||||
const progressCallback = (part, loaded, total) => {
|
||||
totalLoaded += loaded;
|
||||
|
||||
@@ -258,63 +386,31 @@
|
||||
progress(totalLoaded, totalSize);
|
||||
};
|
||||
|
||||
let combinedBuffer = await readTensorFromDb(db, "net.f16");
|
||||
let textModelU8 = await readTensorFromDb(db, "net.text");
|
||||
let textModelFetched = false;
|
||||
let netKeys = ["net_part0", "net_part1", "net_part2", "net_part3", "net_textmodel"];
|
||||
let buffers = await Promise.all(netKeys.map(key => getPart(key)));
|
||||
|
||||
if (combinedBuffer == null) {
|
||||
let dlParts = [
|
||||
getProgressDlForPart(window.MODEL_BASE_URL + '/net_part0.safetensors', progressCallback),
|
||||
getProgressDlForPart(window.MODEL_BASE_URL + '/net_part1.safetensors', progressCallback),
|
||||
getProgressDlForPart(window.MODEL_BASE_URL + '/net_part2.safetensors', progressCallback),
|
||||
getProgressDlForPart(window.MODEL_BASE_URL + '/net_part3.safetensors', progressCallback)
|
||||
];
|
||||
|
||||
if (textModelU8 == null) {
|
||||
dlParts.push(getProgressDlForPart(window.MODEL_BASE_URL + '/net_textmodel.safetensors', progressCallback));
|
||||
// Combine everything except for text model, since that's already f32
|
||||
const totalLength = buffers.reduce((acc, buffer, index, array) => {
|
||||
if (index < 4) {
|
||||
return acc + buffer.byteLength;
|
||||
} else {
|
||||
return acc;
|
||||
}
|
||||
}, 0
|
||||
);
|
||||
|
||||
let buffers = await Promise.all(dlParts);
|
||||
|
||||
// Combine everything except for text model, since that's alreafy f32
|
||||
const totalLength = buffers.reduce((acc, buffer, index, array) => {
|
||||
if (index < 4) {
|
||||
return acc + buffer.byteLength;
|
||||
} else {
|
||||
return acc;
|
||||
}
|
||||
}, 0
|
||||
);
|
||||
|
||||
combinedBuffer = new Uint8Array(totalLength);
|
||||
|
||||
let offset = 0;
|
||||
buffers.forEach((buffer, index) => {
|
||||
if (index < 4) {
|
||||
combinedBuffer.set(new Uint8Array(buffer), offset);
|
||||
offset += buffer.byteLength;
|
||||
buffer = null;
|
||||
}
|
||||
});
|
||||
|
||||
await saveTensorToDb(db, "net.f16", combinedBuffer);
|
||||
|
||||
if (textModelU8 == null) {
|
||||
textModelFetched = true;
|
||||
textModelU8 = new Uint8Array(buffers[4]);
|
||||
await saveTensorToDb(db, "net.text", textModelU8);
|
||||
combinedBuffer = new Uint8Array(totalLength);
|
||||
let offset = 0;
|
||||
buffers.forEach((buffer, index) => {
|
||||
saveTensorToDb(db, netKeys[index], new Uint8Array(buffer));
|
||||
if (index < 4) {
|
||||
combinedBuffer.set(new Uint8Array(buffer), offset);
|
||||
offset += buffer.byteLength;
|
||||
buffer = null;
|
||||
}
|
||||
} else {
|
||||
combinedBuffer = combinedBuffer.content;
|
||||
}
|
||||
|
||||
if (textModelU8 == null) {
|
||||
textModelU8 = new Uint8Array(await getProgressDlForPart(window.MODEL_BASE_URL + '/net_textmodel.safetensors', progressCallback));
|
||||
await saveTensorToDb(db, "net.text", textModelU8);
|
||||
} else if (!textModelFetched) {
|
||||
textModelU8 = textModelU8.content;
|
||||
}
|
||||
});
|
||||
|
||||
let textModelU8 = new Uint8Array(buffers[4]);
|
||||
document.getElementById("modelDlTitle").innerHTML = "Decompressing model";
|
||||
|
||||
const textModelOffset = 3772703308;
|
||||
@@ -322,7 +418,7 @@
|
||||
const metadata = JSON.parse(new TextDecoder("utf8").decode(combinedBuffer.subarray(8, 8 + metadataLength)));
|
||||
|
||||
const allToDecomp = combinedBuffer.byteLength - (8 + metadataLength);
|
||||
const decodeChunkSize = 67107840;
|
||||
const decodeChunkSize = 8388480;
|
||||
const numChunks = Math.ceil(allToDecomp/decodeChunkSize);
|
||||
|
||||
console.log(allToDecomp + " bytes to decompress");
|
||||
@@ -346,16 +442,8 @@
|
||||
let chunkStartF16 = 8 + metadataLength + (decodeChunkSize * i);
|
||||
let chunkEndF16 = chunkStartF16 + decodeChunkSize;
|
||||
let chunk = combinedBuffer.subarray(chunkStartF16, chunkEndF16);
|
||||
|
||||
if (chunk.byteLength %4 != 0) {
|
||||
const paddingBytes = 4 - (chunk.byteLength % 4);
|
||||
const alignedBuffer = new ArrayBuffer(chunk.byteLength + paddingBytes);
|
||||
const alignedView = new Uint8Array(alignedBuffer);
|
||||
alignedView.set(new Uint8Array(chunk));
|
||||
chunk = alignedView;
|
||||
}
|
||||
|
||||
let result = await f16tof32GPU(device, chunk);
|
||||
let uint32Chunk = new Uint32Array(chunk.buffer, chunk.byteOffset, chunk.byteLength / 4);
|
||||
let result = await f16decomp(uint32Chunk);
|
||||
let resultUint8 = new Uint8Array(result.buffer);
|
||||
let chunkStartF32 = 8 + metadataLength + (decodeChunkSize * i * 2);
|
||||
let chunkEndF32 = chunkStartF32 + resultUint8.byteLength;
|
||||
@@ -398,6 +486,7 @@
|
||||
}
|
||||
|
||||
const device = await getDevice();
|
||||
f16decomp = await f16tof32().setup(device, safetensorParts),
|
||||
safetensorParts = await getAndDecompressF16Safetensors(device, progress);
|
||||
|
||||
modelDlTitle.innerHTML = "Compiling model"
|
||||
@@ -421,7 +510,7 @@
|
||||
document.getElementById("btnRunNet").disabled = false;
|
||||
}
|
||||
|
||||
function runStableDiffusion(prompt, steps, guidance) {
|
||||
function runStableDiffusion(prompt, steps, guidance, showStep) {
|
||||
return new Promise(async (resolve, reject) => {
|
||||
let context = await timer(() => nets["textModel"](clipTokenizer.encodeForCLIP(prompt)));
|
||||
let unconditional_context = await timer(() => nets["textModel"](clipTokenizer.encodeForCLIP("")));
|
||||
@@ -456,8 +545,16 @@
|
||||
|
||||
for (let i = timesteps.length - 1; i >= 0; i--) {
|
||||
let timestep = new Float32Array([timesteps[i]]);
|
||||
let x_prev = await timer(() => nets["diffusor"](unconditional_context, context, latent, timestep, new Float32Array([alphas[i]]), new Float32Array([alphas_prev[i]]), new Float32Array([guidance])));
|
||||
let start = performance.now()
|
||||
let x_prev = await nets["diffusor"](unconditional_context, context, latent, timestep, new Float32Array([alphas[i]]), new Float32Array([alphas_prev[i]]), new Float32Array([guidance]));
|
||||
document.getElementById("divStepTime").style.display = "block";
|
||||
document.getElementById("stepTimeValue").innerText = `${(performance.now() - start).toFixed(1)} ms / step`
|
||||
latent = x_prev;
|
||||
|
||||
if (showStep != null) {
|
||||
showStep(await nets["decoder"](latent));
|
||||
}
|
||||
|
||||
document.getElementById("progressBar").value = ((steps - i) / steps) * 100
|
||||
document.getElementById("progressFraction").innerHTML = (steps - i) + "/" + steps
|
||||
}
|
||||
@@ -466,29 +563,53 @@
|
||||
});
|
||||
}
|
||||
|
||||
document.getElementById("btnRunNet").addEventListener("click", function(e) {
|
||||
e.target.disabled = true;
|
||||
function renderImage(image) {
|
||||
let pixels = []
|
||||
let pixelCounter = 0
|
||||
|
||||
runStableDiffusion(document.getElementById("promptText").value, document.getElementById("stepRange").value, document.getElementById("guidanceRange").value).then((image) => {
|
||||
let pixels = []
|
||||
let pixelCounter = 0
|
||||
|
||||
for (let j = 0; j < 512; j++) {
|
||||
for (let k = 0; k < 512; k++) {
|
||||
pixels.push(image[pixelCounter])
|
||||
pixels.push(image[pixelCounter+1])
|
||||
pixels.push(image[pixelCounter+2])
|
||||
pixels.push(255)
|
||||
pixelCounter += 3
|
||||
}
|
||||
for (let j = 0; j < 512; j++) {
|
||||
for (let k = 0; k < 512; k++) {
|
||||
pixels.push(image[pixelCounter])
|
||||
pixels.push(image[pixelCounter+1])
|
||||
pixels.push(image[pixelCounter+2])
|
||||
pixels.push(255)
|
||||
pixelCounter += 3
|
||||
}
|
||||
|
||||
ctx.putImageData(new ImageData(new Uint8ClampedArray(pixels), 512, 512), 0, 0);
|
||||
console.log(image);
|
||||
console.log("Success");
|
||||
e.target.disabled = false;
|
||||
}
|
||||
|
||||
ctx.putImageData(new ImageData(new Uint8ClampedArray(pixels), 512, 512), 0, 0);
|
||||
}
|
||||
|
||||
const handleRunNetAndRenderResult = () => {
|
||||
document.getElementById("btnRunNet").disabled = true;
|
||||
const canvas = document.getElementById("canvas");
|
||||
ctx.clearRect(0, 0, canvas.width, canvas.height);
|
||||
|
||||
const prevTitleValue = document.getElementById("modelDlTitle").innerHTML;
|
||||
document.getElementById("modelDlTitle").innerHTML = "Running model";
|
||||
|
||||
runStableDiffusion(
|
||||
document.getElementById("promptText").value,
|
||||
document.getElementById("stepRange").value,
|
||||
document.getElementById("guidanceRange").value,
|
||||
// Decode at each step
|
||||
null
|
||||
).then((image) => {
|
||||
renderImage(image);
|
||||
}).finally(() => {
|
||||
document.getElementById("modelDlTitle").innerHTML = prevTitleValue;
|
||||
document.getElementById("btnRunNet").disabled = false;
|
||||
});
|
||||
}, false);
|
||||
};
|
||||
|
||||
document.getElementById("btnRunNet").addEventListener("click", handleRunNetAndRenderResult, false);
|
||||
|
||||
document.getElementById("promptForm").addEventListener("submit", function (event) {
|
||||
event.preventDefault();
|
||||
if (document.getElementById("btnRunNet").disabled) return;
|
||||
|
||||
handleRunNetAndRenderResult();
|
||||
})
|
||||
|
||||
const stepSlider = document.getElementById('stepRange');
|
||||
const stepValue = document.getElementById('stepValue');
|
||||
|
||||
Reference in New Issue
Block a user