Skip to content

Commit 9be407f

Browse files
mohnjilesmohnjiles
committed
Give extra python libs for triton/sage/etc with every windows venv & fix pip install requirements parsing when deps have spaces around specifiers
1 parent 9928e32 commit 9be407f

7 files changed

Lines changed: 179 additions & 63 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning 2.0](https://semver.org/spec/v2
1111
### Fixed
1212
- Fixed Inference ControlNet Preprocessors using incorrect resolution and increased maximum of smallest dimension to 16384
1313
- Fixed Triton/Sage install option showing for incompatible GPUs
14+
- Fixed errors from invalid pip specifiers in requirements files
1415

1516
## v2.14.0-pre.1
1617
### Added

‎StabilityMatrix.Avalonia/Helpers/UnixPrerequisiteHelper.cs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -491,4 +491,14 @@ public Task<bool> FixGitLongPaths()
491491
{
492492
throw new PlatformNotSupportedException();
493493
}
494+
495+
[UnsupportedOSPlatform("Linux")]
496+
[UnsupportedOSPlatform("macOS")]
497+
public Task AddMissingLibsToVenv(
498+
DirectoryPath installedPackagePath,
499+
IProgress<ProgressReport>? progress = null
500+
)
501+
{
502+
throw new PlatformNotSupportedException();
503+
}
494504
}

‎StabilityMatrix.Avalonia/Helpers/WindowsPrerequisiteHelper.cs‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ IPyRunner pyRunner
4444

4545
private const string HipSdkDownloadUrl =
4646
"https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-24.Q4-Win10-Win11-For-HIP.exe";
47+
private const string PythonLibsDownloadUrl = "https://cdn.lykos.ai/python_libs_for_sage.zip";
4748

4849
private string HomeDir => settingsManager.LibraryDir;
4950

@@ -633,6 +634,43 @@ public async Task<Process> RunDotnet(
633634
);
634635
}
635636

637+
[SupportedOSPlatform("Windows")]
638+
public async Task AddMissingLibsToVenv(
639+
DirectoryPath installedPackagePath,
640+
IProgress<ProgressReport>? progress = null
641+
)
642+
{
643+
var venvLibsDir = installedPackagePath.JoinDir("venv", "libs");
644+
var venvIncludeDir = installedPackagePath.JoinDir("venv", "include");
645+
if (
646+
venvLibsDir.Exists
647+
&& venvIncludeDir.Exists
648+
&& venvLibsDir.JoinFile("python3.lib").Exists
649+
&& venvLibsDir.JoinFile("python310.lib").Exists
650+
)
651+
{
652+
Logger.Debug("Python libs already installed at {VenvLibsDir}", venvLibsDir);
653+
return;
654+
}
655+
656+
var downloadPath = installedPackagePath.JoinFile("python_libs_for_sage.zip");
657+
var venvDir = installedPackagePath.JoinDir("venv");
658+
await downloadService
659+
.DownloadToFileAsync(PythonLibsDownloadUrl, downloadPath, progress)
660+
.ConfigureAwait(false);
661+
662+
progress?.Report(
663+
new ProgressReport(-1f, message: "Extracting Python libraries", isIndeterminate: true)
664+
);
665+
await ArchiveHelper.Extract7Z(downloadPath, venvDir, progress);
666+
667+
var includeFolder = venvDir.JoinDir("include");
668+
var scriptsIncludeFolder = venvDir.JoinDir("Scripts").JoinDir("include");
669+
await includeFolder.CopyToAsync(scriptsIncludeFolder);
670+
671+
await downloadPath.DeleteAsync();
672+
}
673+
636674
private async Task DownloadAndExtractPrerequisite(
637675
IProgress<ProgressReport>? progress,
638676
string downloadUrl,

‎StabilityMatrix.Core/Helper/IPrerequisiteHelper.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
using System.Diagnostics;
22
using System.Runtime.Versioning;
33
using StabilityMatrix.Core.Models;
4+
using StabilityMatrix.Core.Models.FileInterfaces;
45
using StabilityMatrix.Core.Models.Packages;
56
using StabilityMatrix.Core.Models.Progress;
67
using StabilityMatrix.Core.Processes;
@@ -203,4 +204,5 @@ Task<Process> RunDotnet(
203204
);
204205

205206
Task<bool> FixGitLongPaths();
207+
Task AddMissingLibsToVenv(DirectoryPath installedPackagePath, IProgress<ProgressReport>? progress = null);
206208
}

