Save options to json with image output

Author: saddam213

OnnxStack.WebUI/Hubs/StableDiffusionHub.cs

@@ -2,47 +2,88 @@
 using OnnxStack.StableDiffusion.Common;
 using OnnxStack.StableDiffusion.Config;
 using OnnxStack.WebUI.Models;
-using System;
 using System.Runtime.CompilerServices;
+using System.Text.Json;
+using System.Text.Json.Serialization;
 
 namespace OnnxStack.Web.Hubs
 {
     public class StableDiffusionHub : Hub
     {
         private readonly ILogger<StableDiffusionHub> _logger;
-        private readonly IStableDiffusionService _stableDiffusionService;
         private readonly IWebHostEnvironment _webHostEnvironment;
+        private readonly JsonSerializerOptions _serializerOptions;
+        private readonly IStableDiffusionService _stableDiffusionService;
+
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="StableDiffusionHub"/> class.
+        /// </summary>
+        /// <param name="logger">The logger.</param>
+        /// <param name="stableDiffusionService">The stable diffusion service.</param>
+        /// <param name="webHostEnvironment">The web host environment.</param>
         public StableDiffusionHub(ILogger<StableDiffusionHub> logger, IStableDiffusionService stableDiffusionService, IWebHostEnvironment webHostEnvironment)
         {
             _logger = logger;
             _webHostEnvironment = webHostEnvironment;
             _stableDiffusionService = stableDiffusionService;
+            _serializerOptions = new JsonSerializerOptions { WriteIndented = true, Converters = { new JsonStringEnumConverter() } };
         }
 
+
+        /// <summary>
+        /// Called when a new connection is established with the hub.
+        /// </summary>
         public override async Task OnConnectedAsync()
         {
             _logger.Log(LogLevel.Information, "[OnConnectedAsync], Id: {0}", Context.ConnectionId);
-
             await Clients.Caller.SendAsync("OnMessage", "OnConnectedAsync");
             await base.OnConnectedAsync();
         }
 
 
+        /// <summary>
+        /// Called when a connection with the hub is terminated.
+        /// </summary>
+        /// <param name="exception"></param>
         public override async Task OnDisconnectedAsync(Exception exception)
         {
             _logger.Log(LogLevel.Information, "[OnDisconnectedAsync], Id: {0}", Context.ConnectionId);
-
             await Clients.Caller.SendAsync("OnMessage", "OnDisconnectedAsync");
             await base.OnDisconnectedAsync(exception);
         }
 
 
+        /// <summary>
+        /// Execute Text-To-Image Stable Diffusion
+        /// </summary>
+        /// <param name="options">The options.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
         [HubMethodName("ExecuteTextToImage")]
-        public async IAsyncEnumerable<DiffusionResult> OnExecuteTextToImage(TextToImageOptions options, [EnumeratorCancellation] CancellationToken cancellationToken)
+        public async IAsyncEnumerable<TextToImageResult> OnExecuteTextToImage(TextToImageOptions options, [EnumeratorCancellation] CancellationToken cancellationToken)
         {
-            _logger.Log(LogLevel.Information, "[OnExecuteTextToImage] - New prompt received, Connection: {0}", Context.ConnectionId);
-            var linkedCancellationToken = CancellationTokenSource.CreateLinkedTokenSource(Context.ConnectionAborted, cancellationToken);
+            _logger.Log(LogLevel.Information, "[OnExecuteTextToImage] - New request received, Connection: {0}", Context.ConnectionId);
+            var cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(Context.ConnectionAborted, cancellationToken);
+
+            // TODO: Add support for multiple results
+            var result = await GenerateTextToImageResult(options, cancellationTokenSource.Token);
+            if (result is null)
+                yield break;
 
+            yield return result;
+        }
+
+
+        /// <summary>
+        /// Generates the text to image result.
+        /// </summary>
+        /// <param name="options">The options.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        private async Task<TextToImageResult> GenerateTextToImageResult(TextToImageOptions options, CancellationToken cancellationToken)
+        {
+            options.Seed = GenerateSeed(options.Seed);
             var promptOptions = new PromptOptions
             {
                 Prompt = options.Prompt,
@@ -54,44 +95,103 @@ public async IAsyncEnumerable<DiffusionResult> OnExecuteTextToImage(TextToImageO
             {
                 Width = options.Width,
                 Height = options.Height,
-                Seed = GenerateSeed(options.Seed),
+                Seed = options.Seed,
                 InferenceSteps = options.InferenceSteps,
                 GuidanceScale = options.GuidanceScale,
                 Strength = options.Strength,
                 InitialNoiseLevel = options.InitialNoiseLevel
             };
 
-            // TODO: Add support for multiple results
-            var result = await GenerateTextToImage(promptOptions, schedulerOptions, cancellationToken);
-            if(result is null)
-                yield break;
+            var fileInfo = CreateFileInfo(promptOptions, schedulerOptions);
+            if (!await SaveOptionsFile(fileInfo, options))
+                return null;
 
-            yield return result;
+            if (!await RunStableDiffusion(promptOptions, schedulerOptions, fileInfo, cancellationToken))
+                return null;
+
+            return new TextToImageResult(fileInfo.OutputImage, fileInfo.OutputImageUrl, options, fileInfo.OutputOptionsUrl);
         }
 
-        private async Task<DiffusionResult> GenerateTextToImage(PromptOptions promptOptions, SchedulerOptions schedulerOptions, CancellationToken cancellationToken)
-        {
-            var rand = Path.GetFileNameWithoutExtension(Path.GetRandomFileName());
-            var outputImage = $"{schedulerOptions.Seed}_{promptOptions.SchedulerType}_{rand}.png";
-            var outputImageUrl = CreateOutputImageUrl("TextToImage", outputImage);
-            var outputImageFile = CreateOutputImageFile(outputImageUrl);
 
+        /// <summary>
+        /// Runs the stable diffusion.
+        /// </summary>
+        /// <param name="promptOptions">The prompt options.</param>
+        /// <param name="schedulerOptions">The scheduler options.</param>
+        /// <param name="fileInfo">The file information.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        private async Task<bool> RunStableDiffusion(PromptOptions promptOptions, SchedulerOptions schedulerOptions, FileInfoResult fileInfo, CancellationToken cancellationToken)
+        {
             try
             {
-                await _stableDiffusionService.TextToImageFile(promptOptions, schedulerOptions, outputImageFile, ProgressCallback(), cancellationToken);
-                return new DiffusionResult(outputImage, outputImageUrl);
+                await _stableDiffusionService.TextToImageFile(promptOptions, schedulerOptions, fileInfo.OutputImageFile, ProgressCallback(), cancellationToken);
+                return true;
             }
             catch (OperationCanceledException tex)
             {
                 await Clients.Caller.SendAsync("OnCanceled", tex.Message);
+                _logger.Log(LogLevel.Warning, tex, "[OnExecuteTextToImage] - Operation canceled, Connection: {0}", Context.ConnectionId);
             }
             catch (Exception ex)
             {
                 await Clients.Caller.SendAsync("OnError", ex.Message);
+                _logger.Log(LogLevel.Error, ex, "[OnExecuteTextToImage] - Error generating image, Connection: {0}", Context.ConnectionId);
+            }
+            return false;
+        }
+
+
+        /// <summary>
+        /// Saves the options file.
+        /// </summary>
+        /// <param name="fileInfo">The file information.</param>
+        /// <param name="options">The options.</param>
+        /// <returns></returns>
+        private async Task<bool> SaveOptionsFile(FileInfoResult fileInfo, TextToImageOptions options)
+        {
+            try
+            {
+                using (var stream = File.Create(fileInfo.OutputOptionsFile))
+                {
+                    await JsonSerializer.SerializeAsync(stream, options, _serializerOptions);
+                    return true;
+                }
+            }
+            catch (Exception ex)
+            {
+                _logger.Log(LogLevel.Error, ex, "[SaveOptions] - Error saving model card, Connection: {0}", Context.ConnectionId);
+                return false;
             }
-            return null;
         }
 
+
+        /// <summary>
+        /// Creates the file information.
+        /// </summary>
+        /// <param name="promptOptions">The prompt options.</param>
+        /// <param name="schedulerOptions">The scheduler options.</param>
+        /// <returns></returns>
+        private FileInfoResult CreateFileInfo(PromptOptions promptOptions, SchedulerOptions schedulerOptions)
+        {
+            var rand = Path.GetFileNameWithoutExtension(Path.GetRandomFileName());
+            var output = $"{schedulerOptions.Seed}_{promptOptions.SchedulerType}_{rand}";
+            var outputImage = $"{output}.png";
+            var outputImageUrl = CreateOutputUrl("TextToImage", outputImage);
+            var outputImageFile = UrlToPhysicalPath(outputImageUrl);
+
+            var outputJson = $"{output}.json";
+            var outputJsonUrl = CreateOutputUrl("TextToImage", outputJson);
+            var outputJsonFile = UrlToPhysicalPath(outputJsonUrl);
+            return new FileInfoResult(outputImage, outputImageUrl, outputImageFile, outputJson, outputJsonUrl, outputJsonFile);
+        }
+
+
+        /// <summary>
+        /// Generates the seed.
+        /// </summary>
+        /// <param name="seed">The seed.</param>
+        /// <returns></returns>
         private int GenerateSeed(int seed)
         {
             if (seed > 0)
@@ -100,6 +200,11 @@ private int GenerateSeed(int seed)
             return Random.Shared.Next();
         }
 
+
+        /// <summary>
+        /// Progress callback.
+        /// </summary>
+        /// <returns></returns>
         private Action<int, int> ProgressCallback()
         {
             return async (progress, total) =>
@@ -110,19 +215,32 @@ private Action<int, int> ProgressCallback()
         }
 
 
-        private string CreateOutputImageFile(string url)
+        /// <summary>
+        /// URL path to physical path.
+        /// </summary>
+        /// <param name="url">The URL.</param>
+        /// <returns></returns>
+        private string UrlToPhysicalPath(string url)
         {
             string webRootPath = _webHostEnvironment.WebRootPath;
             string physicalPath = Path.Combine(webRootPath, url.TrimStart('/').Replace('/', '\\'));
             return physicalPath;
         }
 
-        private string CreateOutputImageUrl(string folder, string imageName)
+
+        /// <summary>
+        /// Creates the output URL.
+        /// </summary>
+        /// <param name="folder">The folder.</param>
+        /// <param name="file">The file.</param>
+        /// <returns></returns>
+        private string CreateOutputUrl(string folder, string file)
         {
-            return $"/images/results/{folder}/{imageName}";
+            return $"/images/results/{folder}/{file}";
         }
     }
 
     public record ProgressResult(int Progress, int Total);
-    public record DiffusionResult(string OutputImage, string OutputImageUrl);
+    public record TextToImageResult(string OutputImage, string OutputImageUrl, TextToImageOptions OutputOptions, string OutputOptionsUrl);
+    public record FileInfoResult(string OutputImage, string OutputImageUrl, string OutputImageFile, string OutputOptions, string OutputOptionsUrl, string OutputOptionsFile);
 }

OnnxStack.WebUI/Pages/StableDiffusion/TextToImage.cshtml

@@ -144,7 +144,11 @@
             let diffusionProcess;
 
             const onResponse = (response) => {
-                updateResultImage(response)
+                if (!response)
+                    return;
+
+                console.log(response);
+                updateResultImage(response);
                 processEnd();
             }
 
@@ -157,25 +161,23 @@
             }
 
             const onProgress = (response) => {
-                console.log(response)
+                console.log(response);
                 updateProgress(response);
             }
 
             const onCanceled = (response) => {
-                console.log(response)
+                console.log(response);
                 outputContainer.html('');
                 processEnd();
             }
 
             const executeTextToImage = async () => {
 
-                // Validate form
                 const diffusionParams = serializeFormToJson(optionsForm);
                 if (!validateForm())
                     return;
 
                 processBegin();
-               
                 outputContainer.html(Mustache.render(progressResultTemplate));
                 diffusionProcess = await connection
                     .stream("ExecuteTextToImage", diffusionParams)
@@ -197,8 +199,9 @@
             const updateProgress = (response) => {
                 const increment = Math.max(100 / response.total, 1);
                 const progressPercent = Math.round(Math.min(increment * response.progress, 100), 0);
-                $("#progress-result").css("width", progressPercent + "%");
-                $("#progress-result").text(progressPercent + "%");
+                const progressBar = $("#progress-result");
+                progressBar.css("width", progressPercent + "%");
+                progressBar.text(progressPercent + "%");
             }
 
             const processBegin = () => {