to force using existing checkpoint")
return None
@@ -646,6 +646,14 @@ def detect_pipeline(f: str, op: str = 'model'):
guess = 'Stable Diffusion XL Instruct'
else:
guess = 'Stable Diffusion'
+ if 'LCM_' in f or 'LCM-' in f:
+ if shared.backend == shared.Backend.ORIGINAL:
+ shared.log.warning(f'Model detected as LCM model, but attempting to load using backend=original: {op}={f} size={size} MB')
+ guess = 'Latent Consistency Model'
+ if 'PixArt' in f:
+ if shared.backend == shared.Backend.ORIGINAL:
+ shared.log.warning(f'Model detected as PixArt Alpha model, but attempting to load using backend=original: {op}={f} size={size} MB')
+ guess = 'PixArt Alpha'
pipeline = shared_items.get_pipelines().get(guess, None)
shared.log.info(f'Autodetect: {op}="{guess}" class={pipeline.__name__} file="{f}" size={size}MB')
except Exception as e:
@@ -777,7 +785,7 @@ def set_diffuser_options(sd_model, vae, op: str):
shared.log.debug(f'Setting {op} VAE: name={sd_vae.loaded_vae_file} upcast={sd_model.vae.config.get("force_upcast", None)}')
if shared.opts.cross_attention_optimization == "xFormers" and hasattr(sd_model, 'enable_xformers_memory_efficient_attention'):
sd_model.enable_xformers_memory_efficient_attention()
- if shared.opts.opt_channelslast:
+ if shared.opts.opt_channelslast and hasattr(sd_model, 'unet'):
shared.log.debug(f'Setting {op}: enable channels last')
sd_model.unet.to(memory_format=torch.channels_last)
diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index db9ad937f..62f0c8a73 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -53,14 +53,14 @@ def create_sampler(name, model):
sampler.config = config
sampler.initialize(p=None)
sampler.name = name
- shared.log.debug(f'Sampler: sampler={sampler.name} config={sampler.config.options}')
+ shared.log.debug(f'Sampler: sampler="{sampler.name}" config={sampler.config.options}')
return sampler
elif shared.backend == shared.Backend.DIFFUSERS:
sampler = config.constructor(model)
if not hasattr(model, 'scheduler_config'):
model.scheduler_config = sampler.sampler.config.copy()
model.scheduler = sampler.sampler
- shared.log.debug(f'Sampler: sampler={sampler.name} config={sampler.config}')
+ shared.log.debug(f'Sampler: sampler="{sampler.name}" config={sampler.config}')
return sampler.sampler
else:
return None
diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py
index b613fca99..9cd3ee806 100644
--- a/modules/sd_samplers_diffusers.py
+++ b/modules/sd_samplers_diffusers.py
@@ -17,6 +17,7 @@ try:
UniPCMultistepScheduler,
LMSDiscreteScheduler,
KDPM2AncestralDiscreteScheduler,
+ LCMScheduler,
)
except Exception as e:
import diffusers
@@ -40,6 +41,7 @@ config = {
'LMSD': { 'use_karras_sigmas': False, 'timestep_spacing': 'linspace', 'steps_offset': 0 },
'PNDM': { 'skip_prk_steps': False, 'set_alpha_to_one': False, 'steps_offset': 0 },
'UniPC': { 'solver_order': 2, 'thresholding': False, 'sample_max_value': 1.0, 'predict_x0': 'bh2', 'lower_order_final': True },
+ 'LCM': { 'num_train_timesteps': 1000, 'beta_start': 0.00085, 'beta_end': 0.012, 'beta_schedule': "scaled_linear", 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False },
}
samplers_data_diffusers = [
@@ -58,6 +60,7 @@ samplers_data_diffusers = [
sd_samplers_common.SamplerData('Euler', lambda model: DiffusionSampler('Euler', EulerDiscreteScheduler, model), [], {}),
sd_samplers_common.SamplerData('Euler a', lambda model: DiffusionSampler('Euler a', EulerAncestralDiscreteScheduler, model), [], {}),
sd_samplers_common.SamplerData('Heun', lambda model: DiffusionSampler('Heun', HeunDiscreteScheduler, model), [], {}),
+ sd_samplers_common.SamplerData('LCM', lambda model: DiffusionSampler('LCM', LCMScheduler, model), [], {}),
]
class DiffusionSampler:
@@ -70,8 +73,10 @@ class DiffusionSampler:
return
for key, value in config.get('All', {}).items(): # apply global defaults
self.config[key] = value
+ # shared.log.debug(f'Sampler: name={name} type=all config={self.config}')
for key, value in config.get(name, {}).items(): # apply diffusers per-scheduler defaults
self.config[key] = value
+ # shared.log.debug(f'Sampler: name={name} type=scheduler config={self.config}')
if hasattr(model.scheduler, 'scheduler_config'): # find model defaults
orig_config = model.scheduler.scheduler_config
else:
@@ -79,9 +84,11 @@ class DiffusionSampler:
for key, value in orig_config.items(): # apply model defaults
if key in self.config:
self.config[key] = value
+ # shared.log.debug(f'Sampler: name={name} type=model config={self.config}')
for key, value in kwargs.items(): # apply user args, if any
if key in self.config:
self.config[key] = value
+ # shared.log.debug(f'Sampler: name={name} type=user config={self.config}')
# finally apply user preferences
if shared.opts.schedulers_prediction_type != 'default':
self.config['prediction_type'] = shared.opts.schedulers_prediction_type
diff --git a/modules/shared_items.py b/modules/shared_items.py
index 85fba87a3..49387e56e 100644
--- a/modules/shared_items.py
+++ b/modules/shared_items.py
@@ -27,7 +27,7 @@ def list_crossattention():
def get_pipelines():
import diffusers
from installer import log
- pipelines = {
+ pipelines = { # note: not all pipelines can be used manually as they require prior pipeline next to decoder pipeline
'Autodetect': None,
'Stable Diffusion': getattr(diffusers, 'StableDiffusionPipeline', None),
'Stable Diffusion Img2Img': getattr(diffusers, 'StableDiffusionImg2ImgPipeline', None),
@@ -37,9 +37,15 @@ def get_pipelines():
'Stable Diffusion XL Img2Img': getattr(diffusers, 'StableDiffusionXLImg2ImgPipeline', None),
'Stable Diffusion XL Inpaint': getattr(diffusers, 'StableDiffusionXLInpaintPipeline', None),
'Stable Diffusion XL Instruct': getattr(diffusers, 'StableDiffusionXLInstructPix2PixPipeline', None),
+ 'Latent Consistency Model': getattr(diffusers, 'LatentConsistencyModelPipeline', None),
+ 'PixArt Alpha': getattr(diffusers, 'PixArtAlphaPipeline', None),
+ 'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None),
+ 'Wuerstchen': getattr(diffusers, 'WuerstchenCombinedPipeline', None),
+ 'Kandinsky 2.1': getattr(diffusers, 'KandinskyPipeline', None),
+ 'Kandinsky 2.2': getattr(diffusers, 'KandinskyV22Pipeline', None),
+ 'DeepFloyd IF': getattr(diffusers, 'IFPipeline', None),
'Custom Diffusers Pipeline': getattr(diffusers, 'DiffusionPipeline', None),
- # 'Test': getattr(diffusers, 'TestPipeline', None),
- # 'Kandinsky V1', 'Kandinsky V2', 'DeepFloyd IF', 'Shap-E', 'Kandinsky V1 Img2Img', 'Kandinsky V2 Img2Img', 'DeepFloyd IF Img2Img', 'Shap-E Img2Img',
+ # Segmind SSD-1B, Segmind Tiny
}
for k, v in pipelines.items():
if k != 'Autodetect' and v is None:
diff --git a/modules/styles.py b/modules/styles.py
index 053cfd6fb..427e76ea9 100644
--- a/modules/styles.py
+++ b/modules/styles.py
@@ -5,7 +5,6 @@ import os
import csv
import json
from installer import log
-from modules import paths
class Style():
@@ -66,6 +65,8 @@ def apply_styles_to_extra(p, style: Style):
class StyleDatabase:
def __init__(self, opts):
+ from modules import paths
+
self.no_style = Style("None")
self.styles = {}
self.path = opts.styles_dir
@@ -135,18 +136,33 @@ class StyleDatabase:
return found[0] if len(found) > 0 else self.no_style
def get_style_prompts(self, styles):
+ if styles is None or not isinstance(styles, list):
+ log.error(f'Invalid styles: {styles}')
+ return []
return [self.find_style(x).prompt for x in styles]
def get_negative_style_prompts(self, styles):
+ if styles is None or not isinstance(styles, list):
+ log.error(f'Invalid styles: {styles}')
+ return []
return [self.find_style(x).negative_prompt for x in styles]
def apply_styles_to_prompt(self, prompt, styles):
+ if styles is None or not isinstance(styles, list):
+ log.error(f'Invalid styles: {styles}')
+ return prompt
return apply_styles_to_prompt(prompt, [self.find_style(x).prompt for x in styles])
def apply_negative_styles_to_prompt(self, prompt, styles):
+ if styles is None or not isinstance(styles, list):
+ log.error(f'Invalid styles: {styles}')
+ return prompt
return apply_styles_to_prompt(prompt, [self.find_style(x).negative_prompt for x in styles])
def apply_styles_to_extra(self, p):
+ if p.styles is None or not isinstance(p.styles, list):
+ log.error(f'Invalid styles: {p.styles}')
+ return
for style in p.styles:
s = self.find_style(style)
apply_styles_to_extra(p, s)
@@ -172,19 +188,25 @@ class StyleDatabase:
log.error(f'Failed to save style: name={name} file={path} error={e}')
count = len(list(self.styles))
if count > 0:
- log.debug(f'Saved styles: {path} {count}')
+ log.debug(f'Saved styles: folder="{path}" items={count}')
def load_csv(self, legacy_file):
if not os.path.isfile(legacy_file):
return
with open(legacy_file, "r", encoding="utf-8-sig", newline='') as file:
reader = csv.DictReader(file, skipinitialspace=True)
+ num = 0
for row in reader:
try:
- self.styles[row["name"]] = Style(row["name"], row["prompt"] if "prompt" in row else row["text"], row.get("negative_prompt", ""))
+ name = row["name"]
+ prompt = row["prompt"] if "prompt" in row else row["text"]
+ negative = row.get("negative_prompt", "") if "negative_prompt" in row else row.get("negative", "")
+ self.styles[name] = Style(name, desc=name, prompt=prompt, negative_prompt=negative, extra="")
+ log.debug(f'Migrated style: {self.styles[name].__dict__}')
+ num += 1
except Exception:
- log.error(f'Styles error: file={legacy_file} row={row}')
- log.debug(f'Load legacy styles: file={legacy_file} items={len(self.styles.keys())}')
+ log.error(f'Styles error: file="{legacy_file}" row={row}')
+ log.info(f'Load legacy styles: file="{legacy_file}" loaded={num} created={len(list(self.styles))}')
"""
def save_csv(self, path: str) -> None:
diff --git a/modules/ui.py b/modules/ui.py
index 262507000..ee490a87d 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -511,43 +511,48 @@ def create_ui(startup_timer = None):
txt_prompt_img.change(fn=modules.images.image_data, inputs=[txt_prompt_img], outputs=[txt2img_prompt, txt_prompt_img])
txt2img_paste_fields = [
+ # prompt
(txt2img_prompt, "Prompt"),
(txt2img_negative_prompt, "Negative prompt"),
- # (txt2img_prompt_styles, "Styles"),
- (steps, "Steps"),
- (seed, "Seed"),
- (sampler_index, "Sampler"),
- (cfg_scale, "CFG scale"),
+ # main
(width, "Size-1"),
(height, "Size-2"),
+ # sampler
+ (sampler_index, "Sampler"),
+ (steps, "Steps"),
+ # batch
+ (batch_count, "Batch-1"),
+ (batch_size, "Batch-2"),
+ # seed
+ (seed, "Seed"),
(subseed, "Variation seed"),
(subseed_strength, "Variation strength"),
+ # advanced
+ (cfg_scale, "CFG scale"),
(clip_skip, "Clip skip"),
- (latent_index, "Latent sampler"),
- (latent_index, "Secondary sampler"),
- (denoising_strength, "Denoising strength"),
- (refiner_steps, "Refiner steps"),
- (refiner_start, "Refiner start"),
+ (image_cfg_scale, "Image CFG scale"),
+ (diffusers_guidance_rescale, "CFG rescale"),
(full_quality, "Full quality"),
(restore_faces, "Face restoration"),
- (batch_size, "Batch size"),
- (batch_count, "Batch count"),
- (seed_resize_from_w, "Seed resize from-1"),
- (seed_resize_from_h, "Seed resize from-2"),
+ (tiling, "Tiling"),
+ # second pass
(enable_hr, "Second pass"),
- (hr_force, "Hires force"),
- (hr_scale, "Hires upscale"),
+ (latent_index, "Latent sampler"),
+ (denoising_strength, "Denoising strength"),
(hr_upscaler, "Hires upscaler"),
+ (hr_force, "Hires force"),
(hr_second_pass_steps, "Hires steps"),
+ (hr_scale, "Hires upscale"),
(hr_resize_x, "Hires resize-1"),
(hr_resize_y, "Hires resize-2"),
- (diffusers_guidance_rescale, "CFG rescale"),
- (image_cfg_scale, "Image CFG scale"),
- (refiner_steps, "Refiner steps"),
+ # refiner
(refiner_start, "Refiner start"),
- (tiling, "Tiling"),
- (refiner_negative, "Negative2"),
+ (refiner_steps, "Refiner steps"),
(refiner_prompt, "Prompt2"),
+ (refiner_negative, "Negative2"),
+ # hidden
+ (seed_resize_from_w, "Seed resize from-1"),
+ (seed_resize_from_h, "Seed resize from-2"),
*modules.scripts.scripts_txt2img.infotext_fields
]
parameters_copypaste.add_paste_fields("txt2img", None, txt2img_paste_fields, override_settings)
@@ -838,46 +843,45 @@ def create_ui(startup_timer = None):
ui_extra_networks.setup_ui(extra_networks_ui_img2img, img2img_gallery)
img2img_paste_fields = [
+ # prompt
(img2img_prompt, "Prompt"),
(img2img_negative_prompt, "Negative prompt"),
- # (img2img_prompt_styles, "Styles"),
- (steps, "Steps"),
- (seed, "Seed"),
+ # sampler
(sampler_index, "Sampler"),
- (cfg_scale, "CFG scale"),
+ (steps, "Steps"),
+ # resize
+ (resize_mode, "Resize mode"),
(width, "Size-1"),
(height, "Size-2"),
+ (scale_by, "Resize scale"),
+ # batch
+ (batch_count, "Batch-1"),
+ (batch_size, "Batch-2"),
+ # seed
+ (seed, "Seed"),
(subseed, "Variation seed"),
(subseed_strength, "Variation strength"),
- (full_quality, "Full quality"),
- (clip_skip, "Clip skip"),
- (latent_index, "Latent sampler"),
- (latent_index, "Secondary sampler"),
+ # denoise
(denoising_strength, "Denoising strength"),
+ (refiner_start, "Refiner start"),
+ # advanced
+ (cfg_scale, "CFG scale"),
+ (image_cfg_scale, "Image CFG scale"),
+ (clip_skip, "Clip skip"),
+ (diffusers_guidance_rescale, "CFG rescale"),
+ (full_quality, "Full quality"),
(restore_faces, "Face restoration"),
- (batch_size, "Batch size"),
- (batch_count, "Batch count"),
+ (tiling, "Tiling"),
+ # inpaint
+ (mask_blur, "Mask blur"),
+ (mask_alpha, "Mask alpha"),
+ (inpainting_mask_invert, "Mask invert"),
+ (inpainting_fill, "Masked content"),
+ (inpaint_full_res, "Mask area"),
+ (inpaint_full_res_padding, "Masked padding"),
+ # hidden
(seed_resize_from_w, "Seed resize from-1"),
(seed_resize_from_h, "Seed resize from-2"),
- (resize_mode, "Resize mode"),
- (image_cfg_scale, "Image CFG scale"),
- (diffusers_guidance_rescale, "CFG rescale"),
- (tiling, "Tiling"),
- (mask_blur, "Mask blur"),
- # TODO scale_by add to paste fields
- (scale_by, "UNKNOWN"),
- # from txt2img
- (hr_force, "Hires force"),
- (hr_scale, "Hires upscale"),
- (hr_upscaler, "Hires upscaler"),
- (hr_second_pass_steps, "Hires steps"),
- (hr_second_pass_steps, "Hires steps"),
- (hr_resize_x, "Hires resize-1"),
- (hr_resize_y, "Hires resize-2"),
- (refiner_steps, "Refiner steps"),
- (refiner_start, "Refiner start"),
- (refiner_prompt, "Prompt2"),
- (refiner_negative, "Negative2"),
*modules.scripts.scripts_img2img.infotext_fields
]
parameters_copypaste.add_paste_fields("img2img", init_img, img2img_paste_fields, override_settings)
diff --git a/modules/ui_common.py b/modules/ui_common.py
index 8b4733b73..eff3d91ac 100644
--- a/modules/ui_common.py
+++ b/modules/ui_common.py
@@ -93,6 +93,7 @@ def save_files(js_data, images, html_info, index):
self.index_of_first_image = getattr(self, 'index_of_first_image', 0)
self.infotexts = getattr(self, 'infotexts', [html_info])
self.infotext = self.infotexts[0] if len(self.infotexts) > 0 else html_info
+ self.outpath_grids = shared.opts.outdir_grids or shared.opts.outdir_txt2img_grids
try:
data = json.loads(js_data)
except Exception:
diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py
index 1fafce4bf..6a3e1e6b1 100644
--- a/modules/ui_extensions.py
+++ b/modules/ui_extensions.py
@@ -319,6 +319,7 @@ def create_html(search_text, sort_column):
for ext in sorted(extensions_list, key=sort_function, reverse=sort_reverse):
installed = get_installed(ext)
author = ''
+ updated = datetime.timestamp(datetime.now())
try:
if 'github' in ext['url']:
author = ext['url'].split('/')[-2].split(':')[-1] if '/' in ext['url'] else ext['url'].split(':')[1].split('/')[0]
diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index ff48c82ac..4815e8210 100644
--- a/modules/ui_extra_networks.py
+++ b/modules/ui_extra_networks.py
@@ -40,9 +40,11 @@ card_full = '''
'''
card_list = '''
'''
@@ -222,6 +224,8 @@ class ExtraNetworksPage:
self.items = []
shared.log.error(f'Extra networks error listing items: class={self.__class__.__name__} tab={tabname} {e}')
for item in self.items:
+ if item is None:
+ continue
self.metadata[item["name"]] = item.get("metadata", {})
t1 = time.time()
debug(f'EN create-items: page={self.name} items={len(self.items)} time={t1-t0:.2f}')
@@ -237,7 +241,7 @@ class ExtraNetworksPage:
return f""
subdirs = {}
allowed_folders = [os.path.abspath(x) for x in self.allowed_directories_for_previews()]
- for parentdir, dirs in {d: modelloader.directory_directories(d) for d in allowed_folders}.items():
+ for parentdir, dirs in {d: modelloader.directory_list(d) for d in allowed_folders}.items():
for tgt in dirs.keys():
if shared.backend == shared.Backend.DIFFUSERS:
if os.path.join(paths.models_path, 'Reference') in tgt:
@@ -265,7 +269,8 @@ class ExtraNetworksPage:
self.create_items(tabname)
self.create_xyz_grid()
htmls = []
- self.items.sort(key=lambda x: x["mtime"], reverse=True)
+ if len(self.items) > 0 and self.items[0].get('mtime', None) is not None:
+ self.items.sort(key=lambda x: x["mtime"], reverse=True)
for item in self.items:
htmls.append(self.create_html(item, tabname))
self.html += ''.join(htmls)
@@ -318,7 +323,7 @@ class ExtraNetworksPage:
return 'html/card-no-preview.png'
if shared.opts.diffusers_dir in path:
path = os.path.relpath(path, shared.opts.diffusers_dir)
- ref = os.path.join(paths.models_path, 'Reference')
+ ref = os.path.join('models', 'Reference')
fn = os.path.join(ref, path.replace('models--', '').replace('\\', '/').split('/')[0])
files = listdir(ref)
else:
@@ -539,9 +544,11 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
refresh_time = time.time()
threads = []
for page in get_pages():
- # page.create_items(ui.tabname)
- threads.append(threading.Thread(target=page.create_items, args=[ui.tabname]))
- threads[-1].start()
+ if os.environ.get('SD_EN_DEBUG', None) is not None:
+ threads.append(threading.Thread(target=page.create_items, args=[ui.tabname]))
+ threads[-1].start()
+ else:
+ page.create_items(ui.tabname)
for thread in threads:
thread.join()
for page in get_pages():
diff --git a/modules/ui_extra_networks_checkpoints.py b/modules/ui_extra_networks_checkpoints.py
index f31604079..01902b376 100644
--- a/modules/ui_extra_networks_checkpoints.py
+++ b/modules/ui_extra_networks_checkpoints.py
@@ -1,9 +1,11 @@
+import os
import html
import json
-import os
-from modules import shared, ui_extra_networks, sd_models, paths
+import concurrent
+from modules import shared, ui_extra_networks, sd_models
-reference_dir = os.path.join(paths.models_path, 'Reference')
+
+reference_dir = os.path.join('models', 'Reference')
class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
def __init__(self):
@@ -12,7 +14,7 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
def refresh(self):
shared.refresh_checkpoints()
- def list_reference(self):
+ def list_reference(self): # pylint: disable=inconsistent-return-statements
if shared.backend != shared.Backend.DIFFUSERS:
return []
reference_models = shared.readfile(os.path.join('html', 'reference.json'))
@@ -36,31 +38,38 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
"description": v.get('desc', ''),
}
+ def create_item(self, name):
+ record = None
+ try:
+ checkpoint: sd_models.CheckpointInfo = sd_models.checkpoints_list.get(name)
+ exists = os.path.exists(checkpoint.filename)
+ record = {
+ "type": 'Model',
+ "name": checkpoint.name,
+ "title": checkpoint.title,
+ "filename": checkpoint.filename,
+ "hash": checkpoint.shorthash,
+ "search_term": self.search_terms_from_path(checkpoint.title),
+ "preview": self.find_preview(checkpoint.filename),
+ "local_preview": f"{os.path.splitext(checkpoint.filename)[0]}.{shared.opts.samples_format}",
+ "metadata": checkpoint.metadata,
+ "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"',
+ "mtime": os.path.getmtime(checkpoint.filename) if exists else 0,
+ "size": os.path.getsize(checkpoint.filename) if exists else 0,
+ }
+ record["info"] = self.find_info(checkpoint.filename)
+ record["description"] = self.find_description(checkpoint.filename, record["info"])
+ except Exception as e:
+ shared.log.debug(f"Extra networks error: type=model file={name} {e}")
+ return record
+
def list_items(self):
- checkpoint: sd_models.CheckpointInfo
- checkpoints = sd_models.checkpoints_list.copy()
- for name, checkpoint in checkpoints.items():
- try:
- exists = os.path.exists(checkpoint.filename)
- record = {
- "type": 'Model',
- "name": checkpoint.name,
- "title": checkpoint.title,
- "filename": checkpoint.filename,
- "hash": checkpoint.shorthash,
- "search_term": self.search_terms_from_path(checkpoint.title),
- "preview": self.find_preview(checkpoint.filename),
- "local_preview": f"{os.path.splitext(checkpoint.filename)[0]}.{shared.opts.samples_format}",
- "metadata": checkpoint.metadata,
- "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"',
- "mtime": os.path.getmtime(checkpoint.filename) if exists else 0,
- "size": os.path.getsize(checkpoint.filename) if exists else 0,
- }
- record["info"] = self.find_info(checkpoint.filename)
- record["description"] = self.find_description(checkpoint.filename, record["info"])
- yield record
- except Exception as e:
- shared.log.debug(f"Extra networks error: type=model file={name} {e}")
+ with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+ future_items = {executor.submit(self.create_item, cp): cp for cp in list(sd_models.checkpoints_list.copy())}
+ for future in concurrent.futures.as_completed(future_items):
+ item = future.result()
+ if item is not None:
+ yield item
for record in self.list_reference():
yield record
diff --git a/modules/ui_extra_networks_hypernets.py b/modules/ui_extra_networks_hypernets.py
index 2fde91634..28189dad4 100644
--- a/modules/ui_extra_networks_hypernets.py
+++ b/modules/ui_extra_networks_hypernets.py
@@ -13,7 +13,6 @@ class ExtraNetworksPageHypernetworks(ui_extra_networks.ExtraNetworksPage):
def list_items(self):
for name, path in shared.hypernetworks.items():
try:
- fn = os.path.splitext(path)[0]
name = os.path.relpath(os.path.splitext(path)[0], shared.opts.hypernetwork_dir)
yield {
"type": 'Hypernetwork',
diff --git a/modules/ui_extra_networks_styles.py b/modules/ui_extra_networks_styles.py
index 49a34affb..12f5f0a53 100644
--- a/modules/ui_extra_networks_styles.py
+++ b/modules/ui_extra_networks_styles.py
@@ -1,6 +1,7 @@
import os
import html
import json
+import concurrent
from modules import shared, extra_networks, ui_extra_networks, styles
@@ -62,35 +63,44 @@ class ExtraNetworksPageStyles(ui_extra_networks.ExtraNetworksPage):
}
return item
- def list_items(self):
- for k, style in shared.prompt_styles.styles.items():
- try:
- fn = os.path.splitext(getattr(style, 'filename', ''))[0]
- name = getattr(style, 'name', '')
- if name == '':
- continue
- txt = f'Prompt: {getattr(style, "prompt", "")}'
- if len(getattr(style, 'negative_prompt', '')) > 0:
- txt += f'\nNegative: {style.negative_prompt}'
- yield {
- "type": 'Style',
- "name": name,
- "title": k,
- "filename": style.filename,
- "search_term": f'{txt} {self.search_terms_from_path(name)}',
- "preview": style.preview if getattr(style, 'preview', None) is not None and style.preview.startswith('data:') else self.find_preview(fn),
- "description": style.description if getattr(style, 'description', None) is not None and len(style.description) > 0 else txt,
- "prompt": getattr(style, 'prompt', ''),
- "negative": getattr(style, 'negative_prompt', ''),
- "extra": getattr(style, 'extra', ''),
- "local_preview": f"{fn}.{shared.opts.samples_format}",
- "onclick": '"' + html.escape(f"""return selectStyle({json.dumps(name)})""") + '"',
- "mtime": getattr(style, 'mtime', 0),
- "size": os.path.getsize(style.filename),
- }
- except Exception as e:
- shared.log.debug(f"Extra networks error: type=style file={k} {e}")
+ def create_item(self, k):
+ item = None
+ try:
+ style = shared.prompt_styles.styles.get(k)
+ fn = os.path.splitext(getattr(style, 'filename', ''))[0]
+ name = getattr(style, 'name', '')
+ if name == '':
+ return item
+ txt = f'Prompt: {getattr(style, "prompt", "")}'
+ if len(getattr(style, 'negative_prompt', '')) > 0:
+ txt += f'\nNegative: {style.negative_prompt}'
+ item = {
+ "type": 'Style',
+ "name": name,
+ "title": k,
+ "filename": style.filename,
+ "search_term": f'{txt} {self.search_terms_from_path(name)}',
+ "preview": style.preview if getattr(style, 'preview', None) is not None and style.preview.startswith('data:') else self.find_preview(fn),
+ "description": style.description if getattr(style, 'description', None) is not None and len(style.description) > 0 else txt,
+ "prompt": getattr(style, 'prompt', ''),
+ "negative": getattr(style, 'negative_prompt', ''),
+ "extra": getattr(style, 'extra', ''),
+ "local_preview": f"{fn}.{shared.opts.samples_format}",
+ "onclick": '"' + html.escape(f"""return selectStyle({json.dumps(name)})""") + '"',
+ "mtime": getattr(style, 'mtime', 0),
+ "size": os.path.getsize(style.filename),
+ }
+ except Exception as e:
+ shared.log.debug(f"Extra networks error: type=style file={k} {e}")
+ return item
+ def list_items(self):
+ with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+ future_items = {executor.submit(self.create_item, style): style for style in list(shared.prompt_styles.styles)}
+ for future in concurrent.futures.as_completed(future_items):
+ item = future.result()
+ if item is not None:
+ yield item
def allowed_directories_for_previews(self):
return [v for v in [shared.opts.styles_dir] if v is not None] + ['html']
diff --git a/modules/ui_extra_networks_textual_inversion.py b/modules/ui_extra_networks_textual_inversion.py
index 004e4cecc..f67031029 100644
--- a/modules/ui_extra_networks_textual_inversion.py
+++ b/modules/ui_extra_networks_textual_inversion.py
@@ -1,5 +1,6 @@
import json
import os
+import concurrent
from modules import shared, sd_hijack, sd_models, ui_extra_networks
from modules.textual_inversion.textual_inversion import Embedding
@@ -8,6 +9,7 @@ class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage):
def __init__(self):
super().__init__('Embedding')
self.allow_negative_prompt = True
+ self.embeddings = []
def refresh(self):
if sd_models.model_data.sd_model is None:
@@ -17,51 +19,61 @@ class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage):
elif hasattr(sd_models.model_data.sd_model, 'embedding_db'):
sd_models.model_data.sd_model.embedding_db.load_textual_inversion_embeddings(force_reload=True)
+ def create_item(self, embedding: Embedding):
+ record = None
+ try:
+ path, _ext = os.path.splitext(embedding.filename)
+ tags = {}
+ if embedding.tag is not None:
+ tags[embedding.tag]=1
+ name = os.path.splitext(embedding.basename)[0]
+ record = {
+ "type": 'Embedding',
+ "name": name,
+ "filename": embedding.filename,
+ "preview": self.find_preview(embedding.filename),
+ "search_term": self.search_terms_from_path(name),
+ "prompt": json.dumps(f" {os.path.splitext(embedding.name)[0]}"),
+ "local_preview": f"{path}.{shared.opts.samples_format}",
+ "tags": tags,
+ "mtime": os.path.getmtime(embedding.filename),
+ "size": os.path.getsize(embedding.filename),
+ }
+ record["info"] = self.find_info(embedding.filename)
+ record["description"] = self.find_description(embedding.filename, record["info"])
+ except Exception as e:
+ shared.log.debug(f"Extra networks error: type=embedding file={embedding.filename} {e}")
+ return record
+
def list_items(self):
+
def list_folder(folder):
for filename in os.listdir(folder):
fn = os.path.join(folder, filename)
if os.path.isfile(fn) and (fn.lower().endswith(".pt") or fn.lower().endswith(".safetensors")):
embedding = Embedding(vec=0, name=os.path.basename(fn), filename=fn)
embedding.filename = fn
- embeddings.append(embedding)
+ self.embeddings.append(embedding)
elif os.path.isdir(fn) and not fn.startswith('.'):
list_folder(fn)
if sd_models.model_data.sd_model is None:
- embeddings = []
+ self.embeddings = []
list_folder(shared.opts.embeddings_dir)
elif shared.backend == shared.Backend.ORIGINAL:
- embeddings = list(sd_hijack.model_hijack.embedding_db.word_embeddings.values())
+ self.embeddings = list(sd_hijack.model_hijack.embedding_db.word_embeddings.values())
elif hasattr(sd_models.model_data.sd_model, 'embedding_db'):
- embeddings = list(sd_models.model_data.sd_model.embedding_db.word_embeddings.values())
+ self.embeddings = list(sd_models.model_data.sd_model.embedding_db.word_embeddings.values())
else:
- embeddings = []
- embeddings = sorted(embeddings, key=lambda emb: emb.filename)
- for embedding in embeddings:
- try:
- path, _ext = os.path.splitext(embedding.filename)
- tags = {}
- if embedding.tag is not None:
- tags[embedding.tag]=1
- name = os.path.splitext(embedding.basename)[0]
- record = {
- "type": 'Embedding',
- "name": name,
- "filename": embedding.filename,
- "preview": self.find_preview(embedding.filename),
- "search_term": self.search_terms_from_path(name),
- "prompt": json.dumps(f" {os.path.splitext(embedding.name)[0]}"),
- "local_preview": f"{path}.{shared.opts.samples_format}",
- "tags": tags,
- "mtime": os.path.getmtime(embedding.filename),
- "size": os.path.getsize(embedding.filename),
- }
- record["info"] = self.find_info(embedding.filename)
- record["description"] = self.find_description(embedding.filename, record["info"])
- yield record
- except Exception as e:
- shared.log.debug(f"Extra networks error: type=embedding file={embedding.filename} {e}")
+ self.embeddings = []
+ self.embeddings = sorted(self.embeddings, key=lambda emb: emb.filename)
+
+ with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+ future_items = {executor.submit(self.create_item, net): net for net in self.embeddings}
+ for future in concurrent.futures.as_completed(future_items):
+ item = future.result()
+ if item is not None:
+ yield item
def allowed_directories_for_previews(self):
return list(sd_hijack.model_hijack.embedding_db.embedding_dirs)
diff --git a/modules/ui_interrogate.py b/modules/ui_interrogate.py
index 3adc5b72c..b9bb074ea 100644
--- a/modules/ui_interrogate.py
+++ b/modules/ui_interrogate.py
@@ -35,7 +35,7 @@ class BatchWriter:
def load(clip_model_name):
global ci # pylint: disable=global-statement
if ci is None:
- config = Config(device=devices.get_optimal_device(), cache_path=os.path.join(paths.models_path, 'clip-interrogator'), clip_model_name=clip_model_name, quiet=True)
+ config = Config(device=devices.get_optimal_device(), cache_path=os.path.join(paths.models_path, 'Interrogator'), clip_model_name=clip_model_name, quiet=True)
if low_vram:
config.apply_low_vram_defaults()
shared.log.info(f'Interrogate load: config={config}')
diff --git a/modules/upscaler.py b/modules/upscaler.py
index e60f9a46c..52c50ff45 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -51,6 +51,23 @@ class Upscaler:
except Exception:
pass
+ def find_folder(self, folder, scalers, loaded):
+ for fn in os.listdir(folder): # from folder
+ file_name = os.path.join(folder, fn)
+ if os.path.isdir(file_name):
+ self.find_folder(file_name, scalers, loaded)
+ continue
+ if not file_name.endswith('.pth') and not file_name.endswith('.pt'):
+ continue
+ if file_name not in loaded:
+ model_name = os.path.splitext(fn)[0]
+ scaler = UpscalerData(name=f'{self.name} {model_name}', path=file_name, upscaler=self)
+ scaler.custom = True
+ scalers.append(scaler)
+ loaded.append(file_name)
+ modules.shared.log.debug(f'Upscaler type={self.name} folder="{folder}" model="{model_name}" path="{file_name}"')
+ print(f'Upscaler type={self.name} folder="{folder}" model="{model_name}" path="{file_name}"')
+
def find_scalers(self):
scalers = []
loaded = []
@@ -66,6 +83,8 @@ class Upscaler:
# modules.shared.log.debug(f'Upscaler type={self.name} folder="{self.user_path}" model="{model[0]}" path="{model_path}"')
if not os.path.exists(self.user_path):
return scalers
+ self.find_folder(self.user_path, scalers, loaded)
+ """
for fn in os.listdir(self.user_path): # from folder
if not fn.endswith('.pth') and not fn.endswith('.pt'):
continue
@@ -77,6 +96,7 @@ class Upscaler:
scalers.append(scaler)
loaded.append(file_name)
# modules.shared.log.debug(f'Upscaler type={self.name} folder="{self.user_path}" model="{model_name}" path="{file_name}"')
+ """
return scalers
@abstractmethod
diff --git a/repositories/codeformer/facelib/detection/__init__.py b/repositories/codeformer/facelib/detection/__init__.py
index 5d1f8fc21..1c021d410 100644
--- a/repositories/codeformer/facelib/detection/__init__.py
+++ b/repositories/codeformer/facelib/detection/__init__.py
@@ -1,14 +1,16 @@
import os
+from copy import deepcopy
import torch
from torch import nn
-from copy import deepcopy
-
from facelib.utils import load_file_from_url
from facelib.utils import download_pretrained_models
from facelib.detection.yolov5face.models.common import Conv
-
from .retinaface.retinaface import RetinaFace
from .yolov5face.face_detector import YoloDetector
+from modules import paths
+
+
+model_dir = os.path.join(paths.models_path, 'Codeformer')
def init_detection_model(model_name, half=False, device='cuda'):
@@ -32,7 +34,7 @@ def init_retinaface_model(model_name, half=False, device='cuda'):
else:
raise NotImplementedError(f'{model_name} is not implemented.')
- model_path = load_file_from_url(url=model_url, model_dir='weights/facelib', progress=True, file_name=None)
+ model_path = load_file_from_url(url=model_url, model_dir=model_dir, progress=True, file_name=None)
load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
# remove unnecessary 'module.'
for k, v in deepcopy(load_net).items():
@@ -55,8 +57,8 @@ def init_yolov5face_model(model_name, device='cuda'):
model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5n-face.pth'
else:
raise NotImplementedError(f'{model_name} is not implemented.')
-
- model_path = load_file_from_url(url=model_url, model_dir='weights/facelib', progress=True, file_name=None)
+
+ model_path = load_file_from_url(url=model_url, model_dir=model_dir, progress=True, file_name=None)
load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
model.detector.load_state_dict(load_net, strict=True)
model.detector.eval()
@@ -69,32 +71,3 @@ def init_yolov5face_model(model_name, device='cuda'):
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
return model
-
-
-# Download from Google Drive
-# def init_yolov5face_model(model_name, device='cuda'):
-# if model_name == 'YOLOv5l':
-# model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5l.yaml', device=device)
-# f_id = {'yolov5l-face.pth': '131578zMA6B2x8VQHyHfa6GEPtulMCNzV'}
-# elif model_name == 'YOLOv5n':
-# model = YoloDetector(config_name='facelib/detection/yolov5face/models/yolov5n.yaml', device=device)
-# f_id = {'yolov5n-face.pth': '1fhcpFvWZqghpGXjYPIne2sw1Fy4yhw6o'}
-# else:
-# raise NotImplementedError(f'{model_name} is not implemented.')
-
-# model_path = os.path.join('weights/facelib', list(f_id.keys())[0])
-# if not os.path.exists(model_path):
-# download_pretrained_models(file_ids=f_id, save_path_root='weights/facelib')
-
-# load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
-# model.detector.load_state_dict(load_net, strict=True)
-# model.detector.eval()
-# model.detector = model.detector.to(device).float()
-
-# for m in model.detector.modules():
-# if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
-# m.inplace = True # pytorch 1.7.0 compatibility
-# elif isinstance(m, Conv):
-# m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
-
-# return model
\ No newline at end of file
diff --git a/repositories/codeformer/scripts/download_pretrained_models.py b/repositories/codeformer/scripts/download_pretrained_models.py
index daa6e8ca1..16ff6e170 100644
--- a/repositories/codeformer/scripts/download_pretrained_models.py
+++ b/repositories/codeformer/scripts/download_pretrained_models.py
@@ -1,40 +1,26 @@
-import argparse
import os
-from os import path as osp
-
+from modules import paths
from basicsr.utils.download_util import load_file_from_url
-def download_pretrained_models(method, file_urls):
- save_path_root = f'./weights/{method}'
- os.makedirs(save_path_root, exist_ok=True)
+urls = {
+ 'CodeFormer': {
+ 'codeformer.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
+ },
+ 'facelib': {
+ # 'yolov5l-face.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth',
+ 'detection_Resnet50_Final.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',
+ 'parsing_parsenet.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
+ }
+}
+
+def download_pretrained_models(file_urls):
+ model_dir = os.path.join(paths.models_path, 'Codeformer')
for file_name, file_url in file_urls.items():
- save_path = load_file_from_url(url=file_url, model_dir=save_path_root, progress=True, file_name=file_name)
+ load_file_from_url(url=file_url, model_dir=model_dir, progress=True, file_name=file_name)
if __name__ == '__main__':
- parser = argparse.ArgumentParser()
-
- parser.add_argument(
- 'method',
- type=str,
- help=("Options: 'CodeFormer' 'facelib'. Set to 'all' to download all the models."))
- args = parser.parse_args()
-
- file_urls = {
- 'CodeFormer': {
- 'codeformer.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth'
- },
- 'facelib': {
- # 'yolov5l-face.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth',
- 'detection_Resnet50_Final.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',
- 'parsing_parsenet.pth': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
- }
- }
-
- if args.method == 'all':
- for method in file_urls.keys():
- download_pretrained_models(method, file_urls[method])
- else:
- download_pretrained_models(args.method, file_urls[args.method])
\ No newline at end of file
+ for method in urls.keys():
+ download_pretrained_models(urls[method])
diff --git a/repositories/codeformer/scripts/download_pretrained_models_from_gdrive.py b/repositories/codeformer/scripts/download_pretrained_models_from_gdrive.py
index 7df5be6fc..5a5c6bd44 100644
--- a/repositories/codeformer/scripts/download_pretrained_models_from_gdrive.py
+++ b/repositories/codeformer/scripts/download_pretrained_models_from_gdrive.py
@@ -1,18 +1,16 @@
-import argparse
import os
+from modules import paths
from os import path as osp
-
-# from basicsr.utils.download_util import download_file_from_google_drive
import gdown
-def download_pretrained_models(method, file_ids):
- save_path_root = f'./weights/{method}'
- os.makedirs(save_path_root, exist_ok=True)
+model_dir = os.path.join(paths.models_path, 'Codeformer')
+
+def download_pretrained_models(file_ids):
for file_name, file_id in file_ids.items():
file_url = 'https://drive.google.com/uc?id='+file_id
- save_path = osp.abspath(osp.join(save_path_root, file_name))
+ save_path = osp.abspath(osp.join(model_dir, file_name))
if osp.exists(save_path):
user_response = input(f'{file_name} already exist. Do you want to cover it? Y/N\n')
if user_response.lower() == 'y':
@@ -29,21 +27,13 @@ def download_pretrained_models(method, file_ids):
# download_file_from_google_drive(file_id, save_path)
if __name__ == '__main__':
- parser = argparse.ArgumentParser()
-
- parser.add_argument(
- 'method',
- type=str,
- help=("Options: 'CodeFormer' 'facelib'. Set to 'all' to download all the models."))
- args = parser.parse_args()
-
# file name: file id
# 'dlib': {
# 'mmod_human_face_detector-4cb19393.dat': '1qD-OqY8M6j4PWUP_FtqfwUPFPRMu6ubX',
# 'shape_predictor_5_face_landmarks-c4b1e980.dat': '1vF3WBUApw4662v9Pw6wke3uk1qxnmLdg',
# 'shape_predictor_68_face_landmarks-fbdc2cb8.dat': '1tJyIVdCHaU6IDMDx86BZCxLGZfsWB8yq'
# }
- file_ids = {
+ urls = {
'CodeFormer': {
'codeformer.pth': '1v_E_vZvP-dQPF55Kc5SRCjaKTQXDz-JB'
},
@@ -52,9 +42,5 @@ if __name__ == '__main__':
'parsing_parsenet.pth': '16pkohyZZ8ViHGBk3QtVqxLZKzdo466bK'
}
}
-
- if args.method == 'all':
- for method in file_ids.keys():
- download_pretrained_models(method, file_ids[method])
- else:
- download_pretrained_models(args.method, file_ids[args.method])
\ No newline at end of file
+ for method in urls.keys():
+ download_pretrained_models(urls[method])
diff --git a/repositories/codeformer/web-demos/hugging_face/app.py b/repositories/codeformer/web-demos/hugging_face/app.py
deleted file mode 100644
index 7da0fc947..000000000
--- a/repositories/codeformer/web-demos/hugging_face/app.py
+++ /dev/null
@@ -1,280 +0,0 @@
-"""
-This file is used for deploying hugging face demo:
-https://huggingface.co/spaces/sczhou/CodeFormer
-"""
-
-import sys
-sys.path.append('CodeFormer')
-import os
-import cv2
-import torch
-import torch.nn.functional as F
-import gradio as gr
-
-from torchvision.transforms.functional import normalize
-
-from basicsr.utils import imwrite, img2tensor, tensor2img
-from basicsr.utils.download_util import load_file_from_url
-from facelib.utils.face_restoration_helper import FaceRestoreHelper
-from facelib.utils.misc import is_gray
-from basicsr.archs.rrdbnet_arch import RRDBNet
-from basicsr.utils.realesrgan_utils import RealESRGANer
-
-from basicsr.utils.registry import ARCH_REGISTRY
-
-
-os.system("pip freeze")
-
-pretrain_model_url = {
- 'codeformer': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth',
- 'detection': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth',
- 'parsing': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth',
- 'realesrgan': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth'
-}
-# download weights
-if not os.path.exists('CodeFormer/weights/CodeFormer/codeformer.pth'):
- load_file_from_url(url=pretrain_model_url['codeformer'], model_dir='CodeFormer/weights/CodeFormer', progress=True, file_name=None)
-if not os.path.exists('CodeFormer/weights/facelib/detection_Resnet50_Final.pth'):
- load_file_from_url(url=pretrain_model_url['detection'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)
-if not os.path.exists('CodeFormer/weights/facelib/parsing_parsenet.pth'):
- load_file_from_url(url=pretrain_model_url['parsing'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None)
-if not os.path.exists('CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth'):
- load_file_from_url(url=pretrain_model_url['realesrgan'], model_dir='CodeFormer/weights/realesrgan', progress=True, file_name=None)
-
-# download images
-torch.hub.download_url_to_file(
- 'https://replicate.com/api/models/sczhou/codeformer/files/fa3fe3d1-76b0-4ca8-ac0d-0a925cb0ff54/06.png',
- '01.png')
-torch.hub.download_url_to_file(
- 'https://replicate.com/api/models/sczhou/codeformer/files/a1daba8e-af14-4b00-86a4-69cec9619b53/04.jpg',
- '02.jpg')
-torch.hub.download_url_to_file(
- 'https://replicate.com/api/models/sczhou/codeformer/files/542d64f9-1712-4de7-85f7-3863009a7c3d/03.jpg',
- '03.jpg')
-torch.hub.download_url_to_file(
- 'https://replicate.com/api/models/sczhou/codeformer/files/a11098b0-a18a-4c02-a19a-9a7045d68426/010.jpg',
- '04.jpg')
-torch.hub.download_url_to_file(
- 'https://replicate.com/api/models/sczhou/codeformer/files/7cf19c2c-e0cf-4712-9af8-cf5bdbb8d0ee/012.jpg',
- '05.jpg')
-
-def imread(img_path):
- img = cv2.imread(img_path)
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- return img
-
-# set enhancer with RealESRGAN
-def set_realesrgan():
- half = True if torch.cuda.is_available() else False
- model = RRDBNet(
- num_in_ch=3,
- num_out_ch=3,
- num_feat=64,
- num_block=23,
- num_grow_ch=32,
- scale=2,
- )
- upsampler = RealESRGANer(
- scale=2,
- model_path="CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth",
- model=model,
- tile=400,
- tile_pad=40,
- pre_pad=0,
- half=half,
- )
- return upsampler
-
-upsampler = set_realesrgan()
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-codeformer_net = ARCH_REGISTRY.get("CodeFormer")(
- dim_embd=512,
- codebook_size=1024,
- n_head=8,
- n_layers=9,
- connect_list=["32", "64", "128", "256"],
-).to(device)
-ckpt_path = "CodeFormer/weights/CodeFormer/codeformer.pth"
-checkpoint = torch.load(ckpt_path)["params_ema"]
-codeformer_net.load_state_dict(checkpoint)
-codeformer_net.eval()
-
-os.makedirs('output', exist_ok=True)
-
-def inference(image, background_enhance, face_upsample, upscale, codeformer_fidelity):
- """Run a single prediction on the model"""
- try: # global try
- # take the default setting for the demo
- has_aligned = False
- only_center_face = False
- draw_box = False
- detection_model = "retinaface_resnet50"
- print('Inp:', image, background_enhance, face_upsample, upscale, codeformer_fidelity)
-
- img = cv2.imread(str(image), cv2.IMREAD_COLOR)
- print('\timage size:', img.shape)
-
- upscale = int(upscale) # convert type to int
- if upscale > 4: # avoid memory exceeded due to too large upscale
- upscale = 4
- if upscale > 2 and max(img.shape[:2])>1000: # avoid memory exceeded due to too large img resolution
- upscale = 2
- if max(img.shape[:2]) > 1500: # avoid memory exceeded due to too large img resolution
- upscale = 1
- background_enhance = False
- face_upsample = False
-
- face_helper = FaceRestoreHelper(
- upscale,
- face_size=512,
- crop_ratio=(1, 1),
- det_model=detection_model,
- save_ext="png",
- use_parse=True,
- device=device,
- )
- bg_upsampler = upsampler if background_enhance else None
- face_upsampler = upsampler if face_upsample else None
-
- if has_aligned:
- # the input faces are already cropped and aligned
- img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
- face_helper.is_gray = is_gray(img, threshold=5)
- if face_helper.is_gray:
- print('\tgrayscale input: True')
- face_helper.cropped_faces = [img]
- else:
- face_helper.read_image(img)
- # get face landmarks for each face
- num_det_faces = face_helper.get_face_landmarks_5(
- only_center_face=only_center_face, resize=640, eye_dist_threshold=5
- )
- print(f'\tdetect {num_det_faces} faces')
- # align and warp each face
- face_helper.align_warp_face()
-
- # face restoration for each cropped face
- for idx, cropped_face in enumerate(face_helper.cropped_faces):
- # prepare data
- cropped_face_t = img2tensor(
- cropped_face / 255.0, bgr2rgb=True, float32=True
- )
- normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
- cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
-
- try:
- with torch.no_grad():
- output = codeformer_net(
- cropped_face_t, w=codeformer_fidelity, adain=True
- )[0]
- restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
- del output
- torch.cuda.empty_cache()
- except RuntimeError as error:
- print(f"Failed inference for CodeFormer: {error}")
- restored_face = tensor2img(
- cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
- )
-
- restored_face = restored_face.astype("uint8")
- face_helper.add_restored_face(restored_face)
-
- # paste_back
- if not has_aligned:
- # upsample the background
- if bg_upsampler is not None:
- # Now only support RealESRGAN for upsampling background
- bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
- else:
- bg_img = None
- face_helper.get_inverse_affine(None)
- # paste each restored face to the input image
- if face_upsample and face_upsampler is not None:
- restored_img = face_helper.paste_faces_to_input_image(
- upsample_img=bg_img,
- draw_box=draw_box,
- face_upsampler=face_upsampler,
- )
- else:
- restored_img = face_helper.paste_faces_to_input_image(
- upsample_img=bg_img, draw_box=draw_box
- )
-
- # save restored img
- save_path = f'output/out.png'
- imwrite(restored_img, str(save_path))
-
- restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)
- return restored_img, save_path
- except Exception as error:
- print('Global exception', error)
- return None, None
-
-
-title = "CodeFormer: Robust Face Restoration and Enhancement Network"
-description = r"""
-Official Gradio demo for Towards Robust Blind Face Restoration with Codebook Lookup Transformer (NeurIPS 2022).
-🔥 CodeFormer is a robust face restoration algorithm for old photos or AI-generated faces.
-🤗 Try CodeFormer for improved stable-diffusion generation!
-"""
-article = r"""
-If CodeFormer is helpful, please help to ⭐ the Github Repo. Thanks!
-[](https://github.com/sczhou/CodeFormer)
-
----
-
-📝 **Citation**
-
-If our work is useful for your research, please consider citing:
-```bibtex
-@inproceedings{zhou2022codeformer,
- author = {Zhou, Shangchen and Chan, Kelvin C.K. and Li, Chongyi and Loy, Chen Change},
- title = {Towards Robust Blind Face Restoration with Codebook Lookup TransFormer},
- booktitle = {NeurIPS},
- year = {2022}
-}
-```
-
-📋 **License**
-
-This project is licensed under S-Lab License 1.0.
-Redistribution and use for non-commercial purposes should follow this license.
-
-📧 **Contact**
-
-If you have any questions, please feel free to reach me out at shangchenzhou@gmail.com.
-
-
- 🤗 Find Me:
-
-
-
-"""
-
-demo = gr.Interface(
- inference, [
- gr.inputs.Image(type="filepath", label="Input"),
- gr.inputs.Checkbox(default=True, label="Background_Enhance"),
- gr.inputs.Checkbox(default=True, label="Face_Upsample"),
- gr.inputs.Number(default=2, label="Rescaling_Factor (up to 4)"),
- gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity (0 for better quality, 1 for better identity)')
- ], [
- gr.outputs.Image(type="numpy", label="Output"),
- gr.outputs.File(label="Download the output")
- ],
- title=title,
- description=description,
- article=article,
- examples=[
- ['01.png', True, True, 2, 0.7],
- ['02.jpg', True, True, 2, 0.7],
- ['03.jpg', True, True, 2, 0.7],
- ['04.jpg', True, True, 2, 0.1],
- ['05.jpg', True, True, 2, 0.1]
- ]
- )
-
-demo.queue(concurrency_count=2)
-demo.launch()
\ No newline at end of file
diff --git a/repositories/codeformer/web-demos/replicate/cog.yaml b/repositories/codeformer/web-demos/replicate/cog.yaml
deleted file mode 100644
index 3f4589690..000000000
--- a/repositories/codeformer/web-demos/replicate/cog.yaml
+++ /dev/null
@@ -1,30 +0,0 @@
-"""
-This file is used for deploying replicate demo:
-https://replicate.com/sczhou/codeformer
-"""
-
-build:
- gpu: true
- cuda: "11.3"
- python_version: "3.8"
- system_packages:
- - "libgl1-mesa-glx"
- - "libglib2.0-0"
- python_packages:
- - "ipython==8.4.0"
- - "future==0.18.2"
- - "lmdb==1.3.0"
- - "scikit-image==0.19.3"
- - "torch==1.11.0 --extra-index-url=https://download.pytorch.org/whl/cu113"
- - "torchvision==0.12.0 --extra-index-url=https://download.pytorch.org/whl/cu113"
- - "scipy==1.9.0"
- - "gdown==4.5.1"
- - "pyyaml==6.0"
- - "tb-nightly==2.11.0a20220906"
- - "tqdm==4.64.1"
- - "yapf==0.32.0"
- - "lpips==0.1.4"
- - "Pillow==9.2.0"
- - "opencv-python==4.6.0.66"
-
-predict: "predict.py:Predictor"
diff --git a/repositories/codeformer/web-demos/replicate/predict.py b/repositories/codeformer/web-demos/replicate/predict.py
deleted file mode 100644
index 61935e9e7..000000000
--- a/repositories/codeformer/web-demos/replicate/predict.py
+++ /dev/null
@@ -1,189 +0,0 @@
-"""
-This file is used for deploying replicate demo:
-https://replicate.com/sczhou/codeformer
-running: cog predict -i image=@inputs/whole_imgs/04.jpg -i codeformer_fidelity=0.5 -i upscale=2
-push: cog push r8.im/sczhou/codeformer
-"""
-
-import tempfile
-import cv2
-import torch
-from torchvision.transforms.functional import normalize
-try:
- from cog import BasePredictor, Input, Path
-except Exception:
- print('please install cog package')
-
-from basicsr.utils import imwrite, img2tensor, tensor2img
-from basicsr.archs.rrdbnet_arch import RRDBNet
-from basicsr.utils.realesrgan_utils import RealESRGANer
-from basicsr.utils.registry import ARCH_REGISTRY
-from facelib.utils.face_restoration_helper import FaceRestoreHelper
-
-
-class Predictor(BasePredictor):
- def setup(self):
- """Load the model into memory to make running multiple predictions efficient"""
- self.device = "cuda:0"
- self.upsampler = set_realesrgan()
- self.net = ARCH_REGISTRY.get("CodeFormer")(
- dim_embd=512,
- codebook_size=1024,
- n_head=8,
- n_layers=9,
- connect_list=["32", "64", "128", "256"],
- ).to(self.device)
- ckpt_path = "weights/CodeFormer/codeformer.pth"
- checkpoint = torch.load(ckpt_path)[
- "params_ema"
- ] # update file permission if cannot load
- self.net.load_state_dict(checkpoint)
- self.net.eval()
-
- def predict(
- self,
- image: Path = Input(description="Input image"),
- codeformer_fidelity: float = Input(
- default=0.5,
- ge=0,
- le=1,
- description="Balance the quality (lower number) and fidelity (higher number).",
- ),
- background_enhance: bool = Input(
- description="Enhance background image with Real-ESRGAN", default=True
- ),
- face_upsample: bool = Input(
- description="Upsample restored faces for high-resolution AI-created images",
- default=True,
- ),
- upscale: int = Input(
- description="The final upsampling scale of the image",
- default=2,
- ),
- ) -> Path:
- """Run a single prediction on the model"""
-
- # take the default setting for the demo
- has_aligned = False
- only_center_face = False
- draw_box = False
- detection_model = "retinaface_resnet50"
-
- self.face_helper = FaceRestoreHelper(
- upscale,
- face_size=512,
- crop_ratio=(1, 1),
- det_model=detection_model,
- save_ext="png",
- use_parse=True,
- device=self.device,
- )
-
- bg_upsampler = self.upsampler if background_enhance else None
- face_upsampler = self.upsampler if face_upsample else None
-
- img = cv2.imread(str(image), cv2.IMREAD_COLOR)
-
- if has_aligned:
- # the input faces are already cropped and aligned
- img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
- self.face_helper.cropped_faces = [img]
- else:
- self.face_helper.read_image(img)
- # get face landmarks for each face
- num_det_faces = self.face_helper.get_face_landmarks_5(
- only_center_face=only_center_face, resize=640, eye_dist_threshold=5
- )
- print(f"\tdetect {num_det_faces} faces")
- # align and warp each face
- self.face_helper.align_warp_face()
-
- # face restoration for each cropped face
- for idx, cropped_face in enumerate(self.face_helper.cropped_faces):
- # prepare data
- cropped_face_t = img2tensor(
- cropped_face / 255.0, bgr2rgb=True, float32=True
- )
- normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
- cropped_face_t = cropped_face_t.unsqueeze(0).to(self.device)
-
- try:
- with torch.no_grad():
- output = self.net(
- cropped_face_t, w=codeformer_fidelity, adain=True
- )[0]
- restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
- del output
- torch.cuda.empty_cache()
- except Exception as error:
- print(f"\tFailed inference for CodeFormer: {error}")
- restored_face = tensor2img(
- cropped_face_t, rgb2bgr=True, min_max=(-1, 1)
- )
-
- restored_face = restored_face.astype("uint8")
- self.face_helper.add_restored_face(restored_face)
-
- # paste_back
- if not has_aligned:
- # upsample the background
- if bg_upsampler is not None:
- # Now only support RealESRGAN for upsampling background
- bg_img = bg_upsampler.enhance(img, outscale=upscale)[0]
- else:
- bg_img = None
- self.face_helper.get_inverse_affine(None)
- # paste each restored face to the input image
- if face_upsample and face_upsampler is not None:
- restored_img = self.face_helper.paste_faces_to_input_image(
- upsample_img=bg_img,
- draw_box=draw_box,
- face_upsampler=face_upsampler,
- )
- else:
- restored_img = self.face_helper.paste_faces_to_input_image(
- upsample_img=bg_img, draw_box=draw_box
- )
-
- # save restored img
- out_path = Path(tempfile.mkdtemp()) / 'output.png'
- imwrite(restored_img, str(out_path))
-
- return out_path
-
-
-def imread(img_path):
- img = cv2.imread(img_path)
- img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
- return img
-
-
-def set_realesrgan():
- if not torch.cuda.is_available(): # CPU
- import warnings
-
- warnings.warn(
- "The unoptimized RealESRGAN is slow on CPU. We do not use it. "
- "If you really want to use it, please modify the corresponding codes.",
- category=RuntimeWarning,
- )
- upsampler = None
- else:
- model = RRDBNet(
- num_in_ch=3,
- num_out_ch=3,
- num_feat=64,
- num_block=23,
- num_grow_ch=32,
- scale=2,
- )
- upsampler = RealESRGANer(
- scale=2,
- model_path="./weights/realesrgan/RealESRGAN_x2plus.pth",
- model=model,
- tile=400,
- tile_pad=40,
- pre_pad=0,
- half=True,
- )
- return upsampler
diff --git a/repositories/codeformer/weights/CodeFormer/.gitkeep b/repositories/codeformer/weights/CodeFormer/.gitkeep
deleted file mode 100644
index e69de29bb..000000000
diff --git a/repositories/codeformer/weights/README.md b/repositories/codeformer/weights/README.md
deleted file mode 100644
index 67ad334bd..000000000
--- a/repositories/codeformer/weights/README.md
+++ /dev/null
@@ -1,3 +0,0 @@
-# Weights
-
-Put the downloaded pre-trained models to this folder.
\ No newline at end of file
diff --git a/repositories/codeformer/weights/facelib/.gitkeep b/repositories/codeformer/weights/facelib/.gitkeep
deleted file mode 100644
index e69de29bb..000000000
diff --git a/requirements.txt b/requirements.txt
index c4a0005f5..b62383f18 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -51,7 +51,7 @@ requests==2.31.0
tqdm==4.66.1
accelerate==0.20.3
opencv-python-headless==4.7.0.72
-diffusers==0.22.0
+diffusers==0.23.0
einops==0.4.1
gradio==3.43.2
huggingface_hub==0.18.0
diff --git a/scripts/postprocessing_upscale.py b/scripts/postprocessing_upscale.py
index 0086a3abc..e97ac9b45 100644
--- a/scripts/postprocessing_upscale.py
+++ b/scripts/postprocessing_upscale.py
@@ -16,7 +16,7 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
with FormRow(elem_id="extras_upscale"):
with gr.Tabs(elem_id="extras_resize_mode"):
with gr.TabItem('Scale by', elem_id="extras_scale_by_tab") as tab_scale_by:
- upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4, elem_id="extras_upscaling_resize")
+ upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=2.0, elem_id="extras_upscaling_resize")
with gr.TabItem('Scale to', elem_id="extras_scale_to_tab") as tab_scale_to:
with FormRow():
diff --git a/scripts/sd_upscale.py b/scripts/sd_upscale.py
index 0d35e262b..8e22febbe 100644
--- a/scripts/sd_upscale.py
+++ b/scripts/sd_upscale.py
@@ -18,7 +18,7 @@ class Script(scripts.Script):
info = gr.HTML("Will upscale the image by the selected scale factor; use width and height sliders to set tile size
")
overlap = gr.Slider(minimum=0, maximum=256, step=16, label='Tile overlap', value=64, elem_id=self.elem_id("overlap"))
scale_factor = gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label='Scale Factor', value=2.0, elem_id=self.elem_id("scale_factor"))
- upscaler_index = gr.Radio(label='Upscaler', choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index", elem_id=self.elem_id("upscaler_index"))
+ upscaler_index = gr.Dropdown(label='Upscaler', choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name, type="index", elem_id=self.elem_id("upscaler_index"))
return [info, overlap, upscaler_index, scale_factor]
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index 5e5c00938..912638777 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -86,6 +86,17 @@ def apply_checkpoint(p, x, xs):
p.override_settings['sd_model_checkpoint'] = info.name
+def apply_refiner(p, x, xs):
+ if x == shared.opts.sd_model_refiner:
+ return
+ info = sd_models.get_closet_checkpoint_match(x)
+ if info is None:
+ shared.log.warning(f"XYZ grid: apply refiner unknown checkpoint: {x}")
+ else:
+ sd_models.reload_model_weights(shared.sd_refiner, info)
+ p.override_settings['sd_model_refiner'] = info.name
+
+
def apply_dict(p, x, xs):
if x == shared.opts.sd_model_dict:
return
@@ -240,11 +251,12 @@ axis_options = [
AxisOption("[Second pass] upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]),
AxisOption("[Second pass] sampler", str, apply_latent_sampler, fmt=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
AxisOption("[Second pass] denoising Strength", float, apply_field("denoising_strength")),
- AxisOption("[Second pass] steps", int, apply_field("hr_second_pass_steps")),
+ AxisOption("[Second pass] hires steps", int, apply_field("hr_second_pass_steps")),
AxisOption("[Second pass] CFG scale", float, apply_field("image_cfg_scale")),
AxisOption("[Second pass] guidance rescale", float, apply_field("diffusers_guidance_rescale")),
- AxisOption("[Second pass] refiner start", float, apply_field("refiner_start")),
- AxisOption("[Second pass] refiner start", float, apply_field("refiner_start")),
+ AxisOption("[Refiner] model", str, apply_refiner, fmt=format_value, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list)),
+ AxisOption("[Refiner] refiner start", float, apply_field("refiner_start")),
+ AxisOption("[Refiner] refiner steps", float, apply_field("refiner_steps")),
AxisOption("[TOME] Token merging ratio (txt2img)", float, apply_override('token_merging_ratio')),
AxisOption("[TOME] Token merging ratio (hires)", float, apply_override('token_merging_ratio_hr')),
AxisOption("[FreeU] 1st stage backbone factor", float, apply_setting('freeu_b1')),
diff --git a/webui.py b/webui.py
index 2e3c76f80..8397aced3 100644
--- a/webui.py
+++ b/webui.py
@@ -156,7 +156,7 @@ def initialize():
def load_model():
- if opts.sd_checkpoint_autoload:
+ if opts.sd_checkpoint_autoload and (shared.cmd_opts.ckpt is not None and shared.cmd_opts.ckpt.lower() != 'none'):
shared.state.begin('load')
thread_model = Thread(target=lambda: shared.sd_model)
thread_model.start()
diff --git a/wiki b/wiki
index e999774e3..cd040c02e 160000
--- a/wiki
+++ b/wiki
@@ -1 +1 @@
-Subproject commit e999774e3096ceb89a264548fdfaaa76d891c0df
+Subproject commit cd040c02e4a477135ce08efe4d06672b57456c31
+