使用 C# 和 ONNX Runtime 推理 Stable Diffusion

在本教程中,我们将学习如何在 C# 中对流行的 Stable Diffusion 深度学习模型进行推理。Stable Diffusion 模型接收文本提示并创建代表该文本的图像。请参见下面的示例:

"make a picture of green tree with flowers around it and a red sky" 
Image of browser inferencing on sample images. Image of browser inferencing on sample images.

目录

先决条件

本教程可以在本地运行,也可以通过利用 Azure 机器学习计算在云端运行。

本地运行

使用 Azure 机器学习在云端运行

使用 Hugging Face 下载 Stable Diffusion 模型

Hugging Face 站点拥有丰富的开源模型库。我们将利用并从 Hugging Face 下载 ONNX Stable Diffusion 模型

选择模型版本仓库后,点击 Files and Versions,然后选择 ONNX 分支。如果没有可用的 ONNX 模型分支,请使用 main 分支并将其转换为 ONNX。有关更多信息,请参阅 PyTorch 的 ONNX 转换教程

  • 克隆仓库
    git lfs install
    git clone https://hugging-face.cn/CompVis/stable-diffusion-v1-4 -b onnx
    
  • 将包含 ONNX 文件的文件夹复制到 C# 项目文件夹 \StableDiffusion\StableDiffusion。需要复制的文件夹是:unetvae_decodertext_encodersafety_checker

通过 Hugging Face 的 Diffusers 在 Python 中理解模型

当对预构建模型进行操作化时,花点时间理解此管道中的模型非常有用。此代码基于 Hugging Face Diffusers 库和博客。如果您想了解更多关于其工作原理的信息,请查看这篇精彩的博文以获取更多详情!

使用 C# 进行推理

现在,让我们开始分解如何在 C# 中进行推理!unet 模型接收由连接文本和图像的 CLIP 模型创建的用户提示的文本嵌入。潜在的噪声图像被创建为起点。调度算法和 unet 模型协同工作以对图像进行去噪,从而创建代表文本提示的图像。让我们看看代码。

主函数

主函数设置提示、推理步数和引导比例。然后它调用 UNet.Inference 函数来运行推理。

需要设置的属性有:

  • prompt - 用于图像的文本提示
  • num_inference_steps - 运行推理的步数。步数越多,推理循环运行时间越长,但图像质量应该会提高。
  • guidance_scale - 无分类器引导的比例。数值越高,越会尝试与提示相似,但图像质量可能会下降。
  • batch_size - 创建的图像数量
  • height - 图像的高度。默认为 512,必须是 8 的倍数。
  • width - 图像的宽度。默认为 512,必须是 8 的倍数。

* 注意:有关更多详情,请查看 Hugging Face 博客

//Default args
var prompt = "make a picture of green tree with flowers around it and a red sky";
// Number of steps
var num_inference_steps = 10;

// Scale for classifier-free guidance
var guidance_scale = 7.5;
//num of images requested
var batch_size = 1;
// Load the tokenizer and text encoder to tokenize and encodethe text.
var textTokenized = TextProcessing.TokenizeText(prompt);
var textPromptEmbeddings = TextProcessing.TextEncode(textTokenized).ToArray();
// Create uncond_input of blank tokens
var uncondInputTokens = TextProcessing.CreateUncondInput();
var uncondEmbedding = TextProcessing.TextEncode(uncondInputTokens).ToArray();
// Concat textEmeddings and uncondEmbedding
DenseTensor<float> textEmbeddings = new DenseTensor<float>(ne[] { 2, 77, 768 });
for (var i = 0; i < textPromptEmbeddings.Length; i++)
{
    textEmbeddings[0, i / 768, i % 768] = uncondEmbedding[i];
    textEmbeddings[1, i / 768, i % 768] = textPromptEmbeddings[i];
}
var height = 512;
var width = 512;
// Inference Stable Diff
var image = UNet.Inference(num_inference_steps, textEmbeddings,guidance_scale, batch_size, height, width);
// If image failed or was unsafe it will return null.
if( image == null )
{
    Console.WriteLine("Unable to create image, please try again.");
}

使用 ONNX Runtime Extensions 进行分词

TextProcessing 类包含用于对文本提示进行分词并使用 CLIP 模型文本编码器进行编码的函数。

我们无需在 C# 中重新实现 CLIP 分词器,而是可以利用 ONNX Runtime Extensions 中跨平台的 CLIP 分词器实现。ONNX Runtime Extensions 包含一个 custom_op_cliptok.onnx 文件分词器,用于对文本提示进行分词。该分词器是一个简单的分词器,它将文本分割成单词,然后将单词转换为标记。

  • 文本提示:表示您想要创建的图像的句子或短语。
    make a picture of green tree with flowers aroundit and a red sky
    
  • 文本分词:文本提示被分词为标记列表。每个标记 ID 都是一个数字,代表句子中的一个单词,然后用空白标记填充,以创建 77 个标记的 maxLength。然后将标记 ID 转换为形状为 (1,77) 的张量。

  • 以下是使用 ONNX Runtime Extensions 对文本提示进行分词的代码。
