fast_loader.py

from sdkit import Context
from sdkit.generate import generate_images
from sdkit.utils import load_tensor_file, save_tensor_file

from ldm.util import instantiate_from_config
from omegaconf import OmegaConf

import time

model_path = "/path/to/models/stable-diffusion/sd-v1-4.ckpt"
config_path = "/path/to/sdkit/models/models_db/configs/v1-inference.yaml"


def load_actual_model_from_disk():
    t = time.time()
    sd = load_tensor_file(model_path)
    sd = sd["state_dict"] if "state_dict" in sd else sd

    config = OmegaConf.load(config_path)
    config.model.params.unet_config.params.use_fp16 = True

    model = instantiate_from_config(config.model)
    _, _ = model.load_state_dict(sd, strict=False)

    model = model.half()
    model.eval()
    print("loaded actual model in", time.time() - t, "seconds")

    return model


def load_model_from_cache():
    t = time.time()
    m = load_tensor_file("model_cached.ckpt")
    print("loaded cached model in", time.time() - t, "seconds")
    return m


c = Context()
c.model_paths["stable-diffusion"] = model_path


def save_model_to_cache():
    m = c.models["stable-diffusion"]
    save_tensor_file(m, "model_cached.ckpt")


def generate(filename):
    t = time.time()
    images = generate_images(c, "Astronaut")
    print("generated in", time.time() - t, "seconds")

    images[0].save(f"{filename}.jpg")


def with_optimization():
    c.models["stable-diffusion"] = load_model_from_cache().to("cuda:0")
    generate("with_optimization.jpg")


def without_optimization():
    c.models["stable-diffusion"] = load_actual_model_from_disk().to("cuda:0")
    generate("without_optimization.jpg")

    save_model_to_cache()


## uncomment one of the two below

# with_optimization()
without_optimization()