Skip to content

Commit 3204c05

Browse files
SocketByteSocketByte
committed
GFPGAN integration and many other features
1 parent b2a77da commit 3204c05

File tree

12 files changed

+1117
-313
lines changed

12 files changed

+1117
-313
lines changed

‎README.md

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,55 @@
11
# SDToolkit
22

3-
All in one, batteries-included software to easily generate and upscale AI art using [Stable Diffusion](https://github.com/CompVis/stable-diffusion).
3+
All in one, batteries-included software to easily generate and upscale AI art using [Stable Diffusion](https://github.com/CompVis/stable-diffusion)
44

5-
![Image](https://i.imgur.com/EyjdVXj.png)
6-
![Image](https://i.imgur.com/kB3mL6e.png)
5+
<div style="padding: 1rem; background-color: #ffad59; border-radius: 1rem; color: black; font-size: 1.15rem;">
6+
<p style="padding: 0; margin: 0;"><b>Disclaimer: </b>This software is still in beta, bugs can and will occur at some point. Please use the <i>Issues</i> tab to report a bug. Thank you for using SDToolkit!</p>
7+
</div>
8+
<br/>
9+
10+
![img](https://i.imgur.com/5c2EjT5.png)
11+
![img](https://i.imgur.com/yXfs4YY.png)
12+
13+
## Features
14+
15+
- No setup required, just install and run
16+
- Video2x upscaler
17+
- GFPGAN face restoration
18+
- Ability to choose a context image for AI to try to match
19+
- Result image viewer with zoom
20+
- Optimized for low VRAM GPUs
21+
- Many helpful tooltips that make it very easy to use
22+
- Decent amount of configuration options
23+
24+
## GFPGAN Integration
25+
26+
SDToolkit offers a built-in GFPGAN for face restoration. This is quite a powerful tool to remove face artifacts and beautify the result. It's definitely recommended to use GFPGAN for every prompt that might have faces in it.
27+
28+
![img](https://i.imgur.com/UecbKYL.png)
29+
30+
## Video2x Upscaler
31+
32+
SDToolkit offers a built-in Video2x upscaler. Please remember that it isn't perfect, and upscaling by a very large factor can cause artifacts.
33+
34+
![img](https://i.imgur.com/OQqkZW7.png)
735

836
## Requirements / Info
937

1038
- OS: Windows 64-bit
11-
- At least 25GB of disk space
39+
- At least 35GB of disk space
1240
- At least 8GB of GPU VRAM is recommended
1341
- At least 16GB of RAM is recommended
1442
- CUDA-enabled GPU is recommended (Nvidia)
1543

16-
Since SDToolkit is all-in-one software, the model, upscaling software (video2x) and execution environment (Conda/Python) is included in the setup.
44+
Since SDToolkit is all-in-one software, the model, upscaling software (video2x/GFPGAN) and execution environment (Conda/Python) is included in the setup.
1745

1846
Tested with i7-12700KF and GTX 1080 with full precision. (around 6.7GB of VRAM usage)
1947

2048
Half precision is recommended for RTX cards.
2149

2250
## Installation
2351

24-
Just download and run the setup from the Releases tab. Be reminded that it's a huge file, you should have at least 20-25GB of free disk space.
52+
Just download and run the setup from the Releases tab. Be reminded that it's a huge file, you should have at least 35GB of free disk space.
2553

2654
## License
2755

@@ -35,3 +63,5 @@ You're free to download it yourself at [CompVis's huggingface repository](https:
3563
Scripts and model execution software provided by [basujindal's fork of stable diffusion](https://github.com/basujindal/stable-diffusion/).
3664

3765
Upscaling model and algorithm is provided by [video2x](https://github.com/k4yt3x/video2x).
66+
67+
GFPGAN 1.3 model and algorithm provided by [Tencent's GFPGAN](https://github.com/TencentARC/GFPGAN).

‎SDToolkit/FaceRestoration.cs

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.Diagnostics;
4+
using System.Drawing;
5+
using System.IO;
6+
using System.Linq;
7+
using System.Text;
8+
using System.Threading.Tasks;
9+
using System.Windows.Forms;
10+
11+
namespace SDToolkit
12+
{
13+
class FaceRestoration
14+
{
15+
public static void RestoreWithGFPGAN(Generator.GeneratorConfig config, string[] images)
16+
{
17+
var outputNames = new string[images.Length];
18+
for (var i = 0; i < images.Length; i++)
19+
{
20+
outputNames[i] = DateTimeOffset.Now.ToUnixTimeSeconds().ToString() + "_" + i + ".png";
21+
}
22+
23+
var inputs = "GFPGAN\\inputs\\temp";
24+
if (Directory.Exists(inputs))
25+
{
26+
Directory.Delete(inputs, true);
27+
}
28+
Directory.CreateDirectory(inputs);
29+
for (var i = 0; i < images.Length; i++)
30+
{
31+
var image = images[i];
32+
File.Copy(image, inputs + "\\" + outputNames[i]);
33+
}
34+
35+
var process = new Process
36+
{
37+
StartInfo = new ProcessStartInfo
38+
{
39+
FileName = "cmd.exe",
40+
RedirectStandardInput = true,
41+
UseShellExecute = false,
42+
CreateNoWindow = true,
43+
RedirectStandardOutput = true,
44+
RedirectStandardError = true,
45+
WorkingDirectory = Generator.WorkingDirectory,
46+
},
47+
EnableRaisingEvents = true,
48+
};
49+
50+
process.Exited += (s, e) =>
51+
{
52+
if (config.GenerateButton.InvokeRequired)
53+
{
54+
config.GenerateButton.Invoke(new MethodInvoker(delegate ()
55+
{
56+
if (config.UseGFPGANandVideo2x)
57+
{
58+
config.GenerateButton.Invoke(new MethodInvoker(delegate ()
59+
{
60+
config.GenerateButton.Text = "Upscaling with Video2x...";
61+
config.DebugTextBox.Clear();
62+
}));
63+
64+
Upscaler.Upscale(config, images);
65+
return;
66+
}
67+
68+
for (var i = 0; i < images.Length; i++)
69+
{
70+
var image = images[i];
71+
var pictureBox = config.PictureBoxes[i];
72+
pictureBox.Image = Image.FromFile("GFPGAN\\results\\restored_imgs\\" + outputNames[i]);
73+
}
74+
config.GenerateButton.Text = "Generate";
75+
76+
config.GenerateButton.Enabled = true;
77+
78+
config.DebugTextBox.Clear();
79+
config.ProgressBar.Value = 0;
80+
}));
81+
}
82+
};
83+
84+
process.OutputDataReceived += (s, e) =>
85+
{
86+
Console.WriteLine(e.Data);
87+
Generator.PrintToDebug(config.DebugTextBox, e.Data);
88+
};
89+
process.ErrorDataReceived += (s, e) =>
90+
{
91+
Console.WriteLine(e.Data);
92+
Generator.PrintToDebug(config.DebugTextBox, e.Data);
93+
};
94+
process.Start();
95+
process.BeginErrorReadLine();
96+
process.BeginOutputReadLine();
97+
98+
using (var sw = process.StandardInput)
99+
{
100+
if (sw.BaseStream.CanWrite)
101+
{
102+
var scale = config.Upscale;
103+
if (config.UseGFPGANandVideo2x)
104+
{
105+
scale = 2;
106+
}
107+
sw.WriteLine("conda\\Scripts\\activate.bat");
108+
sw.WriteLine("activate");
109+
sw.WriteLine("cd GFPGAN");
110+
sw.WriteLine("python inference_gfpgan.py -i inputs/temp -o results -v 1.3 -s " + scale);
111+
}
112+
}
113+
}
114+
}
115+
}

‎SDToolkit/Generator.cs

Lines changed: 68 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,18 +15,31 @@ class Generator
1515
{
1616
public static readonly string WorkingDirectory = Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location);
1717

18+
public delegate void OnEnd();
19+
1820
public class GeneratorConfig
1921
{
2022
public string Prompt { get; set; }
2123
public int ResHeight { get; set; }
2224
public int ResWidth { get; set; }
2325
public int Seed { get; set; }
2426
public int Steps { get; set; }
25-
public string UpscalingAlgorithm { get; set; }
2627
public string Precision { get; set; }
28+
public int Upscale { get; set; }
29+
public bool UseContextImage { get; set; }
30+
31+
public string ContextImage { get; set; }
2732

2833
public bool UseTurbo { get; set; }
2934

35+
public bool UseGFPGAN { get; set; }
36+
37+
public bool UseGFPGANandVideo2x { get; set; }
38+
39+
public string Sampler { get; set; }
40+
41+
public OnEnd OnEnd { get; set; }
42+
3043
public TextBox DebugTextBox { get; set; }
3144
public ProgressBar ProgressBar { get; set; }
3245
public Button GenerateButton { get; set; }
@@ -62,7 +75,8 @@ public static void GenerateFromPrompt(GeneratorConfig config)
6275

6376
private static void Run(GeneratorConfig config)
6477
{
65-
var samplesDir = WorkingDirectory + "\\stable-diffusion\\outputs\\txt2img-samples";
78+
var samples = config.UseContextImage ? "img2img-samples" : "txt2img-samples";
79+
var samplesDir = WorkingDirectory + "\\stable-diffusion\\outputs\\" + samples;
6680
if (Directory.Exists(samplesDir))
6781
{
6882
Directory.Delete(samplesDir, true);
@@ -92,28 +106,50 @@ private static void Run(GeneratorConfig config)
92106

93107
process.Exited += (s, e) =>
94108
{
95-
var prompt = config.Prompt;
96-
var converted = string.Join("_", prompt.Split(' '));
109+
try
110+
{
111+
var prompt = config.Prompt;
112+
var converted = string.Join("_", prompt.Split(' '));
97113

98-
var dir = Directory.GetDirectories(WorkingDirectory + @"\stable-diffusion\outputs\txt2img-samples\")[0];
99-
var images = Directory.GetFiles(dir, "*.png");
114+
var dir = Directory.GetDirectories(WorkingDirectory + @"\stable-diffusion\outputs\" + samples)[0];
115+
var images = Directory.GetFiles(dir, "*.png");
100116

101-
config.GenerateButton.Invoke(new MethodInvoker(delegate ()
102-
{
103-
config.GenerateButton.Text = "Upscaling...";
104-
config.DebugTextBox.Clear();
105-
}));
117+
if (config.UseGFPGAN)
118+
{
119+
config.GenerateButton.Invoke(new MethodInvoker(delegate ()
120+
{
121+
config.GenerateButton.Text = "Restoring faces and upscaling with GFPGAN...";
122+
config.DebugTextBox.Clear();
123+
}));
124+
FaceRestoration.RestoreWithGFPGAN(config, images);
125+
return;
126+
}
106127

107-
Upscaler.Upscale(config, images);
128+
config.GenerateButton.Invoke(new MethodInvoker(delegate ()
129+
{
130+
config.GenerateButton.Text = "Upscaling with Video2x...";
131+
config.DebugTextBox.Clear();
132+
}));
133+
134+
Upscaler.Upscale(config, images);
135+
}
136+
catch (Exception ex)
137+
{
138+
MessageBox.Show(ex.Message, "An error occurred", MessageBoxButtons.OK, MessageBoxIcon.Error);
139+
140+
}
108141
};
109142

110143
process.ErrorDataReceived += (s, e) =>
111144
{
145+
Console.WriteLine(e.Data);
146+
PrintToDebug(config.DebugTextBox, e.Data);
147+
112148
if (e.Data == null)
113149
{
114150
return;
115151
}
116-
if (e.Data.Contains("PLMS Sampler:"))
152+
if (e.Data.Contains("Sampler:") || e.Data.Contains("Decoding image:"))
117153
{
118154
var percentage = e.Data.Split(':')[1].Split('%')[0].Trim();
119155
if (percentage.Length == 0)
@@ -134,6 +170,7 @@ private static void Run(GeneratorConfig config)
134170

135171
process.OutputDataReceived += (s, e) =>
136172
{
173+
Console.WriteLine(e.Data);
137174
PrintToDebug(config.DebugTextBox, e.Data);
138175
};
139176

@@ -148,13 +185,29 @@ private static void Run(GeneratorConfig config)
148185
sw.WriteLine("conda\\Scripts\\activate.bat");
149186
sw.WriteLine("activate ldm");
150187
sw.WriteLine("cd stable-diffusion");
151-
sw.WriteLine("python optimizedSD/optimized_txt2img.py --prompt \"" + config.Prompt + "\""
188+
if (!config.UseContextImage)
189+
{
190+
sw.WriteLine("python optimizedSD/optimized_txt2img.py --prompt \"" + config.Prompt + "\""
191+
+ " --sampler " + config.Sampler.ToLower()
152192
+ " --H " + config.ResHeight
153193
+ " --W " + config.ResWidth
154194
+ " --seed " + config.Seed
155195
+ (config.UseTurbo ? " --turbo" : "")
156-
+ (config.Precision == "Full" ? " --precision full" : "")
196+
+ " --precision " + config.Precision
157197
+ " --n_iter 1 --n_samples 8 --ddim_steps " + config.Steps);
198+
}
199+
else
200+
{
201+
sw.WriteLine("python optimizedSD/optimized_img2img.py --strength 0.8 --prompt \"" + config.Prompt + "\""
202+
+ " --init-img " + config.ContextImage
203+
+ " --H " + config.ResHeight
204+
+ " --W " + config.ResWidth
205+
+ " --seed " + config.Seed
206+
+ (config.UseTurbo ? " --turbo" : "")
207+
+ " --precision " + config.Precision
208+
+ " --n_iter 1 --n_samples 8 --ddim_steps " + config.Steps);
209+
210+
}
158211
}
159212
}
160213
}

0 commit comments

Comments
 (0)