‎StabilityMatrix.Core/Models/PackageModification/InstallSageAttentionStep.cs‎

Lines changed: 86 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,56 @@ public async Task ExecuteAsync(IProgress<ProgressReport>? progress = null)
3030
);
3131
}
3232

33+
var venvDir = WorkingDirectory.JoinDir("venv");
34+
35+
await using var venvRunner = PyBaseInstall.Default.CreateVenvRunner(
36+
venvDir,
37+
workingDirectory: WorkingDirectory,
38+
environmentVariables: EnvironmentVariables
39+
);
40+
41+
var torchInfo = await venvRunner.PipShow("torch").ConfigureAwait(false);
42+
var sageWheelUrl = string.Empty;
43+
44+
if (torchInfo == null)
45+
{
46+
sageWheelUrl = string.Empty;
47+
}
48+
else if (torchInfo.Version.Contains("2.5.1") && torchInfo.Version.Contains("cu124"))
49+
{
50+
sageWheelUrl =
51+
"https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu124torch2.5.1-cp310-cp310-win_amd64.whl";
52+
}
53+
else if (torchInfo.Version.Contains("2.6.0") && torchInfo.Version.Contains("cu126"))
54+
{
55+
sageWheelUrl =
56+
"https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu126torch2.6.0-cp310-cp310-win_amd64.whl";
57+
}
58+
else if (torchInfo.Version.Contains("2.7.0") && torchInfo.Version.Contains("cu128"))
59+
{
60+
sageWheelUrl =
61+
"https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu128torch2.7.0-cp310-cp310-win_amd64.whl";
62+
}
63+
64+
var pipArgs = new PipInstallArgs();
65+
if (IsBlackwellGpu)
66+
{
67+
pipArgs = pipArgs.AddArg("--pre");
68+
}
69+
pipArgs = pipArgs.AddArg("triton-windows");
70+
71+
if (!string.IsNullOrWhiteSpace(sageWheelUrl))
72+
{
73+
pipArgs = pipArgs.AddArg(sageWheelUrl);
74+
75+
progress?.Report(
76+
new ProgressReport(-1f, message: "Installing Triton & SageAttention", isIndeterminate: true)
77+
);
78+
await venvRunner.PipInstall(pipArgs, progress.AsProcessOutputHandler()).ConfigureAwait(false);
79+
return;
80+
}
81+
82+
// no wheels, gotta build
3383
if (!prerequisiteHelper.IsVcBuildToolsInstalled)
3484
{
3585
throw new MissingPrerequisiteException(
@@ -63,14 +113,6 @@ public async Task ExecuteAsync(IProgress<ProgressReport>? progress = null)
63113
: cuda126ExpectedPath.JoinFile("nvcc.exe").ToString();
64114
}
65115

66-
var venvDir = WorkingDirectory.JoinDir("venv");
67-
68-
await using var venvRunner = PyBaseInstall.Default.CreateVenvRunner(
69-
venvDir,
70-
workingDirectory: WorkingDirectory,
71-
environmentVariables: EnvironmentVariables
72-
);
73-
74116
venvRunner.UpdateEnvironmentVariables(env =>
75117
{
76118
var cudaBinPath = Path.GetDirectoryName(nvccPath)!;
@@ -88,47 +130,6 @@ public async Task ExecuteAsync(IProgress<ProgressReport>? progress = null)
88130
return env;
89131
});
90132

91-
var torchInfo = await venvRunner.PipShow("torch").ConfigureAwait(false);
92-
var sageWheelUrl = string.Empty;
93-
94-
if (torchInfo == null)
95-
{
96-
sageWheelUrl = string.Empty;
97-
}
98-
else if (torchInfo.Version.Contains("2.5.1") && torchInfo.Version.Contains("cu124"))
99-
{
100-
sageWheelUrl =
101-
"https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu124torch2.5.1-cp310-cp310-win_amd64.whl";
102-
}
103-
else if (torchInfo.Version.Contains("2.6.0") && torchInfo.Version.Contains("cu126"))
104-
{
105-
sageWheelUrl =
106-
"https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu126torch2.6.0-cp310-cp310-win_amd64.whl";
107-
}
108-
else if (torchInfo.Version.Contains("2.7.0") && torchInfo.Version.Contains("cu128"))
109-
{
110-
sageWheelUrl =
111-
"https://github.com/woct0rdho/SageAttention/releases/download/v2.1.1-windows/sageattention-2.1.1+cu128torch2.7.0-cp310-cp310-win_amd64.whl";
112-
}
113-
114-
var pipArgs = new PipInstallArgs();
115-
if (IsBlackwellGpu)
116-
{
117-
pipArgs = pipArgs.AddArg("--pre");
118-
}
119-
pipArgs = pipArgs.AddArg("triton-windows");
120-
121-
if (!string.IsNullOrWhiteSpace(sageWheelUrl))
122-
{
123-
pipArgs = pipArgs.AddArg(sageWheelUrl);
124-
125-
progress?.Report(
126-
new ProgressReport(-1f, message: "Installing Triton & SageAttention", isIndeterminate: true)
127-
);
128-
await venvRunner.PipInstall(pipArgs, progress.AsProcessOutputHandler()).ConfigureAwait(false);
129-
return;
130-
}
131-
132133
progress?.Report(new ProgressReport(-1f, message: "Installing Triton", isIndeterminate: true));
133134

134135
await venvRunner.PipInstall(pipArgs, progress.AsProcessOutputHandler()).ConfigureAwait(false);
@@ -138,19 +139,7 @@ public async Task ExecuteAsync(IProgress<ProgressReport>? progress = null)
138139
progress?.Report(
139140
new ProgressReport(-1f, message: "Downloading Python libraries", isIndeterminate: true)
140141
);
141-
var downloadPath = WorkingDirectory.JoinFile("python_libs_for_sage.zip");
142-
await downloadService
143-
.DownloadToFileAsync(PythonLibsDownloadUrl, downloadPath, progress)
144-
.ConfigureAwait(false);
145-
146-
progress?.Report(
147-
new ProgressReport(-1f, message: "Extracting Python libraries", isIndeterminate: true)
148-
);
149-
await ArchiveHelper.Extract7Z(downloadPath, venvDir, progress).ConfigureAwait(false);
150-
151-
var includeFolder = venvDir.JoinDir("include");
152-
var scriptsIncludeFolder = venvDir.JoinDir("Scripts").JoinDir("include");
153-
await includeFolder.CopyToAsync(scriptsIncludeFolder).ConfigureAwait(false);
142+
await AddMissingLibsToVenv(WorkingDirectory, progress).ConfigureAwait(false);
154143

155144
var sageDir = WorkingDirectory.JoinDir("SageAttention");
156145

@@ -176,5 +165,40 @@ await venvRunner
176165
.ConfigureAwait(false);
177166
}
178167

