StableDiffusionPipeline video stream

Author: saddam213

OnnxStack.Console/Examples/VideoToVideoStreamExample.cs

@@ -0,0 +1,67 @@
+using OnnxStack.Core.Video;
+using OnnxStack.FeatureExtractor.Pipelines;
+using OnnxStack.StableDiffusion.Config;
+using OnnxStack.StableDiffusion.Enums;
+using OnnxStack.StableDiffusion.Pipelines;
+
+namespace OnnxStack.Console.Runner
+{
+    public sealed class VideoToVideoStreamExample : IExampleRunner
+    {
+        private readonly string _outputDirectory;
+        private readonly StableDiffusionConfig _configuration;
+
+        public VideoToVideoStreamExample(StableDiffusionConfig configuration)
+        {
+            _configuration = configuration;
+            _outputDirectory = Path.Combine(Directory.GetCurrentDirectory(), "Examples", nameof(VideoToVideoStreamExample));
+            Directory.CreateDirectory(_outputDirectory);
+        }
+
+        public int Index => 4;
+
+        public string Name => "Video To Video Stream Demo";
+
+        public string Description => "Video Stream Stable Diffusion Inference";
+
+        public async Task RunAsync()
+        {
+
+            // Read Video
+            var videoFile = "C:\\Users\\Deven\\Pictures\\gidsgphy.gif";
+            var videoInfo = await VideoHelper.ReadVideoInfoAsync(videoFile);
+
+            // Loop though the appsettings.json model sets
+            foreach (var modelSet in _configuration.ModelSets)
+            {
+                OutputHelpers.WriteConsole($"Loading Model `{modelSet.Name}`...", ConsoleColor.Cyan);
+
+                // Create Pipeline
+                var pipeline = PipelineBase.CreatePipeline(modelSet);
+
+                // Preload Models (optional)
+                await pipeline.LoadAsync();
+
+                // Add text and video to prompt
+                var promptOptions = new PromptOptions
+                {
+                    Prompt = "Iron Man",
+                    DiffuserType = DiffuserType.ImageToImage
+                };
+
+
+                // Create Video Stream
+                var videoStream = VideoHelper.ReadVideoStreamAsync(videoFile, videoInfo.FrameRate);
+
+                // Create Pipeline Stream
+                var pipelineStream = pipeline.GenerateVideoStreamAsync(videoStream, promptOptions, progressCallback:OutputHelpers.ProgressCallback);
+
+                // Write Video Stream
+                await VideoHelper.WriteVideoStreamAsync(videoInfo, pipelineStream, Path.Combine(_outputDirectory, $"{modelSet.PipelineType}.mp4"));
+
+                //Unload
+                await pipeline.UnloadAsync();
+            }
+        }
+    }
+}

OnnxStack.Core/Image/OnnxImage.cs

@@ -182,6 +182,28 @@ public async Task<Stream> GetImageStreamAsync()
         }
 
 
+        /// <summary>
+        /// Copies the image to stream.
+        /// </summary>
+        /// <param name="destination">The destination.</param>
+        /// <returns></returns>
+        public void CopyToStream(Stream destination)
+        {
+            _imageData.SaveAsPng(destination);
+        }
+
+
+        /// <summary>
+        /// Copies the image to stream.
+        /// </summary>
+        /// <param name="destination">The destination.</param>
+        /// <returns></returns>
+        public Task CopyToStreamAsync(Stream destination)
+        {
+            return _imageData.SaveAsPngAsync(destination);
+        }
+
+
         /// <summary>
         /// Gets the image as tensor.
         /// </summary>

OnnxStack.Core/Video/VideoHelper.cs

