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

Commit 582aeed

Browse files
saddam213saddam213
committed
Remove Medium mode as it has no real benefit
1 parent 9190822 commit 582aeed

27 files changed

+63
-78
lines changed

‎OnnxStack.Core/Model/OnnxModelSession.cs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,10 @@ public async Task LoadAsync()
6969
/// <returns></returns>
7070
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
7274
await Task.Yield();
75+
7376
if (_session is not null)
7477
{
7578
_session.Dispose();

‎OnnxStack.StableDiffusion/Diffusers/DiffuserBase.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ protected virtual async Task<DenseTensor<float>> DecodeLatentsAsync(PromptOption
140140
using (var imageResult = results.First())
141141
{
142142
// Unload if required
143-
if (_memoryMode != MemoryModeType.Maximum)
143+
if (_memoryMode == MemoryModeType.Minimum)
144144
await _vaeDecoder.UnloadAsync();
145145

146146
_logger?.LogEnd("Latents decoded", timestamp);

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

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -148,12 +148,9 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
148148
}
149149

150150
// Unload if required
151-
if (_memoryMode != MemoryModeType.Maximum)
152-
{
153-
await _unet.UnloadAsync();
154-
await _controlNet.UnloadAsync();
155-
}
156-
151+
if (_memoryMode == MemoryModeType.Minimum)
152+
await Task.WhenAll(_controlNet.UnloadAsync(), _unet.UnloadAsync());
153+
157154
// Decode Latents
158155
return await DecodeLatentsAsync(promptOptions, schedulerOptions, latents);
159156
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
107107
}
108108

109109
// Unload if required
110-
if (_memoryMode != MemoryModeType.Maximum)
110+
if (_memoryMode == MemoryModeType.Minimum)
111111
await _unet.UnloadAsync();
112112

113113
// Decode Latents

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -145,11 +145,8 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
145145
}
146146

147147
// Unload if required
148-
if (_memoryMode != MemoryModeType.Maximum)
149-
{
150-
await _unet.UnloadAsync();
151-
await _controlNet.UnloadAsync();
152-
}
148+
if (_memoryMode == MemoryModeType.Minimum)
149+
await Task.WhenAll(_controlNet.UnloadAsync(), _unet.UnloadAsync());
153150

154151
// Decode Latents
155152
return await DecodeLatentsAsync(promptOptions, schedulerOptions, latents);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
7474
using (var result = results.First())
7575
{
7676
// Unload if required
77-
if (_memoryMode != MemoryModeType.Maximum)
77+
if (_memoryMode == MemoryModeType.Minimum)
7878
await _vaeEncoder.UnloadAsync();
7979

8080
var outputResult = result.ToDenseTensor();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
7171
using (var result = results.First())
7272
{
7373
// Unload if required
74-
if (_memoryMode != MemoryModeType.Maximum)
74+
if (_memoryMode == MemoryModeType.Minimum)
7575
await _vaeEncoder.UnloadAsync();
7676

7777
var outputResult = result.ToDenseTensor();

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
142142
}
143143

144144
// Unload if required
145-
if (_memoryMode != MemoryModeType.Maximum)
145+
if (_memoryMode == MemoryModeType.Minimum)
146146
await _unet.UnloadAsync();
147147

148148
// Decode Latents
@@ -173,7 +173,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
173173
using (var result = results.First())
174174
{
175175
// Unload if required
176-
if (_memoryMode != MemoryModeType.Maximum)
176+
if (_memoryMode == MemoryModeType.Minimum)
177177
await _vaeEncoder.UnloadAsync();
178178

179179
var outputResult = result.ToDenseTensor();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
107107
}
108108

109109
// Unload if required
110-
if (_memoryMode != MemoryModeType.Maximum)
110+
if (_memoryMode == MemoryModeType.Minimum)
111111
await _unet.UnloadAsync();
112112

113113
// Decode Latents

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -150,11 +150,8 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
150150
}
151151

152152
// Unload if required
153-
if (_memoryMode != MemoryModeType.Maximum)
154-
{
155-
await _unet.UnloadAsync();
156-
await _controlNet.UnloadAsync();
157-
}
153+
if (_memoryMode == MemoryModeType.Minimum)
154+
await Task.WhenAll(_controlNet.UnloadAsync(), _unet.UnloadAsync());
158155

