diffusers/tools

1#!/usr/bin/env python32from safetensors.torch import load_file as safe_load_file3import time4import sys5 6direct_on_gpu = bool(int(sys.argv[1]))7 8if direct_on_gpu:9    start_time = time.time()10    checkpoint = safe_load_file("/home/patrick_huggingface_co/stable-diffusion-v1-4/unet/diffusion_pytorch_model.safetensors", device=0)11    print("Directly on GPU", time.time() - start_time)12else:13    start_time = time.time()14    checkpoint = safe_load_file("/home/patrick_huggingface_co/stable-diffusion-v1-4/unet/diffusion_pytorch_model.safetensors")15    checkpoint = {k: v.to("cuda:0") for k, v in checkpoint.items()}16    print("On CPU", time.time() - start_time)17