mirror of
https://github.com/Sygil-Dev/sygil-webui.git
synced 2024-12-14 05:58:18 +03:00
Added LDSR to the UI.
This commit is contained in:
ZeroCool940711
parent
a5ddf9f355
commit
0c03cedeb9
@ -50,10 +50,12 @@ def img2img(prompt: str = '', init_info: any = None, init_info_mask: any = None,
|
||||
n_iter: int = 1, cfg_scale: float = 7.5, denoising_strength: float = 0.8,
|
||||
seed: int = -1, noise_mode: int = 0, find_noise_steps: str = "", height: int = 512, width: int = 512, resize_mode: int = 0, fp = None,
|
||||
variant_amount: float = None, variant_seed: int = None, ddim_eta:float = 0.0,
|
||||
write_info_files:bool = True, RealESRGAN_model: str = "RealESRGAN_x4plus_anime_6B",
|
||||
separate_prompts:bool = False, normalize_prompt_weights:bool = True,
|
||||
write_info_files:bool = True, separate_prompts:bool = False, normalize_prompt_weights:bool = True,
|
||||
save_individual_images: bool = True, save_grid: bool = True, group_by_prompt: bool = True,
|
||||
save_as_jpg: bool = True, use_GFPGAN: bool = True, GFPGAN_model: str = 'GFPGANv1.3',use_RealESRGAN: bool = True, loopback: bool = False,
|
||||
save_as_jpg: bool = True, use_GFPGAN: bool = True, GFPGAN_model: str = 'GFPGANv1.3',
|
||||
use_RealESRGAN: bool = True, RealESRGAN_model: str = "RealESRGAN_x4plus_anime_6B",
|
||||
use_LDSR: bool = True, LDSR_model: str = "model",
|
||||
loopback: bool = False,
|
||||
random_seed_loopback: bool = False
|
||||
):
|
||||
|
||||
@ -260,6 +262,8 @@ def img2img(prompt: str = '', init_info: any = None, init_info_mask: any = None,
|
||||
GFPGAN_model=GFPGAN_model,
|
||||
use_RealESRGAN=use_RealESRGAN and is_final_iteration, # Forcefully disable upscaling when using loopback
|
||||
realesrgan_model_name=RealESRGAN_model,
|
||||
use_LDSR=use_LDSR,
|
||||
LDSR_model_name=LDSR_model,
|
||||
normalize_prompt_weights=normalize_prompt_weights,
|
||||
save_individual_images=save_individual_images,
|
||||
init_img=init_img,
|
||||
@ -331,6 +335,8 @@ def img2img(prompt: str = '', init_info: any = None, init_info_mask: any = None,
|
||||
GFPGAN_model=GFPGAN_model,
|
||||
use_RealESRGAN=use_RealESRGAN,
|
||||
realesrgan_model_name=RealESRGAN_model,
|
||||
use_LDSR=use_LDSR,
|
||||
LDSR_model_name=LDSR_model,
|
||||
normalize_prompt_weights=normalize_prompt_weights,
|
||||
save_individual_images=save_individual_images,
|
||||
init_img=init_img,
|
||||
@ -470,6 +476,7 @@ def layout():
|
||||
help="Save a file next to the image with informartion about the generation.")
|
||||
save_as_jpg = st.checkbox("Save samples as jpg", value=st.session_state['defaults'].img2img.save_as_jpg, help="Saves the images as jpg instead of png.")
|
||||
|
||||
#
|
||||
# check if GFPGAN, RealESRGAN and LDSR are available.
|
||||
if "GFPGAN_available" not in st.session_state:
|
||||
GFPGAN_available()
|
||||
@ -477,33 +484,72 @@ def layout():
|
||||
if "RealESRGAN_available" not in st.session_state:
|
||||
RealESRGAN_available()
|
||||
|
||||
if st.session_state["GFPGAN_available"] or st.session_state["RealESRGAN_available"]:
|
||||
if "LDSR_available" not in st.session_state:
|
||||
LDSR_available()
|
||||
|
||||
if st.session_state["GFPGAN_available"] or st.session_state["RealESRGAN_available"] or st.session_state["LDSR_available"]:
|
||||
with st.expander("Post-Processing"):
|
||||
# GFPGAN used for face restoration
|
||||
if st.session_state["GFPGAN_available"]:
|
||||
with st.expander("Face Restoration"):
|
||||
face_restoration_tab, upscaling_tab = st.tabs(["Face Restoration", "Upscaling"])
|
||||
with face_restoration_tab:
|
||||
# GFPGAN used for face restoration
|
||||
if st.session_state["GFPGAN_available"]:
|
||||
#with st.expander("Face Restoration"):
|
||||
#if st.session_state["GFPGAN_available"]:
|
||||
#with st.expander("GFPGAN"):
|
||||
st.session_state["use_GFPGAN"] = st.checkbox("Use GFPGAN", value=st.session_state['defaults'].txt2img.use_GFPGAN,
|
||||
help="Uses the GFPGAN model to improve faces after the generation.\
|
||||
This greatly improve the quality and consistency of faces but uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
help="Uses the GFPGAN model to improve faces after the generation.\
|
||||
This greatly improve the quality and consistency of faces but uses\
|
||||
extra VRAM. Disable if you need the extra VRAM.")
|
||||
|
||||
st.session_state["GFPGAN_model"] = st.selectbox("GFPGAN model", st.session_state["GFPGAN_models"],
|
||||
index=st.session_state["GFPGAN_models"].index(st.session_state['defaults'].general.GFPGAN_model))
|
||||
index=st.session_state["GFPGAN_models"].index(st.session_state['defaults'].general.GFPGAN_model))
|
||||
|
||||
#st.session_state["GFPGAN_strenght"] = st.slider("Effect Strenght", min_value=1, max_value=100, value=1, step=1, help='')
|
||||
else:
|
||||
st.session_state["use_GFPGAN"] = False
|
||||
|
||||
with st.expander("Upscaling"):
|
||||
# RealESRGAN used for upscaling.
|
||||
if st.session_state["RealESRGAN_available"]:
|
||||
st.session_state["use_RealESRGAN"] = st.checkbox("Use RealESRGAN", value=st.session_state['defaults'].txt2img.use_RealESRGAN,
|
||||
help="Uses the RealESRGAN model to upscale the images after the generation.\
|
||||
This greatly improve the quality and lets you have high resolution images but uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
st.session_state["RealESRGAN_model"] = st.selectbox("RealESRGAN model", st.session_state["RealESRGAN_models"],
|
||||
index=st.session_state["RealESRGAN_models"].index(st.session_state['defaults'].general.RealESRGAN_model))
|
||||
else:
|
||||
st.session_state["use_RealESRGAN"] = False
|
||||
st.session_state["RealESRGAN_model"] = "RealESRGAN_x4plus"
|
||||
st.session_state["use_GFPGAN"] = False
|
||||
|
||||
with upscaling_tab:
|
||||
#with st.expander("Upscaling"):
|
||||
# RealESRGAN and LDSR used for upscaling.
|
||||
if st.session_state["RealESRGAN_available"] or st.session_state["LDSR_available"]:
|
||||
|
||||
upscaling_method_list = []
|
||||
if st.session_state["RealESRGAN_available"]:
|
||||
upscaling_method_list.append("RealESRGAN")
|
||||
if st.session_state["LDSR_available"]:
|
||||
upscaling_method_list.append("LDSR")
|
||||
|
||||
st.session_state["upscaling_method"] = st.selectbox("Upscaling Method", upscaling_method_list,
|
||||
index=upscaling_method_list.index(st.session_state['defaults'].general.upscaling_method))
|
||||
|
||||
if st.session_state["RealESRGAN_available"]:
|
||||
# with st.expander("RealESRGAN"):
|
||||
st.session_state["use_RealESRGAN"] = st.checkbox("Use RealESRGAN", value=st.session_state['defaults'].txt2img.use_RealESRGAN,
|
||||
help="Uses the RealESRGAN model to upscale the images after the generation.\
|
||||
This greatly improve the quality and lets you have high resolution images but \
|
||||
uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
|
||||
st.session_state["RealESRGAN_model"] = st.selectbox("RealESRGAN model", st.session_state["RealESRGAN_models"],
|
||||
index=st.session_state["RealESRGAN_models"].index(st.session_state['defaults'].general.RealESRGAN_model))
|
||||
else:
|
||||
st.session_state["use_RealESRGAN"] = False
|
||||
st.session_state["RealESRGAN_model"] = "RealESRGAN_x4plus"
|
||||
|
||||
|
||||
#
|
||||
if st.session_state["LDSR_available"]:
|
||||
#with st.expander("LDSR"):
|
||||
st.session_state["use_LDSR"] = st.checkbox("Use LDSR", value=st.session_state['defaults'].txt2img.use_LDSR,
|
||||
help="Uses the LDSR model to upscale the images after the generation.\
|
||||
This greatly improve the quality and lets you have high resolution images but \
|
||||
uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
|
||||
st.session_state["LDSR_model"] = st.selectbox("LDSR model", st.session_state["LDSR_models"],
|
||||
index=st.session_state["LDSR_models"].index(st.session_state['defaults'].general.LDSR_model))
|
||||
else:
|
||||
st.session_state["use_LDSR"] = False
|
||||
st.session_state["LDSR_model"] = "model"
|
||||
|
||||
with st.expander("Variant"):
|
||||
variant_amount = st.slider("Variant Amount:", value=st.session_state['defaults'].img2img.variant_amount, min_value=0.0, max_value=1.0, step=0.01)
|
||||
@ -594,8 +640,9 @@ def layout():
|
||||
# load the models when we hit the generate button for the first time, it wont be loaded after that so dont worry.
|
||||
with col3_img2img_layout:
|
||||
with hc.HyLoader('Loading Models...', hc.Loaders.standard_loaders,index=[0]):
|
||||
load_models(False, st.session_state["use_GFPGAN"], st.session_state["GFPGAN_model"], st.session_state["use_RealESRGAN"],
|
||||
st.session_state["RealESRGAN_model"], server_state["CustomModel_available"], st.session_state["custom_model"])
|
||||
load_models(st.session_state["use_LDSR"], st.session_state["LDSR_model"], st.session_state["use_GFPGAN"],
|
||||
st.session_state["GFPGAN_model"] , st.session_state["use_RealESRGAN"],
|
||||
st.session_state["RealESRGAN_model"], server_state["CustomModel_available"], st.session_state["custom_model"])
|
||||
|
||||
if uploaded_images:
|
||||
image = Image.open(uploaded_images).convert('RGBA')
|
||||
@ -614,11 +661,13 @@ def layout():
|
||||
seed=seed, noise_mode=noise_mode, find_noise_steps=find_noise_steps, width=width,
|
||||
height=height, variant_amount=variant_amount,
|
||||
ddim_eta=st.session_state.defaults.img2img.ddim_eta, write_info_files=write_info_files,
|
||||
RealESRGAN_model=st.session_state["RealESRGAN_model"],
|
||||
separate_prompts=separate_prompts, normalize_prompt_weights=normalize_prompt_weights,
|
||||
save_individual_images=save_individual_images, save_grid=save_grid,
|
||||
group_by_prompt=group_by_prompt, save_as_jpg=save_as_jpg, use_GFPGAN=st.session_state["use_GFPGAN"],
|
||||
GFPGAN_model=st.session_state["GFPGAN_model"], use_RealESRGAN=st.session_state["use_RealESRGAN"], loopback=loopback
|
||||
GFPGAN_model=st.session_state["GFPGAN_model"],
|
||||
use_RealESRGAN=st.session_state["use_RealESRGAN"], RealESRGAN_model=st.session_state["RealESRGAN_model"],
|
||||
use_LDSR=st.session_state["use_LDSR"], LDSR_model=st.session_state["LDSR_model"],
|
||||
loopback=loopback
|
||||
)
|
||||
|
||||
#show a message when the generation is complete.
|
||||
|
||||
@ -32,6 +32,7 @@ import json
|
||||
|
||||
import base64
|
||||
import os, sys, re, random, datetime, time, math, glob, toml
|
||||
import gc
|
||||
from PIL import Image, ImageFont, ImageDraw, ImageFilter
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
from scipy import integrate
|
||||
@ -41,9 +42,10 @@ import k_diffusion as K
|
||||
import math, requests
|
||||
import mimetypes
|
||||
import numpy as np
|
||||
from numpy import asarray
|
||||
import pynvml
|
||||
import threading
|
||||
import torch
|
||||
import torch, torchvision
|
||||
from torch import autocast
|
||||
from torchvision import transforms
|
||||
import torch.nn as nn
|
||||
@ -51,7 +53,7 @@ from omegaconf import OmegaConf
|
||||
import yaml
|
||||
from pathlib import Path
|
||||
from contextlib import nullcontext
|
||||
from einops import rearrange
|
||||
from einops import rearrange, repeat
|
||||
from ldm.util import instantiate_from_config
|
||||
from retry import retry
|
||||
from slugify import slugify
|
||||
@ -59,8 +61,12 @@ import skimage
|
||||
import piexif
|
||||
import piexif.helper
|
||||
from tqdm import trange
|
||||
from ldm.models.diffusion.ddim import DDIMSampler
|
||||
from ldm.util import ismap
|
||||
|
||||
|
||||
# Temp imports
|
||||
#from basicsr.utils.registry import ARCH_REGISTRY
|
||||
|
||||
|
||||
# end of imports
|
||||
@ -210,7 +216,7 @@ def human_readable_size(size, decimal_places=3):
|
||||
return f"{size:.{decimal_places}f}{unit}"
|
||||
|
||||
|
||||
def load_models(continue_prev_run = False, use_GFPGAN=False, GFPGAN_model='GFPGANv1.3', use_RealESRGAN=False, RealESRGAN_model="RealESRGAN_x4plus",
|
||||
def load_models(use_LDSR = False, LDSR_model='model', use_GFPGAN=False, GFPGAN_model='GFPGANv1.3', use_RealESRGAN=False, RealESRGAN_model="RealESRGAN_x4plus",
|
||||
CustomModel_available=False, custom_model="Stable Diffusion v1.4"):
|
||||
"""Load the different models. We also reuse the models that are already in memory to speed things up instead of loading them again. """
|
||||
|
||||
@ -228,6 +234,26 @@ def load_models(continue_prev_run = False, use_GFPGAN=False, GFPGAN_model='GFPGA
|
||||
).decode("ascii")
|
||||
|
||||
# check what models we want to use and if the they are already loaded.
|
||||
with server_state_lock["LDSR"]:
|
||||
if use_LDSR:
|
||||
if "LDSR" in server_state and server_state["LDSR"].name == LDSR_model:
|
||||
print("LDSR already loaded")
|
||||
else:
|
||||
if "LDSR" in server_state:
|
||||
del server_state["LDSR"]
|
||||
|
||||
# Load GFPGAN
|
||||
if os.path.exists(st.session_state["defaults"].general.LDSR_dir):
|
||||
try:
|
||||
server_state["LDSR"] = load_LDSR(model_name=LDSR_model)
|
||||
print(f"Loaded LDSR")
|
||||
except Exception:
|
||||
import traceback
|
||||
print(f"Error loading LDSR:", file=sys.stderr)
|
||||
print(traceback.format_exc(), file=sys.stderr)
|
||||
else:
|
||||
if "LDSR" in server_state:
|
||||
del server_state["LDSR"]
|
||||
|
||||
with server_state_lock["GFPGAN"]:
|
||||
if use_GFPGAN:
|
||||
@ -735,6 +761,409 @@ class KDiffusionSampler:
|
||||
'cond_scale': unconditional_guidance_scale}, disable=False, callback=generation_callback)
|
||||
#
|
||||
return samples_ddim, None
|
||||
#
|
||||
#create class LDSR
|
||||
class LDSR():
|
||||
#init function
|
||||
def __init__(self, modelPath,yamlPath):
|
||||
self.modelPath = modelPath
|
||||
self.yamlPath = yamlPath
|
||||
#self.model = self.load_model_from_config()
|
||||
#print(self.load_model_from_config(OmegaConf.load(yamlPath), modelPath))
|
||||
#self.print_current_directory()
|
||||
#get currennt directory
|
||||
|
||||
'''
|
||||
def check_model_exists(self):
|
||||
#check if model and yaml exist
|
||||
path = self.pathInput + "/models/ldm/ld_sr".replace('\\',os.sep).replace('/',os.sep)
|
||||
model = self.modelName
|
||||
yaml = self.yamlName
|
||||
if os.path.exists(path):
|
||||
#check if yaml exists
|
||||
if os.path.exists(os.path.join(path,yaml)):
|
||||
print('YAML found')
|
||||
#check if ckpt exists
|
||||
if os.path.exists(os.path.join(path,model)):
|
||||
print('Model found')
|
||||
return os.path.join(path,model), os.path.join(path,yaml)
|
||||
else:
|
||||
return False
|
||||
#return onlyfiles
|
||||
'''
|
||||
def load_model_from_config(self):
|
||||
#print(f"Loading model from {self.modelPath}")
|
||||
pl_sd = torch.load(self.modelPath, map_location="cpu")
|
||||
global_step = pl_sd["global_step"]
|
||||
sd = pl_sd["state_dict"]
|
||||
config = OmegaConf.load(self.yamlPath)
|
||||
model = instantiate_from_config(config.model)
|
||||
m, u = model.load_state_dict(sd, strict=False)
|
||||
model.cuda()
|
||||
model.eval()
|
||||
return {"model": model}#, global_step
|
||||
|
||||
'''
|
||||
def get_model(self):
|
||||
check = self.check_model_exists()
|
||||
if check != False:
|
||||
path_ckpt = check[0]
|
||||
path_conf = check[1]
|
||||
else:
|
||||
print('Model not found, please run the bat file to download the model')
|
||||
config = OmegaConf.load(path_conf)
|
||||
model, step = self.load_model_from_config(config, path_ckpt)
|
||||
return model
|
||||
|
||||
|
||||
def get_custom_cond(mode):
|
||||
dest = "data/example_conditioning"
|
||||
|
||||
if mode == "superresolution":
|
||||
uploaded_img = files.upload()
|
||||
filename = next(iter(uploaded_img))
|
||||
name, filetype = filename.split(".") # todo assumes just one dot in name !
|
||||
os.rename(f"{filename}", f"{dest}/{mode}/custom_{name}.{filetype}")
|
||||
|
||||
elif mode == "text_conditional":
|
||||
#w = widgets.Text(value='A cake with cream!', disabled=True)
|
||||
w = 'Empty Test'
|
||||
display.display(w)
|
||||
|
||||
with open(f"{dest}/{mode}/custom_{w.value[:20]}.txt", 'w') as f:
|
||||
f.write(w.value)
|
||||
|
||||
elif mode == "class_conditional":
|
||||
#w = widgets.IntSlider(min=0, max=1000)
|
||||
w = 1000
|
||||
display.display(w)
|
||||
with open(f"{dest}/{mode}/custom.txt", 'w') as f:
|
||||
f.write(w.value)
|
||||
|
||||
else:
|
||||
raise NotImplementedError(f"cond not implemented for mode{mode}")
|
||||
'''
|
||||
|
||||
def get_cond_options(self,mode):
|
||||
path = "data/example_conditioning"
|
||||
path = os.path.join(path, mode)
|
||||
onlyfiles = [f for f in sorted(os.listdir(path))]
|
||||
return path, onlyfiles
|
||||
|
||||
'''
|
||||
def select_cond_path(mode):
|
||||
path = "data/example_conditioning" # todo
|
||||
path = os.path.join(path, mode)
|
||||
onlyfiles = [f for f in sorted(os.listdir(path))]
|
||||
|
||||
selected = widgets.RadioButtons(
|
||||
options=onlyfiles,
|
||||
description='Select conditioning:',
|
||||
disabled=False
|
||||
)
|
||||
display.display(selected)
|
||||
selected_path = os.path.join(path, selected.value)
|
||||
return selected_path
|
||||
'''
|
||||
|
||||
|
||||
|
||||
'''
|
||||
# Google Collab stuff
|
||||
def visualize_cond_img(path):
|
||||
display.display(ipyimg(filename=path))
|
||||
'''
|
||||
|
||||
def run(self,model, selected_path, task, custom_steps, eta, resize_enabled=False, classifier_ckpt=None, global_step=None):
|
||||
def make_convolutional_sample(batch, model, mode="vanilla", custom_steps=None, eta=1.0, swap_mode=False, masked=False,
|
||||
invert_mask=True, quantize_x0=False, custom_schedule=None, decode_interval=1000,
|
||||
resize_enabled=False, custom_shape=None, temperature=1., noise_dropout=0., corrector=None,
|
||||
corrector_kwargs=None, x_T=None, save_intermediate_vid=False, make_progrow=True,ddim_use_x0_pred=False):
|
||||
log = dict()
|
||||
|
||||
z, c, x, xrec, xc = model.get_input(batch, model.first_stage_key,
|
||||
return_first_stage_outputs=True,
|
||||
force_c_encode=not (hasattr(model, 'split_input_params')
|
||||
and model.cond_stage_key == 'coordinates_bbox'),
|
||||
return_original_cond=True)
|
||||
|
||||
log_every_t = 1 if save_intermediate_vid else None
|
||||
|
||||
if custom_shape is not None:
|
||||
z = torch.randn(custom_shape)
|
||||
# print(f"Generating {custom_shape[0]} samples of shape {custom_shape[1:]}")
|
||||
|
||||
z0 = None
|
||||
|
||||
log["input"] = x
|
||||
log["reconstruction"] = xrec
|
||||
|
||||
if ismap(xc):
|
||||
log["original_conditioning"] = model.to_rgb(xc)
|
||||
if hasattr(model, 'cond_stage_key'):
|
||||
log[model.cond_stage_key] = model.to_rgb(xc)
|
||||
|
||||
else:
|
||||
log["original_conditioning"] = xc if xc is not None else torch.zeros_like(x)
|
||||
if model.cond_stage_model:
|
||||
log[model.cond_stage_key] = xc if xc is not None else torch.zeros_like(x)
|
||||
if model.cond_stage_key =='class_label':
|
||||
log[model.cond_stage_key] = xc[model.cond_stage_key]
|
||||
|
||||
with model.ema_scope("Plotting"):
|
||||
t0 = time.time()
|
||||
img_cb = None
|
||||
|
||||
sample, intermediates = convsample_ddim(model, c, steps=custom_steps, shape=z.shape,
|
||||
eta=eta,
|
||||
quantize_x0=quantize_x0, img_callback=img_cb, mask=None, x0=z0,
|
||||
temperature=temperature, noise_dropout=noise_dropout,
|
||||
score_corrector=corrector, corrector_kwargs=corrector_kwargs,
|
||||
x_T=x_T, log_every_t=log_every_t)
|
||||
t1 = time.time()
|
||||
|
||||
if ddim_use_x0_pred:
|
||||
sample = intermediates['pred_x0'][-1]
|
||||
|
||||
x_sample = model.decode_first_stage(sample)
|
||||
|
||||
try:
|
||||
x_sample_noquant = model.decode_first_stage(sample, force_not_quantize=True)
|
||||
log["sample_noquant"] = x_sample_noquant
|
||||
log["sample_diff"] = torch.abs(x_sample_noquant - x_sample)
|
||||
except:
|
||||
pass
|
||||
|
||||
log["sample"] = x_sample
|
||||
log["time"] = t1 - t0
|
||||
|
||||
return log
|
||||
def convsample_ddim(model, cond, steps, shape, eta=1.0, callback=None, normals_sequence=None,
|
||||
mask=None, x0=None, quantize_x0=False, img_callback=None,
|
||||
temperature=1., noise_dropout=0., score_corrector=None,
|
||||
corrector_kwargs=None, x_T=None, log_every_t=None
|
||||
):
|
||||
|
||||
ddim = DDIMSampler(model)
|
||||
bs = shape[0] # dont know where this comes from but wayne
|
||||
shape = shape[1:] # cut batch dim
|
||||
print(f"Sampling with eta = {eta}; steps: {steps}")
|
||||
samples, intermediates = ddim.sample(steps, batch_size=bs, shape=shape, conditioning=cond, callback=callback,
|
||||
normals_sequence=normals_sequence, quantize_x0=quantize_x0, eta=eta,
|
||||
mask=mask, x0=x0, temperature=temperature, verbose=False,
|
||||
score_corrector=score_corrector,
|
||||
corrector_kwargs=corrector_kwargs, x_T=x_T)
|
||||
|
||||
return samples, intermediates
|
||||
# global stride
|
||||
def get_cond(mode, selected_path):
|
||||
example = dict()
|
||||
if mode == "superresolution":
|
||||
up_f = 4
|
||||
#visualize_cond_img(selected_path)
|
||||
|
||||
c = selected_path.convert('RGB')
|
||||
c = torch.unsqueeze(torchvision.transforms.ToTensor()(c), 0)
|
||||
c_up = torchvision.transforms.functional.resize(c, size=[up_f * c.shape[2], up_f * c.shape[3]], antialias=True)
|
||||
c_up = rearrange(c_up, '1 c h w -> 1 h w c')
|
||||
c = rearrange(c, '1 c h w -> 1 h w c')
|
||||
c = 2. * c - 1.
|
||||
|
||||
c = c.to(torch.device("cuda"))
|
||||
example["LR_image"] = c
|
||||
example["image"] = c_up
|
||||
|
||||
return example
|
||||
example = get_cond(task, selected_path)
|
||||
|
||||
save_intermediate_vid = False
|
||||
n_runs = 1
|
||||
masked = False
|
||||
guider = None
|
||||
ckwargs = None
|
||||
mode = 'ddim'
|
||||
ddim_use_x0_pred = False
|
||||
temperature = 1.
|
||||
eta = eta
|
||||
make_progrow = True
|
||||
custom_shape = None
|
||||
|
||||
height, width = example["image"].shape[1:3]
|
||||
split_input = height >= 128 and width >= 128
|
||||
|
||||
if split_input:
|
||||
ks = 128
|
||||
stride = 64
|
||||
vqf = 4 #
|
||||
model.split_input_params = {"ks": (ks, ks), "stride": (stride, stride),
|
||||
"vqf": vqf,
|
||||
"patch_distributed_vq": True,
|
||||
"tie_braker": False,
|
||||
"clip_max_weight": 0.5,
|
||||
"clip_min_weight": 0.01,
|
||||
"clip_max_tie_weight": 0.5,
|
||||
"clip_min_tie_weight": 0.01}
|
||||
else:
|
||||
if hasattr(model, "split_input_params"):
|
||||
delattr(model, "split_input_params")
|
||||
|
||||
invert_mask = False
|
||||
|
||||
x_T = None
|
||||
for n in range(n_runs):
|
||||
if custom_shape is not None:
|
||||
x_T = torch.randn(1, custom_shape[1], custom_shape[2], custom_shape[3]).to(model.device)
|
||||
x_T = repeat(x_T, '1 c h w -> b c h w', b=custom_shape[0])
|
||||
|
||||
logs = make_convolutional_sample(example, model,
|
||||
mode=mode, custom_steps=custom_steps,
|
||||
eta=eta, swap_mode=False , masked=masked,
|
||||
invert_mask=invert_mask, quantize_x0=False,
|
||||
custom_schedule=None, decode_interval=10,
|
||||
resize_enabled=resize_enabled, custom_shape=custom_shape,
|
||||
temperature=temperature, noise_dropout=0.,
|
||||
corrector=guider, corrector_kwargs=ckwargs, x_T=x_T, save_intermediate_vid=save_intermediate_vid,
|
||||
make_progrow=make_progrow,ddim_use_x0_pred=ddim_use_x0_pred
|
||||
)
|
||||
return logs
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
|
||||
def superResolution(self,image,ddimSteps=100,preDownScale='None',postDownScale='None'):
|
||||
diffMode = 'superresolution'
|
||||
model = self.load_model_from_config()
|
||||
#@title Import location
|
||||
#@markdown ***File height and width should be multiples of 64, or image will be padded.***
|
||||
|
||||
#@markdown *To change upload settings without adding more, run and cancel upload*
|
||||
#import_method = 'Directory' #@param ['Google Drive', 'Upload']
|
||||
#output_subfolder_name = 'processed' #@param {type: 'string'}
|
||||
|
||||
#@markdown Drive method options:
|
||||
#drive_directory = '/content/drive/MyDrive/upscaleTest' #@param {type: 'string'}
|
||||
|
||||
#@markdown Upload method options:
|
||||
#remove_previous_uploads = False #@param {type: 'boolean'}
|
||||
#save_output_to_drive = False #@param {type: 'boolean'}
|
||||
#zip_if_not_drive = False #@param {type: 'boolean'}
|
||||
'''
|
||||
os.makedirs(pathInput+'/content/input'.replace('\\',os.sep).replace('/',os.sep), exist_ok=True)
|
||||
output_directory = os.getcwd()+f'/content/output/{output_subfolder_name}'.replace('\\',os.sep).replace('/',os.sep)
|
||||
os.makedirs(output_directory, exist_ok=True)
|
||||
uploaded_img = pathInput+'/content/input/'.replace('\\',os.sep).replace('/',os.sep)
|
||||
pathInput, dirsInput, filesInput = next(os.walk(pathInput+'/content/input').replace('\\',os.sep).replace('/',os.sep))
|
||||
file_count = len(filesInput)
|
||||
print(f'Found {file_count} files total')
|
||||
'''
|
||||
|
||||
|
||||
#Run settings
|
||||
|
||||
diffusion_steps = int(ddimSteps) #@param [25, 50, 100, 250, 500, 1000]
|
||||
eta = 1.0 #@param {type: 'raw'}
|
||||
stride = 0 #not working atm
|
||||
|
||||
# ####Scaling options:
|
||||
# Downsampling to 256px first will often improve the final image and runs faster.
|
||||
|
||||
# You can improve sharpness without upscaling by upscaling and then downsampling to the original size (i.e. Super Resolution)
|
||||
pre_downsample = preDownScale #@param ['None', '1/2', '1/4']
|
||||
|
||||
post_downsample = postDownScale #@param ['None', 'Original Size', '1/2', '1/4']
|
||||
|
||||
# Nearest gives sharper results, but may look more pixellated. Lancoz is much higher quality, but result may be less crisp.
|
||||
downsample_method = 'Lanczos' #@param ['Nearest', 'Lanczos']
|
||||
|
||||
|
||||
overwrite_prior_runs = True #@param {type: 'boolean'}
|
||||
|
||||
#pathProcessed, dirsProcessed, filesProcessed = next(os.walk(output_directory))
|
||||
|
||||
#for img in filesInput:
|
||||
# if img in filesProcessed and overwrite_prior_runs is False:
|
||||
# print(f'Skipping {img}: Already processed')
|
||||
# continue
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
#dir = pathInput
|
||||
#filepath = os.path.join(dir, img).replace('\\',os.sep).replace('/',os.sep)
|
||||
|
||||
im_og = image
|
||||
width_og, height_og = im_og.size
|
||||
|
||||
#Downsample Pre
|
||||
if pre_downsample == '1/2':
|
||||
downsample_rate = 2
|
||||
elif pre_downsample == '1/4':
|
||||
downsample_rate = 4
|
||||
else:
|
||||
downsample_rate = 1
|
||||
# get system temp directory
|
||||
#dir = tempfile.gettempdir()
|
||||
width_downsampled_pre = width_og//downsample_rate
|
||||
height_downsampled_pre = height_og//downsample_rate
|
||||
if downsample_rate != 1:
|
||||
print(f'Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]')
|
||||
im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS)
|
||||
#os.makedirs(dir, exist_ok=True)
|
||||
#im_og.save(dir + '/ldsr/temp.png'.replace('\\',os.sep).replace('/',os.sep))
|
||||
#filepath = dir + '/ldsr/temp.png'.replace('\\',os.sep).replace('/',os.sep)
|
||||
|
||||
logs = self.run(model["model"], im_og, diffMode, diffusion_steps, eta)
|
||||
|
||||
sample = logs["sample"]
|
||||
sample = sample.detach().cpu()
|
||||
sample = torch.clamp(sample, -1., 1.)
|
||||
sample = (sample + 1.) / 2. * 255
|
||||
sample = sample.numpy().astype(np.uint8)
|
||||
sample = np.transpose(sample, (0, 2, 3, 1))
|
||||
#print(sample.shape)
|
||||
a = Image.fromarray(sample[0])
|
||||
|
||||
#Downsample Post
|
||||
if post_downsample == '1/2':
|
||||
downsample_rate = 2
|
||||
elif post_downsample == '1/4':
|
||||
downsample_rate = 4
|
||||
else:
|
||||
downsample_rate = 1
|
||||
|
||||
width, height = a.size
|
||||
width_downsampled_post = width//downsample_rate
|
||||
height_downsampled_post = height//downsample_rate
|
||||
|
||||
if downsample_method == 'Lanczos':
|
||||
aliasing = Image.LANCZOS
|
||||
else:
|
||||
aliasing = Image.NEAREST
|
||||
|
||||
if downsample_rate != 1:
|
||||
print(f'Downsampling from [{width}, {height}] to [{width_downsampled_post}, {height_downsampled_post}]')
|
||||
a = a.resize((width_downsampled_post, height_downsampled_post), aliasing)
|
||||
elif post_downsample == 'Original Size':
|
||||
print(f'Downsampling from [{width}, {height}] to Original Size [{width_og}, {height_og}]')
|
||||
a = a.resize((width_og, height_og), aliasing)
|
||||
|
||||
#display.display(a)
|
||||
#a.save(f'{output_directory}/{img}')
|
||||
del model
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
'''
|
||||
if import_method != 'Google Drive' and zip_if_not_drive is True:
|
||||
print('Zipping files')
|
||||
current_time = datetime.now().strftime('%y%m%d-%H%M%S_%f')
|
||||
output_zip_name = 'output'+str(current_time)+'.zip'
|
||||
#!zip -r {output_zip_name} {output_directory}
|
||||
print(f'Zipped outputs in {output_zip_name}')
|
||||
'''
|
||||
print(f'Processing finished!')
|
||||
return a
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@ -796,13 +1225,15 @@ def load_GFPGAN(model_name='GFPGANv1.3'):
|
||||
|
||||
with server_state_lock['GFPGAN']:
|
||||
if st.session_state['defaults'].general.gfpgan_cpu or st.session_state['defaults'].general.extra_models_cpu:
|
||||
server_state['GFPGAN'] = GFPGANer(model_path=model_path, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=torch.device('cpu'))
|
||||
server_state['GFPGAN'] = GFPGANer(model_path=model_path, upscale=2, arch='clean',
|
||||
channel_multiplier=2, bg_upsampler=None, device=torch.device('cpu'))
|
||||
|
||||
elif st.session_state['defaults'].general.extra_models_gpu:
|
||||
server_state['GFPGAN'] = GFPGANer(model_path=model_path, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None,
|
||||
server_state['GFPGAN'] = GFPGANer(model_path=model_path, upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=None,
|
||||
device=torch.device(f"cuda:{st.session_state['defaults'].general.gfpgan_gpu}"))
|
||||
else:
|
||||
server_state['GFPGAN'] = GFPGANer(model_path=model_path, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None,
|
||||
server_state['GFPGAN'] = GFPGANer(model_path=model_path, upscale=2, arch='clean',
|
||||
channel_multiplier=2, bg_upsampler=None,
|
||||
device=torch.device(f"cuda:{st.session_state['defaults'].general.gpu}"))
|
||||
|
||||
# Add the model_name to model loaded so we can later
|
||||
@ -853,24 +1284,27 @@ def load_RealESRGAN(model_name: str):
|
||||
|
||||
#
|
||||
@retry(tries=5)
|
||||
def load_LDSR(checking=False):
|
||||
model_name = 'model'
|
||||
yaml_name = 'project'
|
||||
model_path = os.path.join(st.session_state['defaults'].general.LDSR_dir, 'experiments/pretrained_models', model_name + '.ckpt')
|
||||
yaml_path = os.path.join(st.session_state['defaults'].general.LDSR_dir, 'experiments/pretrained_models', yaml_name + '.yaml')
|
||||
def load_LDSR(model_name="model", config="project", checking=False):
|
||||
#model_name = 'model'
|
||||
#yaml_name = 'project'
|
||||
|
||||
model_path = os.path.join(st.session_state['defaults'].general.LDSR_dir, model_name + '.ckpt')
|
||||
yaml_path = os.path.join(st.session_state['defaults'].general.LDSR_dir, config + '.yaml')
|
||||
|
||||
if not os.path.isfile(model_path):
|
||||
raise Exception("LDSR model not found at path "+model_path)
|
||||
raise Exception("LDSR model not found at path " + model_path)
|
||||
if not os.path.isfile(yaml_path):
|
||||
raise Exception("LDSR model not found at path "+yaml_path)
|
||||
raise Exception("LDSR model not found at path " + yaml_path)
|
||||
if checking == True:
|
||||
return True
|
||||
|
||||
sys.path.append(os.path.abspath(st.session_state['defaults'].general.LDSR_dir))
|
||||
from LDSR import LDSR
|
||||
LDSRObject = LDSR(model_path, yaml_path)
|
||||
#sys.path.append(os.path.abspath(st.session_state['defaults'].general.LDSR_dir))
|
||||
#from LDSR import LDSR
|
||||
server_state['LDSR'] = LDSR(model_path, yaml_path)
|
||||
|
||||
return LDSRObject
|
||||
server_state['LDSR'].name = model_name
|
||||
|
||||
return server_state['LDSR']
|
||||
|
||||
#
|
||||
|
||||
@ -1382,6 +1816,27 @@ def RealESRGAN_available():
|
||||
st.session_state["RealESRGAN_available"] = True
|
||||
else:
|
||||
st.session_state["RealESRGAN_available"] = False
|
||||
#
|
||||
def LDSR_available():
|
||||
#with server_state_lock["RealESRGAN_models"]:
|
||||
#
|
||||
# Allow for custom models to be used instead of the default one,
|
||||
# an example would be Waifu-Diffusion or any other fine tune of stable diffusion
|
||||
st.session_state["LDSR_models"]:sorted = []
|
||||
|
||||
for root, dirs, files in os.walk(st.session_state['defaults'].general.LDSR_dir):
|
||||
for file in files:
|
||||
if os.path.splitext(file)[1] == '.ckpt':
|
||||
st.session_state["LDSR_models"].append(os.path.splitext(file)[0])
|
||||
|
||||
#print (st.session_state['defaults'].general.LDSR_dir)
|
||||
#print (st.session_state["LDSR_models"])
|
||||
#with server_state_lock["LDSR_available"]:
|
||||
if len(st.session_state["LDSR_models"]) > 0:
|
||||
st.session_state["LDSR_available"] = True
|
||||
else:
|
||||
st.session_state["LDSR_available"] = False
|
||||
|
||||
|
||||
def save_sample(image, sample_path_i, filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
|
||||
normalize_prompt_weights, use_GFPGAN, write_info_files, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback,
|
||||
@ -1564,7 +2019,7 @@ def oxlamon_matrix(prompt, seed, n_iter, batch_size):
|
||||
def process_images(
|
||||
outpath, func_init, func_sample, prompt, seed, sampler_name, save_grid, batch_size,
|
||||
n_iter, steps, cfg_scale, width, height, prompt_matrix, use_GFPGAN, GFPGAN_model, use_RealESRGAN, realesrgan_model_name,
|
||||
ddim_eta=0.0, normalize_prompt_weights=True, init_img=None, init_mask=None,
|
||||
use_LDSR, LDSR_model_name, ddim_eta=0.0, normalize_prompt_weights=True, init_img=None, init_mask=None,
|
||||
mask_blur_strength=3, mask_restore=False, denoising_strength=0.75, noise_mode=0, find_noise_steps=1, resize_mode=None, uses_loopback=False,
|
||||
uses_random_seed_loopback=False, sort_samples=True, write_info_files=True, jpg_sample=False,
|
||||
variant_amount=0.0, variant_seed=None, save_individual_images: bool = True):
|
||||
@ -1791,11 +2246,16 @@ def process_images(
|
||||
|
||||
if use_GFPGAN and server_state["GFPGAN"] is not None and not use_RealESRGAN:
|
||||
st.session_state["progress_bar_text"].text("Running GFPGAN on image %d of %d..." % (i+1, len(x_samples_ddim)))
|
||||
#skip_save = True # #287 >_>
|
||||
|
||||
torch_gc()
|
||||
cropped_faces, restored_faces, restored_img = server_state["GFPGAN"].enhance(x_sample[:,:,::-1], has_aligned=False, only_center_face=False, paste_back=True)
|
||||
|
||||
gfpgan_sample = restored_img[:,:,::-1]
|
||||
gfpgan_image = Image.fromarray(gfpgan_sample)
|
||||
|
||||
#if st.session_state["GFPGAN_strenght"]:
|
||||
#gfpgan_sample = Image.blend(image, gfpgan_image, st.session_state["GFPGAN_strenght"])
|
||||
|
||||
gfpgan_filename = original_filename + '-gfpgan'
|
||||
|
||||
save_sample(gfpgan_image, sample_path_i, gfpgan_filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
|
||||
@ -1808,6 +2268,32 @@ def process_images(
|
||||
|
||||
if simple_templating:
|
||||
grid_captions.append( captions[i] + "\ngfpgan" )
|
||||
|
||||
#
|
||||
elif use_GFPGAN and server_state["GFPGAN"] is not None and not use_LDSR:
|
||||
st.session_state["progress_bar_text"].text("Running GFPGAN on image %d of %d..." % (i+1, len(x_samples_ddim)))
|
||||
|
||||
torch_gc()
|
||||
cropped_faces, restored_faces, restored_img = server_state["GFPGAN"].enhance(x_sample[:,:,::-1], has_aligned=False, only_center_face=False, paste_back=True)
|
||||
|
||||
gfpgan_sample = restored_img[:,:,::-1]
|
||||
gfpgan_image = Image.fromarray(gfpgan_sample)
|
||||
|
||||
#if st.session_state["GFPGAN_strenght"]:
|
||||
#gfpgan_sample = Image.blend(image, gfpgan_image, st.session_state["GFPGAN_strenght"])
|
||||
|
||||
gfpgan_filename = original_filename + '-gfpgan'
|
||||
|
||||
save_sample(gfpgan_image, sample_path_i, gfpgan_filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
|
||||
normalize_prompt_weights, use_GFPGAN, write_info_files, prompt_matrix, init_img, uses_loopback,
|
||||
uses_random_seed_loopback, save_grid, sort_samples, sampler_name, ddim_eta,
|
||||
n_iter, batch_size, i, denoising_strength, resize_mode, False, server_state["loaded_model"])
|
||||
|
||||
output_images.append(gfpgan_image) #287
|
||||
run_images.append(gfpgan_image)
|
||||
|
||||
if simple_templating:
|
||||
grid_captions.append( captions[i] + "\ngfpgan" )
|
||||
|
||||
elif use_RealESRGAN and server_state["RealESRGAN"] is not None and not use_GFPGAN:
|
||||
st.session_state["progress_bar_text"].text("Running RealESRGAN on image %d of %d..." % (i+1, len(x_samples_ddim)))
|
||||
@ -1836,6 +2322,36 @@ def process_images(
|
||||
|
||||
if simple_templating:
|
||||
grid_captions.append( captions[i] + "\nesrgan" )
|
||||
|
||||
#
|
||||
elif use_LDSR and server_state["LDSR"] is not None and not use_GFPGAN:
|
||||
print ("Running LDSR on image %d of %d..." % (i+1, len(x_samples_ddim)))
|
||||
st.session_state["progress_bar_text"].text("Running LDSR on image %d of %d..." % (i+1, len(x_samples_ddim)))
|
||||
#skip_save = True # #287 >_>
|
||||
torch_gc()
|
||||
|
||||
if server_state["LDSR"].name != LDSR_model_name:
|
||||
#try_loading_RealESRGAN(realesrgan_model_name)
|
||||
load_models(use_LDSR=use_LDSR, LDSR_model=LDSR_model_name, use_GFPGAN=use_GFPGAN, use_RealESRGAN=use_RealESRGAN, RealESRGAN_model=realesrgan_model_name)
|
||||
|
||||
result = server_state["LDSR"].superResolution(image, 2, 2, 2)
|
||||
ldsr_filename = original_filename + '-ldsr4x'
|
||||
ldsr_sample = result[:,:,::-1]
|
||||
ldsr_image = Image.fromarray(ldsr_sample)
|
||||
|
||||
#save_sample(image, sample_path_i, original_filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
|
||||
#normalize_prompt_weights, use_GFPGAN, write_info_files, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback, skip_save,
|
||||
#save_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoising_strength, resize_mode)
|
||||
|
||||
save_sample(esrgan_image, sample_path_i, ldsr_filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
|
||||
normalize_prompt_weights, use_GFPGAN, write_info_files, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback,
|
||||
save_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoising_strength, resize_mode, False, server_state["loaded_model"])
|
||||
|
||||
output_images.append(ldsr_image) #287
|
||||
run_images.append(ldsr_image)
|
||||
|
||||
if simple_templating:
|
||||
grid_captions.append( captions[i] + "\nldsr" )
|
||||
|
||||
elif use_RealESRGAN and server_state["RealESRGAN"] is not None and use_GFPGAN and server_state["GFPGAN"] is not None:
|
||||
st.session_state["progress_bar_text"].text("Running GFPGAN+RealESRGAN on image %d of %d..." % (i+1, len(x_samples_ddim)))
|
||||
@ -1862,6 +2378,34 @@ def process_images(
|
||||
|
||||
if simple_templating:
|
||||
grid_captions.append( captions[i] + "\ngfpgan_esrgan" )
|
||||
|
||||
#
|
||||
elif use_LDSR and server_state["LDSR"] is not None and use_GFPGAN and server_state["GFPGAN"] is not None:
|
||||
st.session_state["progress_bar_text"].text("Running GFPGAN+LDSR on image %d of %d..." % (i+1, len(x_samples_ddim)))
|
||||
#skip_save = True # #287 >_>
|
||||
torch_gc()
|
||||
cropped_faces, restored_faces, restored_img = server_state["LDSR"].enhance(x_sample[:,:,::-1], has_aligned=False, only_center_face=False, paste_back=True)
|
||||
gfpgan_sample = restored_img[:,:,::-1]
|
||||
|
||||
if server_state["LDSR"].model.name != ldsr_model_name:
|
||||
#try_loading_RealESRGAN(realesrgan_model_name)
|
||||
load_models(use_LDSR=use_LDSR, LDSR_model=LDSR_model_name,use_GFPGAN=use_GFPGAN, use_RealESRGAN=use_RealESRGAN, RealESRGAN_model=realesrgan_model_name)
|
||||
|
||||
output, img_mode = server_state["LDSR"].enhance(gfpgan_sample[:,:,::-1])
|
||||
gfpgan_ldsr_filename = original_filename + '-gfpgan-ldsr4x'
|
||||
gfpgan_ldsr_sample = output[:,:,::-1]
|
||||
gfpgan_ldsr_image = Image.fromarray(gfpgan_ldsr_sample)
|
||||
|
||||
save_sample(gfpgan_ldsr_image, sample_path_i, gfpgan_ldsr_filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
|
||||
normalize_prompt_weights, False, write_info_files, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback,
|
||||
save_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoising_strength, resize_mode, False, server_state["loaded_model"])
|
||||
|
||||
output_images.append(gfpgan_ldsr_image) #287
|
||||
run_images.append(gfpgan_ldsr_image)
|
||||
|
||||
if simple_templating:
|
||||
grid_captions.append( captions[i] + "\ngfpgan_ldsr" )
|
||||
|
||||
else:
|
||||
output_images.append(image)
|
||||
run_images.append(image)
|
||||
|
||||
@ -96,7 +96,9 @@ def txt2img(prompt: str, ddim_steps: int, sampler_name: str, realesrgan_model_na
|
||||
height: int, width: int, separate_prompts:bool = False, normalize_prompt_weights:bool = True,
|
||||
save_individual_images: bool = True, save_grid: bool = True, group_by_prompt: bool = True,
|
||||
save_as_jpg: bool = True, use_GFPGAN: bool = True, GFPGAN_model: str = 'GFPGANv1.3', use_RealESRGAN: bool = True,
|
||||
RealESRGAN_model: str = "RealESRGAN_x4plus_anime_6B", fp = None, variant_amount: float = None,
|
||||
RealESRGAN_model: str = "RealESRGAN_x4plus_anime_6B", use_LDSR: bool = True,
|
||||
LDSR_model: str = "model",
|
||||
fp = None, variant_amount: float = None,
|
||||
variant_seed: int = None, ddim_eta:float = 0.0, write_info_files:bool = True):
|
||||
|
||||
outpath = st.session_state['defaults'].general.outdir_txt2img or st.session_state['defaults'].general.outdir or "outputs/txt2img-samples"
|
||||
@ -152,6 +154,8 @@ def txt2img(prompt: str, ddim_steps: int, sampler_name: str, realesrgan_model_na
|
||||
GFPGAN_model=st.session_state["GFPGAN_model"],
|
||||
use_RealESRGAN=st.session_state["use_RealESRGAN"],
|
||||
realesrgan_model_name=realesrgan_model_name,
|
||||
use_LDSR=st.session_state["use_LDSR"],
|
||||
LDSR_model_name=LDSR_model,
|
||||
ddim_eta=ddim_eta,
|
||||
normalize_prompt_weights=normalize_prompt_weights,
|
||||
save_individual_images=save_individual_images,
|
||||
@ -171,7 +175,8 @@ def txt2img(prompt: str, ddim_steps: int, sampler_name: str, realesrgan_model_na
|
||||
#err_msg = f'CRASHED:<br><textarea rows="5" style="color:white;background: black;width: -webkit-fill-available;font-family: monospace;font-size: small;font-weight: bold;">{str(e)}</textarea><br><br>Please wait while the program restarts.'
|
||||
#stats = err_msg
|
||||
#return [], seed, 'err', stats
|
||||
|
||||
|
||||
#
|
||||
def layout():
|
||||
with st.form("txt2img-inputs"):
|
||||
st.session_state["generation_mode"] = "txt2img"
|
||||
@ -286,35 +291,76 @@ def layout():
|
||||
# check if GFPGAN, RealESRGAN and LDSR are available.
|
||||
if "GFPGAN_available" not in st.session_state:
|
||||
GFPGAN_available()
|
||||
|
||||
|
||||
if "RealESRGAN_available" not in st.session_state:
|
||||
RealESRGAN_available()
|
||||
|
||||
if st.session_state["GFPGAN_available"] or st.session_state["RealESRGAN_available"]:
|
||||
with st.expander("Post-Processing"):
|
||||
# GFPGAN used for face restoration
|
||||
if st.session_state["GFPGAN_available"]:
|
||||
with st.expander("Face Restoration"):
|
||||
st.session_state["use_GFPGAN"] = st.checkbox("Use GFPGAN", value=st.session_state['defaults'].txt2img.use_GFPGAN,
|
||||
help="Uses the GFPGAN model to improve faces after the generation.\
|
||||
This greatly improve the quality and consistency of faces but uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
|
||||
st.session_state["GFPGAN_model"] = st.selectbox("GFPGAN model", st.session_state["GFPGAN_models"],
|
||||
index=st.session_state["GFPGAN_models"].index(st.session_state['defaults'].general.GFPGAN_model))
|
||||
else:
|
||||
st.session_state["use_GFPGAN"] = False
|
||||
if "LDSR_available" not in st.session_state:
|
||||
LDSR_available()
|
||||
|
||||
with st.expander("Upscaling"):
|
||||
# RealESRGAN used for upscaling.
|
||||
if st.session_state["RealESRGAN_available"]:
|
||||
st.session_state["use_RealESRGAN"] = st.checkbox("Use RealESRGAN", value=st.session_state['defaults'].txt2img.use_RealESRGAN,
|
||||
help="Uses the RealESRGAN model to upscale the images after the generation.\
|
||||
This greatly improve the quality and lets you have high resolution images but uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
st.session_state["RealESRGAN_model"] = st.selectbox("RealESRGAN model", st.session_state["RealESRGAN_models"],
|
||||
index=st.session_state["RealESRGAN_models"].index(st.session_state['defaults'].general.RealESRGAN_model))
|
||||
if st.session_state["GFPGAN_available"] or st.session_state["RealESRGAN_available"] or st.session_state["LDSR_available"]:
|
||||
with st.expander("Post-Processing"):
|
||||
face_restoration_tab, upscaling_tab = st.tabs(["Face Restoration", "Upscaling"])
|
||||
with face_restoration_tab:
|
||||
# GFPGAN used for face restoration
|
||||
if st.session_state["GFPGAN_available"]:
|
||||
#with st.expander("Face Restoration"):
|
||||
#if st.session_state["GFPGAN_available"]:
|
||||
#with st.expander("GFPGAN"):
|
||||
st.session_state["use_GFPGAN"] = st.checkbox("Use GFPGAN", value=st.session_state['defaults'].txt2img.use_GFPGAN,
|
||||
help="Uses the GFPGAN model to improve faces after the generation.\
|
||||
This greatly improve the quality and consistency of faces but uses\
|
||||
extra VRAM. Disable if you need the extra VRAM.")
|
||||
|
||||
st.session_state["GFPGAN_model"] = st.selectbox("GFPGAN model", st.session_state["GFPGAN_models"],
|
||||
index=st.session_state["GFPGAN_models"].index(st.session_state['defaults'].general.GFPGAN_model))
|
||||
|
||||
#st.session_state["GFPGAN_strenght"] = st.slider("Effect Strenght", min_value=1, max_value=100, value=1, step=1, help='')
|
||||
|
||||
else:
|
||||
st.session_state["use_RealESRGAN"] = False
|
||||
st.session_state["RealESRGAN_model"] = "RealESRGAN_x4plus"
|
||||
st.session_state["use_GFPGAN"] = False
|
||||
|
||||
with upscaling_tab:
|
||||
#with st.expander("Upscaling"):
|
||||
# RealESRGAN and LDSR used for upscaling.
|
||||
if st.session_state["RealESRGAN_available"] or st.session_state["LDSR_available"]:
|
||||
|
||||
upscaling_method_list = []
|
||||
if st.session_state["RealESRGAN_available"]:
|
||||
upscaling_method_list.append("RealESRGAN")
|
||||
if st.session_state["LDSR_available"]:
|
||||
upscaling_method_list.append("LDSR")
|
||||
|
||||
st.session_state["upscaling_method"] = st.selectbox("Upscaling Method", upscaling_method_list,
|
||||
index=upscaling_method_list.index(st.session_state['defaults'].general.upscaling_method))
|
||||
|
||||
if st.session_state["RealESRGAN_available"]:
|
||||
# with st.expander("RealESRGAN"):
|
||||
st.session_state["use_RealESRGAN"] = st.checkbox("Use RealESRGAN", value=st.session_state['defaults'].txt2img.use_RealESRGAN,
|
||||
help="Uses the RealESRGAN model to upscale the images after the generation.\
|
||||
This greatly improve the quality and lets you have high resolution images but \
|
||||
uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
|
||||
st.session_state["RealESRGAN_model"] = st.selectbox("RealESRGAN model", st.session_state["RealESRGAN_models"],
|
||||
index=st.session_state["RealESRGAN_models"].index(st.session_state['defaults'].general.RealESRGAN_model))
|
||||
else:
|
||||
st.session_state["use_RealESRGAN"] = False
|
||||
st.session_state["RealESRGAN_model"] = "RealESRGAN_x4plus"
|
||||
|
||||
|
||||
#
|
||||
if st.session_state["LDSR_available"]:
|
||||
#with st.expander("LDSR"):
|
||||
st.session_state["use_LDSR"] = st.checkbox("Use LDSR", value=st.session_state['defaults'].txt2img.use_LDSR,
|
||||
help="Uses the LDSR model to upscale the images after the generation.\
|
||||
This greatly improve the quality and lets you have high resolution images but \
|
||||
uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
|
||||
st.session_state["LDSR_model"] = st.selectbox("LDSR model", st.session_state["LDSR_models"],
|
||||
index=st.session_state["LDSR_models"].index(st.session_state['defaults'].general.LDSR_model))
|
||||
else:
|
||||
st.session_state["use_LDSR"] = False
|
||||
st.session_state["LDSR_model"] = "model"
|
||||
|
||||
with st.expander("Variant"):
|
||||
variant_amount = st.slider("Variant Amount:", value=st.session_state['defaults'].txt2img.variant_amount.value,
|
||||
@ -330,67 +376,65 @@ def layout():
|
||||
generate_col1.write("")
|
||||
generate_button = generate_col1.form_submit_button("Generate")
|
||||
|
||||
if generate_button:
|
||||
#print("Loading models")
|
||||
# load the models when we hit the generate button for the first time, it wont be loaded after that so dont worry.
|
||||
#print (server_state['CustomModel_available'])
|
||||
#print (st.session_state['custom_model'])
|
||||
#
|
||||
if generate_button:
|
||||
|
||||
with col2:
|
||||
with hc.HyLoader('Loading Models...', hc.Loaders.standard_loaders,index=[0]):
|
||||
load_models(False, st.session_state["use_GFPGAN"], st.session_state["GFPGAN_model"] , st.session_state["use_RealESRGAN"],
|
||||
load_models(st.session_state["use_LDSR"], st.session_state["LDSR_model"], st.session_state["use_GFPGAN"], st.session_state["GFPGAN_model"] , st.session_state["use_RealESRGAN"],
|
||||
st.session_state["RealESRGAN_model"], server_state["CustomModel_available"], st.session_state["custom_model"])
|
||||
|
||||
try:
|
||||
#
|
||||
output_images, seeds, info, stats = txt2img(prompt, st.session_state.sampling_steps, sampler_name, st.session_state["RealESRGAN_model"], batch_count, batch_size,
|
||||
cfg_scale, seed, height, width, separate_prompts, normalize_prompt_weights, save_individual_images,
|
||||
save_grid, group_by_prompt, save_as_jpg, st.session_state["use_GFPGAN"], st.session_state['GFPGAN_model'],
|
||||
st.session_state["use_RealESRGAN"], st.session_state["RealESRGAN_model"],
|
||||
variant_amount=variant_amount, variant_seed=variant_seed, write_info_files=write_info_files)
|
||||
#try:
|
||||
#
|
||||
output_images, seeds, info, stats = txt2img(prompt, st.session_state.sampling_steps, sampler_name, st.session_state["RealESRGAN_model"], batch_count, batch_size,
|
||||
cfg_scale, seed, height, width, separate_prompts, normalize_prompt_weights, save_individual_images,
|
||||
save_grid, group_by_prompt, save_as_jpg, st.session_state["use_GFPGAN"], st.session_state['GFPGAN_model'],
|
||||
use_LDSR=st.session_state["use_LDSR"], LDSR_model=st.session_state["LDSR_model"],
|
||||
variant_amount=variant_amount, variant_seed=variant_seed, write_info_files=write_info_files)
|
||||
|
||||
message.success('Render Complete: ' + info + '; Stats: ' + stats, icon="✅")
|
||||
message.success('Render Complete: ' + info + '; Stats: ' + stats, icon="✅")
|
||||
|
||||
#history_tab,col1,col2,col3,PlaceHolder,col1_cont,col2_cont,col3_cont = st.session_state['historyTab']
|
||||
#history_tab,col1,col2,col3,PlaceHolder,col1_cont,col2_cont,col3_cont = st.session_state['historyTab']
|
||||
|
||||
#if 'latestImages' in st.session_state:
|
||||
#for i in output_images:
|
||||
##push the new image to the list of latest images and remove the oldest one
|
||||
##remove the last index from the list\
|
||||
#st.session_state['latestImages'].pop()
|
||||
##add the new image to the start of the list
|
||||
#st.session_state['latestImages'].insert(0, i)
|
||||
#PlaceHolder.empty()
|
||||
#with PlaceHolder.container():
|
||||
#col1, col2, col3 = st.columns(3)
|
||||
#col1_cont = st.container()
|
||||
#col2_cont = st.container()
|
||||
#col3_cont = st.container()
|
||||
#images = st.session_state['latestImages']
|
||||
#with col1_cont:
|
||||
#with col1:
|
||||
#[st.image(images[index]) for index in [0, 3, 6] if index < len(images)]
|
||||
#with col2_cont:
|
||||
#with col2:
|
||||
#[st.image(images[index]) for index in [1, 4, 7] if index < len(images)]
|
||||
#with col3_cont:
|
||||
#with col3:
|
||||
#[st.image(images[index]) for index in [2, 5, 8] if index < len(images)]
|
||||
#historyGallery = st.empty()
|
||||
#if 'latestImages' in st.session_state:
|
||||
#for i in output_images:
|
||||
##push the new image to the list of latest images and remove the oldest one
|
||||
##remove the last index from the list\
|
||||
#st.session_state['latestImages'].pop()
|
||||
##add the new image to the start of the list
|
||||
#st.session_state['latestImages'].insert(0, i)
|
||||
#PlaceHolder.empty()
|
||||
#with PlaceHolder.container():
|
||||
#col1, col2, col3 = st.columns(3)
|
||||
#col1_cont = st.container()
|
||||
#col2_cont = st.container()
|
||||
#col3_cont = st.container()
|
||||
#images = st.session_state['latestImages']
|
||||
#with col1_cont:
|
||||
#with col1:
|
||||
#[st.image(images[index]) for index in [0, 3, 6] if index < len(images)]
|
||||
#with col2_cont:
|
||||
#with col2:
|
||||
#[st.image(images[index]) for index in [1, 4, 7] if index < len(images)]
|
||||
#with col3_cont:
|
||||
#with col3:
|
||||
#[st.image(images[index]) for index in [2, 5, 8] if index < len(images)]
|
||||
#historyGallery = st.empty()
|
||||
|
||||
## check if output_images length is the same as seeds length
|
||||
#with gallery_tab:
|
||||
#st.markdown(createHTMLGallery(output_images,seeds), unsafe_allow_html=True)
|
||||
## check if output_images length is the same as seeds length
|
||||
#with gallery_tab:
|
||||
#st.markdown(createHTMLGallery(output_images,seeds), unsafe_allow_html=True)
|
||||
|
||||
|
||||
#st.session_state['historyTab'] = [history_tab,col1,col2,col3,PlaceHolder,col1_cont,col2_cont,col3_cont]
|
||||
|
||||
with gallery_tab:
|
||||
print(seeds)
|
||||
sdGallery(output_images)
|
||||
|
||||
#st.session_state['historyTab'] = [history_tab,col1,col2,col3,PlaceHolder,col1_cont,col2_cont,col3_cont]
|
||||
|
||||
with gallery_tab:
|
||||
print(seeds)
|
||||
sdGallery(output_images)
|
||||
|
||||
|
||||
except (StopException, KeyError):
|
||||
print(f"Received Streamlit StopException")
|
||||
#except (StopException, KeyError):
|
||||
#print(f"Received Streamlit StopException")
|
||||
|
||||
# this will render all the images at the end of the generation but its better if its moved to a second tab inside col2 and shown as a gallery.
|
||||
# use the current col2 first tab to show the preview_img and update it as its generated.
|
||||
|
||||
@ -715,8 +715,11 @@ def layout():
|
||||
st.session_state["normalize_prompt_weights"] = st.checkbox("Normalize Prompt Weights.",
|
||||
value=st.session_state['defaults'].txt2vid.normalize_prompt_weights, help="Ensure the sum of all weights add up to 1.0")
|
||||
st.session_state["save_individual_images"] = st.checkbox("Save individual images.",
|
||||
value=st.session_state['defaults'].txt2vid.save_individual_images, help="Save each image generated before any filter or enhancement is applied.")
|
||||
st.session_state["save_video"] = st.checkbox("Save video",value=st.session_state['defaults'].txt2vid.save_video, help="Save a video with all the images generated as frames at the end of the generation.")
|
||||
value=st.session_state['defaults'].txt2vid.save_individual_images,
|
||||
help="Save each image generated before any filter or enhancement is applied.")
|
||||
st.session_state["save_video"] = st.checkbox("Save video",value=st.session_state['defaults'].txt2vid.save_video,
|
||||
help="Save a video with all the images generated as frames at the end of the generation.")
|
||||
|
||||
st.session_state["group_by_prompt"] = st.checkbox("Group results by prompt", value=st.session_state['defaults'].txt2vid.group_by_prompt,
|
||||
help="Saves all the images with the same prompt into the same folder. When using a prompt matrix each prompt combination will have its own folder.")
|
||||
st.session_state["write_info_files"] = st.checkbox("Write Info file", value=st.session_state['defaults'].txt2vid.write_info_files,
|
||||
@ -729,13 +732,17 @@ def layout():
|
||||
st.session_state["save_as_jpg"] = st.checkbox("Save samples as jpg", value=st.session_state['defaults'].txt2vid.save_as_jpg, help="Saves the images as jpg instead of png.")
|
||||
|
||||
if server_state["GFPGAN_available"]:
|
||||
st.session_state["use_GFPGAN"] = st.checkbox("Use GFPGAN", value=st.session_state['defaults'].txt2vid.use_GFPGAN, help="Uses the GFPGAN model to improve faces after the generation. This greatly improve the quality and consistency of faces but uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
st.session_state["use_GFPGAN"] = st.checkbox("Use GFPGAN", value=st.session_state['defaults'].txt2vid.use_GFPGAN,
|
||||
help="Uses the GFPGAN model to improve faces after the generation. This greatly improve the quality and consistency \
|
||||
of faces but uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
else:
|
||||
st.session_state["use_GFPGAN"] = False
|
||||
|
||||
if server_state["RealESRGAN_available"]:
|
||||
st.session_state["use_RealESRGAN"] = st.checkbox("Use RealESRGAN", value=st.session_state['defaults'].txt2vid.use_RealESRGAN,
|
||||
help="Uses the RealESRGAN model to upscale the images after the generation. This greatly improve the quality and lets you have high resolution images but uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
help="Uses the RealESRGAN model to upscale the images after the generation. \
|
||||
This greatly improve the quality and lets you have high resolution images but \
|
||||
uses extra VRAM. Disable if you need the extra VRAM.")
|
||||
st.session_state["RealESRGAN_model"] = st.selectbox("RealESRGAN model", ["RealESRGAN_x4plus", "RealESRGAN_x4plus_anime_6B"], index=0)
|
||||
else:
|
||||
st.session_state["use_RealESRGAN"] = False
|
||||
|
||||
Loading…
Reference in New Issue
Block a user