Skip to content
This repository was archived by the owner on Dec 9, 2025. It is now read-only.

Commit d83a59c

Browse files
saddam213saddam213
authored
Merge pull request #111 from saddam213/LowVram
Support for low memory devices
2 parents da3a66d + 582aeed commit d83a59c

47 files changed

Lines changed: 379 additions & 145 deletions

Some content is hidden

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

‎OnnxStack.Core/Model/OnnxModelSession.cs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,14 +67,18 @@ public async Task LoadAsync()
6767
/// Unloads the model session.
6868
/// </summary>
6969
/// <returns></returns>
70-
public Task UnloadAsync()
70+
public async Task UnloadAsync()
7171
{
72+
// TODO: deadlock on model dispose when no synchronization context exists(console app)
73+
// Task.Yield seems to force a context switch resolving any issues, revist this
74+
await Task.Yield();
75+
7276
if (_session is not null)
7377
{
74-
_metadata = null;
7578
_session.Dispose();
79+
_metadata = null;
80+
_session = null;
7681
}
77-
return Task.CompletedTask;
7882
}
7983

8084

‎OnnxStack.StableDiffusion/Config/PipelineOptions.cs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
using OnnxStack.StableDiffusion.Enums;
2+
3+
namespace OnnxStack.StableDiffusion.Config
4+
{
5+
public record PipelineOptions(string Name, MemoryModeType MemoryMode);
6+
7+
}

‎OnnxStack.StableDiffusion/Config/StableDiffusionModelSet.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ public record StableDiffusionModelSet : IOnnxModelSetConfig
1414
public int SampleSize { get; set; } = 512;
1515
public DiffuserPipelineType PipelineType { get; set; }
1616
public List<DiffuserType> Diffusers { get; set; } = new List<DiffuserType>();
17-
17+
public MemoryModeType MemoryMode { get; set; }
1818
public int DeviceId { get; set; }
1919
public int InterOpNumThreads { get; set; }
2020
public int IntraOpNumThreads { get; set; }

‎OnnxStack.StableDiffusion/Diffusers/DiffuserBase.cs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ public abstract class DiffuserBase : IDiffuser
2222
protected readonly UNetConditionModel _unet;
2323
protected readonly AutoEncoderModel _vaeDecoder;
2424
protected readonly AutoEncoderModel _vaeEncoder;
25+
protected readonly MemoryModeType _memoryMode;
2526

2627
/// <summary>
2728
/// Initializes a new instance of the <see cref="DiffuserBase"/> class.
@@ -31,12 +32,13 @@ public abstract class DiffuserBase : IDiffuser
3132
/// <param name="vaeDecoder">The vae decoder.</param>
3233
/// <param name="vaeEncoder">The vae encoder.</param>
3334
/// <param name="logger">The logger.</param>
34-
public DiffuserBase(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, ILogger logger = default)
35+
public DiffuserBase(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, MemoryModeType memoryMode, ILogger logger = default)
3536
{
3637
_logger = logger;
3738
_unet = unet;
3839
_vaeDecoder = vaeDecoder;
3940
_vaeEncoder = vaeEncoder;
41+
_memoryMode = memoryMode;
4042
}
4143

4244
/// <summary>
@@ -137,10 +139,15 @@ protected virtual async Task<DenseTensor<float>> DecodeLatentsAsync(PromptOption
137139
var results = await _vaeDecoder.RunInferenceAsync(inferenceParameters);
138140
using (var imageResult = results.First())
139141
{
142+
// Unload if required
143+
if (_memoryMode == MemoryModeType.Minimum)
144+
await _vaeDecoder.UnloadAsync();
145+
140146
_logger?.LogEnd("Latents decoded", timestamp);
141147
return imageResult.ToDenseTensor();
142148
}
143149
}
150+
144151
}
145152

146153

