github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

progress

Browse Source
This commit is contained in:
George Hotz
2025-10-10 22:44:41 +08:00
parent 79393bddb4
commit dd2ff2ddb9
3 changed files with 11 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
tinygrad/apps/img.py
View File

@@ -1,23 +0,0 @@
# challenge, implement image generation in under 150 lines. can import from tinygrad.apps.llm
from tinygrad import Tensor, nn
if __name__ == "__main__":
vae_state_dict = nn.state.safe_load(Tensor.from_url("https://huggingface.co/stabilityai/sdxl-vae/resolve/main/diffusion_pytorch_model.safetensors"))
for k,v in vae_state_dict.items():
print(k, v.shape, v.dtype)
state_dict = nn.state.safe_load(Tensor.from_url("https://huggingface.co/deepseek-ai/JanusFlow-1.3B/resolve/main/model.safetensors"))
# language_model = Llama variant, apps.llm should just work
#
for k,v in state_dict.items():
print(k, v.shape, v.dtype)
# 1. Tokenise the prompt
# 2. Run language_model.model.layers.... to embed
# 3. Rectified flow
# 4. VAE decode
# 5. Save png

17
tinygrad/apps/resnet.py
View File

@@ -1,18 +1,17 @@
# classification in 50 lines?
# used for verification of img.py
# classification in 50 lines
import sys
from tinygrad import nn, Tensor
class Bottleneck:
expansion = 4
def __init__(self, in_c, mid_c, stride=1):
super().__init__()
out_c = mid_c * self.expansion
self.conv1, self.bn1 = nn.Conv2d(in_c, mid_c, 1, bias=False), nn.BatchNorm2d(mid_c)
self.conv2, self.bn2 = nn.Conv2d(mid_c, mid_c, 3, stride, 1, bias=False), nn.BatchNorm2d(mid_c)
self.conv3, self.bn3 = nn.Conv2d(mid_c, out_c, 1, bias=False), nn.BatchNorm2d(out_c)
self.downsample = (stride != 1 or in_c != out_c) and [nn.Conv2d(in_c, out_c, 1, stride, bias=False), nn.BatchNorm2d(out_c)] or []
def forward(self, x:Tensor) -> Tensor:
def __call__(self, x:Tensor) -> Tensor:
id = x.sequential(self.downsample)
x = self.bn1(self.conv1(x)).relu()
x = self.bn2(self.conv2(x)).relu()
@@ -21,7 +20,6 @@ class Bottleneck:
class ResNet50:
def __init__(self, num_classes=1000):
super().__init__()
self.conv1, self.bn1 = nn.Conv2d(3, 64, 7, 2, 3, bias=False), nn.BatchNorm2d(64)
self.layer1 = self._make_layer(64, 64, 3, 1) # 256 out
self.layer2 = self._make_layer(256, 128, 4, 2) # 512 out
@@ -34,7 +32,7 @@ class ResNet50:
for _ in range(1, blocks): layers.append(Bottleneck(mid_c * Bottleneck.expansion, mid_c))
return layers
def forward(self, x):
def __call__(self, x:Tensor) -> Tensor:
x = self.bn1(self.conv1(x)).relu()
x = x.max_pool2d()
x = self.layer1(x); x = self.layer2(x); x = self.layer3(x); x = self.layer4(x)
@@ -42,9 +40,8 @@ class ResNet50:
return self.fc(x)
if __name__ == "__main__":
# TODO: tiny png library
state_dict = nn.state.safe_load(Tensor.from_url("https://huggingface.co/timm/resnet50.a1_in1k/resolve/main/model.safetensors"))
for k,v in state_dict.items():
print(k, v.shape)
model = ResNet50()
nn.state.load_state_dict(model, state_dict)
nn.state.load_state_dict(model, state_dict)
img = nn.state.png_load(Tensor.from_url(sys.argv[1]))
print(model(img).argmax())

4
tinygrad/nn/state.py
View File

@@ -349,3 +349,7 @@ def gguf_load(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
for name, dims, typ, off in t_infos: state_dict[name] = ggml_data_to_tensor(tensor[data_start + off:], prod(dims), typ).reshape(*reversed(dims))
return kv_data, state_dict
@accept_filename
def png_load(t:Tensor) -> Tensor:
assert t[0:8].tolist() == [0x89,0x50,0x4E,0x47,0x0D,0x0A,0x1A,0x0A], "not a PNG"