public static int[] TokenizeText(string text)
{
            // Create Tokenizer and tokenize the sentence.
            var tokenizerOnnxPath = Directory.GetCurrentDirectory().ToString() + ("\\text_tokenizer\\custom_op_cliptok.onnx");

            // Create session options for custom op of extensions
            using var sessionOptions = new SessionOptions();
            var customOp = "ortextensions.dll";
            sessionOptions.RegisterCustomOpLibraryV2(customOp, out var libraryHandle);

            // Create an InferenceSession from the onnx clip tokenizer.
            using var tokenizeSession = new InferenceSession(tokenizerOnnxPath, sessionOptions);

            // Create input tensor from text
            using var inputTensor = OrtValue.CreateTensorWithEmptyStrings(OrtAllocator.DefaultInstance, new long[] { 1 });
            inputTensor.StringTensorSetElementAt(text.AsSpan(), 0);

            var inputs = new Dictionary<string, OrtValue>
            {
                {  "string_input", inputTensor }
            };

            // Run session and send the input data in to get inference output. 
            using var runOptions = new RunOptions();
            using var tokens = tokenizeSession.Run(runOptions, inputs, tokenizeSession.OutputNames);

            var inputIds = tokens[0].GetTensorDataAsSpan<long>();

            // Cast inputIds to Int32
            var InputIdsInt = new int[inputIds.Length];
            for(int i = 0; i < inputIds.Length; i++)
            {
                InputIdsInt[i] = (int)inputIds[i];
            }

            Console.WriteLine(String.Join(" ", InputIdsInt));

            var modelMaxLength = 77;
            // Pad array with 49407 until length is modelMaxLength
            if (InputIdsInt.Length < modelMaxLength)
            {
                var pad = Enumerable.Repeat(49407, 77 - InputIdsInt.Length).ToArray();
                InputIdsInt = InputIdsInt.Concat(pad).ToArray();
            }
            return InputIdsInt;
}

tensor([[49406,  1078,   320,  1674,   539,  1901,  2677,   593,  4023,  1630,
           585,   537,   320,   736,  2390, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,
         49407, 49407, 49407, 49407, 49407, 49407, 49407]])

使用 CLIP 文本编码器模型进行文本嵌入

这些标记被发送到文本编码器模型,并转换为形状为 (1, 77, 768) 的张量,其中第一维是批大小,第二维是标记数量,第三维是嵌入大小。文本编码器是一个 OpenAI CLIP 模型,它连接文本到图像。

文本编码器创建文本嵌入,该嵌入经过训练将文本提示编码为用于引导图像生成的向量。然后将文本嵌入与 uncond 嵌入连接起来,创建发送到 unet 模型进行推理的文本嵌入。

  • 文本嵌入:从分词结果创建的代表文本提示的数字向量。文本嵌入由 text_encoder 模型创建。
        public static float[] TextEncoder(int[] tokenizedInput)
        {
            // Create input tensor. OrtValue will not copy, will read from managed memory
            using var input_ids = OrtValue.CreateTensorValueFromMemory<int>(tokenizedInput,
                new long[] { 1, tokenizedInput.Count() });

            var textEncoderOnnxPath = Directory.GetCurrentDirectory().ToString() + ("\\text_encoder\\model.onnx");

            using var encodeSession = new InferenceSession(textEncoderOnnxPath);

            // Pre-allocate the output so it goes to a managed buffer
            // we know the shape
            var lastHiddenState = new float[1 * 77 * 768];
            using var outputOrtValue = OrtValue.CreateTensorValueFromMemory<float>(lastHiddenState, new long[] { 1, 77, 768 });

            string[] input_names = { "input_ids" };
            OrtValue[] inputs = { input_ids };

            string[] output_names = { encodeSession.OutputNames[0] };
            OrtValue[] outputs = { outputOrtValue };

            // Run inference.
            using var runOptions = new RunOptions();
            encodeSession.Run(runOptions, input_names, inputs, output_names, outputs);

            return lastHiddenState;
        }
