Update text to image service to generate actual image files

The text to image generation functionality of the service is now functional, as opposed to previously throwing a NotImplementedException. It pulls a directory from the service's attributes collection to use as output location for the generated images. Logging has also been improved within the service, registering status updates and error messages. Test cases are updated to match the new changes.
frankhaugen · Jan 9, 2024 · 830422e · 830422e
1 parent ca23b39
commit 830422e
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 31 deletions.
diff --git a/...untime.StableDiffusion/Frank.SemanticKernel.Connectors.OnnxRuntime.StableDiffusion.csproj b/...untime.StableDiffusion/Frank.SemanticKernel.Connectors.OnnxRuntime.StableDiffusion.csproj
@@ -3,14 +3,29 @@
     <PropertyGroup>
     </PropertyGroup>
 
-    <ItemGroup>
-        <PackageReference Include="Microsoft.SemanticKernel" Version="1.0.1"/>
-        <PackageReference Include="Microsoft.ML.OnnxRuntime" Version="1.16.3"/>
-        <PackageReference Include="Microsoft.ML.OnnxRuntime.Gpu" Version="1.16.3"/>
-        <PackageReference Include="StableDiffusion.ML.OnnxRuntime" Version="1.1.3" />
+    <ItemGroup><!--        <PackageReference Include="Microsoft.AI.DirectML" Version="1.13.0" />-->
+<!--        <PackageReference Include="Microsoft.ML.OnnxRuntime.Managed" Version="1.16.3" />-->
+<!--        <PackageReference Include="Microsoft.SemanticKernel" Version="1.0.1"/>-->
+<!--        <PackageReference Include="Microsoft.ML.OnnxRuntime" Version="1.16.3"/>-->
+<!--        <PackageReference Include="Microsoft.ML.OnnxRuntime.Gpu" Version="1.16.3"/>-->
+<!--        <PackageReference Include="StableDiffusion.ML.OnnxRuntime" Version="1.1.3" />-->
+<!--        <PackageReference Include="Microsoft.ML" Version="3.0.0" />-->
+<!--        <PackageReference Include="Microsoft.ML.OnnxRuntime.DirectML" Version="1.16.3" />-->
+
+      <PackageReference Include="Microsoft.ML.OnnxRuntime" Version="1.16.3" />
+
+      <PackageReference Include="Microsoft.ML.OnnxRuntime.DirectML" Version="1.16.3" />
+
+      <PackageReference Include="Microsoft.ML.OnnxRuntime.Gpu" Version="1.16.3" />
+
+      <PackageReference Include="Microsoft.SemanticKernel.Abstractions" Version="1.0.1" />
+
+      <PackageReference Include="StableDiffusion.ML.OnnxRuntime" Version="1.1.3" />
     </ItemGroup>
 
     <ItemGroup>
-        <None Include=".\models\*.onnx" CopyToOutputDirectory="PreserveNewest" />
+        <None Update="models\*\model.onnx">
+          <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
+        </None>
     </ItemGroup>
 </Project>
diff --git a/...emanticKernel.Connectors.OnnxRuntime.StableDiffusion/StableDiffusionTextToImageService.cs b/...emanticKernel.Connectors.OnnxRuntime.StableDiffusion/StableDiffusionTextToImageService.cs
@@ -1,5 +1,8 @@
-using Microsoft.SemanticKernel;
+using System.Text;
+using Microsoft.Extensions.Logging;
+using Microsoft.SemanticKernel;
 using Microsoft.SemanticKernel.TextToImage;
+using SixLabors.ImageSharp.Formats.Png;
 using StableDiffusion.ML.OnnxRuntime;
 
 namespace Frank.SemanticKernel.Connectors.OnnxRuntime.StableDiffusion;
@@ -15,33 +18,33 @@ public class StableDiffusionTextToImageService : ITextToImageService
     public const string UNET_PATH = "unet/model.onnx";
     public const string TEXT_ENCODER = "text_encoder/model.onnx";
     public const string PERSISTENT_OUTPUT_DIRECTORY_KEY = "PersistentOutputDirectory";
+
+    private readonly ILogger<StableDiffusionTextToImageService> _logger;
+
+    public StableDiffusionTextToImageService(ILogger<StableDiffusionTextToImageService> logger)
+    {
+        _logger = logger;
+    }
 
     internal Dictionary<string, object?> InternalAttributes { get; } = new Dictionary<string, object?>();
 
     public IReadOnlyDictionary<string, object?> Attributes => InternalAttributes;
 
