forked from tinygrad/tinygrad
Bring back WebGPU (#7063)
* Start from andredaprato:webgpu-clean * Fix infs * inf wgsl function is not needed * Emulated ulong for threefry, more tests passing * Randomness tests passing * Update model export to support new changes in webgpu, efficientnet export works again * Simplify shift emulation in wgsl * Delete test file * Fix bigger than u32 u32 literal * Why was skip copies added here? * Python3.12 for webgpu tests * Fix model export syntax error * Get test ops passing with some skips * Fix lint * Much simpler shift * Run more tests * Timestamp queries are not supported in CI, so skip search tests * All fancy indexing passing * r is ctx * Run more dtype tests by using is_dtype_supported * Cleanup ulong shift rendering * UPat -> Pat, UOps -> Ops * Pat -> UPat * Refactor render_ushift if-else * Pattern to avoid ulong mul * Remove vals_dtype * is_nan trick + rewrite, test_isnan passing * Rewrite a * select(1, nan, gate) -> select(a, nan, gate) * No arg, just op * Support char, uchar, short, ushort * Run test_index_mnis now that we have uint8 * Fix pyling * Save 3 lines by using base Compiler * No more long emulation * Remove fixup_binops * No more external_local_bufx wgsl specific cstyle modif, use base extra_pm * Simpler, faster copyin/out * Skip some new tests that use long * Fix typo * copyout touchup * Save lines by using render_cast * WebGL is not supported in core, delete it from is_dtype_supported * More narrow test skips for some unary tests * TernaryOps, UnaryOps -> Ops * TinyGrad supports WebGPU * StableDiffusion demo: f16tof32 gpu is a lib, update UI * Packed load/store, no more scale_size, no core tinygrad changes * Rename copyin, copyout * Device -> dev * Fix lint * Pattern matcher rule for packed load/store * Refactor * Shorter packed load/store * this should fix lint * Fix mypy * SD compile script working * New SD webgpu UI * New default prompt * New SD weights * Fix title when webgpu not available * Run symbolic tests, simplify is_nan, use round_up * Show step time on UI * Bump minimum wgpu version to v0.19 * Fix latent --------- Co-authored-by: George Hotz <[email protected]>
This commit is contained in:
Ahmed Harmouche
co-authored by
George Hotz
George Hotz
parent
ff3f2a9c1a
commit
10618aba98
@@ -6,7 +6,7 @@ from tinygrad.tensor import Tensor
|
||||
from tinygrad import Device
|
||||
from tinygrad.helpers import fetch
|
||||
from typing import NamedTuple, Any, List
|
||||
from pathlib import Path
|
||||
import requests
|
||||
import argparse
|
||||
import numpy as np
|
||||
|
||||
@@ -60,16 +60,22 @@ def split_safetensor(fn):
|
||||
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:
|
||||
with open(os.path.join(os.path.dirname(__file__), f'./net_textmodel.safetensors'), "wb+") as f:
|
||||
f.write(net_bytes[text_model_start+8+json_len:])
|
||||
|
||||
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()
|
||||
@@ -94,12 +100,21 @@ if __name__ == "__main__":
|
||||
|
||||
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.base.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()])
|
||||
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])
|
||||
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()])
|
||||
@@ -148,14 +163,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"""
|
||||
@@ -185,20 +200,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);
|
||||
}}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user