‎OnnxStack.StableDiffusion/Diffusers/InstaFlow/ControlNetDiffuser.cs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,8 @@ public class ControlNetDiffuser : InstaFlowDiffuser
3030
/// <param name="vaeDecoder">The vae decoder.</param>
3131
/// <param name="vaeEncoder">The vae encoder.</param>
3232
/// <param name="logger">The logger.</param>
33-
public ControlNetDiffuser(ControlNetModel controlNet, UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, ILogger logger = default)
34-
: base(unet, vaeDecoder, vaeEncoder, logger)
33+
public ControlNetDiffuser(ControlNetModel controlNet, UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, MemoryModeType memoryMode, ILogger logger = default)
34+
: base(unet, vaeDecoder, vaeEncoder, memoryMode, logger)
3535
{
3636
_controlNet = controlNet;
3737
}
@@ -144,9 +144,13 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
144144
}
145145

146146
ReportProgress(progressCallback, step, timesteps.Count, latents);
147-
_logger?.LogEnd($"Step {step}/{timesteps.Count}", stepTime);
147+
_logger?.LogEnd(LogLevel.Debug, $"Step {step}/{timesteps.Count}", stepTime);
148148
}
149149

150+
// Unload if required
151+
if (_memoryMode == MemoryModeType.Minimum)
152+
await Task.WhenAll(_controlNet.UnloadAsync(), _unet.UnloadAsync());
153+
150154
// Decode Latents
151155
return await DecodeLatentsAsync(promptOptions, schedulerOptions, latents);
152156
}

‎OnnxStack.StableDiffusion/Diffusers/InstaFlow/InstaFlowDiffuser.cs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@ public abstract class InstaFlowDiffuser : DiffuserBase
2626
/// <param name="vaeDecoder">The vae decoder.</param>
2727
/// <param name="vaeEncoder">The vae encoder.</param>
2828
/// <param name="logger">The logger.</param>
29-
public InstaFlowDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, ILogger logger = default)
30-
: base(unet, vaeDecoder, vaeEncoder, logger) { }
29+
public InstaFlowDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, MemoryModeType memoryMode, ILogger logger = default)
30+
: base(unet, vaeDecoder, vaeEncoder, memoryMode, logger) { }
3131

3232
/// <summary>
3333
/// Gets the type of the pipeline.
@@ -103,9 +103,13 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
103103
}
104104

105105
ReportProgress(progressCallback, step, timesteps.Count, latents);
106-
_logger?.LogEnd($"Step {step}/{timesteps.Count}", stepTime);
106+
_logger?.LogEnd(LogLevel.Debug, $"Step {step}/{timesteps.Count}", stepTime);
107107
}
108108

109+
// Unload if required
110+
if (_memoryMode == MemoryModeType.Minimum)
111+
await _unet.UnloadAsync();
112+
109113
// Decode Latents
110114
return await DecodeLatentsAsync(promptOptions, schedulerOptions, latents);
111115
}

‎OnnxStack.StableDiffusion/Diffusers/InstaFlow/TextDiffuser.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,8 @@ public sealed class TextDiffuser : InstaFlowDiffuser
2020
/// <param name="vaeDecoder">The vae decoder.</param>
2121
/// <param name="vaeEncoder">The vae encoder.</param>
2222
/// <param name="logger">The logger.</param>
23-
public TextDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, ILogger logger = default)
24-
: base(unet, vaeDecoder, vaeEncoder, logger) { }
23+
public TextDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, MemoryModeType memoryMode, ILogger logger = default)
24+
: base(unet, vaeDecoder, vaeEncoder, memoryMode, logger) { }
2525

2626

2727
/// <summary>