-    public async Task<string> GenerateImageAsync(string description, int width, int height, Kernel? kernel = null, CancellationToken cancellationToken = new CancellationToken())
-    {
-        throw new NotImplementedException();
-    }
-
-    public void Run(DirectoryInfo? executionDirectory = null)
+    public async Task<string> GenerateImageAsync(string description, int width, int height, Kernel? kernel = null, CancellationToken cancellationToken = default)
     {
-        if (executionDirectory == null)
-        {
-            executionDirectory = new DirectoryInfo(AppContext.BaseDirectory);
-        }
-
-        var modelsDirectory = executionDirectory.CreateSubdirectory("models");
-
+        var directory = InternalAttributes[PERSISTENT_OUTPUT_DIRECTORY_KEY] as DirectoryInfo;
+        if (directory == null)
+            throw new InvalidOperationException($"The '{PERSISTENT_OUTPUT_DIRECTORY_KEY}' attribute must be set before calling this method.");
+
+        var modelsDirectory = directory.CreateSubdirectory("models");
         EnsureModelsExist(modelsDirectory);
 
         //test how long this takes to execute
         var watch = System.Diagnostics.Stopwatch.StartNew();
 
         //Default args
         var prompt = "A cat with a lightsaber fighting a dog with a two-sided lightsaber-staff.";
-        Console.WriteLine(prompt);
+        _logger.LogInformation(prompt);
 
         var config = new StableDiffusionConfig
         {
@@ -69,12 +72,16 @@ public void Run(DirectoryInfo? executionDirectory = null)
         // If image failed or was unsafe it will return null.
         if (image == null)
         {
-            Console.WriteLine("Unable to create image, please try again.");
+            _logger.LogError("Unable to create image, please try again.");
         }
         // Stop the timer
         watch.Stop();
         var elapsedMs = watch.ElapsedMilliseconds;
-        Console.WriteLine("Time taken: " + elapsedMs + "ms");
+        _logger.LogInformation("Time taken: " + elapsedMs + "ms");
+
+        using var imageStream = new MemoryStream(); 
+        await image!.SaveAsync(imageStream, new PngEncoder(), cancellationToken);
+        return Encoding.UTF8.GetString(imageStream.ToArray());
     }
 
     /// <summary>

diff --git a/Frank.SemanticKernel.Tests/Connectors/OnnxRuntime/TextToImagesServiceExtensionsTests.cs b/Frank.SemanticKernel.Tests/Connectors/OnnxRuntime/TextToImagesServiceExtensionsTests.cs
@@ -1,5 +1,6 @@
 using FluentAssertions;
 using Frank.SemanticKernel.Connectors.OnnxRuntime.StableDiffusion;
+using Frank.Testing.Logging;
 using Xunit.Abstractions;
 
 namespace Frank.SemanticKernel.Tests.Connectors.OnnxRuntime;
@@ -17,7 +18,7 @@ public TextToImagesServiceExtensionsTests(ITestOutputHelper testOutputHelper)
     public void Test1()
     {
         // Arrange
-        var service = new StableDiffusionTextToImageService();
+        var service = new StableDiffusionTextToImageService(_testOutputHelper.CreateTestLogger<StableDiffusionTextToImageService>());
         var value = new DirectoryInfo(Directory.GetCurrentDirectory());
         service.AddPersistentOutputDirectory(value);
 

diff --git a/Frank.SemanticKernel.Tests/Connectors/OnnxRuntime/TextToImagesServiceTests.cs b/Frank.SemanticKernel.Tests/Connectors/OnnxRuntime/TextToImagesServiceTests.cs
@@ -1,5 +1,6 @@
-using FluentAssertions;
+using System.Text;
 using Frank.SemanticKernel.Connectors.OnnxRuntime.StableDiffusion;
+using Frank.Testing.Logging;
 using Xunit.Abstractions;
 
 namespace Frank.SemanticKernel.Tests.Connectors.OnnxRuntime;
@@ -14,19 +15,18 @@ public TextToImagesServiceTests(ITestOutputHelper testOutputHelper)
     }
 
     [Fact]
-    public void Test1()
+    public async Task TestAsync()
     {
         // Arrange
-        var service = new StableDiffusionTextToImageService();
+        var service = new StableDiffusionTextToImageService(_testOutputHelper.CreateTestLogger<StableDiffusionTextToImageService>());
         var value = new DirectoryInfo(Directory.GetCurrentDirectory());
         service.AddPersistentOutputDirectory(value);
 
         // Act
-        var result = service.Attributes[StableDiffusionTextToImageService.PERSISTENT_OUTPUT_DIRECTORY_KEY] as DirectoryInfo;
+        var result = await service.GenerateImageAsync("Hello world!", 512, 512);
+        await File.WriteAllBytesAsync(Path.Combine(value.FullName, "test.png"), Encoding.UTF8.GetBytes(result));
 
         // Assert
-        Assert.NotNull(result);
-        _testOutputHelper.WriteLine(result.FullName);
-        result.Should().Be(value);
+        _testOutputHelper.WriteLine(Path.Combine(value.FullName, "test.png"));
     }
 }