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

Commit f40ec99

Browse files
committed
Model downloader prototype
1 parent 850b32a commit f40ec99

27 files changed

+2257
-610
lines changed

OnnxStack.Console/appsettings.json

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
"IsEnabled": true,
1414
"PadTokenId": 49407,
1515
"BlankTokenId": 49407,
16-
"InputTokenLimit": 2048,
1716
"TokenizerLimit": 77,
1817
"EmbeddingsLength": 768,
1918
"ScaleFactor": 0.18215,

OnnxStack.Core/Config/OnnxModelType.cs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,5 @@ public enum OnnxModelType
77
TextEncoder = 20,
88
VaeEncoder = 30,
99
VaeDecoder = 40,
10-
SafetyChecker = 100,
1110
}
1211
}

OnnxStack.StableDiffusion/Common/IModelOptions.cs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ public interface IModelOptions : IOnnxModel
1212
int BlankTokenId { get; set; }
1313
float ScaleFactor { get; set; }
1414
int TokenizerLimit { get; set; }
15-
int InputTokenLimit { get; set; }
1615
int EmbeddingsLength { get; set; }
1716
DiffuserPipelineType PipelineType { get; set; }
1817
List<DiffuserType> Diffusers { get; set; }

OnnxStack.StableDiffusion/Config/ModelOptions.cs

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ public class ModelOptions : IModelOptions, IOnnxModelSetConfig
1515
public bool IsEnabled { get; set; }
1616
public int PadTokenId { get; set; }
1717
public int BlankTokenId { get; set; }
18-
public int InputTokenLimit { get; set; } = 2048;
1918
public int TokenizerLimit { get; set; }
2019
public int EmbeddingsLength { get; set; }
2120
public float ScaleFactor { get; set; }
@@ -34,7 +33,7 @@ public class ModelOptions : IModelOptions, IOnnxModelSetConfig
3433

3534
public void InitBlankTokenArray()
3635
{
37-
BlankTokenValueArray = Enumerable.Repeat(BlankTokenId, InputTokenLimit).ToImmutableArray();
36+
BlankTokenValueArray = Enumerable.Repeat(BlankTokenId, 20480).ToImmutableArray();
3837
}
3938
}
4039
}

OnnxStack.StableDiffusion/README.md

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,6 @@ Each model can be assigned to its own device, which is handy if you have only a
153153
"Name": "StableDiffusion 1.5",
154154
"PadTokenId": 49407,
155155
"BlankTokenId": 49407,
156-
"InputTokenLimit": 512,
157156
"TokenizerLimit": 77,
158157
"EmbeddingsLength": 768,
159158
"ScaleFactor": 0.18215,

OnnxStack.UI/App.xaml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,9 @@
1212
xmlns:converters="clr-namespace:OnnxStack.UI.Converters">
1313
<Application.Resources>
1414

15+
<BitmapImage x:Key="PlaceholderIcon" UriSource="/Images/Icon.png" />
1516
<BitmapImage x:Key="PlaceholderImage" UriSource="/Images/placeholder.png" />
17+
<BitmapImage x:Key="PlaceholderImageSmall" UriSource="/Images/placeholder_sm.png" />
1618
<Storyboard x:Key="LoadingAnimation">
1719
<DoubleAnimation Storyboard.TargetProperty="(Image.RenderTransform).(RotateTransform.Angle)" To="360" Duration="0:0:1" RepeatBehavior="Forever" />
1820
</Storyboard>

OnnxStack.UI/App.xaml.cs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414
using System.Diagnostics;
1515
using System.Windows.Controls;
1616
using System.Windows.Threading;
17+
using OnnxStack.UI.Views;
18+
using System.IO;
1719