@@ -103,7 +103,7 @@ public static async Task WriteVideoStreamAsync(VideoInfo videoInfo, IAsyncEnumer
                 await foreach (var frame in videoStream)
                 {
                     // Write each frame to the input stream of FFMPEG
-                    await videoWriter.StandardInput.BaseStream.WriteAsync(frame.GetImageBytes(), cancellationToken);
+                    await frame.CopyToStreamAsync(videoWriter.StandardInput.BaseStream);
                 }
 
                 // Done close stream and wait for app to process

OnnxStack.StableDiffusion/Config/PromptOptions.cs

@@ -5,7 +5,7 @@
 
 namespace OnnxStack.StableDiffusion.Config
 {
-    public class PromptOptions
+    public record PromptOptions
     {
         public DiffuserType DiffuserType { get; set; }
 

OnnxStack.StableDiffusion/Pipelines/Base/IPipeline.cs

@@ -14,6 +14,13 @@ namespace OnnxStack.StableDiffusion.Pipelines
 {
     public interface IPipeline
     {
+
+        /// <summary>
+        /// Gets the name.
+        /// </summary>
+        string Name { get; }
+
+
         /// <summary>
         /// Gets the pipelines supported diffusers.
         /// </summary>
@@ -127,5 +134,17 @@ public interface IPipeline
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
         IAsyncEnumerable<BatchVideoResult> GenerateVideoBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+
+        /// <summary>
+        /// Runs the video stream pipeline returning each frame as an OnnxImage.
+        /// </summary>
+        /// <param name="videoFrames">The video frames.</param>
+        /// <param name="promptOptions">The prompt options.</param>
+        /// <param name="schedulerOptions">The scheduler options.</param>
+        /// <param name="controlNet">The control net.</param>
+        /// <param name="progressCallback">The progress callback.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        IAsyncEnumerable<OnnxImage> GenerateVideoStreamAsync(IAsyncEnumerable<OnnxImage> videoFrames, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
     }
 }

OnnxStack.StableDiffusion/Pipelines/Base/PipelineBase.cs

@@ -161,6 +161,19 @@ protected PipelineBase(PipelineOptions pipelineOptions, ILogger logger)
         public abstract IAsyncEnumerable<BatchVideoResult> GenerateVideoBatchAsync(BatchOptions batchOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
 
 
+        /// <summary>
+        /// Runs the video stream pipeline returning each frame as an OnnxImage.
+        /// </summary>
+        /// <param name="videoFrames">The video frames.</param>
+        /// <param name="promptOptions">The prompt options.</param>
+        /// <param name="schedulerOptions">The scheduler options.</param>
+        /// <param name="controlNet">The control net.</param>
+        /// <param name="progressCallback">The progress callback.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        public abstract IAsyncEnumerable<OnnxImage> GenerateVideoStreamAsync(IAsyncEnumerable<OnnxImage> videoFrames, PromptOptions promptOptions, SchedulerOptions schedulerOptions = default, ControlNetModel controlNet = default, Action<DiffusionProgress> progressCallback = null, CancellationToken cancellationToken = default);
+
+
         /// <summary>
         /// Creates the diffuser.
         /// </summary>

OnnxStack.StableDiffusion/Pipelines/StableDiffusionPipeline.cs

@@ -266,6 +266,49 @@ public override async IAsyncEnumerable<BatchVideoResult> GenerateVideoBatchAsync
         }
 
 
+        /// <summary>
+        /// Runs the video stream pipeline returning each frame as an OnnxImage.
+        /// </summary>
+        /// <param name="videoFrames">The video frames.</param>
+        /// <param name="promptOptions">The prompt options.</param>
+        /// <param name="schedulerOptions">The scheduler options.</param>
+        /// <param name="controlNet">The control net.</param>
+        /// <param name="progressCallback">The progress callback.</param>
+        /// <param name="cancellationToken">The cancellation token.</param>
+        /// <returns></returns>
+        public override async IAsyncEnumerable<OnnxImage> GenerateVideoStreamAsync(IAsyncEnumerable<OnnxImage> videoFrames, PromptOptions promptOptions, SchedulerOptions schedulerOptions = null, ControlNetModel controlNet = null, Action<DiffusionProgress> progressCallback = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            var diffuseTime = _logger?.LogBegin("Diffuser starting...");
+            var options = GetSchedulerOptionsOrDefault(schedulerOptions);
+            _logger?.Log($"Model: {Name}, Pipeline: {PipelineType}, Diffuser: {promptOptions.DiffuserType}, Scheduler: {options.SchedulerType}");
+
+            // Check guidance
+            var performGuidance = ShouldPerformGuidance(options);
+
+            // Process prompts
+            var promptEmbeddings = await CreatePromptEmbedsAsync(promptOptions, performGuidance);
+
+            // Create Diffuser
+            var diffuser = CreateDiffuser(promptOptions.DiffuserType, controlNet);
+
+            // Diffuse
+            await foreach (var videoFrame in videoFrames)
+            {
+                var frameOptions = promptOptions with
+                {
+                    InputImage = promptOptions.DiffuserType == DiffuserType.ImageToImage || promptOptions.DiffuserType == DiffuserType.ControlNetImage
+                        ? videoFrame : default,
+                    InputContolImage = promptOptions.DiffuserType == DiffuserType.ControlNet || promptOptions.DiffuserType == DiffuserType.ControlNetImage
+                        ? videoFrame : default,
+                };
+
+                yield return new OnnxImage(await DiffuseImageAsync(diffuser, frameOptions, options, promptEmbeddings, performGuidance, progressCallback, cancellationToken));
+            }
+
+            _logger?.LogEnd($"Diffuser complete", diffuseTime);
+        }
+
+
         /// <summary>
         /// Runs the pipeline
         /// </summary>