StableDiffusionService initial video support

Author: saddam213

OnnxStack.Core/Services/IVideoService.cs

@@ -76,7 +76,7 @@ public interface IVideoService
         /// <param name="videoFPS">The video FPS.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        Task<VideoResult> CreateVideoAsync(IEnumerable<byte[]> videoFrames, float videoFPS, CancellationToken cancellationToken = default);
+        Task<VideoOutput> CreateVideoAsync(IEnumerable<byte[]> videoFrames, float videoFPS, CancellationToken cancellationToken = default);
 
 
         /// <summary>
@@ -85,7 +85,7 @@ public interface IVideoService
         /// <param name="videoFrames">The video frames.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        Task<VideoResult> CreateVideoAsync(VideoFrames videoFrames, CancellationToken cancellationToken = default);
+        Task<VideoOutput> CreateVideoAsync(VideoFrames videoFrames, CancellationToken cancellationToken = default);
 
 
         /// <summary>

OnnxStack.Core/Services/VideoService.cs

@@ -88,9 +88,9 @@ public async Task<VideoInfo> GetVideoInfoAsync(byte[] videoBytes, CancellationTo
         /// <param name="videoFrames">The video frames.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        public async Task<VideoResult> CreateVideoAsync(VideoFrames videoFrames, CancellationToken cancellationToken = default)
+        public async Task<VideoOutput> CreateVideoAsync(VideoFrames videoFrames, CancellationToken cancellationToken = default)
         {
-            return await CreateVideoInternalAsync(videoFrames.Frames, videoFrames.FPS, cancellationToken);
+            return await CreateVideoInternalAsync(videoFrames.Frames, videoFrames.Info.FPS, cancellationToken);
         }
 
 
@@ -101,7 +101,7 @@ public async Task<VideoResult> CreateVideoAsync(VideoFrames videoFrames, Cancell
         /// <param name="videoFPS">The video FPS.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        public async Task<VideoResult> CreateVideoAsync(IEnumerable<byte[]> videoFrames, float videoFPS, CancellationToken cancellationToken = default)
+        public async Task<VideoOutput> CreateVideoAsync(IEnumerable<byte[]> videoFrames, float videoFPS, CancellationToken cancellationToken = default)
         {
             return await CreateVideoInternalAsync(videoFrames, videoFPS, cancellationToken);
         }
@@ -139,7 +139,9 @@ public async Task<VideoFrames> CreateFramesAsync(VideoInput videoInput, float vi
         /// <returns></returns>
         public async Task<VideoFrames> CreateFramesAsync(byte[] videoBytes, float videoFPS, CancellationToken cancellationToken = default)
         {
-            return new VideoFrames(videoFPS, await CreateFramesInternalAsync(videoBytes, videoFPS, cancellationToken).ToListAsync(cancellationToken));
+            var videoInfo = await GetVideoInfoAsync(videoBytes, cancellationToken);
+            var videoFrames = await CreateFramesInternalAsync(videoBytes, videoFPS, cancellationToken).ToListAsync(cancellationToken);
+            return new VideoFrames(videoInfo, videoFrames);
         }
 
 
@@ -155,7 +157,10 @@ public async Task<VideoFrames> CreateFramesAsync(Stream videoStream, float video
             using (var memoryStream = new MemoryStream())
             {
                 await memoryStream.CopyToAsync(videoStream, cancellationToken).ConfigureAwait(false);
-                return new VideoFrames(videoFPS, await CreateFramesInternalAsync(memoryStream.ToArray(), videoFPS, cancellationToken).ToListAsync(cancellationToken));
+                var videoBytes = memoryStream.ToArray();
+                var videoInfo = await GetVideoInfoAsync(videoBytes, cancellationToken);
+                var videoFrames = await CreateFramesInternalAsync(videoBytes, videoFPS, cancellationToken).ToListAsync(cancellationToken);
+                return new VideoFrames(videoInfo, videoFrames);
             }
         }
 
@@ -197,7 +202,7 @@ private async Task<VideoInfo> GetVideoInfoInternalAsync(MemoryStream videoStream
         /// <param name="fps">The FPS.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        private async Task<VideoResult> CreateVideoInternalAsync(IEnumerable<byte[]> imageData, float fps = 15, CancellationToken cancellationToken = default)
+        private async Task<VideoOutput> CreateVideoInternalAsync(IEnumerable<byte[]> imageData, float fps = 15, CancellationToken cancellationToken = default)
         {
             string tempVideoPath = GetTempFilename();
             try
@@ -224,7 +229,7 @@ private async Task<VideoResult> CreateVideoInternalAsync(IEnumerable<byte[]> ima
 
                     // Analyze the result
                     var videoInfo = await GetVideoInfoAsync(videoResult);
-                    return new VideoResult(videoResult, videoInfo);
+                    return new VideoOutput(videoResult, videoInfo);
                 }
             }
             finally

OnnxStack.Core/Video/VideoFrames.cs

@@ -2,5 +2,5 @@
 
 namespace OnnxStack.Core.Video
 {
-    public record VideoFrames(float FPS, List<byte[]> Frames);
+    public record VideoFrames(VideoInfo Info, IReadOnlyList<byte[]> Frames);
 }

OnnxStack.Core/Video/VideoInfo.cs

@@ -2,5 +2,5 @@
 
 namespace OnnxStack.Core.Video
 {
-    public record VideoInfo(int Width, int Height, TimeSpan Duration, int Fps);
+    public record VideoInfo(int Width, int Height, TimeSpan Duration, int FPS);
 }

OnnxStack.Core/Video/VideoOutput.cs

@@ -0,0 +1,4 @@
+namespace OnnxStack.Core.Video
+{
+    public record VideoOutput(byte[] Data, VideoInfo Info);
+}

OnnxStack.Core/Video/VideoResult.cs

@@ -1,4 +1,5 @@
 using OnnxStack.Core.Image;
+using OnnxStack.Core.Video;
 using OnnxStack.StableDiffusion.Enums;
 using System.ComponentModel.DataAnnotations;
 
@@ -19,6 +20,9 @@ public class PromptOptions
 
         public InputImage InputImageMask { get; set; }
 
+        public VideoFrames InputVideo { get; set; }
+
+        public bool HasInputVideo => InputVideo?.Frames?.Count > 0;
         public bool HasInputImage => InputImage?.HasImage ?? false;
         public bool HasInputImageMask => InputImageMask?.HasImage ?? false;
     }

OnnxStack.StableDiffusion/Config/PromptOptions.cs

@@ -84,6 +84,8 @@ public record SchedulerOptions
         public float AestheticScore { get; set; } = 6f;
         public float AestheticNegativeScore { get; set; } = 2.5f;
 
+        public float VideoFPS { get; set; }
+
         public bool IsKarrasScheduler
         {
             get

OnnxStack.StableDiffusion/Config/SchedulerOptions.cs

@@ -17,6 +17,9 @@ public enum DiffuserType
         ImageInpaintLegacy = 3,
 
         [Description("Image To Animation")]
-        ImageToAnimation = 4
+        ImageToAnimation = 4,
+
+        [Description("Video To Video")]
+        VideoToVideo = 5
     }
 }

OnnxStack.StableDiffusion/Enums/DiffuserType.cs

@@ -1,6 +1,5 @@
 using Microsoft.ML.OnnxRuntime.Tensors;
 using OnnxStack.Core;
-using OnnxStack.Core.Config;
 using OnnxStack.Core.Services;
 using OnnxStack.StableDiffusion.Common;
 using OnnxStack.StableDiffusion.Config;
@@ -26,6 +25,7 @@ namespace OnnxStack.StableDiffusion.Services
     /// <seealso cref="OnnxStack.StableDiffusion.Common.IStableDiffusionService" />
     public sealed class StableDiffusionService : IStableDiffusionService
     {
+        private readonly IVideoService _videoService;
         private readonly IOnnxModelService _modelService;
         private readonly StableDiffusionConfig _configuration;
         private readonly ConcurrentDictionary<DiffuserPipelineType, IPipeline> _pipelines;
@@ -34,10 +34,11 @@ public sealed class StableDiffusionService : IStableDiffusionService
         /// Initializes a new instance of the <see cref="StableDiffusionService"/> class.
         /// </summary>
         /// <param name="schedulerService">The scheduler service.</param>
-        public StableDiffusionService(StableDiffusionConfig configuration, IOnnxModelService onnxModelService, IEnumerable<IPipeline> pipelines)
+        public StableDiffusionService(StableDiffusionConfig configuration, IOnnxModelService onnxModelService, IVideoService videoService, IEnumerable<IPipeline> pipelines)
         {
             _configuration = configuration;
             _modelService = onnxModelService;
+            _videoService = videoService;
             _pipelines = pipelines.ToConcurrentDictionary(k => k.PipelineType, k => k);
         }
 
@@ -115,9 +116,11 @@ public async Task<Image<Rgba32>> GenerateAsImageAsync(StableDiffusionModelSet mo
         /// <returns>The diffusion result as <see cref="byte[]"/></returns>
         public async Task<byte[]> GenerateAsBytesAsync(StableDiffusionModelSet model, PromptOptions prompt, SchedulerOptions options, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
         {
-            return await GenerateAsync(model, prompt, options, progressCallback, cancellationToken)
-                .ContinueWith(t => t.Result.ToImageBytes(), cancellationToken)
-                .ConfigureAwait(false);
+            var generateResult = await GenerateAsync(model, prompt, options, progressCallback, cancellationToken).ConfigureAwait(false);
+            if (!prompt.HasInputVideo)
+                return generateResult.ToImageBytes();
+
+            return await GetVideoResultAsBytesAsync(options, generateResult, cancellationToken).ConfigureAwait(false);
         }
 
 
@@ -131,9 +134,11 @@ public async Task<byte[]> GenerateAsBytesAsync(StableDiffusionModelSet model, Pr
         /// <returns>The diffusion result as <see cref="System.IO.Stream"/></returns>
         public async Task<Stream> GenerateAsStreamAsync(StableDiffusionModelSet model, PromptOptions prompt, SchedulerOptions options, Action<int, int> progressCallback = null, CancellationToken cancellationToken = default)
         {
-            return await GenerateAsync(model, prompt, options, progressCallback, cancellationToken)
-                .ContinueWith(t => t.Result.ToImageStream(), cancellationToken)
-                .ConfigureAwait(false);
+            var generateResult = await GenerateAsync(model, prompt, options, progressCallback, cancellationToken).ConfigureAwait(false);
+            if (!prompt.HasInputVideo)
+                return generateResult.ToImageStream();
+
+            return await GetVideoResultAsStreamAsync(options, generateResult, cancellationToken).ConfigureAwait(false);
         }
 
 
@@ -183,7 +188,12 @@ public async IAsyncEnumerable<Image<Rgba32>> GenerateBatchAsImageAsync(StableDif
         public async IAsyncEnumerable<byte[]> GenerateBatchAsBytesAsync(StableDiffusionModelSet modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions, Action<int, int, int, int> progressCallback = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
             await foreach (var result in GenerateBatchAsync(modelOptions, promptOptions, schedulerOptions, batchOptions, progressCallback, cancellationToken))
-                yield return result.ImageResult.ToImageBytes();
+            {
+                if (!promptOptions.HasInputVideo)
+                    yield return result.ImageResult.ToImageBytes();
+
+                yield return await GetVideoResultAsBytesAsync(schedulerOptions, result.ImageResult, cancellationToken).ConfigureAwait(false);
+            }
         }
 
 
@@ -200,7 +210,12 @@ public async IAsyncEnumerable<byte[]> GenerateBatchAsBytesAsync(StableDiffusionM
         public async IAsyncEnumerable<Stream> GenerateBatchAsStreamAsync(StableDiffusionModelSet modelOptions, PromptOptions promptOptions, SchedulerOptions schedulerOptions, BatchOptions batchOptions, Action<int, int, int, int> progressCallback = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
             await foreach (var result in GenerateBatchAsync(modelOptions, promptOptions, schedulerOptions, batchOptions, progressCallback, cancellationToken))
-                yield return result.ImageResult.ToImageStream();
+            {
+                if (!promptOptions.HasInputVideo)
+                    yield return result.ImageResult.ToImageStream();
+
+                yield return await GetVideoResultAsStreamAsync(schedulerOptions, result.ImageResult, cancellationToken).ConfigureAwait(false);
+            }
         }
 
 
@@ -237,6 +252,21 @@ private IAsyncEnumerable<BatchResult> DiffuseBatchAsync(StableDiffusionModelSet
             return diffuser.DiffuseBatchAsync(modelOptions, promptOptions, schedulerOptions, batchOptions, progress, cancellationToken);
         }
 
+        private async Task<byte[]> GetVideoResultAsBytesAsync(SchedulerOptions options, DenseTensor<float> tensorResult, CancellationToken cancellationToken = default)
+        {
+            var frameTensors = tensorResult
+                .Split(tensorResult.Dimensions[0])
+                .Select(x => x.ToImageBytes());
+
+            var videoResult = await _videoService.CreateVideoAsync(frameTensors, options.VideoFPS, cancellationToken);
+            return videoResult.Data;
+        }
+
+        private async Task<MemoryStream> GetVideoResultAsStreamAsync(SchedulerOptions options, DenseTensor<float> tensorResult, CancellationToken cancellationToken = default)
+        {
+            return new MemoryStream(await GetVideoResultAsBytesAsync(options, tensorResult, cancellationToken));
+        }
+
 
     }
 }