Skip to content
This repository was archived by the owner on Nov 27, 2024. It is now read-only.

Commit e96fe8e

Browse files
authored
Merge pull request #85 from saddam213/Video
Video Support
2 parents 656bf3f + ddfb9c8 commit e96fe8e

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

49 files changed

+2257
-178
lines changed

OnnxStack.Console/Examples/StableDiffusionBatch.cs

+7-5
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
1-
using OnnxStack.StableDiffusion.Common;
1+
using OnnxStack.Core.Image;
2+
using OnnxStack.StableDiffusion.Common;
23
using OnnxStack.StableDiffusion.Config;
34
using OnnxStack.StableDiffusion.Enums;
45
using OnnxStack.StableDiffusion.Helpers;
6+
using OnnxStack.StableDiffusion.Models;
57
using SixLabors.ImageSharp;
68

79
namespace OnnxStack.Console.Runner
@@ -58,13 +60,13 @@ public async Task RunAsync()
5860
await _stableDiffusionService.LoadModelAsync(model);
5961

6062
var batchIndex = 0;
61-
var callback = (int batch, int batchCount, int step, int steps) =>
63+
var callback = (DiffusionProgress progress) =>
6264
{
63-
batchIndex = batch;
64-
OutputHelpers.WriteConsole($"Image: {batch}/{batchCount} - Step: {step}/{steps}", ConsoleColor.Cyan);
65+
batchIndex = progress.BatchValue;
66+
OutputHelpers.WriteConsole($"Image: {progress.BatchValue}/{progress.BatchMax} - Step: {progress.StepValue}/{progress.StepMax}", ConsoleColor.Cyan);
6567
};
6668

