aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--.gitignore7
-rw-r--r--README.md57
-rw-r--r--SwinIR/put_swinir_models_here.txt1
-rw-r--r--artists.csv2
-rw-r--r--embeddings/Place Textual Inversion embeddings here.txt (renamed from ESRGAN/Put ESRGAN models here.txt)0
-rw-r--r--javascript/hints.js5
-rw-r--r--javascript/ui.js10
-rw-r--r--launch.py9
-rw-r--r--models/Stable-diffusion/Put Stable Diffusion checkpoints here.txt (renamed from models/Put Stable Diffusion checkpoints here.txt)0
-rw-r--r--modules/bsrgan_model.py78
-rw-r--r--modules/bsrgan_model_arch.py102
-rw-r--r--modules/codeformer_model.py44
-rw-r--r--modules/esrgan_model.py106
-rw-r--r--modules/extras.py11
-rw-r--r--modules/gfpgan_model.py90
-rw-r--r--modules/images.py75
-rw-r--r--modules/ldsr_model.py103
-rw-r--r--modules/ldsr_model_arch.py222
-rw-r--r--modules/modelloader.py140
-rw-r--r--modules/paths.py22
-rw-r--r--modules/processing.py3
-rw-r--r--modules/prompt_parser.py90
-rw-r--r--modules/realesrgan_model.py192
-rw-r--r--modules/scripts.py2
-rw-r--r--modules/sd_hijack.py118
-rw-r--r--modules/sd_models.py60
-rw-r--r--modules/sd_samplers.py16
-rw-r--r--modules/shared.py46
-rw-r--r--modules/styles.py6
-rw-r--r--modules/swinir.py123
-rw-r--r--modules/swinir_model.py139
-rw-r--r--modules/swinir_model_arch.py (renamed from modules/swinir_arch.py)1734
-rw-r--r--modules/ui.py75
-rw-r--r--modules/upscaler.py121
-rw-r--r--requirements.txt2
-rw-r--r--scripts/outpainting_mk_2.py40
-rw-r--r--scripts/sd_upscale.py2
-rw-r--r--scripts/xy_grid.py10
-rw-r--r--style.css14
-rw-r--r--webui-user.sh3
-rw-r--r--webui.py51
-rwxr-xr-xwebui.sh7
42 files changed, 2487 insertions, 1451 deletions
diff --git a/.gitignore b/.gitignore
index b71e1875..3532dab3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,10 +1,13 @@
__pycache__
-/ESRGAN
+*.ckpt
+*.pth
+/ESRGAN/*
+/SwinIR/*
/repositories
/venv
/tmp
/model.ckpt
-/models/**/*.ckpt
+/models/**/*
/GFPGANv1.3.pth
/gfpgan/weights/*.pth
/ui-config.json
diff --git a/README.md b/README.md
index 50e222ee..5ded94f9 100644
--- a/README.md
+++ b/README.md
@@ -3,50 +3,64 @@ A browser interface based on Gradio library for Stable Diffusion.
![](txt2img_Screenshot.png)
+Check the [custom scripts](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Custom-Scripts) wiki page for extra scripts developed by users.
+
## Features
[Detailed feature showcase with images](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features):
- Original txt2img and img2img modes
- One click install and run script (but you still must install python and git)
- Outpainting
- Inpainting
-- Prompt matrix
+- Prompt
- Stable Diffusion upscale
-- Attention
-- Loopback
-- X/Y plot
+- Attention, specify parts of text that the model should pay more attention to
+ - a man in a ((txuedo)) - will pay more attentinoto tuxedo
+ - a man in a (txuedo:1.21) - alternative syntax
+- Loopback, run img2img procvessing multiple times
+- X/Y plot, a way to draw a 2 dimensional plot of images with different parameters
- Textual Inversion
+ - have as many embeddings as you want and use any names you like for them
+ - use multiple embeddings with different numbers of vectors per token
+ - works with half precision floating point numbers
- Extras tab with:
- GFPGAN, neural network that fixes faces
- CodeFormer, face restoration tool as an alternative to GFPGAN
- RealESRGAN, neural network upscaler
- - ESRGAN, neural network with a lot of third party models
+ - ESRGAN, neural network upscaler with a lot of third party models
- SwinIR, neural network upscaler
- LDSR, Latent diffusion super resolution upscaling
- Resizing aspect ratio options
- Sampling method selection
- Interrupt processing at any time
-- 4GB video card support
-- Correct seeds for batches
+- 4GB video card support (also reports of 2GB working)
+- Correct seeds for batches
- Prompt length validation
-- Generation parameters added as text to PNG
-- Tab to view an existing picture's generation parameters
+ - get length of prompt in tokensas you type
+ - get a warning after geenration if some text was truncated
+- Generation parameters
+ - parameters you used to generate images are saved with that image
+ - in PNG chunks for PNG, in EXIF for JPEG
+ - can drag the image to PNG info tab to restore generation parameters and automatically copy them into UI
+ - can be disabled in settings
- Settings page
-- Running custom code from UI
+- Running arbitrary python code from UI (must run with commandline flag to enable)
- Mouseover hints for most UI elements
- Possible to change defaults/mix/max/step values for UI elements via text config
- Random artist button
-- Tiling support: UI checkbox to create images that can be tiled like textures
+- Tiling support, a checkbox to create images that can be tiled like textures
- Progress bar and live image generation preview
-- Negative prompt
-- Styles
-- Variations
-- Seed resizing
-- CLIP interrogator
-- Prompt Editing
-- Batch Processing
+- Negative prompt, an extra text field that allows you to list what you don't want to see in generated image
+- Styles, a way to save part of prompt and easily apply them via dropdown later
+- Variations, a way to generate same image but with tiny differences
+- Seed resizing, a way to generate same image but at slightly different resolution
+- CLIP interrogator, a button that tries to guess prompt from an image
+- Prompt Editing, a way to change prompt mid-generation, say to start making a watermelon and switch to anime girl midway
+- Batch Processing, process a group of files using img2img
- Img2img Alternative
-- Highres Fix
-- LDSR Upscaling
+- Highres Fix, a convenience option to produce high resolution pictures in one click without usual distortions
+- Reloading checkpoints on the fly
+- Checkpoint Merger, a tab that allows you to merge two checkpoints into one
+- [Custom scripts](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Custom-Scripts) with many extensions from community
## Installation and Running
Make sure the required [dependencies](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Dependencies) are met and follow the instructions available for both [NVidia](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-NVidia-GPUs) (recommended) and [AMD](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-AMD-GPUs) GPUs.
@@ -83,6 +97,9 @@ bash <(wget -qO- https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusio
Find the instructions [here](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Installation-on-Apple-Silicon).
+## Contributing
+Here's how to add code to this repo: [Contributing](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing)
+
## Documentation
The documentation was moved from this README over to the project's [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki).
diff --git a/SwinIR/put_swinir_models_here.txt b/SwinIR/put_swinir_models_here.txt
deleted file mode 100644
index 8b137891..00000000
--- a/SwinIR/put_swinir_models_here.txt
+++ /dev/null
@@ -1 +0,0 @@
-
diff --git a/artists.csv b/artists.csv
index c92d08f5..14ba2022 100644
--- a/artists.csv
+++ b/artists.csv
@@ -359,7 +359,6 @@ Antanas Sutkus,0.7369492,black-white
Leonora Carrington,0.73726475,scribbles
Hieronymus Bosch,0.7369955,scribbles
A. J. Casson,0.73666203,scribbles
-A.J.Casson,0.73666203,scribbles
Chaim Soutine,0.73662066,scribbles
Artur Bordalo,0.7364549,weird
Thomas Allom,0.68792284,fineart
@@ -1907,7 +1906,6 @@ Alex Schomburg,0.46614102,digipa-low-impact
Bastien L. Deharme,0.583349,special
František Jakub Prokyš,0.58782333,fineart
Jesper Ejsing,0.58782053,fineart
-Jesper Ejsing,0.58782053,fineart
Odd Nerdrum,0.53551745,digipa-high-impact
Tom Lovell,0.5877577,fineart
Ayami Kojima,0.5877416,fineart
diff --git a/ESRGAN/Put ESRGAN models here.txt b/embeddings/Place Textual Inversion embeddings here.txt
index e69de29b..e69de29b 100644
--- a/ESRGAN/Put ESRGAN models here.txt
+++ b/embeddings/Place Textual Inversion embeddings here.txt
diff --git a/javascript/hints.js b/javascript/hints.js
index 59dd770c..84694eeb 100644
--- a/javascript/hints.js
+++ b/javascript/hints.js
@@ -15,6 +15,7 @@ titles = {
"\u267b\ufe0f": "Reuse seed from last generation, mostly useful if it was randomed",
"\u{1f3a8}": "Add a random artist to the prompt.",
"\u2199\ufe0f": "Read generation parameters from prompt into user interface.",
+ "\uD83D\uDCC2": "Open images output directory",
"Inpaint a part of image": "Draw a mask over an image, and the script will regenerate the masked area with content according to prompt",
"SD upscale": "Upscale image normally, split result into tiles, improve each tile using img2img, merge whole image back",
@@ -57,8 +58,8 @@ titles = {
"Interrogate": "Reconstruct prompt from existing image and put it into the prompt field.",
- "Images filename pattern": "Use following tags to define how filenames for images are chosen: [steps], [cfg], [prompt], [prompt_spaces], [width], [height], [styles], [sampler], [seed], [model_hash], [prompt_words], [date], [datetime], [job_timestamp]; leave empty for default.",
- "Directory name pattern": "Use following tags to define how subdirectories for images and grids are chosen: [steps], [cfg], [prompt], [prompt_spaces], [width], [height], [styles], [sampler], [seed], [model_hash], [prompt_words], [date], [datetime], [job_timestamp]; leave empty for default.",
+ "Images filename pattern": "Use following tags to define how filenames for images are chosen: [steps], [cfg], [prompt], [prompt_no_styles], [prompt_spaces], [width], [height], [styles], [sampler], [seed], [model_hash], [prompt_words], [date], [datetime], [job_timestamp]; leave empty for default.",
+ "Directory name pattern": "Use following tags to define how subdirectories for images and grids are chosen: [steps], [cfg], [prompt], [prompt_no_styles], [prompt_spaces], [width], [height], [styles], [sampler], [seed], [model_hash], [prompt_words], [date], [datetime], [job_timestamp]; leave empty for default.",
"Max prompt words": "Set the maximum number of words to be used in the [prompt_words] option; ATTENTION: If the words are too long, they may exceed the maximum length of the file path that the system can handle",
"Loopback": "Process an image, use it as an input, repeat.",
diff --git a/javascript/ui.js b/javascript/ui.js
index 562d2552..bfe02410 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -186,10 +186,12 @@ onUiUpdate(function(){
if (!txt2img_textarea) {
txt2img_textarea = gradioApp().querySelector("#txt2img_prompt > label > textarea");
txt2img_textarea?.addEventListener("input", () => update_token_counter("txt2img_token_button"));
+ txt2img_textarea?.addEventListener("keyup", (event) => submit_prompt(event, "txt2img_generate"));
}
if (!img2img_textarea) {
img2img_textarea = gradioApp().querySelector("#img2img_prompt > label > textarea");
img2img_textarea?.addEventListener("input", () => update_token_counter("img2img_token_button"));
+ img2img_textarea?.addEventListener("keyup", (event) => submit_prompt(event, "img2img_generate"));
}
})
@@ -197,6 +199,14 @@ let txt2img_textarea, img2img_textarea = undefined;
let wait_time = 800
let token_timeout;
+function submit_prompt(event, generate_button_id) {
+ if (event.altKey && event.keyCode === 13) {
+ event.preventDefault();
+ gradioApp().getElementById(generate_button_id).click();
+ return;
+ }
+}
+
function update_token_counter(button_id) {
if (token_timeout)
clearTimeout(token_timeout);
diff --git a/launch.py b/launch.py
index 0e6b64ab..d2793ed2 100644
--- a/launch.py
+++ b/launch.py
@@ -1,5 +1,4 @@
# this scripts installs necessary requirements and launches main program in webui.py
-
import subprocess
import os
import sys
@@ -19,10 +18,9 @@ gfpgan_package = os.environ.get('GFPGAN_PACKAGE', "git+https://github.com/Tencen
stable_diffusion_commit_hash = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', "69ae4b35e0a0f6ee1af8bb9a5d0016ccb27e36dc")
taming_transformers_commit_hash = os.environ.get('TAMING_TRANSFORMERS_COMMIT_HASH', "24268930bf1dce879235a7fddd0b2355b84d7ea6")
-k_diffusion_commit_hash = os.environ.get('K_DIFFUSION_COMMIT_HASH', "9e3002b7cd64df7870e08527b7664eb2f2f5f3f5")
+k_diffusion_commit_hash = os.environ.get('K_DIFFUSION_COMMIT_HASH', "a7ec1974d4ccb394c2dca275f42cd97490618924")
codeformer_commit_hash = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")
blip_commit_hash = os.environ.get('BLIP_COMMIT_HASH', "48211a1594f1321b00f14c9f7a5b4813144b2fb9")
-ldsr_commit_hash = os.environ.get('LDSR_COMMIT_HASH', "abf33e7002d59d9085081bce93ec798dcabd49af")
args = shlex.split(commandline_args)
@@ -120,8 +118,6 @@ git_clone("https://github.com/CompVis/taming-transformers.git", repo_dir('taming
git_clone("https://github.com/crowsonkb/k-diffusion.git", repo_dir('k-diffusion'), "K-diffusion", k_diffusion_commit_hash)
git_clone("https://github.com/sczhou/CodeFormer.git", repo_dir('CodeFormer'), "CodeFormer", codeformer_commit_hash)
git_clone("https://github.com/salesforce/BLIP.git", repo_dir('BLIP'), "BLIP", blip_commit_hash)
-# Using my repo until my changes are merged, as this makes interfacing with our version of SD-web a lot easier
-git_clone("https://github.com/Hafiidz/latent-diffusion", repo_dir('latent-diffusion'), "LDSR", ldsr_commit_hash)
if not is_installed("lpips"):
run_pip(f"install -r {os.path.join(repo_dir('CodeFormer'), 'requirements.txt')}", "requirements for CodeFormer")
@@ -130,6 +126,9 @@ run_pip(f"install -r {requirements_file}", "requirements for Web UI")
sys.argv += args
+if "--exit" in args:
+ print("Exiting because of --exit argument")
+ exit(0)
def start_webui():
print(f"Launching Web UI with arguments: {' '.join(sys.argv[1:])}")
diff --git a/models/Put Stable Diffusion checkpoints here.txt b/models/Stable-diffusion/Put Stable Diffusion checkpoints here.txt
index e69de29b..e69de29b 100644
--- a/models/Put Stable Diffusion checkpoints here.txt
+++ b/models/Stable-diffusion/Put Stable Diffusion checkpoints here.txt
diff --git a/modules/bsrgan_model.py b/modules/bsrgan_model.py
new file mode 100644
index 00000000..e62c6657
--- /dev/null
+++ b/modules/bsrgan_model.py
@@ -0,0 +1,78 @@
+import os.path
+import sys
+import traceback
+
+import PIL.Image
+import numpy as np
+import torch
+from basicsr.utils.download_util import load_file_from_url
+
+import modules.upscaler
+from modules import shared, modelloader
+from modules.bsrgan_model_arch import RRDBNet
+from modules.paths import models_path
+
+
+class UpscalerBSRGAN(modules.upscaler.Upscaler):
+ def __init__(self, dirname):
+ self.name = "BSRGAN"
+ self.model_path = os.path.join(models_path, self.name)
+ self.model_name = "BSRGAN 4x"
+ self.model_url = "https://github.com/cszn/KAIR/releases/download/v1.0/BSRGAN.pth"
+ self.user_path = dirname
+ super().__init__()
+ model_paths = self.find_models(ext_filter=[".pt", ".pth"])
+ scalers = []
+ if len(model_paths) == 0:
+ scaler_data = modules.upscaler.UpscalerData(self.model_name, self.model_url, self, 4)
+ scalers.append(scaler_data)
+ for file in model_paths:
+ if "http" in file:
+ name = self.model_name
+ else:
+ name = modelloader.friendly_name(file)
+ try:
+ scaler_data = modules.upscaler.UpscalerData(name, file, self, 4)
+ scalers.append(scaler_data)
+ except Exception:
+ print(f"Error loading BSRGAN model: {file}", file=sys.stderr)
+ print(traceback.format_exc(), file=sys.stderr)
+ self.scalers = scalers
+
+ def do_upscale(self, img: PIL.Image, selected_file):
+ torch.cuda.empty_cache()
+ model = self.load_model(selected_file)
+ if model is None:
+ return img
+ model.to(shared.device)
+ torch.cuda.empty_cache()
+ img = np.array(img)
+ img = img[:, :, ::-1]
+ img = np.moveaxis(img, 2, 0) / 255
+ img = torch.from_numpy(img).float()
+ img = img.unsqueeze(0).to(shared.device)
+ with torch.no_grad():
+ output = model(img)
+ output = output.squeeze().float().cpu().clamp_(0, 1).numpy()
+ output = 255. * np.moveaxis(output, 0, 2)
+ output = output.astype(np.uint8)
+ output = output[:, :, ::-1]
+ torch.cuda.empty_cache()
+ return PIL.Image.fromarray(output, 'RGB')
+
+ def load_model(self, path: str):
+ if "http" in path:
+ filename = load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="%s.pth" % self.name,
+ progress=True)
+ else:
+ filename = path
+ if not os.path.exists(filename) or filename is None:
+ print(f"BSRGAN: Unable to load model from {filename}", file=sys.stderr)
+ return None
+ model = RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=4) # define network
+ model.load_state_dict(torch.load(filename), strict=True)
+ model.eval()
+ for k, v in model.named_parameters():
+ v.requires_grad = False
+ return model
+
diff --git a/modules/bsrgan_model_arch.py b/modules/bsrgan_model_arch.py
new file mode 100644
index 00000000..cb4d1c13
--- /dev/null
+++ b/modules/bsrgan_model_arch.py
@@ -0,0 +1,102 @@
+import functools
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+
+
+def initialize_weights(net_l, scale=1):
+ if not isinstance(net_l, list):
+ net_l = [net_l]
+ for net in net_l:
+ for m in net.modules():
+ if isinstance(m, nn.Conv2d):
+ init.kaiming_normal_(m.weight, a=0, mode='fan_in')
+ m.weight.data *= scale # for residual block
+ if m.bias is not None:
+ m.bias.data.zero_()
+ elif isinstance(m, nn.Linear):
+ init.kaiming_normal_(m.weight, a=0, mode='fan_in')
+ m.weight.data *= scale
+ if m.bias is not None:
+ m.bias.data.zero_()
+ elif isinstance(m, nn.BatchNorm2d):
+ init.constant_(m.weight, 1)
+ init.constant_(m.bias.data, 0.0)
+
+
+def make_layer(block, n_layers):
+ layers = []
+ for _ in range(n_layers):
+ layers.append(block())
+ return nn.Sequential(*layers)
+
+
+class ResidualDenseBlock_5C(nn.Module):
+ def __init__(self, nf=64, gc=32, bias=True):
+ super(ResidualDenseBlock_5C, self).__init__()
+ # gc: growth channel, i.e. intermediate channels
+ self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)
+ self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)
+ self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)
+ self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)
+ self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)
+ self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+ # initialization
+ initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+ def forward(self, x):
+ x1 = self.lrelu(self.conv1(x))
+ x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+ x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+ x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+ x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+ return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+ '''Residual in Residual Dense Block'''
+
+ def __init__(self, nf, gc=32):
+ super(RRDB, self).__init__()
+ self.RDB1 = ResidualDenseBlock_5C(nf, gc)
+ self.RDB2 = ResidualDenseBlock_5C(nf, gc)
+ self.RDB3 = ResidualDenseBlock_5C(nf, gc)
+
+ def forward(self, x):
+ out = self.RDB1(x)
+ out = self.RDB2(out)
+ out = self.RDB3(out)
+ return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+ def __init__(self, in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=4):
+ super(RRDBNet, self).__init__()
+ RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)
+ self.sf = sf
+
+ self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True)
+ self.RRDB_trunk = make_layer(RRDB_block_f, nb)
+ self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+ #### upsampling
+ self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+ if self.sf==4:
+ self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+ self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
+ self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)
+
+ self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+ def forward(self, x):
+ fea = self.conv_first(x)
+ trunk = self.trunk_conv(self.RRDB_trunk(fea))
+ fea = fea + trunk
+
+ fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest')))
+ if self.sf==4:
+ fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))
+ out = self.conv_last(self.lrelu(self.HRconv(fea)))
+
+ return out \ No newline at end of file
diff --git a/modules/codeformer_model.py b/modules/codeformer_model.py
index 2177291a..8769e1db 100644
--- a/modules/codeformer_model.py
+++ b/modules/codeformer_model.py
@@ -5,31 +5,31 @@ import traceback
import cv2
import torch
-from modules import shared, devices
-from modules.paths import script_path
-import modules.shared
import modules.face_restoration
-from importlib import reload
+import modules.shared
+from modules import shared, devices, modelloader
+from modules.paths import script_path, models_path
-# codeformer people made a choice to include modified basicsr librry to their projectwhich makes
-# it utterly impossiblr to use it alongside with other libraries that also use basicsr, like GFPGAN.
+# codeformer people made a choice to include modified basicsr library to their project which makes
+# it utterly impossible to use it alongside with other libraries that also use basicsr, like GFPGAN.
# I am making a choice to include some files from codeformer to work around this issue.
-
-pretrain_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
+model_dir = "Codeformer"
+model_path = os.path.join(models_path, model_dir)
+model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
have_codeformer = False
codeformer = None
-def setup_codeformer():
+
+def setup_model(dirname):
+ global model_path
+ if not os.path.exists(model_path):
+ os.makedirs(model_path)
+
path = modules.paths.paths.get("CodeFormer", None)
if path is None:
return
-
- # both GFPGAN and CodeFormer use bascisr, one has it installed from pip the other uses its own
- #stored_sys_path = sys.path
- #sys.path = [path] + sys.path
-
try:
from torchvision.transforms.functional import normalize
from modules.codeformer.codeformer_arch import CodeFormer
@@ -44,18 +44,23 @@ def setup_codeformer():
def name(self):
return "CodeFormer"
- def __init__(self):
+ def __init__(self, dirname):
self.net = None
self.face_helper = None
+ self.cmd_dir = dirname
def create_models(self):
if self.net is not None and self.face_helper is not None:
self.net.to(devices.device_codeformer)
return self.net, self.face_helper
-
+ model_paths = modelloader.load_models(model_path, model_url, self.cmd_dir, download_name='codeformer-v0.1.0.pth')
+ if len(model_paths) != 0:
+ ckpt_path = model_paths[0]
+ else:
+ print("Unable to load codeformer model.")
+ return None, None
net = net_class(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256']).to(devices.device_codeformer)
- ckpt_path = load_file_from_url(url=pretrain_model_url, model_dir=os.path.join(path, 'weights/CodeFormer'), progress=True)
checkpoint = torch.load(ckpt_path)['params_ema']
net.load_state_dict(checkpoint)
net.eval()
@@ -74,6 +79,9 @@ def setup_codeformer():
original_resolution = np_image.shape[0:2]
self.create_models()
+ if self.net is None or self.face_helper is None:
+ return np_image
+
self.face_helper.clean_all()
self.face_helper.read_image(np_image)
self.face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
@@ -116,7 +124,7 @@ def setup_codeformer():
have_codeformer = True
global codeformer
- codeformer = FaceRestorerCodeFormer()
+ codeformer = FaceRestorerCodeFormer(dirname)
shared.face_restorers.append(codeformer)
except Exception:
diff --git a/modules/esrgan_model.py b/modules/esrgan_model.py
index 7f3baf31..ea91abfe 100644
--- a/modules/esrgan_model.py
+++ b/modules/esrgan_model.py
@@ -1,26 +1,22 @@
import os
-import sys
-import traceback
import numpy as np
import torch
from PIL import Image
+from basicsr.utils.download_util import load_file_from_url
import modules.esrgam_model_arch as arch
-from modules import shared
-from modules.shared import opts
+from modules import shared, modelloader, images
from modules.devices import has_mps
-import modules.images
+from modules.paths import models_path
+from modules.upscaler import Upscaler, UpscalerData
+from modules.shared import opts
-def load_model(filename):
+def fix_model_layers(crt_model, pretrained_net):
# this code is adapted from https://github.com/xinntao/ESRGAN
- pretrained_net = torch.load(filename, map_location='cpu' if has_mps else None)
- crt_model = arch.RRDBNet(3, 3, 64, 23, gc=32)
-
if 'conv_first.weight' in pretrained_net:
- crt_model.load_state_dict(pretrained_net)
- return crt_model
+ return pretrained_net
if 'model.0.weight' not in pretrained_net:
is_realesrgan = "params_ema" in pretrained_net and 'body.0.rdb1.conv1.weight' in pretrained_net["params_ema"]
@@ -72,9 +68,59 @@ def load_model(filename):
crt_net['conv_last.weight'] = pretrained_net['model.10.weight']
crt_net['conv_last.bias'] = pretrained_net['model.10.bias']