Wrap video bytes in class with ControlImage

saddam213 · saddam213 · commit d5414a3d5e8e · 2024-01-29T14:40:49.000+13:00
diff --git a/OnnxStack.Core/Services/IVideoService.cs b/OnnxStack.Core/Services/IVideoService.cs
@@ -104,6 +104,6 @@ public interface IVideoService
         /// <param name="targetFPS">The target FPS.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        IAsyncEnumerable<byte[]> StreamFramesAsync(byte[] videoBytes, float targetFPS, CancellationToken cancellationToken = default);
+        IAsyncEnumerable<VideoFrame> StreamFramesAsync(byte[] videoBytes, float targetFPS, CancellationToken cancellationToken = default);
     }
 }
diff --git a/OnnxStack.Core/Services/VideoService.cs b/OnnxStack.Core/Services/VideoService.cs
@@ -104,7 +104,10 @@ public async Task<VideoOutput> CreateVideoAsync(VideoFrames videoFrames, Cancell
         /// <returns></returns>
         public async Task<VideoOutput> CreateVideoAsync(DenseTensor<float> videoTensor, float videoFPS, CancellationToken cancellationToken = default)
         {
-            var videoFrames = await videoTensor.ToVideoFramesAsBytesAsync().ToListAsync(cancellationToken);
+            var videoFrames = await videoTensor
+                .ToVideoFramesAsBytesAsync()
+                .Select(x => new VideoFrame(x))
+                .ToListAsync(cancellationToken);
             return await CreateVideoInternalAsync(videoFrames, videoFPS, cancellationToken);
         }
 
@@ -118,7 +121,8 @@ public async Task<VideoOutput> CreateVideoAsync(DenseTensor<float> videoTensor,
         /// <returns></returns>
         public async Task<VideoOutput> CreateVideoAsync(IEnumerable<byte[]> videoFrames, float videoFPS, CancellationToken cancellationToken = default)
         {
-            return await CreateVideoInternalAsync(videoFrames, videoFPS, cancellationToken);
+            var frames = videoFrames.Select(x => new VideoFrame(x));
+            return await CreateVideoInternalAsync(frames, videoFPS, cancellationToken);
         }
 
 
@@ -190,7 +194,7 @@ public async Task<VideoFrames> CreateFramesAsync(Stream videoStream, float video
         /// <param name="targetFPS">The target FPS.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
-        public IAsyncEnumerable<byte[]> StreamFramesAsync(byte[] videoBytes, float targetFPS, CancellationToken cancellationToken = default)
+        public IAsyncEnumerable<VideoFrame> StreamFramesAsync(byte[] videoBytes, float targetFPS, CancellationToken cancellationToken = default)
         {
             return CreateFramesInternalAsync(videoBytes, targetFPS, cancellationToken);
         }
@@ -220,13 +224,13 @@ 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<VideoOutput> CreateVideoInternalAsync(IEnumerable<byte[]> imageData, float fps = 15, CancellationToken cancellationToken = default)
+        private async Task<VideoOutput> CreateVideoInternalAsync(IEnumerable<VideoFrame> imageData, float fps = 15, CancellationToken cancellationToken = default)
         {
             string tempVideoPath = GetTempFilename();
             try
             {
                 // Analyze first fram to get some details
-                var frameInfo = await GetVideoInfoAsync(imageData.First());
+                var frameInfo = await GetVideoInfoAsync(imageData.First().Frame);
                 var aspectRatio = (double)frameInfo.Width / frameInfo.Height;
                 using (var videoWriter = CreateWriter(tempVideoPath, fps, aspectRatio))
                 {
@@ -235,7 +239,7 @@ private async Task<VideoOutput> CreateVideoInternalAsync(IEnumerable<byte[]> ima
                     foreach (var image in imageData)
                     {
                         // Write each frame to the input stream of FFMPEG
-                        await videoWriter.StandardInput.BaseStream.WriteAsync(image, cancellationToken);
+                        await videoWriter.StandardInput.BaseStream.WriteAsync(image.Frame, cancellationToken);
                     }
 
                     // Done close stream and wait for app to process
@@ -265,7 +269,7 @@ private async Task<VideoOutput> CreateVideoInternalAsync(IEnumerable<byte[]> ima
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <returns></returns>
         /// <exception cref="Exception">Invalid PNG header</exception>
-        private async IAsyncEnumerable<byte[]> CreateFramesInternalAsync(byte[] videoData, float fps = 15, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        private async IAsyncEnumerable<VideoFrame> CreateFramesInternalAsync(byte[] videoData, float fps = 15, [EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
             string tempVideoPath = GetTempFilename();
             try
@@ -325,7 +329,7 @@ private async IAsyncEnumerable<byte[]> CreateFramesInternalAsync(byte[] videoDat
                                 break;
                         }
 
-                        yield return buffer[..currentIndex];
+                        yield return new VideoFrame(buffer[..currentIndex]);
                     }
 
                     if (cancellationToken.IsCancellationRequested)
diff --git a/OnnxStack.Core/Video/VideoFrame.cs b/OnnxStack.Core/Video/VideoFrame.cs
@@ -0,0 +1,6 @@
+﻿using OnnxStack.Core.Image;
+
+namespace OnnxStack.Core.Video
+{
+    public record VideoFrame(byte[] Frame, InputImage ControlImage = default);
+}
diff --git a/OnnxStack.Core/Video/VideoFrames.cs b/OnnxStack.Core/Video/VideoFrames.cs
@@ -2,5 +2,5 @@
 
 namespace OnnxStack.Core.Video
 {
-    public record VideoFrames(VideoInfo Info, IReadOnlyList<byte[]> Frames);
+    public record VideoFrames(VideoInfo Info, IReadOnlyList<VideoFrame> Frames);
 }
diff --git a/OnnxStack.ImageUpscaler/Services/UpscaleService.cs b/OnnxStack.ImageUpscaler/Services/UpscaleService.cs
@@ -265,7 +265,7 @@ public async Task<List<DenseTensor<float>>> GenerateInternalAsync(UpscaleModelSe
             var outputTensors = new List<DenseTensor<float>>();
             foreach (var frame in videoFrames.Frames)
             {
-                using (var imageFrame = Image.Load<Rgba32>(frame))
+                using (var imageFrame = Image.Load<Rgba32>(frame.Frame))
                 {
                     var input = CreateInputParams(imageFrame, modelSession.SampleSize, modelSession.ScaleFactor);
                     var outputDimension = new[] { 1, modelSession.Channels, 0, 0 };
diff --git a/OnnxStack.StableDiffusion/Helpers/ModelFactory.cs b/OnnxStack.StableDiffusion/Helpers/ModelFactory.cs
@@ -50,7 +50,6 @@ public static StableDiffusionModelSet CreateModelSet(string modelFolder, Diffuse
             var textEncoder2Config = default(TextEncoderModelConfig);
             var tokenizerConfig = new TokenizerModelConfig
             {
-                TokenizerLength = modelType == ModelType.Turbo ? 1024 : 768,
                 OnnxModelPath = tokenizerPath
             };
 
@@ -106,6 +105,10 @@ public static StableDiffusionModelSet CreateModelSet(string modelFolder, Diffuse
                 };
             }
 
+            // SD-Turbo has TokenizerLength 1024
+            if (pipeline == DiffuserPipelineType.StableDiffusion && modelType == ModelType.Turbo)
+                tokenizerConfig.TokenizerLength = 1024;
+
             var configuration = new StableDiffusionModelSet
             {
                 IsEnabled = true,
diff --git a/OnnxStack.StableDiffusion/Pipelines/Base/PipelineBase.cs b/OnnxStack.StableDiffusion/Pipelines/Base/PipelineBase.cs
@@ -155,8 +155,20 @@ protected async Task<DenseTensor<float>> DiffuseVideoAsync(IDiffuser diffuser, P
             foreach (var videoFrame in videoFrames)
             {
                 frameIndex++;
-                promptOptions.InputImage = promptOptions.DiffuserType == DiffuserType.ControlNet ? default : new InputImage(videoFrame);
-                promptOptions.InputContolImage = promptOptions.DiffuserType == DiffuserType.ImageToImage ? default : new InputImage(videoFrame);
+                //  byte[] videoFrame = videoFrames[i].Frame;
+                if (promptOptions.DiffuserType == DiffuserType.ControlNet || promptOptions.DiffuserType == DiffuserType.ControlNetImage)
+                {
+                    // ControlNetImage uses frame as input image
+                    if (promptOptions.DiffuserType == DiffuserType.ControlNetImage)
+                        promptOptions.InputImage = new InputImage(videoFrame.Frame);
+
+                    promptOptions.InputContolImage = videoFrame.ControlImage;
+                }
+                else
+                {
+                    promptOptions.InputImage = new InputImage(videoFrame.Frame);
+                }
+
                 var frameResultTensor = await diffuser.DiffuseAsync(promptOptions, schedulerOptions, promptEmbeddings, performGuidance, schedulerFrameCallback, cancellationToken);
 
                 // Frame Progress
@@ -223,5 +235,24 @@ protected Action<DiffusionProgress> CreateBatchCallback(Action<DiffusionProgress
                 BatchValue = batchIndex()
             });
         }
+
+
+        /// <summary>
+        /// Creates the pipeline from a ModelSet configuration.
+        /// </summary>
+        /// <param name="modelSet">The model set.</param>
+        /// <param name="logger">The logger.</param>
+        /// <returns></returns>
+        public static IPipeline CreatePipeline(StableDiffusionModelSet modelSet, ILogger logger = default)
+        {
+            return modelSet.PipelineType switch
+            {
+                DiffuserPipelineType.StableDiffusionXL => StableDiffusionXLPipeline.CreatePipeline(modelSet, logger),
+                DiffuserPipelineType.LatentConsistency => LatentConsistencyPipeline.CreatePipeline(modelSet, logger),
+                DiffuserPipelineType.LatentConsistencyXL => LatentConsistencyXLPipeline.CreatePipeline(modelSet, logger),
+                DiffuserPipelineType.InstaFlow => InstaFlowPipeline.CreatePipeline(modelSet, logger),
+                _ => StableDiffusionPipeline.CreatePipeline(modelSet, logger)
+            };
+        }
     }
 }
diff --git a/OnnxStack.StableDiffusion/Pipelines/StableDiffusionPipeline.cs b/OnnxStack.StableDiffusion/Pipelines/StableDiffusionPipeline.cs
@@ -123,14 +123,14 @@ public override Task LoadAsync()
         /// Unloads the pipeline.
         /// </summary>
         /// <returns></returns>
-        public override Task UnloadAsync()
+        public override async Task UnloadAsync()
         {
+            await Task.Yield();
             _unet?.Dispose();
             _tokenizer?.Dispose();
             _textEncoder?.Dispose();
             _vaeDecoder?.Dispose();
             _vaeEncoder?.Dispose();
-            return Task.CompletedTask;
         }
 
 
@@ -371,7 +371,7 @@ protected IEnumerable<int> PadWithBlankTokens(IEnumerable<int> inputs, int requi
         /// <param name="modelSet">The model set.</param>
         /// <param name="logger">The logger.</param>
         /// <returns></returns>
-        public static StableDiffusionPipeline CreatePipeline(StableDiffusionModelSet modelSet, ILogger logger = default)
+        public static new StableDiffusionPipeline CreatePipeline(StableDiffusionModelSet modelSet, ILogger logger = default)
         {
             var unet = new UNetConditionModel(modelSet.UnetConfig.ApplyDefaults(modelSet));
             var tokenizer = new TokenizerModel(modelSet.TokenizerConfig.ApplyDefaults(modelSet));
diff --git a/OnnxStack.StableDiffusion/Pipelines/StableDiffusionXLPipeline.cs b/OnnxStack.StableDiffusion/Pipelines/StableDiffusionXLPipeline.cs
@@ -43,7 +43,7 @@ public StableDiffusionXLPipeline(string name, TokenizerModel tokenizer, Tokenize
             {
                 SchedulerType.Euler,
                 SchedulerType.EulerAncestral,
-                SchedulerType.DDIM,
+                SchedulerType.DDPM,
                 SchedulerType.KDPM2
             };
             _defaultSchedulerOptions = defaultSchedulerOptions ?? new SchedulerOptions
diff --git a/OnnxStack.StableDiffusion/Registration.cs b/OnnxStack.StableDiffusion/Registration.cs
@@ -3,8 +3,6 @@
 using OnnxStack.Core.Services;
 using OnnxStack.StableDiffusion.Common;
 using OnnxStack.StableDiffusion.Config;
-using OnnxStack.StableDiffusion.Diffusers;
-using OnnxStack.StableDiffusion.Pipelines;
 using OnnxStack.StableDiffusion.Services;
 using SixLabors.ImageSharp;
 using SixLabors.ImageSharp.Memory;
diff --git a/OnnxStack.StableDiffusion/Services/StableDiffusionService.cs b/OnnxStack.StableDiffusion/Services/StableDiffusionService.cs
@@ -70,13 +70,13 @@ public async Task<bool> LoadModelAsync(StableDiffusionModelSet model)
         /// </summary>
         /// <param name="model">The model.</param>
         /// <returns></returns>
-        public Task<bool> UnloadModelAsync(StableDiffusionModelSet model)
+        public async Task<bool> UnloadModelAsync(StableDiffusionModelSet model)
         {
             if (_pipelines.Remove(model, out var pipeline))
             {
-                pipeline?.UnloadAsync();
+                await pipeline?.UnloadAsync();
             }
-            return Task.FromResult(true);
+            return true;
         }
 
 

Original file line number	Diff line number	Diff line change
`@@ -104,6 +104,6 @@ public interface IVideoService`
`104`	`104`	`/// <param name="targetFPS">The target FPS.</param>`
`105`	`105`	`/// <param name="cancellationToken">The cancellation token.</param>`
`106`	`106`	`/// <returns></returns>`
`107`		`- IAsyncEnumerable<byte[]> StreamFramesAsync(byte[] videoBytes, float targetFPS, CancellationToken cancellationToken = default);`
	`107`	`+ IAsyncEnumerable<VideoFrame> StreamFramesAsync(byte[] videoBytes, float targetFPS, CancellationToken cancellationToken = default);`
`108`	`108`	`}`
`109`	`109`	`}`
Original file line number	Diff line number	Diff line change
`@@ -104,7 +104,10 @@ public async Task<VideoOutput> CreateVideoAsync(VideoFrames videoFrames, Cancell`
`104`	`104`	`/// <returns></returns>`
`105`	`105`	`public async Task<VideoOutput> CreateVideoAsync(DenseTensor<float> videoTensor, float videoFPS, CancellationToken cancellationToken = default)`
`106`	`106`	`{`
`107`		`- var videoFrames = await videoTensor.ToVideoFramesAsBytesAsync().ToListAsync(cancellationToken);`
	`107`	`+ var videoFrames = await videoTensor`
	`108`	`+ .ToVideoFramesAsBytesAsync()`
	`109`	`+ .Select(x => new VideoFrame(x))`
	`110`	`+ .ToListAsync(cancellationToken);`
`108`	`111`	`return await CreateVideoInternalAsync(videoFrames, videoFPS, cancellationToken);`
`109`	`112`	`}`
`110`	`113`
`@@ -118,7 +121,8 @@ public async Task<VideoOutput> CreateVideoAsync(DenseTensor<float> videoTensor,`
`118`	`121`	`/// <returns></returns>`
`119`	`122`	`public async Task<VideoOutput> CreateVideoAsync(IEnumerable<byte[]> videoFrames, float videoFPS, CancellationToken cancellationToken = default)`
`120`	`123`	`{`
`121`		`- return await CreateVideoInternalAsync(videoFrames, videoFPS, cancellationToken);`
	`124`	`+ var frames = videoFrames.Select(x => new VideoFrame(x));`
	`125`	`+ return await CreateVideoInternalAsync(frames, videoFPS, cancellationToken);`
`122`	`126`	`}`
`123`	`127`
`124`	`128`
`@@ -190,7 +194,7 @@ public async Task<VideoFrames> CreateFramesAsync(Stream videoStream, float video`
`190`	`194`	`/// <param name="targetFPS">The target FPS.</param>`
`191`	`195`	`/// <param name="cancellationToken">The cancellation token.</param>`
`192`	`196`	`/// <returns></returns>`
`193`		`- public IAsyncEnumerable<byte[]> StreamFramesAsync(byte[] videoBytes, float targetFPS, CancellationToken cancellationToken = default)`
	`197`	`+ public IAsyncEnumerable<VideoFrame> StreamFramesAsync(byte[] videoBytes, float targetFPS, CancellationToken cancellationToken = default)`
`194`	`198`	`{`
`195`	`199`	`return CreateFramesInternalAsync(videoBytes, targetFPS, cancellationToken);`
`196`	`200`	`}`
`@@ -220,13 +224,13 @@ private async Task<VideoInfo> GetVideoInfoInternalAsync(MemoryStream videoStream`
`220`	`224`	`/// <param name="fps">The FPS.</param>`
`221`	`225`	`/// <param name="cancellationToken">The cancellation token.</param>`
`222`	`226`	`/// <returns></returns>`
`223`		`- private async Task<VideoOutput> CreateVideoInternalAsync(IEnumerable<byte[]> imageData, float fps = 15, CancellationToken cancellationToken = default)`
	`227`	`+ private async Task<VideoOutput> CreateVideoInternalAsync(IEnumerable<VideoFrame> imageData, float fps = 15, CancellationToken cancellationToken = default)`
`224`	`228`	`{`
`225`	`229`	`string tempVideoPath = GetTempFilename();`
`226`	`230`	`try`
`227`	`231`	`{`
`228`	`232`	`// Analyze first fram to get some details`
`229`		`- var frameInfo = await GetVideoInfoAsync(imageData.First());`
	`233`	`+ var frameInfo = await GetVideoInfoAsync(imageData.First().Frame);`
`230`	`234`	`var aspectRatio = (double)frameInfo.Width / frameInfo.Height;`
`231`	`235`	`using (var videoWriter = CreateWriter(tempVideoPath, fps, aspectRatio))`
`232`	`236`	`{`
`@@ -235,7 +239,7 @@ private async Task<VideoOutput> CreateVideoInternalAsync(IEnumerable<byte[]> ima`
`235`	`239`	`foreach (var image in imageData)`
`236`	`240`	`{`
`237`	`241`	`// Write each frame to the input stream of FFMPEG`
`238`		`- await videoWriter.StandardInput.BaseStream.WriteAsync(image, cancellationToken);`
	`242`	`+ await videoWriter.StandardInput.BaseStream.WriteAsync(image.Frame, cancellationToken);`
`239`	`243`	`}`
`240`	`244`
`241`	`245`	`// Done close stream and wait for app to process`
`@@ -265,7 +269,7 @@ private async Task<VideoOutput> CreateVideoInternalAsync(IEnumerable<byte[]> ima`
`265`	`269`	`/// <param name="cancellationToken">The cancellation token.</param>`
`266`	`270`	`/// <returns></returns>`
`267`	`271`	`/// <exception cref="Exception">Invalid PNG header</exception>`
`268`		`- private async IAsyncEnumerable<byte[]> CreateFramesInternalAsync(byte[] videoData, float fps = 15, [EnumeratorCancellation] CancellationToken cancellationToken = default)`
	`272`	`+ private async IAsyncEnumerable<VideoFrame> CreateFramesInternalAsync(byte[] videoData, float fps = 15, [EnumeratorCancellation] CancellationToken cancellationToken = default)`
`269`	`273`	`{`
`270`	`274`	`string tempVideoPath = GetTempFilename();`
`271`	`275`	`try`
`@@ -325,7 +329,7 @@ private async IAsyncEnumerable<byte[]> CreateFramesInternalAsync(byte[] videoDat`
`325`	`329`	`break;`
`326`	`330`	`}`
`327`	`331`
`328`		`- yield return buffer[..currentIndex];`
	`332`	`+ yield return new VideoFrame(buffer[..currentIndex]);`
`329`	`333`	`}`
`330`	`334`
`331`	`335`	`if (cancellationToken.IsCancellationRequested)`
Original file line number	Diff line number	Diff line change
`@@ -2,5 +2,5 @@`
`2`	`2`
`3`	`3`	`namespace OnnxStack.Core.Video`
`4`	`4`	`{`
`5`		`- public record VideoFrames(VideoInfo Info, IReadOnlyList<byte[]> Frames);`
	`5`	`+ public record VideoFrames(VideoInfo Info, IReadOnlyList<VideoFrame> Frames);`
`6`	`6`	`}`
Original file line number	Diff line number	Diff line change
`@@ -265,7 +265,7 @@ public async Task<List<DenseTensor<float>>> GenerateInternalAsync(UpscaleModelSe`
`265`	`265`	`var outputTensors = new List<DenseTensor<float>>();`
`266`	`266`	`foreach (var frame in videoFrames.Frames)`
`267`	`267`	`{`
`268`		`- using (var imageFrame = Image.Load<Rgba32>(frame))`
	`268`	`+ using (var imageFrame = Image.Load<Rgba32>(frame.Frame))`
`269`	`269`	`{`
`270`	`270`	`var input = CreateInputParams(imageFrame, modelSession.SampleSize, modelSession.ScaleFactor);`
`271`	`271`	`var outputDimension = new[] { 1, modelSession.Channels, 0, 0 };`
Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,6 @@ public static StableDiffusionModelSet CreateModelSet(string modelFolder, Diffuse`
`50`	`50`	`var textEncoder2Config = default(TextEncoderModelConfig);`
`51`	`51`	`var tokenizerConfig = new TokenizerModelConfig`
`52`	`52`	`{`
`53`		`- TokenizerLength = modelType == ModelType.Turbo ? 1024 : 768,`
`54`	`53`	`OnnxModelPath = tokenizerPath`
`55`	`54`	`};`
`56`	`55`
`@@ -106,6 +105,10 @@ public static StableDiffusionModelSet CreateModelSet(string modelFolder, Diffuse`
`106`	`105`	`};`
`107`	`106`	`}`
`108`	`107`
	`108`	`+ // SD-Turbo has TokenizerLength 1024`
	`109`	`+ if (pipeline == DiffuserPipelineType.StableDiffusion && modelType == ModelType.Turbo)`
	`110`	`+ tokenizerConfig.TokenizerLength = 1024;`
	`111`	`+`
`109`	`112`	`var configuration = new StableDiffusionModelSet`
`110`	`113`	`{`
`111`	`114`	`IsEnabled = true,`
Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ public StableDiffusionXLPipeline(string name, TokenizerModel tokenizer, Tokenize`
`43`	`43`	`{`
`44`	`44`	`SchedulerType.Euler,`
`45`	`45`	`SchedulerType.EulerAncestral,`
`46`		`- SchedulerType.DDIM,`
	`46`	`+ SchedulerType.DDPM,`
`47`	`47`	`SchedulerType.KDPM2`
`48`	`48`	`};`
`49`	`49`	`_defaultSchedulerOptions = defaultSchedulerOptions ?? new SchedulerOptions`