Files
onepilot/tinygrad_repo/test/external/external_benchmark_load_stable_diffusion.py
T
Vehicle Researcher 33d5cfc393 openpilot v0.10.3 release
date: 2025-12-18T23:23:16
master commit: 3cdee7b54718ee14bd85befd6c5bad3d699c5479
2025-12-18 23:23:21 -08:00

14 lines
565 B
Python

from tinygrad.helpers import fetch, Timing
from tinygrad.device import Device
from tinygrad.nn.state import torch_load, load_state_dict
from examples.stable_diffusion import StableDiffusion
# run "sudo purge" before testing on OS X to avoid the memory cache
if __name__ == "__main__":
fn = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
model = StableDiffusion()
with Timing():
load_state_dict(model, torch_load(fn)['state_dict'], strict=False)
Device[Device.DEFAULT].synchronize()