67-
await foreach (var result in _stableDiffusionService.GenerateBatchAsync(model, promptOptions, schedulerOptions, batchOptions, callback))
69+
await foreach (var result in _stableDiffusionService.GenerateBatchAsync(model, promptOptions, schedulerOptions, batchOptions, default))
6870
{
6971
var outputFilename = Path.Combine(_outputDirectory, $"{batchIndex}_{result.SchedulerOptions.Seed}.png");
7072
var image = result.ImageResult.ToImage();

OnnxStack.Console/OnnxStack.Console.csproj

+6-7
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
<ImplicitUsings>enable</ImplicitUsings>
77
<Nullable>disable</Nullable>
88
<PlatformTarget>x64</PlatformTarget>
9-
<Configurations>Debug;Release;Debug-DirectML;Debug-Cuda;Debug-TensorRT;Release-DirectML;Release-Cuda;Release-TensorRT</Configurations>
9+
<Configurations>Debug;Release;Debug-Cuda;Debug-TensorRT;Release-Cuda;Release-TensorRT</Configurations>
1010
</PropertyGroup>
1111

1212
<ItemGroup>
@@ -15,17 +15,16 @@
1515
</ItemGroup>
1616

1717
<ItemGroup>
18-
<PackageReference Include="OnnxStack.StableDiffusion" Version="0.14.0" Condition=" '$(Configuration)' == 'Release' OR '$(Configuration)' == 'Release-DirectML' OR '$(Configuration)' == 'Release-Cuda' OR '$(Configuration)' == 'Release-TensorRT'" />
19-
<PackageReference Include="OnnxStack.ImageUpscaler" Version="0.14.0" Condition=" '$(Configuration)' == 'Release' OR '$(Configuration)' == 'Release-DirectML' OR '$(Configuration)' == 'Release-Cuda' OR '$(Configuration)' == 'Release-TensorRT'" />
20-
<ProjectReference Include="..\OnnxStack.StableDiffusion\OnnxStack.StableDiffusion.csproj" Condition=" '$(Configuration)' == 'Debug' OR '$(Configuration)' == 'Debug-DirectML' OR '$(Configuration)' == 'Debug-Cuda' OR '$(Configuration)' == 'Debug-TensorRT'" />
21-
<ProjectReference Include="..\OnnxStack.ImageUpscaler\OnnxStack.ImageUpscaler.csproj" Condition=" '$(Configuration)' == 'Debug' OR '$(Configuration)' == 'Debug-DirectML' OR '$(Configuration)' == 'Debug-Cuda' OR '$(Configuration)' == 'Debug-TensorRT'" />
18+
<PackageReference Include="OnnxStack.StableDiffusion" Version="0.14.0" Condition=" '$(Configuration)' == 'Release' OR '$(Configuration)' == 'Release-Cuda' OR '$(Configuration)' == 'Release-TensorRT'" />
19+
<PackageReference Include="OnnxStack.ImageUpscaler" Version="0.14.0" Condition=" '$(Configuration)' == 'Release' OR '$(Configuration)' == 'Release-Cuda' OR '$(Configuration)' == 'Release-TensorRT'" />
20+
<ProjectReference Include="..\OnnxStack.StableDiffusion\OnnxStack.StableDiffusion.csproj" Condition=" '$(Configuration)' == 'Debug' OR '$(Configuration)' == 'Debug-Cuda' OR '$(Configuration)' == 'Debug-TensorRT'" />
21+
<ProjectReference Include="..\OnnxStack.ImageUpscaler\OnnxStack.ImageUpscaler.csproj" Condition=" '$(Configuration)' == 'Debug' OR '$(Configuration)' == 'Debug-Cuda' OR '$(Configuration)' == 'Debug-TensorRT'" />
2222
</ItemGroup>
2323

2424
<ItemGroup>
25-
<PackageReference Include="Microsoft.ML.OnnxRuntime" Version="1.16.3" Condition=" '$(Configuration)' == 'Debug' OR '$(Configuration)' == 'Release' " />
25+
<PackageReference Include="Microsoft.ML.OnnxRuntime.DirectML" Version="1.16.3" Condition=" '$(Configuration)' == 'Debug' OR '$(Configuration)' == 'Release' " />
2626
<PackageReference Include="Microsoft.ML.OnnxRuntime.Gpu" Version="1.16.3" Condition=" '$(Configuration)' == 'Debug-TensorRT' OR '$(Configuration)' == 'Release-TensorRT'" />
2727
<PackageReference Include="Microsoft.ML.OnnxRuntime.Gpu" Version="1.16.3" Condition=" '$(Configuration)' == 'Debug-Cuda' OR '$(Configuration)' == 'Release-Cuda'" />
28-
<PackageReference Include="Microsoft.ML.OnnxRuntime.DirectML" Version="1.16.3" Condition=" '$(Configuration)' == 'Debug-DirectML' OR '$(Configuration)' == 'Release-DirectML'" />
2928
</ItemGroup>
3029

3130
<ItemGroup>

OnnxStack.Core/Config/OnnxStackConfig.cs

+4
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,10 @@ namespace OnnxStack.Core.Config
44
{
55
public class OnnxStackConfig : IConfigSection
66
{
7+
public string TempPath { get; set; } = ".temp";
8+
public string FFmpegPath { get; set; } = "ffmpeg.exe";
9+
public string FFprobePath { get; set; } = "ffprobe.exe";
10+
711
public void Initialize()
812
{
913
}

OnnxStack.Core/Image/Extensions.cs

+73
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
using Microsoft.ML.OnnxRuntime.Tensors;
2+
using SixLabors.ImageSharp.PixelFormats;
3+
using SixLabors.ImageSharp;
4+
using System;
5+
using System.Collections.Generic;
6+
using System.Linq;
7+
using System.Text;
8+
using System.Threading.Tasks;
9+
using System.IO;
10+
11+
namespace OnnxStack.Core.Image
12+
{
13+
public static class Extensions
14+
{
15+
public static Image<Rgba32> ToImage(this DenseTensor<float> imageTensor)
16+
{
17+
var height = imageTensor.Dimensions[2];
18+
var width = imageTensor.Dimensions[3];
19+
var hasAlpha = imageTensor.Dimensions[1] == 4;
20+
var result = new Image<Rgba32>(width, height);
21+
for (var y = 0; y < height; y++)
22+
{
23+
for (var x = 0; x < width; x++)
24+
{
25+
result[x, y] = new Rgba32(
26+
CalculateByte(imageTensor, 0, y, x),
27+
CalculateByte(imageTensor, 1, y, x),
28+
CalculateByte(imageTensor, 2, y, x),
29+
hasAlpha ? CalculateByte(imageTensor, 3, y, x) : byte.MaxValue
30+
);
31+
}
32+
}
33+
return result;
34+
}
35+
36+
/// <summary>
37+
/// Converts to image byte array.
38+
/// </summary>
39+
/// <param name="imageTensor">The image tensor.</param>
40+
/// <returns></returns>
41+
public static byte[] ToImageBytes(this DenseTensor<float> imageTensor)
42+
{
43+
using (var image = imageTensor.ToImage())
44+
using (var memoryStream = new MemoryStream())
45+
{
46+
image.SaveAsPng(memoryStream);
47+
return memoryStream.ToArray();
48+
}
49+
}
50+
51+
/// <summary>
52+
/// Converts to image byte array.
53+
/// </summary>
54+
/// <param name="imageTensor">The image tensor.</param>
55+
/// <returns></returns>
56+
public static async Task<byte[]> ToImageBytesAsync(this DenseTensor<float> imageTensor)
57+
{
58+
using (var image = imageTensor.ToImage())
59+
using (var memoryStream = new MemoryStream())
60+
{
61+
await image.SaveAsPngAsync(memoryStream);
62+
return memoryStream.ToArray();
63+
}
64+
}
65+
66+
67+
private static byte CalculateByte(Tensor<float> imageTensor, int index, int y, int x)
68+
{
69+
return (byte)Math.Round(Math.Clamp(imageTensor[0, index, y, x] / 2 + 0.5, 0, 1) * 255);
70+
}
71+
72+
}
73+
}

OnnxStack.Core/OnnxStack.Core.csproj

+2
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,15 @@
3535
</ItemGroup>
3636

3737
<ItemGroup>
38+
<PackageReference Include="FFMpegCore" Version="5.1.0" />
3839
<PackageReference Include="Microsoft.Extensions.DependencyInjection.Abstractions" Version="8.0.0" />
3940
<PackageReference Include="Microsoft.Extensions.Hosting.Abstractions" Version="8.0.0" />
4041
<PackageReference Include="Microsoft.Extensions.Logging.Abstractions" Version="8.0.0" />
4142
<PackageReference Include="Microsoft.ML" Version="3.0.0" />
4243
<PackageReference Include="Microsoft.ML.OnnxRuntime.Extensions" Version="0.9.0" />
4344
<PackageReference Include="Microsoft.ML.OnnxRuntime.Managed" Version="1.16.3" />
4445
<PackageReference Include="SixLabors.ImageSharp" Version="3.1.1" />
46+
<PackageReference Include="System.Linq.Async" Version="6.0.1" />
4547
</ItemGroup>
4648

4749
<ItemGroup>
+109
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
using Microsoft.ML.OnnxRuntime.Tensors;
2+
using OnnxStack.Core.Video;
3+
using System.Collections.Generic;
4+
using System.IO;
5+
using System.Threading;
6+
using System.Threading.Tasks;
7+
8+
namespace OnnxStack.Core.Services
9+
{
10+
/// <summary>
11+
/// Service with basic handling of video for use in OnnxStack, Frame->Video and Video->Frames
12+
/// </summary>
13+
public interface IVideoService
14+
{
15+
/// <summary>
16+
/// Gets the video information asynchronous.
17+
/// </summary>
18+
/// <param name="videoBytes">The video bytes.</param>
19+
/// <param name="cancellationToken">The cancellation token.</param>
20+
/// <returns></returns>
21+
Task<VideoInfo> GetVideoInfoAsync(byte[] videoBytes, CancellationToken cancellationToken = default);
22+
23+
/// <summary>
24+
/// Gets the video information asynchronous.
25+
/// </summary>
26+
/// <param name="videoStream">The video stream.</param>
27+
/// <param name="cancellationToken">The cancellation token.</param>
28+
/// <returns></returns>
29+
Task<VideoInfo> GetVideoInfoAsync(Stream videoStream, CancellationToken cancellationToken = default);
30+
31+
/// <summary>
32+
/// Gets the video information, Size, FPS, Duration etc.
33+
/// </summary>
34+
/// <param name="videoInput">The video input.</param>
35+
/// <param name="cancellationToken">The cancellation token.</param>
36+
/// <returns></returns>
37+
/// <exception cref="ArgumentException">No video data found</exception>
38+
Task<VideoInfo> GetVideoInfoAsync(VideoInput videoInput, CancellationToken cancellationToken = default);
39+
40+
41+
/// <summary>
42+
/// Creates a collection of PNG frames from a video source
43+
/// </summary>
44+
/// <param name="videoBytes">The video bytes.</param>
45+
/// <param name="videoFPS">The video FPS.</param>
46+
/// <param name="cancellationToken">The cancellation token.</param>
47+
/// <returns></returns>
48+
Task<VideoFrames> CreateFramesAsync(byte[] videoBytes, float videoFPS, CancellationToken cancellationToken = default);
49+
50+
51+
/// <summary>
52+
/// Creates a collection of PNG frames from a video source
53+
/// </summary>
54+
/// <param name="videoStream">The video stream.</param>
55+
/// <param name="videoFPS">The video FPS.</param>
56+
/// <param name="cancellationToken">The cancellation token.</param>
57+
/// <returns></returns>
58+
Task<VideoFrames> CreateFramesAsync(Stream videoStream, float videoFPS, CancellationToken cancellationToken = default);
59+
60+
61+
/// <summary>
62+
/// Creates a collection of PNG frames from a video source
63+
/// </summary>
64+
/// <param name="videoInput">The video input.</param>
65+
/// <param name="videoFPS">The video FPS.</param>
66+
/// <param name="cancellationToken">The cancellation token.</param>
67+
/// <returns></returns>
68+
/// <exception cref="NotSupportedException">VideoTensor not supported</exception>
69+
/// <exception cref="ArgumentException">No video data found</exception>
70+
Task<VideoFrames> CreateFramesAsync(VideoInput videoInput, float videoFPS, CancellationToken cancellationToken = default);
71+
72+
73+
/// <summary>
74+
/// Creates and MP4 video from a collection of PNG images.
75+
/// </summary>
76+
/// <param name="videoFrames">The video frames.</param>
77+
/// <param name="videoFPS">The video FPS.</param>
78+
/// <param name="cancellationToken">The cancellation token.</param>
79+
/// <returns></returns>
80+
Task<VideoOutput> CreateVideoAsync(IEnumerable<byte[]> videoFrames, float videoFPS, CancellationToken cancellationToken = default);
81+
82+
83+
/// <summary>
84+
/// Creates and MP4 video from a collection of PNG images.
85+
/// </summary>
86+
/// <param name="videoFrames">The video frames.</param>
87+
/// <param name="cancellationToken">The cancellation token.</param>
88+
/// <returns></returns>
89+
Task<VideoOutput> CreateVideoAsync(VideoFrames videoFrames, CancellationToken cancellationToken = default);
90+
91+
// <summary>
92+
/// Creates and MP4 video from a collection of PNG images.
93+
/// </summary>
94+
/// <param name="videoTensor">The video frames.</param>
95+
/// <param name="videoFPS">The video FPS.</param>
96+
/// <param name="cancellationToken">The cancellation token.</param>
97+
/// <returns></returns>
98+
Task<VideoOutput> CreateVideoAsync(DenseTensor<float> videoTensor, float videoFPS, CancellationToken cancellationToken = default);
99+
100+
/// <summary>
101+
/// Streams frames as PNG as they are processed from a video source
102+
/// </summary>
103+
/// <param name="videoBytes">The video bytes.</param>
104+
/// <param name="targetFPS">The target FPS.</param>
105+
/// <param name="cancellationToken">The cancellation token.</param>
106+
/// <returns></returns>
107+
IAsyncEnumerable<byte[]> StreamFramesAsync(byte[] videoBytes, float targetFPS, CancellationToken cancellationToken = default);
108+
}
109+
}

0 commit comments

Comments
 (0)