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

Commit 000ee0d

Browse files
saddam213saddam213
committed
Video Upscale process
1 parent 72473e8 commit 000ee0d

File tree

6 files changed

+149
-32
lines changed

6 files changed

+149
-32
lines changed

‎OnnxStack.Core/Extensions/TensorExtension.cs‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
using Microsoft.ML.OnnxRuntime.Tensors;
2+
using System;
3+
4+
namespace OnnxStack.Core
5+
{
6+
public static class TensorExtension
7+
{
8+
/// <summary>
9+
/// Concatenates the specified tensors along the 0 axis.
10+
/// </summary>
11+
/// <param name="tensor1">The tensor1.</param>
12+
/// <param name="tensor2">The tensor2.</param>
13+
/// <param name="axis">The axis.</param>
14+
/// <returns></returns>
15+
/// <exception cref="System.NotImplementedException">Only axis 0 is supported</exception>
16+
public static DenseTensor<float> Concatenate(this DenseTensor<float> tensor1, DenseTensor<float> tensor2, int axis = 0)
17+
{
18+
if (tensor1 == null)
19+
return tensor2.ToDenseTensor();
20+
21+
if (axis != 0 && axis != 2)
22+
throw new NotImplementedException("Only axis 0, 2 is supported");
23+
24+
if (axis == 2)
25+
return Concatenate(tensor1, tensor2);
26+
27+
var dimensions = tensor1.Dimensions.ToArray();
28+
dimensions[0] += tensor2.Dimensions[0];
29+
30+
var buffer = new float[tensor1.Length + tensor2.Length].AsMemory();
31+
tensor1.Buffer.CopyTo(buffer[..(int)tensor1.Length]);
32+
tensor2.Buffer.CopyTo(buffer[(int)tensor1.Length..]);
33+
return new DenseTensor<float>(buffer, dimensions);
34+
}
35+
}
36+
}