‎OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ControlNetDiffuser.cs‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ public class ControlNetDiffuser : LatentConsistencyDiffuser
2929
/// <param name="vaeDecoder">The vae decoder.</param>
3030
/// <param name="vaeEncoder">The vae encoder.</param>
3131
/// <param name="logger">The logger.</param>
32-
public ControlNetDiffuser(ControlNetModel controlNet, UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, ILogger logger = default)
33-
: base(unet, vaeDecoder, vaeEncoder, logger)
32+
public ControlNetDiffuser(ControlNetModel controlNet, UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, MemoryModeType memoryMode, ILogger logger = default)
33+
: base(unet, vaeDecoder, vaeEncoder, memoryMode, logger)
3434
{
3535
_controlNet = controlNet;
3636
}
@@ -141,9 +141,13 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
141141
}
142142

143143
ReportProgress(progressCallback, step, timesteps.Count, latents);
144-
_logger?.LogEnd($"Step {step}/{timesteps.Count}", stepTime);
144+
_logger?.LogEnd(LogLevel.Debug, $"Step {step}/{timesteps.Count}", stepTime);
145145
}
146146

147+
// Unload if required
148+
if (_memoryMode == MemoryModeType.Minimum)
149+
await Task.WhenAll(_controlNet.UnloadAsync(), _unet.UnloadAsync());
150+
147151
// Decode Latents
148152
return await DecodeLatentsAsync(promptOptions, schedulerOptions, latents);
149153
}

‎OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ControlNetImageDiffuser.cs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ public sealed class ControlNetImageDiffuser : ControlNetDiffuser
2727
/// <param name="vaeDecoder">The vae decoder.</param>
2828
/// <param name="vaeEncoder">The vae encoder.</param>
2929
/// <param name="logger">The logger.</param>
30-
public ControlNetImageDiffuser(ControlNetModel controlNet, UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, ILogger logger = default)
31-
: base(controlNet, unet, vaeDecoder, vaeEncoder, logger) { }
30+
public ControlNetImageDiffuser(ControlNetModel controlNet, UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, MemoryModeType memoryMode, ILogger logger = default)
31+
: base(controlNet, unet, vaeDecoder, vaeEncoder, memoryMode, logger) { }
3232

3333

3434
/// <summary>
@@ -73,6 +73,10 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
7373
var results = await _vaeEncoder.RunInferenceAsync(inferenceParameters);
7474
using (var result = results.First())
7575
{
76+
// Unload if required
77+
if (_memoryMode == MemoryModeType.Minimum)
78+
await _vaeEncoder.UnloadAsync();
79+
7680
var outputResult = result.ToDenseTensor();
7781
var scaledSample = outputResult.MultiplyBy(_vaeEncoder.ScaleFactor);
7882
return scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);

‎OnnxStack.StableDiffusion/Diffusers/LatentConsistency/ImageDiffuser.cs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@ public sealed class ImageDiffuser : LatentConsistencyDiffuser
2626
/// <param name="vaeDecoder">The vae decoder.</param>
2727
/// <param name="vaeEncoder">The vae encoder.</param>
2828
/// <param name="logger">The logger.</param>
29-
public ImageDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, ILogger logger = default)
30-
: base(unet, vaeDecoder, vaeEncoder, logger) { }
29+
public ImageDiffuser(UNetConditionModel unet, AutoEncoderModel vaeDecoder, AutoEncoderModel vaeEncoder, MemoryModeType memoryMode, ILogger logger = default)
30+
: base(unet, vaeDecoder, vaeEncoder, memoryMode, logger) { }
3131

3232

3333
/// <summary>
@@ -70,6 +70,10 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
7070
var results = await _vaeEncoder.RunInferenceAsync(inferenceParameters);
7171
using (var result = results.First())
7272
{
73+
// Unload if required
74+
if (_memoryMode == MemoryModeType.Minimum)
75+
await _vaeEncoder.UnloadAsync();
76+
7377
var outputResult = result.ToDenseTensor();
7478
var scaledSample = outputResult.MultiplyBy(_vaeEncoder.ScaleFactor);
7579
return scheduler.AddNoise(scaledSample, scheduler.CreateRandomSample(scaledSample.Dimensions), timesteps);

0 commit comments

Comments
 (0)