Merge pull request #1109 from d8ahazard/ModelLoader (25414bcd) · Commits · github_fork / Stable Diffusion Webui

.gitignore

+5 −2

Original line number	Diff line number	Diff line
		__pycache__
		/ESRGAN
		*.ckpt
		*.pth
		/ESRGAN/*
		/SwinIR/*
		/repositories
		/venv
		/tmp
		/model.ckpt
		/models/*/.ckpt
		/models/*/
		/GFPGANv1.3.pth
		/gfpgan/weights/*.pth
		/ui-config.json

ESRGAN/Put ESRGAN models here.txt

deleted100644 → 0

+0 −0

Empty file deleted.

SwinIR/put_swinir_models_here.txt

deleted100644 → 0

+0 −1

Original line number	Diff line number	Diff line

launch.py

+6 −5

Original line number	Diff line number	Diff line
		# this scripts installs necessary requirements and launches main program in webui.py

		import shutil
		import subprocess
		import os
		import sys
		@@ -22,7 +22,6 @@ taming_transformers_commit_hash = os.environ.get('TAMING_TRANSFORMERS_COMMIT_HAS
		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,9 +119,11 @@ 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 os.path.isdir(repo_dir('latent-diffusion')):
		try:
		shutil.rmtree(repo_dir('latent-diffusion'))
		except:
		pass
		if not is_installed("lpips"):
		run_pip(f"install -r {os.path.join(repo_dir('CodeFormer'), 'requirements.txt')}", "requirements for CodeFormer")

modules/bsrgan_model.py

0 → 100644

+79 −0

Original line number	Diff line number	Diff line
		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("Unable to load %s from %s" % (self.model_dir, filename))
		return None
		print("Loading %s from %s" % (self.model_dir, filename))
		model = RRDBNet(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=2) # 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

Admin message