‎OnnxStack.Core/Services/VideoService.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,8 @@ public async Task<VideoFrames> CreateFramesAsync(VideoInput videoInput, float vi
138138
return await CreateFramesAsync(videoInput.VideoBytes, videoFPS, cancellationToken);
139139
if (videoInput.VideoStream is not null)
140140
return await CreateFramesAsync(videoInput.VideoStream, videoFPS, cancellationToken);
141+
if (videoInput.VideoFrames is not null)
142+
return videoInput.VideoFrames;
141143
if (videoInput.VideoTensor is not null)
142144
throw new NotSupportedException("VideoTensor not supported");
143145

‎OnnxStack.ImageUpscaler/Registration.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
using Microsoft.Extensions.DependencyInjection;
22
using OnnxStack.Core;
33
using OnnxStack.Core.Config;
4+
using OnnxStack.Core.Services;
45
using OnnxStack.ImageUpscaler.Config;
5-
using OnnxStack.ImageUpscaler.Services;
66

77
namespace OnnxStack.ImageUpscaler
88
{
@@ -38,7 +38,7 @@ public static void AddOnnxStackImageUpscaler(this IServiceCollection serviceColl
3838
/// <param name="serviceCollection">The service collection.</param>
3939
private static void RegisterServices(this IServiceCollection serviceCollection)
4040
{
41-
serviceCollection.AddSingleton<IUpscaleService, UpscaleService>();
41+
serviceCollection.AddSingleton<IVideoService, VideoService>();
4242
}
4343

4444

‎OnnxStack.ImageUpscaler/Services/IUpscaleService.cs‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using Microsoft.ML.OnnxRuntime.Tensors;
22
using OnnxStack.Core.Image;
3+
using OnnxStack.Core.Video;
34
using OnnxStack.StableDiffusion.Config;
45
using SixLabors.ImageSharp;
56
using SixLabors.ImageSharp.PixelFormats;
@@ -67,5 +68,32 @@ public interface IUpscaleService
6768
/// <param name="inputImage">The input image.</param>
6869
/// <returns></returns>
6970
Task<Stream> GenerateAsStreamAsync(UpscaleModelSet modelOptions, InputImage inputImage);
71+
72+
73+
/// <summary>
74+
/// Generates the upscaled video.
75+
/// </summary>
76+
/// <param name="modelOptions">The model options.</param>
77+
/// <param name="videoInput">The video input.</param>
78+
/// <returns></returns>
79+
Task<DenseTensor<float>> GenerateAsync(UpscaleModelSet modelOptions, VideoInput videoInput);
80+
81+
82+
/// <summary>
83+
/// Generates the upscaled video.
84+
/// </summary>
85+
/// <param name="modelOptions">The model options.</param>
86+
/// <param name="videoInput">The video input.</param>
87+
/// <returns></returns>
88+
Task<byte[]> GenerateAsByteAsync(UpscaleModelSet modelOptions, VideoInput videoInput);
89+
90+
91+
/// <summary>
92+
/// Generates the upscaled video.
93+
/// </summary>
94+
/// <param name="modelOptions">The model options.</param>
95+
/// <param name="videoInput">The video input.</param>
96+
/// <returns></returns>
97+
Task<Stream> GenerateAsStreamAsync(UpscaleModelSet modelOptions, VideoInput videoInput);
7098
}
7199
}

‎OnnxStack.ImageUpscaler/Services/UpscaleService.cs‎

Lines changed: 80 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,14 @@
44
using OnnxStack.Core.Image;
55
using OnnxStack.Core.Model;
66
using OnnxStack.Core.Services;
7+
using OnnxStack.Core.Video;
78
using OnnxStack.ImageUpscaler.Config;
89
using OnnxStack.ImageUpscaler.Extensions;
910
using OnnxStack.ImageUpscaler.Models;
1011
using OnnxStack.StableDiffusion.Config;
1112
using SixLabors.ImageSharp;
1213
using SixLabors.ImageSharp.PixelFormats;
1314
using SixLabors.ImageSharp.Processing;
14-
using System;
1515
using System.Collections.Generic;
1616
using System.IO;
1717
using System.Linq;
@@ -23,17 +23,19 @@ public class UpscaleService : IUpscaleService
2323
{
2424
private readonly IOnnxModelService _modelService;
2525
private readonly ImageUpscalerConfig _configuration;
26+
private readonly IVideoService _videoService;
2627

2728
/// <summary>
2829
/// Initializes a new instance of the <see cref="UpscaleService"/> class.
2930
/// </summary>
3031
/// <param name="configuration">The configuration.</param>
3132
/// <param name="modelService">The model service.</param>
3233
/// <param name="imageService">The image service.</param>
33-
public UpscaleService(ImageUpscalerConfig configuration, IOnnxModelService modelService)
34+
public UpscaleService(ImageUpscalerConfig configuration, IOnnxModelService modelService, IVideoService videoService)
3435
{
3536
_configuration = configuration;
3637
_modelService = modelService;
38+
_videoService = videoService;
3739
}
3840

3941

@@ -129,6 +131,63 @@ public async Task<Stream> GenerateAsStreamAsync(UpscaleModelSet modelOptions, In
129131
image.SaveAsPng(memoryStream);
130132
return memoryStream;
131133
}
134+
135+
136+
/// <summary>
137+
/// Generates the upscaled video.
138+
/// </summary>
139+
/// <param name="modelOptions">The model options.</param>
140+
/// <param name="videoInput">The video input.</param>
141+
/// <returns></returns>
142+
public async Task<DenseTensor<float>> GenerateAsync(UpscaleModelSet modelOptions, VideoInput videoInput)
143+
{
144+
DenseTensor<float> output = default;
145+
var videoInfo = await _videoService.GetVideoInfoAsync(videoInput);
146+
var videoFrames = await _videoService.CreateFramesAsync(videoInput, videoInfo.FPS);
147+
foreach (var frame in videoFrames.Frames)
148+
{
149+
var image = await GenerateInternalAsync(modelOptions, new InputImage(frame));
150+
output = output.Concatenate(image.ToDenseTensor(new[] { 1, 3, image.Height, image.Width }));
151+
}
152+
return output;
153+
}
154+
155+
156+
/// <summary>
157+
/// Generates the upscaled video.
158+
/// </summary>
159+
/// <param name="modelOptions">The model options.</param>
160+
/// <param name="videoInput">The video input.</param>
161+
/// <returns></returns>
162+
public async Task<byte[]> GenerateAsByteAsync(UpscaleModelSet modelOptions, VideoInput videoInput)
163+
{
164+
List<byte[]> output = new List<byte[]>();
165+
var videoInfo = await _videoService.GetVideoInfoAsync(videoInput);
166+
var videoFrames = await _videoService.CreateFramesAsync(videoInput, videoInfo.FPS);
167+
foreach (var frame in videoFrames.Frames)
168+
{
169+
170+
var image = await GenerateInternalAsync(modelOptions, new InputImage(frame));
171+
var ms = new MemoryStream();
172+
await image.SaveAsPngAsync(ms);
173+
output.Add(ms.ToArray());
174+
}
175+
176+
var videoResult = await _videoService.CreateVideoAsync(output, videoInfo.FPS);
177+
return videoResult.Data;
178+
}
179+
180+
181+
/// <summary>
182+
/// Generates the upscaled video.
183+
/// </summary>
184+
/// <param name="modelOptions">The model options.</param>
185+
/// <param name="videoInput">The video input.</param>
186+
/// <returns></returns>
187+
public async Task<Stream> GenerateAsStreamAsync(UpscaleModelSet modelOptions, VideoInput videoInput)
188+
{
189+
return new MemoryStream(await GenerateAsByteAsync(modelOptions, videoInput));
190+
}
132191

133192

134193
/// <summary>
@@ -167,6 +226,25 @@ private async Task<Image<Rgba32>> GenerateInternalAsync(UpscaleModelSet modelSet
167226
}
168227

169228

229+
/// <summary>
230+
/// Generates an upscaled video of the source provided.
231+
/// </summary>
232+
/// <param name="modelOptions">The model options.</param>
233+
/// <param name="videoInput">The video input.</param>
234+
/// <returns></returns>
235+
public async Task<IEnumerable<Image<Rgba32>>> GenerateInternalAsync(UpscaleModelSet modelOptions, VideoInput videoInput)
236+
{
237+
var output = new List<Image<Rgba32>>();
238+
var videoInfo = await _videoService.GetVideoInfoAsync(videoInput);
239+
var videoFrames = await _videoService.CreateFramesAsync(videoInput, videoInfo.FPS);
240+
foreach (var frame in videoFrames.Frames)
241+
{
242+
output.Add(await GenerateInternalAsync(modelOptions, new InputImage(frame)));
243+
}
244+
return output;
245+
}
246+
247+
170248
/// <summary>
171249
/// Creates the input parameters.
172250
/// </summary>
@@ -181,6 +259,5 @@ private UpscaleInput CreateInputParams(Image<Rgba32> imageSource, int maxTileSiz
181259
var height = imageSource.Height * scaleFactor;
182260
return new UpscaleInput(tiles, width, height);
183261
}
184-
185262
}
186263
}