168+
private async Task AddMissingLibsToVenv(
169+
DirectoryPath installedPackagePath,
170+
IProgress<ProgressReport>? progress = null
171+
)
172+
{
173+
var venvLibsDir = installedPackagePath.JoinDir("venv", "libs");
174+
var venvIncludeDir = installedPackagePath.JoinDir("venv", "include");
175+
if (
176+
venvLibsDir.Exists
177+
&& venvIncludeDir.Exists
178+
&& venvLibsDir.JoinFile("python3.lib").Exists
179+
&& venvLibsDir.JoinFile("python310.lib").Exists
180+
)
181+
{
182+
return;
183+
}
184+
185+
var downloadPath = installedPackagePath.JoinFile("python_libs_for_sage.zip");
186+
var venvDir = installedPackagePath.JoinDir("venv");
187+
await downloadService
188+
.DownloadToFileAsync(PythonLibsDownloadUrl, downloadPath, progress)
189+
.ConfigureAwait(false);
190+
191+
progress?.Report(
192+
new ProgressReport(-1f, message: "Extracting Python libraries", isIndeterminate: true)
193+
);
194+
await ArchiveHelper.Extract7Z(downloadPath, venvDir, progress).ConfigureAwait(false);
195+
196+
var includeFolder = venvDir.JoinDir("include");
197+
var scriptsIncludeFolder = venvDir.JoinDir("Scripts").JoinDir("include");
198+
await includeFolder.CopyToAsync(scriptsIncludeFolder).ConfigureAwait(false);
199+
200+
await downloadPath.DeleteAsync().ConfigureAwait(false);
201+
}
202+
179203
public string ProgressTitle => "Installing Triton and SageAttention";
180204
}