1820
namespace OnnxStack.UI
1921
{
@@ -33,13 +35,15 @@ public App()
3335

3436
// Add OnnxStackStableDiffusion
3537
builder.Services.AddOnnxStackStableDiffusion();
38+
builder.Services.AddOnnxStackConfig<OnnxStackUIConfig>();
3639

3740
// Add Windows
3841
builder.Services.AddSingleton<MainWindow>();
3942
builder.Services.AddTransient<MessageDialog>();
4043
builder.Services.AddTransient<TextInputDialog>();
4144
builder.Services.AddTransient<CropImageDialog>();
4245
builder.Services.AddSingleton<IDialogService, DialogService>();
46+
builder.Services.AddSingleton<IModelDownloadService, ModelDownloadService>();
4347

4448
// Build App
4549
_applicationHost = builder.Build();
Lines changed: 239 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,239 @@
1+
using LibGit2Sharp;
2+
using System;
3+
using System.Collections.Generic;
4+
using System.Diagnostics;
5+
using System.IO;
6+
using System.Linq;
7+
using System.Threading;
8+
using System.Threading.Tasks;
9+
10+
namespace OnnxStack.UI.Helpers
11+
{
12+
public class RepositoryDownloader : IDisposable
13+
{
14+
private readonly LFSFilter _lfsFilter;
15+
private readonly string _repositoryUrl;
16+
private readonly string _destinationPath;
17+
private readonly CloneOptions _cloneOptions;
18+
private readonly Action<string, double> _progressCallback;
19+
private readonly CancellationTokenSource _cancellationTokenSource;
20+
21+
public RepositoryDownloader(string repositoryUrl, string destinationPath, Action<string, double> progressCallback = null, CancellationToken cancellationToken = default)
22+
{
23+
_repositoryUrl = repositoryUrl;
24+
_destinationPath = destinationPath;
25+
_progressCallback = progressCallback;
26+
_cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
27+
_lfsFilter = new LFSFilter("lfs", new[] { new FilterAttributeEntry("lfs") });
28+
_cloneOptions = new CloneOptions { OnCheckoutProgress = (f, p, t) => _progressCallback?.Invoke(f, (p * 100.0 / t)) };
29+
if (!GlobalSettings.GetRegisteredFilters().Any(x => x.Name == "lfs"))
30+
GlobalSettings.RegisterFilter(_lfsFilter);
31+
}
32+
33+
34+
public Task DownloadAsync()
35+
{
36+
if (!Directory.Exists(_destinationPath))
37+
Directory.CreateDirectory(_destinationPath);
38+
39+
try
40+
{
41+
// Perform the clone operation.
42+
return Task.Factory.StartNew(() =>
43+
Repository.Clone(_repositoryUrl, _destinationPath, _cloneOptions),
44+
_cancellationTokenSource.Token,
45+
TaskCreationOptions.LongRunning,
46+
TaskScheduler.Default);
47+
}
48+
catch (Exception)
49+
{
50+
_lfsFilter.KillProcess();
51+
throw;
52+
}
53+
}
54+
55+
public void Cancel()
56+
{
57+
_cancellationTokenSource.Cancel();
58+
_lfsFilter.KillProcess();
59+
}
60+
61+
public void Dispose()
62+
{
63+
Cancel();
64+
}
65+
}
66+
67+
public class LFSFilter : Filter
68+
{
69+
private Process _process;
70+
private FilterMode _mode;
71+
72+
73+
/// <summary>
74+
/// Initializes a new instance of the <see cref="LFSFilter"/> class.
75+
/// </summary>
76+
/// <param name="name"></param>
77+
/// <param name="attributes"></param>
78+
public LFSFilter(string name, IEnumerable<FilterAttributeEntry> attributes)
79+
: base(name, attributes)
80+
{
81+
// Kill the downloads if the app exists or crashes
82+
AppDomain.CurrentDomain.DomainUnload += (s, e) => KillProcess();
83+
AppDomain.CurrentDomain.ProcessExit += (s, e) => KillProcess();
84+
AppDomain.CurrentDomain.UnhandledException += (s, e) => KillProcess();
85+
}
86+
87+
88+
/// <summary>
89+
/// Kills the process.
90+
/// </summary>
91+
public void KillProcess()
92+
{
93+
try
94+
{
95+
if (_process is not null)
96+
{
97+
_process.Kill(true);
98+
_process.WaitForExit();
99+
}
100+
}
101+
catch (Exception)
102+
{
103+
104+
}
105+
}
106+
107+
108+
/// <summary>
109+
/// Initialize callback on filter
110+
/// </summary>
111+
protected override void Initialize()
112+
{
113+
base.Initialize();
114+
}
115+
116+
117+
/// <summary>
118+
/// Clean the input stream and write to the output stream.
119+
/// </summary>
120+
/// <param name="path">The path of the file being filtered</param>
121+
/// <param name="root">The path of the working directory for the owning repository</param>
122+
/// <param name="input">Input from the upstream filter or input reader</param>
123+
/// <param name="output">Output to the downstream filter or output writer</param>
124+
protected override void Clean(string path, string root, Stream input, Stream output)
125+
{
126+
try
127+
{
128+
// write file data to stdin
129+
input.CopyTo(_process.StandardInput.BaseStream);
130+
input.Flush();
131+
}
132+
catch (Exception)
133+
{
134+
135+
}
136+
}
137+
138+
139+
/// <summary>
140+
/// Complete callback on filter
141+
/// This optional callback will be invoked when the upstream filter is
142+
/// closed. Gives the filter a chance to perform any final actions or
143+
/// necissary clean up.
144+
/// </summary>
145+
/// <param name="path">The path of the file being filtered</param>
146+
/// <param name="root">The path of the working directory for the owning repository</param>
147+
/// <param name="output">Output to the downstream filter or output writer</param>
148+
protected override void Complete(string path, string root, Stream output)
149+
{
150+
try
151+
{
152+
// finalize stdin and wait for git-lfs to finish
153+
_process.StandardInput.Flush();
154+
_process.StandardInput.Close();
155+
if (_mode == FilterMode.Clean)
156+
{
157+
_process.WaitForExit();
158+
159+
// write git-lfs pointer for 'clean' to git or file data for 'smudge' to working copy
160+
_process.StandardOutput.BaseStream.CopyTo(output);
161+
_process.StandardOutput.BaseStream.Flush();
162+
_process.StandardOutput.Close();
163+
output.Flush();
164+
output.Close();
165+
}
166+
else if (_mode == FilterMode.Smudge)
167+
{
168+
// write git-lfs pointer for 'clean' to git or file data for 'smudge' to working copy
169+
_process.StandardOutput.BaseStream.CopyTo(output);
170+
_process.StandardOutput.BaseStream.Flush();
171+
_process.StandardOutput.Close();
172+
output.Flush();
173+
output.Close();
174+
175+
_process.WaitForExit();
176+
}
177+
178+
_process.Dispose();
179+
}
180+
catch (Exception)
181+
{
182+
183+
}
184+
}
185+
186+
187+
/// <summary>
188+
/// Indicates that a filter is going to be applied for the given file for
189+
/// the given mode.
190+
/// </summary>
191+
/// <param name="path">The path of the file being filtered</param>
192+
/// <param name="root">The path of the working directory for the owning repository</param>
193+
/// <param name="mode">The filter mode</param>
194+
protected override void Create(string path, string root, FilterMode mode)
195+
{
196+
try
197+
{
198+
_mode = mode;
199+
// launch git-lfs
200+
_process = new Process();
201+
_process.StartInfo.FileName = "git-lfs";
202+
_process.StartInfo.Arguments = string.Format("{0} {1}", mode == FilterMode.Clean ? "clean" : "smudge", path);
203+
_process.StartInfo.WorkingDirectory = root;
204+
_process.StartInfo.RedirectStandardInput = true;
205+
_process.StartInfo.RedirectStandardOutput = true;
206+
_process.StartInfo.RedirectStandardError = true;
207+
_process.StartInfo.CreateNoWindow = true;
208+
_process.StartInfo.UseShellExecute = false;
209+
_process.Start();
210+
}
211+
catch
212+
{
213+
throw new Exception("Git-LFS is not installed or is not initialized");
214+
}
215+
}
216+
217+
218+
/// <summary>
219+
/// Smudge the input stream and write to the output stream.
220+
/// </summary>
221+
/// <param name="path">The path of the file being filtered</param>
222+
/// <param name="root">The path of the working directory for the owning repository</param>
223+
/// <param name="input">Input from the upstream filter or input reader</param>
224+
/// <param name="output">Output to the downstream filter or output writer</param>
225+
protected override void Smudge(string path, string root, Stream input, Stream output)
226+
{
227+
try
228+
{
229+
// write git-lfs pointer to stdin
230+
input.CopyTo(_process.StandardInput.BaseStream);
231+
input.Flush();
232+
}
233+
catch (Exception)
234+
{
235+
236+
}
237+
}
238+
}
239+
}

OnnxStack.UI/Images/Icon.png

-6.92 KB
Loading
-9.26 KB
Binary file not shown.

0 commit comments

Comments
 (0)