‎OnnxStack.StableDiffusion/Helpers/TensorHelper.cs‎

Lines changed: 1 addition & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -247,33 +247,7 @@ public static DenseTensor<float> Clip(this DenseTensor<float> tensor, float minV
247247
}
248248

249249

250-
/// <summary>
251-
/// Concatenates the specified tensors along the 0 axis.
252-
/// </summary>
253-
/// <param name="tensor1">The tensor1.</param>
254-
/// <param name="tensor2">The tensor2.</param>
255-
/// <param name="axis">The axis.</param>
256-
/// <returns></returns>
257-
/// <exception cref="System.NotImplementedException">Only axis 0 is supported</exception>
258-
public static DenseTensor<float> Concatenate(this DenseTensor<float> tensor1, DenseTensor<float> tensor2, int axis = 0)
259-
{
260-
if (tensor1 == null)
261-
return tensor2.ToDenseTensor();
262-
263-
if (axis != 0 && axis != 2)
264-
throw new NotImplementedException("Only axis 0, 2 is supported");
265-
266-
if (axis == 2)
267-
return Concatenate(tensor1, tensor2);
268-
269-
var dimensions = tensor1.Dimensions.ToArray();
270-
dimensions[0] += tensor2.Dimensions[0];
271-
272-
var buffer = new float[tensor1.Length + tensor2.Length].AsMemory();
273-
tensor1.Buffer.CopyTo(buffer[..(int)tensor1.Length]);
274-
tensor2.Buffer.CopyTo(buffer[(int)tensor1.Length..]);
275-
return new DenseTensor<float>(buffer, dimensions);
276-
}
250+
277251

278252

279253
private static DenseTensor<float> Concatenate(DenseTensor<float> tensor1, DenseTensor<float> tensor2)

0 commit comments

Comments
 (0)