Skip to content

Commit

Permalink
GFPGAN integration and many other features
Browse files Browse the repository at this point in the history
  • Loading branch information
SocketByte committed Oct 14, 2022
1 parent b2a77da commit 3204c05
Show file tree
Hide file tree
Showing 12 changed files with 1,117 additions and 313 deletions.
42 changes: 36 additions & 6 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,27 +1,55 @@
# SDToolkit

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

![Image](https://i.imgur.com/EyjdVXj.png)
![Image](https://i.imgur.com/kB3mL6e.png)
<div style="padding: 1rem; background-color: #ffad59; border-radius: 1rem; color: black; font-size: 1.15rem;">
<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>
</div>
<br/>

![img](https://i.imgur.com/5c2EjT5.png)
![img](https://i.imgur.com/yXfs4YY.png)

## Features

- No setup required, just install and run
- Video2x upscaler
- GFPGAN face restoration
- Ability to choose a context image for AI to try to match
- Result image viewer with zoom
- Optimized for low VRAM GPUs
- Many helpful tooltips that make it very easy to use
- Decent amount of configuration options

## GFPGAN Integration

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.

![img](https://i.imgur.com/UecbKYL.png)

## Video2x Upscaler

SDToolkit offers a built-in Video2x upscaler. Please remember that it isn't perfect, and upscaling by a very large factor can cause artifacts.

![img](https://i.imgur.com/OQqkZW7.png)

## Requirements / Info

- OS: Windows 64-bit
- At least 25GB of disk space
- At least 35GB of disk space
- At least 8GB of GPU VRAM is recommended
- At least 16GB of RAM is recommended
- CUDA-enabled GPU is recommended (Nvidia)

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

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

Half precision is recommended for RTX cards.

## Installation

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.
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.

## License

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

Upscaling model and algorithm is provided by [video2x](https://github.com/k4yt3x/video2x).

GFPGAN 1.3 model and algorithm provided by [Tencent's GFPGAN](https://github.com/TencentARC/GFPGAN).
115 changes: 115 additions & 0 deletions SDToolkit/FaceRestoration.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,115 @@
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Drawing;
using System.IO;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using System.Windows.Forms;

namespace SDToolkit
{
class FaceRestoration
{
public static void RestoreWithGFPGAN(Generator.GeneratorConfig config, string[] images)
{
var outputNames = new string[images.Length];
for (var i = 0; i < images.Length; i++)
{
outputNames[i] = DateTimeOffset.Now.ToUnixTimeSeconds().ToString() + "_" + i + ".png";
}

var inputs = "GFPGAN\\inputs\\temp";
if (Directory.Exists(inputs))
{
Directory.Delete(inputs, true);
}
Directory.CreateDirectory(inputs);
for (var i = 0; i < images.Length; i++)
{
var image = images[i];
File.Copy(image, inputs + "\\" + outputNames[i]);
}

var process = new Process
{
StartInfo = new ProcessStartInfo
{
FileName = "cmd.exe",
RedirectStandardInput = true,
UseShellExecute = false,
CreateNoWindow = true,
RedirectStandardOutput = true,
RedirectStandardError = true,
WorkingDirectory = Generator.WorkingDirectory,
},
EnableRaisingEvents = true,
};

process.Exited += (s, e) =>
{
if (config.GenerateButton.InvokeRequired)
{
config.GenerateButton.Invoke(new MethodInvoker(delegate ()
{
if (config.UseGFPGANandVideo2x)
{
config.GenerateButton.Invoke(new MethodInvoker(delegate ()
{
config.GenerateButton.Text = "Upscaling with Video2x...";
config.DebugTextBox.Clear();
}));

Upscaler.Upscale(config, images);
return;
}

for (var i = 0; i < images.Length; i++)
{
var image = images[i];
var pictureBox = config.PictureBoxes[i];
pictureBox.Image = Image.FromFile("GFPGAN\\results\\restored_imgs\\" + outputNames[i]);
}
config.GenerateButton.Text = "Generate";

config.GenerateButton.Enabled = true;

config.DebugTextBox.Clear();
config.ProgressBar.Value = 0;
}));
}
};

process.OutputDataReceived += (s, e) =>
{
Console.WriteLine(e.Data);
Generator.PrintToDebug(config.DebugTextBox, e.Data);
};
process.ErrorDataReceived += (s, e) =>
{
Console.WriteLine(e.Data);
Generator.PrintToDebug(config.DebugTextBox, e.Data);
};
process.Start();
process.BeginErrorReadLine();
process.BeginOutputReadLine();

using (var sw = process.StandardInput)
{
if (sw.BaseStream.CanWrite)
{
var scale = config.Upscale;
if (config.UseGFPGANandVideo2x)
{
scale = 2;
}
sw.WriteLine("conda\\Scripts\\activate.bat");
sw.WriteLine("activate");
sw.WriteLine("cd GFPGAN");
sw.WriteLine("python inference_gfpgan.py -i inputs/temp -o results -v 1.3 -s " + scale);
}
}
}
}
}
83 changes: 68 additions & 15 deletions SDToolkit/Generator.cs
Original file line number Diff line number Diff line change
Expand Up @@ -15,18 +15,31 @@ class Generator
{
public static readonly string WorkingDirectory = Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location);

public delegate void OnEnd();

public class GeneratorConfig
{
public string Prompt { get; set; }
public int ResHeight { get; set; }
public int ResWidth { get; set; }
public int Seed { get; set; }
public int Steps { get; set; }
public string UpscalingAlgorithm { get; set; }
public string Precision { get; set; }
public int Upscale { get; set; }
public bool UseContextImage { get; set; }

public string ContextImage { get; set; }

public bool UseTurbo { get; set; }

public bool UseGFPGAN { get; set; }

public bool UseGFPGANandVideo2x { get; set; }

public string Sampler { get; set; }

public OnEnd OnEnd { get; set; }

public TextBox DebugTextBox { get; set; }
public ProgressBar ProgressBar { get; set; }
public Button GenerateButton { get; set; }
Expand Down Expand Up @@ -62,7 +75,8 @@ public static void GenerateFromPrompt(GeneratorConfig config)

private static void Run(GeneratorConfig config)
{
var samplesDir = WorkingDirectory + "\\stable-diffusion\\outputs\\txt2img-samples";
var samples = config.UseContextImage ? "img2img-samples" : "txt2img-samples";
var samplesDir = WorkingDirectory + "\\stable-diffusion\\outputs\\" + samples;
if (Directory.Exists(samplesDir))
{
Directory.Delete(samplesDir, true);
Expand Down Expand Up @@ -92,28 +106,50 @@ private static void Run(GeneratorConfig config)

process.Exited += (s, e) =>
{
var prompt = config.Prompt;
var converted = string.Join("_", prompt.Split(' '));
try
{
var prompt = config.Prompt;
var converted = string.Join("_", prompt.Split(' '));

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

config.GenerateButton.Invoke(new MethodInvoker(delegate ()
{
config.GenerateButton.Text = "Upscaling...";
config.DebugTextBox.Clear();
}));
if (config.UseGFPGAN)
{
config.GenerateButton.Invoke(new MethodInvoker(delegate ()
{
config.GenerateButton.Text = "Restoring faces and upscaling with GFPGAN...";
config.DebugTextBox.Clear();
}));
FaceRestoration.RestoreWithGFPGAN(config, images);
return;
}

Upscaler.Upscale(config, images);
config.GenerateButton.Invoke(new MethodInvoker(delegate ()
{
config.GenerateButton.Text = "Upscaling with Video2x...";
config.DebugTextBox.Clear();
}));

Upscaler.Upscale(config, images);
}
catch (Exception ex)
{
MessageBox.Show(ex.Message, "An error occurred", MessageBoxButtons.OK, MessageBoxIcon.Error);

}
};

process.ErrorDataReceived += (s, e) =>
{
Console.WriteLine(e.Data);
PrintToDebug(config.DebugTextBox, e.Data);

if (e.Data == null)
{
return;
}
if (e.Data.Contains("PLMS Sampler:"))
if (e.Data.Contains("Sampler:") || e.Data.Contains("Decoding image:"))
{
var percentage = e.Data.Split(':')[1].Split('%')[0].Trim();
if (percentage.Length == 0)
Expand All @@ -134,6 +170,7 @@ private static void Run(GeneratorConfig config)

process.OutputDataReceived += (s, e) =>
{
Console.WriteLine(e.Data);
PrintToDebug(config.DebugTextBox, e.Data);
};

Expand All @@ -148,13 +185,29 @@ private static void Run(GeneratorConfig config)
sw.WriteLine("conda\\Scripts\\activate.bat");
sw.WriteLine("activate ldm");
sw.WriteLine("cd stable-diffusion");
sw.WriteLine("python optimizedSD/optimized_txt2img.py --prompt \"" + config.Prompt + "\""
if (!config.UseContextImage)
{
sw.WriteLine("python optimizedSD/optimized_txt2img.py --prompt \"" + config.Prompt + "\""
+ " --sampler " + config.Sampler.ToLower()
+ " --H " + config.ResHeight
+ " --W " + config.ResWidth
+ " --seed " + config.Seed
+ (config.UseTurbo ? " --turbo" : "")
+ (config.Precision == "Full" ? " --precision full" : "")
+ " --precision " + config.Precision
+ " --n_iter 1 --n_samples 8 --ddim_steps " + config.Steps);
}
else
{
sw.WriteLine("python optimizedSD/optimized_img2img.py --strength 0.8 --prompt \"" + config.Prompt + "\""
+ " --init-img " + config.ContextImage
+ " --H " + config.ResHeight
+ " --W " + config.ResWidth
+ " --seed " + config.Seed
+ (config.UseTurbo ? " --turbo" : "")
+ " --precision " + config.Precision
+ " --n_iter 1 --n_samples 8 --ddim_steps " + config.Steps);

}
}
}
}
Expand Down
Loading

0 comments on commit 3204c05

Please sign in to comment.