159156
// Decode Latents
160157
return await DecodeLatentsAsync(promptOptions, schedulerOptions, latents);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
7272
using (var result = results.First())
7373
{
7474
// Unload if required
75-
if (_memoryMode != MemoryModeType.Maximum)
75+
if (_memoryMode == MemoryModeType.Minimum)
7676
await _vaeEncoder.UnloadAsync();
7777

7878
var outputResult = result.ToDenseTensor();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
7373
using (var result = results.First())
7474
{
7575
// Unload if required
76-
if (_memoryMode != MemoryModeType.Maximum)
76+
if (_memoryMode == MemoryModeType.Minimum)
7777
await _vaeEncoder.UnloadAsync();
7878

7979
var outputResult = result.ToDenseTensor();

‎OnnxStack.StableDiffusion/Diffusers/LatentConsistencyXL/InpaintLegacyDiffuser.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
124124
}
125125

126126
// Unload if required
127-
if (_memoryMode != MemoryModeType.Maximum)
127+
if (_memoryMode == MemoryModeType.Minimum)
128128
await _unet.UnloadAsync();
129129

130130
// Decode Latents
@@ -168,7 +168,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
168168
using (var result = results.First())
169169
{
170170
// Unload if required
171-
if (_memoryMode != MemoryModeType.Maximum)
171+
if (_memoryMode == MemoryModeType.Minimum)
172172
await _vaeEncoder.UnloadAsync();
173173

174174
var outputResult = result.ToDenseTensor();

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -141,11 +141,8 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
141141
}
142142

143143
// Unload if required
144-
if (_memoryMode != MemoryModeType.Maximum)
145-
{
146-
await _unet.UnloadAsync();
147-
await _controlNet.UnloadAsync();
148-
}
144+
if (_memoryMode == MemoryModeType.Minimum)
145+
await Task.WhenAll(_controlNet.UnloadAsync(), _unet.UnloadAsync());
149146

150147
// Decode Latents
151148
return await DecodeLatentsAsync(promptOptions, schedulerOptions, latents);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
7272
using (var result = results.First())
7373
{
7474
// Unload if required
75-
if (_memoryMode != MemoryModeType.Maximum)
75+
if (_memoryMode == MemoryModeType.Minimum)
7676
await _vaeEncoder.UnloadAsync();
7777

7878
var outputResult = result.ToDenseTensor();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
7272
using (var result = results.First())
7373
{
7474
// Unload if required
75-
if (_memoryMode != MemoryModeType.Maximum)
75+
if (_memoryMode == MemoryModeType.Minimum)
7676
await _vaeEncoder.UnloadAsync();
7777

7878
var outputResult = result.ToDenseTensor();

‎OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintDiffuser.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
112112
}
113113

114114
// Unload if required
115-
if (_memoryMode != MemoryModeType.Maximum)
115+
if (_memoryMode == MemoryModeType.Minimum)
116116
await _unet.UnloadAsync();
117117

118118
// Decode Latents
@@ -229,7 +229,7 @@ private async Task<DenseTensor<float>> PrepareImageMask(PromptOptions promptOpti
229229
using (var result = results.First())
230230
{
231231
// Unload if required
232-
if (_memoryMode != MemoryModeType.Maximum)
232+
if (_memoryMode == MemoryModeType.Minimum)
233233
await _vaeEncoder.UnloadAsync();
234234

235235
var sample = result.ToDenseTensor();

‎OnnxStack.StableDiffusion/Diffusers/StableDiffusion/InpaintLegacyDiffuser.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
118118
}
119119

120120
// Unload if required
121-
if (_memoryMode != MemoryModeType.Maximum)
121+
if (_memoryMode == MemoryModeType.Minimum)
122122
await _unet.UnloadAsync();
123123

124124
// Decode Latents
@@ -162,7 +162,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
162162
using (var result = results.First())
163163
{
164164
// Unload if required
165-
if (_memoryMode != MemoryModeType.Maximum)
165+
if (_memoryMode == MemoryModeType.Minimum)
166166
await _vaeEncoder.UnloadAsync();
167167

168168
var outputResult = result.ToDenseTensor();

‎OnnxStack.StableDiffusion/Diffusers/StableDiffusion/StableDiffusionDiffuser.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
100100
}
101101

102102
// Unload if required
103-
if (_memoryMode != MemoryModeType.Maximum)
103+
if (_memoryMode == MemoryModeType.Minimum)
104104
await _unet.UnloadAsync();
105105

106106

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -151,11 +151,8 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
151151
}
152152

153153
// Unload if required
154-
if (_memoryMode != MemoryModeType.Maximum)
155-
{
156-
await _unet.UnloadAsync();
157-
await _controlNet.UnloadAsync();
158-
}
154+
if (_memoryMode == MemoryModeType.Minimum)
155+
await Task.WhenAll(_controlNet.UnloadAsync(), _unet.UnloadAsync());
159156

160157
// Decode Latents
161158
return await DecodeLatentsAsync(promptOptions, schedulerOptions, latents);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
7575
using (var result = results.First())
7676
{
7777
// Unload if required
78-
if (_memoryMode != MemoryModeType.Maximum)
78+
if (_memoryMode == MemoryModeType.Minimum)
7979
await _vaeEncoder.UnloadAsync();
8080

8181
var outputResult = result.ToDenseTensor();

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
7575
using (var result = results.First())
7676
{
7777
// Unload if required
78-
if (_memoryMode != MemoryModeType.Maximum)
78+
if (_memoryMode == MemoryModeType.Minimum)
7979
await _vaeEncoder.UnloadAsync();
8080

8181
var outputResult = result.ToDenseTensor();

‎OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/InpaintLegacyDiffuser.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,7 +125,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
125125
}
126126

127127
// Unload if required
128-
if (_memoryMode != MemoryModeType.Maximum)
128+
if (_memoryMode == MemoryModeType.Minimum)
129129
await _unet.UnloadAsync();
130130

131131
// Decode Latents
@@ -172,7 +172,7 @@ protected override async Task<DenseTensor<float>> PrepareLatentsAsync(PromptOpti
172172
using (var result = results.First())
173173
{
174174
// Unload if required
175-
if (_memoryMode != MemoryModeType.Maximum)
175+
if (_memoryMode == MemoryModeType.Minimum)
176176
await _vaeEncoder.UnloadAsync();
177177

178178
var outputResult = result.ToDenseTensor();

‎OnnxStack.StableDiffusion/Diffusers/StableDiffusionXL/StableDiffusionXLDiffuser.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ public override async Task<DenseTensor<float>> DiffuseAsync(PromptOptions prompt
108108
}
109109

110110
// Unload if required
111-
if (_memoryMode != MemoryModeType.Maximum)
111+
if (_memoryMode == MemoryModeType.Minimum)
112112
await _unet.UnloadAsync();
113113

114114
// Decode Latents

‎OnnxStack.StableDiffusion/Enums/MemoryModeType.cs

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,6 @@
33
public enum MemoryModeType
44
{
55
Maximum = 0,
6-
Medium = 1,
7-
Minimum = 2
6+
Minimum = 10
87
}
9-
108
}

‎OnnxStack.StableDiffusion/Pipelines/StableDiffusionPipeline.cs

Lines changed: 15 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -108,24 +108,18 @@ public StableDiffusionPipeline(PipelineOptions pipelineOptions, TokenizerModel t
108108
/// </summary>
109109
public override Task LoadAsync()
110110
{
111-
if (_pipelineOptions.MemoryMode == MemoryModeType.Maximum)
112-
{
113-
return Task.WhenAll(
114-
_unet.LoadAsync(),
115-
_tokenizer.LoadAsync(),
116-
_textEncoder.LoadAsync(),
117-
_vaeDecoder.LoadAsync(),
118-
_vaeEncoder.LoadAsync());
119-
}
120-
121-
if (_pipelineOptions.MemoryMode == MemoryModeType.Medium)
122-
{
123-
return Task.WhenAll(
124-
_unet.LoadAsync(),
125-
_tokenizer.LoadAsync(),
126-
_textEncoder.LoadAsync());
127-
}
128-
return Task.CompletedTask;
111+
if (_pipelineOptions.MemoryMode == MemoryModeType.Minimum)
112+
return Task.CompletedTask;
113+
114+
// Preload all models into VRAM
115+
return Task.WhenAll
116+
(
117+
_unet.LoadAsync(),
118+
_tokenizer.LoadAsync(),
119+
_textEncoder.LoadAsync(),
120+
_vaeDecoder.LoadAsync(),
121+
_vaeEncoder.LoadAsync()
122+
);
129123
}
130124

131125

@@ -135,7 +129,10 @@ public override Task LoadAsync()
135129
/// <returns></returns>
136130
public override async Task UnloadAsync()
137131
{
132+
// TODO: deadlock on model dispose when no synchronization context exists(console app)
133+
// Task.Yield seems to force a context switch resolving any issues, revist this
138134
await Task.Yield();
135+
139136
_unet?.Dispose();
140137
_tokenizer?.Dispose();
141138
_textEncoder?.Dispose();

‎OnnxStack.StableDiffusion/Pipelines/StableDiffusionXLPipeline.cs

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -68,14 +68,16 @@ public StableDiffusionXLPipeline(PipelineOptions pipelineOptions, TokenizerModel
6868
/// </summary>
6969
public override Task LoadAsync()
7070
{
71-
if (_pipelineOptions.MemoryMode != MemoryModeType.Minimum)
72-
{
73-
return Task.WhenAll(
74-
_tokenizer2.LoadAsync(),
75-
_textEncoder2.LoadAsync(),
76-
base.LoadAsync());
77-
}
78-
return base.LoadAsync();
71+
if (_pipelineOptions.MemoryMode == MemoryModeType.Minimum)
72+
return base.LoadAsync();
73+
74+
// Preload all models into VRAM
75+
return Task.WhenAll
76+
(
77+
_tokenizer2.LoadAsync(),
78+
_textEncoder2.LoadAsync(),
79+
base.LoadAsync()
80+
);
7981
}
8082

8183

0 commit comments

Comments
 (0)