‎StabilityMatrix.Core/Models/Packages/BaseGitPackage.cs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,18 @@ public async Task<PyVenvRunner> SetupVenvPure(
210210
await venvRunner.Setup(true, onConsoleOutput).ConfigureAwait(false);
211211
}
212212

213+
if (!Compat.IsWindows)
214+
return venvRunner;
215+
216+
try
217+
{
218+
await PrerequisiteHelper.AddMissingLibsToVenv(installedPackagePath).ConfigureAwait(false);
219+
}
220+
catch (Exception e)
221+
{
222+
Logger.Warn(e, "Failed to add missing libs to venv");
223+
}
224+
213225
return venvRunner;
214226
}
215227

‎StabilityMatrix.Core/Python/PipInstallArgs.cs‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,8 @@ public PipInstallArgs WithParsedFromRequirementsTxt(
4444
.Where(s => !s.StartsWith('#'))
4545
.Select(s => s.Contains('#') ? s.Substring(0, s.IndexOf('#')) : s)
4646
.Select(s => s.Trim())
47-
.Where(s => !string.IsNullOrWhiteSpace(s));
47+
.Where(s => !string.IsNullOrWhiteSpace(s))
48+
.Select(NormalizePackageSpecifier);
4849

4950
if (excludePattern is not null)
5051
{
@@ -56,6 +57,34 @@ public PipInstallArgs WithParsedFromRequirementsTxt(
5657
return this.AddArgs(requirementsEntries.Select(Argument.Quoted).ToArray());
5758
}
5859

60+
/// <summary>
61+
/// Normalizes a package specifier by removing spaces around version constraint operators.
62+
/// </summary>
63+
/// <param name="specifier">The package specifier to normalize.</param>
64+
/// <returns>The normalized package specifier.</returns>
65+
private static string NormalizePackageSpecifier(string specifier)
66+
{
67+
// Skip normalization for special pip commands that start with a hyphen
68+
if (specifier.StartsWith('-'))
69+
return specifier;
70+
71+
// Regex to match common version constraint patterns with spaces
72+
// Matches: package >= 1.0.0, package <= 1.0.0, package == 1.0.0, etc.
73+
var versionConstraintPattern = new Regex(@"^([a-zA-Z0-9\-_.]+)\s*(>=|<=|==|>|<|!=|~=)\s*(.+)$");
74+
75+
var match = versionConstraintPattern.Match(specifier);
76+
if (match.Success)
77+
{
78+
var packageName = match.Groups[1].Value;
79+
var versionOperator = match.Groups[2].Value;
80+
var version = match.Groups[3].Value;
81+
82+
return $"{packageName}{versionOperator}{version}";
83+
}
84+
85+
return specifier;
86+
}
87+
5988
public PipInstallArgs WithUserOverrides(List<PipPackageSpecifierOverride> overrides)
6089
{
6190
var newArgs = this;

0 commit comments

Comments
 (0)