diffusers/tools

1#!/usr/bin/env python32import torch3from warnings import warn4from diffusers import (5    AutoencoderKL,6    DiffusionPipeline,7)8import hashlib9 10base = "stabilityai/stable-diffusion-xl-base-1.0"11adapter1 = 'nerijs/pixel-art-xl'12weightname1 = 'pixel-art-xl.safetensors'13 14adapter2 = 'Alexzyx/lora-trained-xl-colab'15weightname2 = None16 17inputs = "elephant"18kwargs = {}19 20if torch.cuda.is_available():21    kwargs["torch_dtype"] = torch.float1622 23#vae = AutoencoderKL.from_pretrained(24#    "madebyollin/sdxl-vae-fp16-fix",25#    torch_dtype=torch.float16,  # load fp16 fix VAE26#)27#kwargs["vae"] = vae28#kwargs["variant"] = "fp16"29#30 31model = DiffusionPipeline.from_pretrained(32    base, **kwargs33)34 35if torch.cuda.is_available():36    model.to("cuda")37 38 39def inference(adapter, weightname):40    model.load_lora_weights(adapter, weight_name=weightname)41    try:42        model.fuse_lora(safe_fusing=True)43    except ValueError:44        warn(f"{adapter} and {weightname} is broken. LoRA is not fused.")45        model.unload_lora_weights()46 47    data = model(inputs, num_inference_steps=1).images[0]48    model.unfuse_lora()49    model.unload_lora_weights()50    filename = '/tmp/hello.jpg'51    data.save(filename, format='jpeg')52    with open(filename, 'rb') as f:53        md5 = hashlib.md5(f.read()).hexdigest()54    print("Adapter %s, md5sum %s" % (adapter, md5))55    if md5 == '40c78c9fd4daeff01c988c3532fdd51b':56        print("BLACK SCREEN IMAGE for adapter %s" % adapter)57 58 59inference(adapter1, weightname1)60inference(adapter2, weightname2)61inference(adapter1, weightname1)62inference(adapter1, weightname1)63