torch.Size([1, 77, 768])
tensor([[[-0.3884,  0.0229, -0.0522,  ..., -0.4899, -0.3066,  0.0675],
         [ 0.0520, -0.6046,  1.9268,  ..., -0.3985,  0.9645, -0.4424],
         [-0.8027, -0.4533,  1.7525,  ..., -1.0365,  0.6296,  1.0712],
         ...,
         [-0.6833,  0.3571, -1.1353,  ..., -1.4067,  0.0142,  0.3566],
         [-0.7049,  0.3517, -1.1524,  ..., -1.4381,  0.0090,  0.3777],
         [-0.6155,  0.4283, -1.1282,  ..., -1.4256, -0.0285,  0.3206]]],

推理循环:UNet 模型、时间步和 LMS 调度器

调度器

调度算法和 unet 模型协同工作,对图像进行去噪,从而创建代表文本提示的图像。可以使用不同的调度算法,要了解更多信息,请查看 Hugging Face 的这篇博客。在此示例中,我们将使用 `LMSDiscreteScheduler`,它是基于 HuggingFace scheduling_lms_discrete.py 创建的。

时间步

推理循环是运行调度算法和 unet 模型的主循环。该循环运行 timesteps 的次数,这些时间步是根据推理步数和其他参数由调度算法计算得出的。

在此示例中,我们有 10 个推理步,计算出以下时间步:

// Get path to model to create inference session.
var modelPath = Directory.GetCurrentDirectory().ToString() + ("\\unet\\model.onnx");
var scheduler = new LMSDiscreteScheduler();
var timesteps = scheduler.SetTimesteps(numInferenceSteps);
tensor([999., 888., 777., 666., 555., 444., 333., 222., 111.,   0.])

潜在变量

latents 是模型输入中使用的噪声图像张量。它通过 GenerateLatentSample 函数创建,生成一个形状为 (1,4,64,64) 的随机张量。seed 可以设置为随机数或固定数字。如果 seed 设置为固定数字,每次都将使用相同的潜在张量。这对于调试或每次都想创建相同图像时非常有用。

var seed = new Random().Next();
var latents = GenerateLatentSample(batchSize, height, width,seed, scheduler.InitNoiseSigma);

Image of browser inferencing on sample images.

推理循环

对于每个推理步骤,潜在图像被复制以创建形状为 (2,4,64,64) 的张量,然后将其缩放并与 unet 模型进行推理。输出张量 (2,4,64,64) 被分割并应用引导。然后将结果张量作为去噪过程的一部分发送到 LMSDiscreteScheduler 步骤,并返回调度步骤的结果张量,循环再次完成,直到达到 num_inference_steps

var modelPath = Directory.GetCurrentDirectory().ToString() + ("\\unet\\model.onnx");
var scheduler = new LMSDiscreteScheduler();
var timesteps = scheduler.SetTimesteps(numInferenceSteps);

var seed = new Random().Next();
var latents = GenerateLatentSample(batchSize, height, width, seed, scheduler.InitNoiseSigma);

// Create Inference Session
using var options = new SessionOptions();
using var unetSession = new InferenceSession(modelPath, options);

var latentInputShape = new int[] { 2, 4, height / 8, width / 8 };
var splitTensorsShape = new int[] { 1, 4, height / 8, width / 8 };

for (int t = 0; t < timesteps.Length; t++)
{
    // torch.cat([latents] * 2)
    var latentModelInput = TensorHelper.Duplicate(latents.ToArray(), latentInputShape);

    // Scale the input
    latentModelInput = scheduler.ScaleInput(latentModelInput, timesteps[t]);

    // Create model input of text embeddings, scaled latent image and timestep
    var input = CreateUnetModelInput(textEmbeddings, latentModelInput, timesteps[t]);

    // Run Inference
    using var output = unetSession.Run(input);
    var outputTensor = output[0].Value as DenseTensor<float>;

    // Split tensors from 2,4,64,64 to 1,4,64,64
    var splitTensors = TensorHelper.SplitTensor(outputTensor, splitTensorsShape);
    var noisePred = splitTensors.Item1;
    var noisePredText = splitTensors.Item2;

    // Perform guidance
    noisePred = performGuidance(noisePred, noisePredText, guidanceScale);

    // LMS Scheduler Step
    latents = scheduler.Step(noisePred, timesteps[t], latents);
}

使用 VAEDecoder 对 output 进行后处理

推理循环完成后,将结果张量进行缩放,然后发送到 vae_decoder 模型以解码图像。最后,将解码后的图像张量转换为图像并保存到磁盘。

public static Tensor<float> Decoder(List<NamedOnnxValue> input)
{
    // Load the model which will be used to decode the latents into image space. 
   var vaeDecoderModelPath = Directory.GetCurrentDirectory().ToString() + ("\\vae_decoder\\model.onnx");
    
    // Create an InferenceSession from the Model Path.
    var vaeDecodeSession = new InferenceSession(vaeDecoderModelPath);

   // Run session and send the input data in to get inference output. 
    var output = vaeDecodeSession.Run(input);
    var result = (output.ToList().First().Value as Tensor<float>);
    return result;
}

public static Image<Rgba32> ConvertToImage(Tensor<float> output, int width = 512, int height = 512, string imageName = "sample")
{
    var result = new Image<Rgba32>(width, height);
    for (var y = 0; y < height; y++)
    {
        for (var x = 0; x < width; x++)
        {
            result[x, y] = new Rgba32(
                (byte)(Math.Round(Math.Clamp((output[0, 0, y, x] / 2 + 0.5), 0, 1) * 255)),
                (byte)(Math.Round(Math.Clamp((output[0, 1, y, x] / 2 + 0.5), 0, 1) * 255)),
                (byte)(Math.Round(Math.Clamp((output[0, 2, y, x] / 2 + 0.5), 0, 1) * 255))
            );
        }
    }
    result.Save($@"C:/code/StableDiffusion/{imageName}.png");
    return result;
}

结果图像

image

结论

这是关于如何在 C# 中运行 Stable Diffusion 的高级概述。它涵盖了主要概念并提供了实现示例。要获取完整代码,请查看 Stable Diffusion C# 示例

资源