", "exec"), module.__dict__)
36 |
37 |
38 | class Script(scripts.Script):
39 |
40 | def title(self):
41 | return "Custom code"
42 |
43 | def show(self, is_img2img):
44 | return cmd_opts.allow_code
45 |
46 | def ui(self, is_img2img):
47 | example = """from modules.processing import process_images
48 |
49 | p.width = 768
50 | p.height = 768
51 | p.batch_size = 2
52 | p.steps = 10
53 |
54 | return process_images(p)
55 | """
56 |
57 |
58 | code = gr.Code(value=example, language="python", label="Python code", elem_id=self.elem_id("code"))
59 | indent_level = gr.Number(label='Indent level', value=2, precision=0, elem_id=self.elem_id("indent_level"))
60 |
61 | return [code, indent_level]
62 |
63 | def run(self, p, code, indent_level):
64 | assert cmd_opts.allow_code, '--allow-code option must be enabled'
65 |
66 | display_result_data = [[], -1, ""]
67 |
68 | def display(imgs, s=display_result_data[1], i=display_result_data[2]):
69 | display_result_data[0] = imgs
70 | display_result_data[1] = s
71 | display_result_data[2] = i
72 |
73 | from types import ModuleType
74 | module = ModuleType("testmodule")
75 | module.__dict__.update(globals())
76 | module.p = p
77 | module.display = display
78 |
79 | indent = " " * indent_level
80 | indented = code.replace('\n', f"\n{indent}")
81 | body = f"""def __webuitemp__():
82 | {indent}{indented}
83 | __webuitemp__()"""
84 |
85 | result = exec_with_return(body, module)
86 |
87 | if isinstance(result, Processed):
88 | return result
89 |
90 | return Processed(p, *display_result_data)
91 |
--------------------------------------------------------------------------------
/modules/ui_postprocessing.py:
--------------------------------------------------------------------------------
1 | import gradio as gr
2 | from modules import scripts, shared, ui_common, postprocessing, call_queue
3 | import modules.generation_parameters_copypaste as parameters_copypaste
4 |
5 |
6 | def create_ui():
7 | tab_index = gr.State(value=0)
8 |
9 | with gr.Row().style(equal_height=False, variant='compact'):
10 | with gr.Column(variant='compact'):
11 | with gr.Tabs(elem_id="mode_extras"):
12 | with gr.TabItem('Single Image', id="single_image", elem_id="extras_single_tab") as tab_single:
13 | extras_image = gr.Image(label="Source", source="upload", interactive=True, type="pil", elem_id="extras_image")
14 |
15 | with gr.TabItem('Batch Process', id="batch_process", elem_id="extras_batch_process_tab") as tab_batch:
16 | image_batch = gr.Files(label="Batch Process", interactive=True, elem_id="extras_image_batch")
17 |
18 | with gr.TabItem('Batch from Directory', id="batch_from_directory", elem_id="extras_batch_directory_tab") as tab_batch_dir:
19 | extras_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs, placeholder="A directory on the same machine where the server is running.", elem_id="extras_batch_input_dir")
20 | extras_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs, placeholder="Leave blank to save images to the default path.", elem_id="extras_batch_output_dir")
21 | show_extras_results = gr.Checkbox(label='Show result images', value=True, elem_id="extras_show_extras_results")
22 |
23 | submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')
24 |
25 | script_inputs = scripts.scripts_postproc.setup_ui()
26 |
27 | with gr.Column():
28 | result_images, html_info_x, html_info, html_log = ui_common.create_output_panel("extras", shared.opts.outdir_extras_samples)
29 |
30 | tab_single.select(fn=lambda: 0, inputs=[], outputs=[tab_index])
31 | tab_batch.select(fn=lambda: 1, inputs=[], outputs=[tab_index])
32 | tab_batch_dir.select(fn=lambda: 2, inputs=[], outputs=[tab_index])
33 |
34 | submit.click(
35 | fn=call_queue.wrap_gradio_gpu_call(postprocessing.run_postprocessing, extra_outputs=[None, '']),
36 | inputs=[
37 | tab_index,
38 | extras_image,
39 | image_batch,
40 | extras_batch_input_dir,
41 | extras_batch_output_dir,
42 | show_extras_results,
43 | *script_inputs
44 | ],
45 | outputs=[
46 | result_images,
47 | html_info_x,
48 | html_info,
49 | ]
50 | )
51 |
52 | parameters_copypaste.add_paste_fields("extras", extras_image, None)
53 |
54 | extras_image.change(
55 | fn=scripts.scripts_postproc.image_changed,
56 | inputs=[], outputs=[]
57 | )
58 |
--------------------------------------------------------------------------------
/modules/ui_tempdir.py:
--------------------------------------------------------------------------------
1 | import os
2 | import tempfile
3 | from collections import namedtuple
4 | from pathlib import Path
5 |
6 | import gradio.components
7 |
8 | from PIL import PngImagePlugin
9 |
10 | from modules import shared
11 |
12 |
13 | Savedfile = namedtuple("Savedfile", ["name"])
14 |
15 |
16 | def register_tmp_file(gradio, filename):
17 | if hasattr(gradio, 'temp_file_sets'): # gradio 3.15
18 | gradio.temp_file_sets[0] = gradio.temp_file_sets[0] | {os.path.abspath(filename)}
19 |
20 | if hasattr(gradio, 'temp_dirs'): # gradio 3.9
21 | gradio.temp_dirs = gradio.temp_dirs | {os.path.abspath(os.path.dirname(filename))}
22 |
23 |
24 | def check_tmp_file(gradio, filename):
25 | if hasattr(gradio, 'temp_file_sets'):
26 | return any(filename in fileset for fileset in gradio.temp_file_sets)
27 |
28 | if hasattr(gradio, 'temp_dirs'):
29 | return any(Path(temp_dir).resolve() in Path(filename).resolve().parents for temp_dir in gradio.temp_dirs)
30 |
31 | return False
32 |
33 |
34 | def save_pil_to_file(self, pil_image, dir=None):
35 | already_saved_as = getattr(pil_image, 'already_saved_as', None)
36 | if already_saved_as and os.path.isfile(already_saved_as):
37 | register_tmp_file(shared.demo, already_saved_as)
38 | filename = already_saved_as
39 |
40 | if not shared.opts.save_images_add_number:
41 | filename += f'?{os.path.getmtime(already_saved_as)}'
42 |
43 | return filename
44 |
45 | if shared.opts.temp_dir != "":
46 | dir = shared.opts.temp_dir
47 |
48 | use_metadata = False
49 | metadata = PngImagePlugin.PngInfo()
50 | for key, value in pil_image.info.items():
51 | if isinstance(key, str) and isinstance(value, str):
52 | metadata.add_text(key, value)
53 | use_metadata = True
54 |
55 | file_obj = tempfile.NamedTemporaryFile(delete=False, suffix=".png", dir=dir)
56 | pil_image.save(file_obj, pnginfo=(metadata if use_metadata else None))
57 | return file_obj.name
58 |
59 |
60 | # override save to file function so that it also writes PNG info
61 | gradio.components.IOComponent.pil_to_temp_file = save_pil_to_file
62 |
63 |
64 | def on_tmpdir_changed():
65 | if shared.opts.temp_dir == "" or shared.demo is None:
66 | return
67 |
68 | os.makedirs(shared.opts.temp_dir, exist_ok=True)
69 |
70 | register_tmp_file(shared.demo, os.path.join(shared.opts.temp_dir, "x"))
71 |
72 |
73 | def cleanup_tmpdr():
74 | temp_dir = shared.opts.temp_dir
75 | if temp_dir == "" or not os.path.isdir(temp_dir):
76 | return
77 |
78 | for root, _, files in os.walk(temp_dir, topdown=False):
79 | for name in files:
80 | _, extension = os.path.splitext(name)
81 | if extension != ".png":
82 | continue
83 |
84 | filename = os.path.join(root, name)
85 | os.remove(filename)
86 |
--------------------------------------------------------------------------------
/configs/instruct-pix2pix.yaml:
--------------------------------------------------------------------------------
1 | # File modified by authors of InstructPix2Pix from original (https://github.com/CompVis/stable-diffusion).
2 | # See more details in LICENSE.
3 |
4 | model:
5 | base_learning_rate: 1.0e-04
6 | target: modules.models.diffusion.ddpm_edit.LatentDiffusion
7 | params:
8 | linear_start: 0.00085
9 | linear_end: 0.0120
10 | num_timesteps_cond: 1
11 | log_every_t: 200
12 | timesteps: 1000
13 | first_stage_key: edited
14 | cond_stage_key: edit
15 | # image_size: 64
16 | # image_size: 32
17 | image_size: 16
18 | channels: 4
19 | cond_stage_trainable: false # Note: different from the one we trained before
20 | conditioning_key: hybrid
21 | monitor: val/loss_simple_ema
22 | scale_factor: 0.18215
23 | use_ema: false
24 |
25 | scheduler_config: # 10000 warmup steps
26 | target: ldm.lr_scheduler.LambdaLinearScheduler
27 | params:
28 | warm_up_steps: [ 0 ]
29 | cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
30 | f_start: [ 1.e-6 ]
31 | f_max: [ 1. ]
32 | f_min: [ 1. ]
33 |
34 | unet_config:
35 | target: ldm.modules.diffusionmodules.openaimodel.UNetModel
36 | params:
37 | image_size: 32 # unused
38 | in_channels: 8
39 | out_channels: 4
40 | model_channels: 320
41 | attention_resolutions: [ 4, 2, 1 ]
42 | num_res_blocks: 2
43 | channel_mult: [ 1, 2, 4, 4 ]
44 | num_heads: 8
45 | use_spatial_transformer: True
46 | transformer_depth: 1
47 | context_dim: 768
48 | use_checkpoint: True
49 | legacy: False
50 |
51 | first_stage_config:
52 | target: ldm.models.autoencoder.AutoencoderKL
53 | params:
54 | embed_dim: 4
55 | monitor: val/rec_loss
56 | ddconfig:
57 | double_z: true
58 | z_channels: 4
59 | resolution: 256
60 | in_channels: 3
61 | out_ch: 3
62 | ch: 128
63 | ch_mult:
64 | - 1
65 | - 2
66 | - 4
67 | - 4
68 | num_res_blocks: 2
69 | attn_resolutions: []
70 | dropout: 0.0
71 | lossconfig:
72 | target: torch.nn.Identity
73 |
74 | cond_stage_config:
75 | target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
76 |
77 | data:
78 | target: main.DataModuleFromConfig
79 | params:
80 | batch_size: 128
81 | num_workers: 1
82 | wrap: false
83 | validation:
84 | target: edit_dataset.EditDataset
85 | params:
86 | path: data/clip-filtered-dataset
87 | cache_dir: data/
88 | cache_name: data_10k
89 | split: val
90 | min_text_sim: 0.2
91 | min_image_sim: 0.75
92 | min_direction_sim: 0.2
93 | max_samples_per_prompt: 1
94 | min_resize_res: 512
95 | max_resize_res: 512
96 | crop_res: 512
97 | output_as_edit: False
98 | real_input: True
99 |
--------------------------------------------------------------------------------
/modules/styles.py:
--------------------------------------------------------------------------------
1 | import csv
2 | import os
3 | import os.path
4 | import typing
5 | import shutil
6 |
7 |
8 | class PromptStyle(typing.NamedTuple):
9 | name: str
10 | prompt: str
11 | negative_prompt: str
12 |
13 |
14 | def merge_prompts(style_prompt: str, prompt: str) -> str:
15 | if "{prompt}" in style_prompt:
16 | res = style_prompt.replace("{prompt}", prompt)
17 | else:
18 | parts = filter(None, (prompt.strip(), style_prompt.strip()))
19 | res = ", ".join(parts)
20 |
21 | return res
22 |
23 |
24 | def apply_styles_to_prompt(prompt, styles):
25 | for style in styles:
26 | prompt = merge_prompts(style, prompt)
27 |
28 | return prompt
29 |
30 |
31 | class StyleDatabase:
32 | def __init__(self, path: str):
33 | self.no_style = PromptStyle("None", "", "")
34 | self.styles = {}
35 | self.path = path
36 |
37 | self.reload()
38 |
39 | def reload(self):
40 | self.styles.clear()
41 |
42 | if not os.path.exists(self.path):
43 | return
44 |
45 | with open(self.path, "r", encoding="utf-8-sig", newline='') as file:
46 | reader = csv.DictReader(file, skipinitialspace=True)
47 | for row in reader:
48 | # Support loading old CSV format with "name, text"-columns
49 | prompt = row["prompt"] if "prompt" in row else row["text"]
50 | negative_prompt = row.get("negative_prompt", "")
51 | self.styles[row["name"]] = PromptStyle(row["name"], prompt, negative_prompt)
52 |
53 | def get_style_prompts(self, styles):
54 | return [self.styles.get(x, self.no_style).prompt for x in styles]
55 |
56 | def get_negative_style_prompts(self, styles):
57 | return [self.styles.get(x, self.no_style).negative_prompt for x in styles]
58 |
59 | def apply_styles_to_prompt(self, prompt, styles):
60 | return apply_styles_to_prompt(prompt, [self.styles.get(x, self.no_style).prompt for x in styles])
61 |
62 | def apply_negative_styles_to_prompt(self, prompt, styles):
63 | return apply_styles_to_prompt(prompt, [self.styles.get(x, self.no_style).negative_prompt for x in styles])
64 |
65 | def save_styles(self, path: str) -> None:
66 | # Always keep a backup file around
67 | if os.path.exists(path):
68 | shutil.copy(path, f"{path}.bak")
69 |
70 | fd = os.open(path, os.O_RDWR|os.O_CREAT)
71 | with os.fdopen(fd, "w", encoding="utf-8-sig", newline='') as file:
72 | # _fields is actually part of the public API: typing.NamedTuple is a replacement for collections.NamedTuple,
73 | # and collections.NamedTuple has explicit documentation for accessing _fields. Same goes for _asdict()
74 | writer = csv.DictWriter(file, fieldnames=PromptStyle._fields)
75 | writer.writeheader()
76 | writer.writerows(style._asdict() for k, style in self.styles.items())
77 |
--------------------------------------------------------------------------------
/modules/sd_vae_taesd.py:
--------------------------------------------------------------------------------
1 | """
2 | Tiny AutoEncoder for Stable Diffusion
3 | (DNN for encoding / decoding SD's latent space)
4 |
5 | https://github.com/madebyollin/taesd
6 | """
7 | import os
8 | import torch
9 | import torch.nn as nn
10 |
11 | from modules import devices, paths_internal
12 |
13 | sd_vae_taesd = None
14 |
15 |
16 | def conv(n_in, n_out, **kwargs):
17 | return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs)
18 |
19 |
20 | class Clamp(nn.Module):
21 | @staticmethod
22 | def forward(x):
23 | return torch.tanh(x / 3) * 3
24 |
25 |
26 | class Block(nn.Module):
27 | def __init__(self, n_in, n_out):
28 | super().__init__()
29 | self.conv = nn.Sequential(conv(n_in, n_out), nn.ReLU(), conv(n_out, n_out), nn.ReLU(), conv(n_out, n_out))
30 | self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity()
31 | self.fuse = nn.ReLU()
32 |
33 | def forward(self, x):
34 | return self.fuse(self.conv(x) + self.skip(x))
35 |
36 |
37 | def decoder():
38 | return nn.Sequential(
39 | Clamp(), conv(4, 64), nn.ReLU(),
40 | Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
41 | Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
42 | Block(64, 64), Block(64, 64), Block(64, 64), nn.Upsample(scale_factor=2), conv(64, 64, bias=False),
43 | Block(64, 64), conv(64, 3),
44 | )
45 |
46 |
47 | class TAESD(nn.Module):
48 | latent_magnitude = 3
49 | latent_shift = 0.5
50 |
51 | def __init__(self, decoder_path="taesd_decoder.pth"):
52 | """Initialize pretrained TAESD on the given device from the given checkpoints."""
53 | super().__init__()
54 | self.decoder = decoder()
55 | self.decoder.load_state_dict(
56 | torch.load(decoder_path, map_location='cpu' if devices.device.type != 'cuda' else None))
57 |
58 | @staticmethod
59 | def unscale_latents(x):
60 | """[0, 1] -> raw latents"""
61 | return x.sub(TAESD.latent_shift).mul(2 * TAESD.latent_magnitude)
62 |
63 |
64 | def download_model(model_path):
65 | model_url = 'https://github.com/madebyollin/taesd/raw/main/taesd_decoder.pth'
66 |
67 | if not os.path.exists(model_path):
68 | os.makedirs(os.path.dirname(model_path), exist_ok=True)
69 |
70 | print(f'Downloading TAESD decoder to: {model_path}')
71 | torch.hub.download_url_to_file(model_url, model_path)
72 |
73 |
74 | def model():
75 | global sd_vae_taesd
76 |
77 | if sd_vae_taesd is None:
78 | model_path = os.path.join(paths_internal.models_path, "VAE-taesd", "taesd_decoder.pth")
79 | download_model(model_path)
80 |
81 | if os.path.exists(model_path):
82 | sd_vae_taesd = TAESD(model_path)
83 | sd_vae_taesd.eval()
84 | sd_vae_taesd.to(devices.device, devices.dtype)
85 | else:
86 | raise FileNotFoundError('TAESD model not found')
87 |
88 | return sd_vae_taesd.decoder
89 |
--------------------------------------------------------------------------------
/modules/memmon.py:
--------------------------------------------------------------------------------
1 | import threading
2 | import time
3 | from collections import defaultdict
4 |
5 | import torch
6 |
7 |
8 | class MemUsageMonitor(threading.Thread):
9 | run_flag = None
10 | device = None
11 | disabled = False
12 | opts = None
13 | data = None
14 |
15 | def __init__(self, name, device, opts):
16 | threading.Thread.__init__(self)
17 | self.name = name
18 | self.device = device
19 | self.opts = opts
20 |
21 | self.daemon = True
22 | self.run_flag = threading.Event()
23 | self.data = defaultdict(int)
24 |
25 | try:
26 | self.cuda_mem_get_info()
27 | torch.cuda.memory_stats(self.device)
28 | except Exception as e: # AMD or whatever
29 | print(f"Warning: caught exception '{e}', memory monitor disabled")
30 | self.disabled = True
31 |
32 | def cuda_mem_get_info(self):
33 | index = self.device.index if self.device.index is not None else torch.cuda.current_device()
34 | return torch.cuda.mem_get_info(index)
35 |
36 | def run(self):
37 | if self.disabled:
38 | return
39 |
40 | while True:
41 | self.run_flag.wait()
42 |
43 | torch.cuda.reset_peak_memory_stats()
44 | self.data.clear()
45 |
46 | if self.opts.memmon_poll_rate <= 0:
47 | self.run_flag.clear()
48 | continue
49 |
50 | self.data["min_free"] = self.cuda_mem_get_info()[0]
51 |
52 | while self.run_flag.is_set():
53 | free, total = self.cuda_mem_get_info()
54 | self.data["min_free"] = min(self.data["min_free"], free)
55 |
56 | time.sleep(1 / self.opts.memmon_poll_rate)
57 |
58 | def dump_debug(self):
59 | print(self, 'recorded data:')
60 | for k, v in self.read().items():
61 | print(k, -(v // -(1024 ** 2)))
62 |
63 | print(self, 'raw torch memory stats:')
64 | tm = torch.cuda.memory_stats(self.device)
65 | for k, v in tm.items():
66 | if 'bytes' not in k:
67 | continue
68 | print('\t' if 'peak' in k else '', k, -(v // -(1024 ** 2)))
69 |
70 | print(torch.cuda.memory_summary())
71 |
72 | def monitor(self):
73 | self.run_flag.set()
74 |
75 | def read(self):
76 | if not self.disabled:
77 | free, total = self.cuda_mem_get_info()
78 | self.data["free"] = free
79 | self.data["total"] = total
80 |
81 | torch_stats = torch.cuda.memory_stats(self.device)
82 | self.data["active"] = torch_stats["active.all.current"]
83 | self.data["active_peak"] = torch_stats["active_bytes.all.peak"]
84 | self.data["reserved"] = torch_stats["reserved_bytes.all.current"]
85 | self.data["reserved_peak"] = torch_stats["reserved_bytes.all.peak"]
86 | self.data["system_peak"] = total - self.data["min_free"]
87 |
88 | return self.data
89 |
90 | def stop(self):
91 | self.run_flag.clear()
92 | return self.read()
93 |
--------------------------------------------------------------------------------
/modules/textual_inversion/learn_schedule.py:
--------------------------------------------------------------------------------
1 | import tqdm
2 |
3 |
4 | class LearnScheduleIterator:
5 | def __init__(self, learn_rate, max_steps, cur_step=0):
6 | """
7 | specify learn_rate as "0.001:100, 0.00001:1000, 1e-5:10000" to have lr of 0.001 until step 100, 0.00001 until 1000, and 1e-5 until 10000
8 | """
9 |
10 | pairs = learn_rate.split(',')
11 | self.rates = []
12 | self.it = 0
13 | self.maxit = 0
14 | try:
15 | for pair in pairs:
16 | if not pair.strip():
17 | continue
18 | tmp = pair.split(':')
19 | if len(tmp) == 2:
20 | step = int(tmp[1])
21 | if step > cur_step:
22 | self.rates.append((float(tmp[0]), min(step, max_steps)))
23 | self.maxit += 1
24 | if step > max_steps:
25 | return
26 | elif step == -1:
27 | self.rates.append((float(tmp[0]), max_steps))
28 | self.maxit += 1
29 | return
30 | else:
31 | self.rates.append((float(tmp[0]), max_steps))
32 | self.maxit += 1
33 | return
34 | assert self.rates
35 | except (ValueError, AssertionError) as e:
36 | raise Exception('Invalid learning rate schedule. It should be a number or, for example, like "0.001:100, 0.00001:1000, 1e-5:10000" to have lr of 0.001 until step 100, 0.00001 until 1000, and 1e-5 until 10000.') from e
37 |
38 |
39 | def __iter__(self):
40 | return self
41 |
42 | def __next__(self):
43 | if self.it < self.maxit:
44 | self.it += 1
45 | return self.rates[self.it - 1]
46 | else:
47 | raise StopIteration
48 |
49 |
50 | class LearnRateScheduler:
51 | def __init__(self, learn_rate, max_steps, cur_step=0, verbose=True):
52 | self.schedules = LearnScheduleIterator(learn_rate, max_steps, cur_step)
53 | (self.learn_rate, self.end_step) = next(self.schedules)
54 | self.verbose = verbose
55 |
56 | if self.verbose:
57 | print(f'Training at rate of {self.learn_rate} until step {self.end_step}')
58 |
59 | self.finished = False
60 |
61 | def step(self, step_number):
62 | if step_number < self.end_step:
63 | return False
64 |
65 | try:
66 | (self.learn_rate, self.end_step) = next(self.schedules)
67 | except StopIteration:
68 | self.finished = True
69 | return False
70 | return True
71 |
72 | def apply(self, optimizer, step_number):
73 | if not self.step(step_number):
74 | return
75 |
76 | if self.verbose:
77 | tqdm.tqdm.write(f'Training at rate of {self.learn_rate} until step {self.end_step}')
78 |
79 | for pg in optimizer.param_groups:
80 | pg['lr'] = self.learn_rate
81 |
82 |
--------------------------------------------------------------------------------
/modules/txt2img.py:
--------------------------------------------------------------------------------
1 | import modules.scripts
2 | from modules import sd_samplers, processing
3 | from modules.generation_parameters_copypaste import create_override_settings_dict
4 | from modules.shared import opts, cmd_opts
5 | import modules.shared as shared
6 | from modules.ui import plaintext_to_html
7 |
8 |
9 |
10 | def txt2img(id_task: str, prompt: str, negative_prompt: str, prompt_styles, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, hr_scale: float, hr_upscaler: str, hr_second_pass_steps: int, hr_resize_x: int, hr_resize_y: int, hr_sampler_index: int, hr_prompt: str, hr_negative_prompt, override_settings_texts, *args):
11 | override_settings = create_override_settings_dict(override_settings_texts)
12 |
13 | p = processing.StableDiffusionProcessingTxt2Img(
14 | sd_model=shared.sd_model,
15 | outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,
16 | outpath_grids=opts.outdir_grids or opts.outdir_txt2img_grids,
17 | prompt=prompt,
18 | styles=prompt_styles,
19 | negative_prompt=negative_prompt,
20 | seed=seed,
21 | subseed=subseed,
22 | subseed_strength=subseed_strength,
23 | seed_resize_from_h=seed_resize_from_h,
24 | seed_resize_from_w=seed_resize_from_w,
25 | seed_enable_extras=seed_enable_extras,
26 | sampler_name=sd_samplers.samplers[sampler_index].name,
27 | batch_size=batch_size,
28 | n_iter=n_iter,
29 | steps=steps,
30 | cfg_scale=cfg_scale,
31 | width=width,
32 | height=height,
33 | restore_faces=restore_faces,
34 | tiling=tiling,
35 | enable_hr=enable_hr,
36 | denoising_strength=denoising_strength if enable_hr else None,
37 | hr_scale=hr_scale,
38 | hr_upscaler=hr_upscaler,
39 | hr_second_pass_steps=hr_second_pass_steps,
40 | hr_resize_x=hr_resize_x,
41 | hr_resize_y=hr_resize_y,
42 | hr_sampler_name=sd_samplers.samplers_for_img2img[hr_sampler_index - 1].name if hr_sampler_index != 0 else None,
43 | hr_prompt=hr_prompt,
44 | hr_negative_prompt=hr_negative_prompt,
45 | override_settings=override_settings,
46 | )
47 |
48 | p.scripts = modules.scripts.scripts_txt2img
49 | p.script_args = args
50 |
51 | if cmd_opts.enable_console_prompts:
52 | print(f"\ntxt2img: {prompt}", file=shared.progress_print_out)
53 |
54 | processed = modules.scripts.scripts_txt2img.run(p, *args)
55 |
56 | if processed is None:
57 | processed = processing.process_images(p)
58 |
59 | p.close()
60 |
61 | shared.total_tqdm.clear()
62 |
63 | generation_info_js = processed.js()
64 | if opts.samples_log_stdout:
65 | print(generation_info_js)
66 |
67 | if opts.do_not_show_images:
68 | processed.images = []
69 |
70 | return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments)
71 |
--------------------------------------------------------------------------------
/modules/hashes.py:
--------------------------------------------------------------------------------
1 | import hashlib
2 | import json
3 | import os.path
4 |
5 | import filelock
6 |
7 | from modules import shared
8 | from modules.paths import data_path
9 |
10 |
11 | cache_filename = os.path.join(data_path, "cache.json")
12 | cache_data = None
13 |
14 |
15 | def dump_cache():
16 | with filelock.FileLock(f"{cache_filename}.lock"):
17 | with open(cache_filename, "w", encoding="utf8") as file:
18 | json.dump(cache_data, file, indent=4)
19 |
20 |
21 | def cache(subsection):
22 | global cache_data
23 |
24 | if cache_data is None:
25 | with filelock.FileLock(f"{cache_filename}.lock"):
26 | if not os.path.isfile(cache_filename):
27 | cache_data = {}
28 | else:
29 | with open(cache_filename, "r", encoding="utf8") as file:
30 | cache_data = json.load(file)
31 |
32 | s = cache_data.get(subsection, {})
33 | cache_data[subsection] = s
34 |
35 | return s
36 |
37 |
38 | def calculate_sha256(filename):
39 | hash_sha256 = hashlib.sha256()
40 | blksize = 1024 * 1024
41 |
42 | with open(filename, "rb") as f:
43 | for chunk in iter(lambda: f.read(blksize), b""):
44 | hash_sha256.update(chunk)
45 |
46 | return hash_sha256.hexdigest()
47 |
48 |
49 | def sha256_from_cache(filename, title, use_addnet_hash=False):
50 | hashes = cache("hashes-addnet") if use_addnet_hash else cache("hashes")
51 | ondisk_mtime = os.path.getmtime(filename)
52 |
53 | if title not in hashes:
54 | return None
55 |
56 | cached_sha256 = hashes[title].get("sha256", None)
57 | cached_mtime = hashes[title].get("mtime", 0)
58 |
59 | if ondisk_mtime > cached_mtime or cached_sha256 is None:
60 | return None
61 |
62 | return cached_sha256
63 |
64 |
65 | def sha256(filename, title, use_addnet_hash=False):
66 | hashes = cache("hashes-addnet") if use_addnet_hash else cache("hashes")
67 |
68 | sha256_value = sha256_from_cache(filename, title, use_addnet_hash)
69 | if sha256_value is not None:
70 | return sha256_value
71 |
72 | if shared.cmd_opts.no_hashing:
73 | return None
74 |
75 | print(f"Calculating sha256 for {filename}: ", end='')
76 | if use_addnet_hash:
77 | with open(filename, "rb") as file:
78 | sha256_value = addnet_hash_safetensors(file)
79 | else:
80 | sha256_value = calculate_sha256(filename)
81 | print(f"{sha256_value}")
82 |
83 | hashes[title] = {
84 | "mtime": os.path.getmtime(filename),
85 | "sha256": sha256_value,
86 | }
87 |
88 | dump_cache()
89 |
90 | return sha256_value
91 |
92 |
93 | def addnet_hash_safetensors(b):
94 | """kohya-ss hash for safetensors from https://github.com/kohya-ss/sd-scripts/blob/main/library/train_util.py"""
95 | hash_sha256 = hashlib.sha256()
96 | blksize = 1024 * 1024
97 |
98 | b.seek(0)
99 | header = b.read(8)
100 | n = int.from_bytes(header, "little")
101 |
102 | offset = n + 8
103 | b.seek(offset)
104 | for chunk in iter(lambda: b.read(blksize), b""):
105 | hash_sha256.update(chunk)
106 |
107 | return hash_sha256.hexdigest()
108 |
109 |
--------------------------------------------------------------------------------
/modules/deepbooru.py:
--------------------------------------------------------------------------------
1 | import os
2 | import re
3 |
4 | import torch
5 | import numpy as np
6 |
7 | from modules import modelloader, paths, deepbooru_model, devices, images, shared
8 |
9 | re_special = re.compile(r'([\\()])')
10 |
11 |
12 | class DeepDanbooru:
13 | def __init__(self):
14 | self.model = None
15 |
16 | def load(self):
17 | if self.model is not None:
18 | return
19 |
20 | files = modelloader.load_models(
21 | model_path=os.path.join(paths.models_path, "torch_deepdanbooru"),
22 | model_url='https://github.com/AUTOMATIC1111/TorchDeepDanbooru/releases/download/v1/model-resnet_custom_v3.pt',
23 | ext_filter=[".pt"],
24 | download_name='model-resnet_custom_v3.pt',
25 | )
26 |
27 | self.model = deepbooru_model.DeepDanbooruModel()
28 | self.model.load_state_dict(torch.load(files[0], map_location="cpu"))
29 |
30 | self.model.eval()
31 | self.model.to(devices.cpu, devices.dtype)
32 |
33 | def start(self):
34 | self.load()
35 | self.model.to(devices.device)
36 |
37 | def stop(self):
38 | if not shared.opts.interrogate_keep_models_in_memory:
39 | self.model.to(devices.cpu)
40 | devices.torch_gc()
41 |
42 | def tag(self, pil_image):
43 | self.start()
44 | res = self.tag_multi(pil_image)
45 | self.stop()
46 |
47 | return res
48 |
49 | def tag_multi(self, pil_image, force_disable_ranks=False):
50 | threshold = shared.opts.interrogate_deepbooru_score_threshold
51 | use_spaces = shared.opts.deepbooru_use_spaces
52 | use_escape = shared.opts.deepbooru_escape
53 | alpha_sort = shared.opts.deepbooru_sort_alpha
54 | include_ranks = shared.opts.interrogate_return_ranks and not force_disable_ranks
55 |
56 | pic = images.resize_image(2, pil_image.convert("RGB"), 512, 512)
57 | a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
58 |
59 | with torch.no_grad(), devices.autocast():
60 | x = torch.from_numpy(a).to(devices.device)
61 | y = self.model(x)[0].detach().cpu().numpy()
62 |
63 | probability_dict = {}
64 |
65 | for tag, probability in zip(self.model.tags, y):
66 | if probability < threshold:
67 | continue
68 |
69 | if tag.startswith("rating:"):
70 | continue
71 |
72 | probability_dict[tag] = probability
73 |
74 | if alpha_sort:
75 | tags = sorted(probability_dict)
76 | else:
77 | tags = [tag for tag, _ in sorted(probability_dict.items(), key=lambda x: -x[1])]
78 |
79 | res = []
80 |
81 | filtertags = {x.strip().replace(' ', '_') for x in shared.opts.deepbooru_filter_tags.split(",")}
82 |
83 | for tag in [x for x in tags if x not in filtertags]:
84 | probability = probability_dict[tag]
85 | tag_outformat = tag
86 | if use_spaces:
87 | tag_outformat = tag_outformat.replace('_', ' ')
88 | if use_escape:
89 | tag_outformat = re.sub(re_special, r'\\\1', tag_outformat)
90 | if include_ranks:
91 | tag_outformat = f"({tag_outformat}:{probability:.3f})"
92 |
93 | res.append(tag_outformat)
94 |
95 | return ", ".join(res)
96 |
97 |
98 | model = DeepDanbooru()
99 |
--------------------------------------------------------------------------------
/.eslintrc.js:
--------------------------------------------------------------------------------
1 | /* global module */
2 | module.exports = {
3 | env: {
4 | browser: true,
5 | es2021: true,
6 | },
7 | extends: "eslint:recommended",
8 | parserOptions: {
9 | ecmaVersion: "latest",
10 | },
11 | rules: {
12 | "arrow-spacing": "error",
13 | "block-spacing": "error",
14 | "brace-style": "error",
15 | "comma-dangle": ["error", "only-multiline"],
16 | "comma-spacing": "error",
17 | "comma-style": ["error", "last"],
18 | "curly": ["error", "multi-line", "consistent"],
19 | "eol-last": "error",
20 | "func-call-spacing": "error",
21 | "function-call-argument-newline": ["error", "consistent"],
22 | "function-paren-newline": ["error", "consistent"],
23 | "indent": ["error", 4],
24 | "key-spacing": "error",
25 | "keyword-spacing": "error",
26 | "linebreak-style": ["error", "unix"],
27 | "no-extra-semi": "error",
28 | "no-mixed-spaces-and-tabs": "error",
29 | "no-multi-spaces": "error",
30 | "no-redeclare": ["error", {builtinGlobals: false}],
31 | "no-trailing-spaces": "error",
32 | "no-unused-vars": "off",
33 | "no-whitespace-before-property": "error",
34 | "object-curly-newline": ["error", {consistent: true, multiline: true}],
35 | "object-curly-spacing": ["error", "never"],
36 | "operator-linebreak": ["error", "after"],
37 | "quote-props": ["error", "consistent-as-needed"],
38 | "semi": ["error", "always"],
39 | "semi-spacing": "error",
40 | "semi-style": ["error", "last"],
41 | "space-before-blocks": "error",
42 | "space-before-function-paren": ["error", "never"],
43 | "space-in-parens": ["error", "never"],
44 | "space-infix-ops": "error",
45 | "space-unary-ops": "error",
46 | "switch-colon-spacing": "error",
47 | "template-curly-spacing": ["error", "never"],
48 | "unicode-bom": "error",
49 | },
50 | globals: {
51 | //script.js
52 | gradioApp: "readonly",
53 | onUiLoaded: "readonly",
54 | onUiUpdate: "readonly",
55 | onOptionsChanged: "readonly",
56 | uiCurrentTab: "writable",
57 | uiElementIsVisible: "readonly",
58 | uiElementInSight: "readonly",
59 | executeCallbacks: "readonly",
60 | //ui.js
61 | opts: "writable",
62 | all_gallery_buttons: "readonly",
63 | selected_gallery_button: "readonly",
64 | selected_gallery_index: "readonly",
65 | switch_to_txt2img: "readonly",
66 | switch_to_img2img_tab: "readonly",
67 | switch_to_img2img: "readonly",
68 | switch_to_sketch: "readonly",
69 | switch_to_inpaint: "readonly",
70 | switch_to_inpaint_sketch: "readonly",
71 | switch_to_extras: "readonly",
72 | get_tab_index: "readonly",
73 | create_submit_args: "readonly",
74 | restart_reload: "readonly",
75 | updateInput: "readonly",
76 | //extraNetworks.js
77 | requestGet: "readonly",
78 | popup: "readonly",
79 | // from python
80 | localization: "readonly",
81 | // progrssbar.js
82 | randomId: "readonly",
83 | requestProgress: "readonly",
84 | // imageviewer.js
85 | modalPrevImage: "readonly",
86 | modalNextImage: "readonly",
87 | }
88 | };
89 |
--------------------------------------------------------------------------------
/test/test_txt2img.py:
--------------------------------------------------------------------------------
1 |
2 | import pytest
3 | import requests
4 |
5 |
6 | @pytest.fixture()
7 | def url_txt2img(base_url):
8 | return f"{base_url}/sdapi/v1/txt2img"
9 |
10 |
11 | @pytest.fixture()
12 | def simple_txt2img_request():
13 | return {
14 | "batch_size": 1,
15 | "cfg_scale": 7,
16 | "denoising_strength": 0,
17 | "enable_hr": False,
18 | "eta": 0,
19 | "firstphase_height": 0,
20 | "firstphase_width": 0,
21 | "height": 64,
22 | "n_iter": 1,
23 | "negative_prompt": "",
24 | "prompt": "example prompt",
25 | "restore_faces": False,
26 | "s_churn": 0,
27 | "s_noise": 1,
28 | "s_tmax": 0,
29 | "s_tmin": 0,
30 | "sampler_index": "Euler a",
31 | "seed": -1,
32 | "seed_resize_from_h": -1,
33 | "seed_resize_from_w": -1,
34 | "steps": 3,
35 | "styles": [],
36 | "subseed": -1,
37 | "subseed_strength": 0,
38 | "tiling": False,
39 | "width": 64,
40 | }
41 |
42 |
43 | def test_txt2img_simple_performed(url_txt2img, simple_txt2img_request):
44 | assert requests.post(url_txt2img, json=simple_txt2img_request).status_code == 200
45 |
46 |
47 | def test_txt2img_with_negative_prompt_performed(url_txt2img, simple_txt2img_request):
48 | simple_txt2img_request["negative_prompt"] = "example negative prompt"
49 | assert requests.post(url_txt2img, json=simple_txt2img_request).status_code == 200
50 |
51 |
52 | def test_txt2img_with_complex_prompt_performed(url_txt2img, simple_txt2img_request):
53 | simple_txt2img_request["prompt"] = "((emphasis)), (emphasis1:1.1), [to:1], [from::2], [from:to:0.3], [alt|alt1]"
54 | assert requests.post(url_txt2img, json=simple_txt2img_request).status_code == 200
55 |
56 |
57 | def test_txt2img_not_square_image_performed(url_txt2img, simple_txt2img_request):
58 | simple_txt2img_request["height"] = 128
59 | assert requests.post(url_txt2img, json=simple_txt2img_request).status_code == 200
60 |
61 |
62 | def test_txt2img_with_hrfix_performed(url_txt2img, simple_txt2img_request):
63 | simple_txt2img_request["enable_hr"] = True
64 | assert requests.post(url_txt2img, json=simple_txt2img_request).status_code == 200
65 |
66 |
67 | def test_txt2img_with_tiling_performed(url_txt2img, simple_txt2img_request):
68 | simple_txt2img_request["tiling"] = True
69 | assert requests.post(url_txt2img, json=simple_txt2img_request).status_code == 200
70 |
71 |
72 | def test_txt2img_with_restore_faces_performed(url_txt2img, simple_txt2img_request):
73 | simple_txt2img_request["restore_faces"] = True
74 | assert requests.post(url_txt2img, json=simple_txt2img_request).status_code == 200
75 |
76 |
77 | @pytest.mark.parametrize("sampler", ["PLMS", "DDIM", "UniPC"])
78 | def test_txt2img_with_vanilla_sampler_performed(url_txt2img, simple_txt2img_request, sampler):
79 | simple_txt2img_request["sampler_index"] = sampler
80 | assert requests.post(url_txt2img, json=simple_txt2img_request).status_code == 200
81 |
82 |
83 | def test_txt2img_multiple_batches_performed(url_txt2img, simple_txt2img_request):
84 | simple_txt2img_request["n_iter"] = 2
85 | assert requests.post(url_txt2img, json=simple_txt2img_request).status_code == 200
86 |
87 |
88 | def test_txt2img_batch_performed(url_txt2img, simple_txt2img_request):
89 | simple_txt2img_request["batch_size"] = 2
90 | assert requests.post(url_txt2img, json=simple_txt2img_request).status_code == 200
91 |
--------------------------------------------------------------------------------
/script.js:
--------------------------------------------------------------------------------
1 | function gradioApp() {
2 | const elems = document.getElementsByTagName('gradio-app');
3 | const elem = elems.length == 0 ? document : elems[0];
4 |
5 | if (elem !== document) {
6 | elem.getElementById = function(id) {
7 | return document.getElementById(id);
8 | };
9 | }
10 | return elem.shadowRoot ? elem.shadowRoot : elem;
11 | }
12 |
13 | function get_uiCurrentTab() {
14 | return gradioApp().querySelector('#tabs button.selected');
15 | }
16 |
17 | function get_uiCurrentTabContent() {
18 | return gradioApp().querySelector('.tabitem[id^=tab_]:not([style*="display: none"])');
19 | }
20 |
21 | var uiUpdateCallbacks = [];
22 | var uiLoadedCallbacks = [];
23 | var uiTabChangeCallbacks = [];
24 | var optionsChangedCallbacks = [];
25 | var uiCurrentTab = null;
26 |
27 | function onUiUpdate(callback) {
28 | uiUpdateCallbacks.push(callback);
29 | }
30 | function onUiLoaded(callback) {
31 | uiLoadedCallbacks.push(callback);
32 | }
33 | function onUiTabChange(callback) {
34 | uiTabChangeCallbacks.push(callback);
35 | }
36 | function onOptionsChanged(callback) {
37 | optionsChangedCallbacks.push(callback);
38 | }
39 |
40 | function runCallback(x, m) {
41 | try {
42 | x(m);
43 | } catch (e) {
44 | (console.error || console.log).call(console, e.message, e);
45 | }
46 | }
47 | function executeCallbacks(queue, m) {
48 | queue.forEach(function(x) {
49 | runCallback(x, m);
50 | });
51 | }
52 |
53 | var executedOnLoaded = false;
54 |
55 | document.addEventListener("DOMContentLoaded", function() {
56 | var mutationObserver = new MutationObserver(function(m) {
57 | if (!executedOnLoaded && gradioApp().querySelector('#txt2img_prompt')) {
58 | executedOnLoaded = true;
59 | executeCallbacks(uiLoadedCallbacks);
60 | }
61 |
62 | executeCallbacks(uiUpdateCallbacks, m);
63 | const newTab = get_uiCurrentTab();
64 | if (newTab && (newTab !== uiCurrentTab)) {
65 | uiCurrentTab = newTab;
66 | executeCallbacks(uiTabChangeCallbacks);
67 | }
68 | });
69 | mutationObserver.observe(gradioApp(), {childList: true, subtree: true});
70 | });
71 |
72 | /**
73 | * Add a ctrl+enter as a shortcut to start a generation
74 | */
75 | document.addEventListener('keydown', function(e) {
76 | var handled = false;
77 | if (e.key !== undefined) {
78 | if ((e.key == "Enter" && (e.metaKey || e.ctrlKey || e.altKey))) handled = true;
79 | } else if (e.keyCode !== undefined) {
80 | if ((e.keyCode == 13 && (e.metaKey || e.ctrlKey || e.altKey))) handled = true;
81 | }
82 | if (handled) {
83 | var button = get_uiCurrentTabContent().querySelector('button[id$=_generate]');
84 | if (button) {
85 | button.click();
86 | }
87 | e.preventDefault();
88 | }
89 | });
90 |
91 | /**
92 | * checks that a UI element is not in another hidden element or tab content
93 | */
94 | function uiElementIsVisible(el) {
95 | if (el === document) {
96 | return true;
97 | }
98 |
99 | const computedStyle = getComputedStyle(el);
100 | const isVisible = computedStyle.display !== 'none';
101 |
102 | if (!isVisible) return false;
103 | return uiElementIsVisible(el.parentNode);
104 | }
105 |
106 | function uiElementInSight(el) {
107 | const clRect = el.getBoundingClientRect();
108 | const windowHeight = window.innerHeight;
109 | const isOnScreen = clRect.bottom > 0 && clRect.top < windowHeight;
110 |
111 | return isOnScreen;
112 | }
113 |
--------------------------------------------------------------------------------
/extensions-builtin/LDSR/scripts/ldsr_model.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | import traceback
4 |
5 | from basicsr.utils.download_util import load_file_from_url
6 |
7 | from modules.upscaler import Upscaler, UpscalerData
8 | from ldsr_model_arch import LDSR
9 | from modules import shared, script_callbacks
10 | import sd_hijack_autoencoder # noqa: F401
11 | import sd_hijack_ddpm_v1 # noqa: F401
12 |
13 |
14 | class UpscalerLDSR(Upscaler):
15 | def __init__(self, user_path):
16 | self.name = "LDSR"
17 | self.user_path = user_path
18 | self.model_url = "https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1"
19 | self.yaml_url = "https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1"
20 | super().__init__()
21 | scaler_data = UpscalerData("LDSR", None, self)
22 | self.scalers = [scaler_data]
23 |
24 | def load_model(self, path: str):
25 | # Remove incorrect project.yaml file if too big
26 | yaml_path = os.path.join(self.model_path, "project.yaml")
27 | old_model_path = os.path.join(self.model_path, "model.pth")
28 | new_model_path = os.path.join(self.model_path, "model.ckpt")
29 |
30 | local_model_paths = self.find_models(ext_filter=[".ckpt", ".safetensors"])
31 | local_ckpt_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("model.ckpt")]), None)
32 | local_safetensors_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("model.safetensors")]), None)
33 | local_yaml_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("project.yaml")]), None)
34 |
35 | if os.path.exists(yaml_path):
36 | statinfo = os.stat(yaml_path)
37 | if statinfo.st_size >= 10485760:
38 | print("Removing invalid LDSR YAML file.")
39 | os.remove(yaml_path)
40 |
41 | if os.path.exists(old_model_path):
42 | print("Renaming model from model.pth to model.ckpt")
43 | os.rename(old_model_path, new_model_path)
44 |
45 | if local_safetensors_path is not None and os.path.exists(local_safetensors_path):
46 | model = local_safetensors_path
47 | else:
48 | model = local_ckpt_path if local_ckpt_path is not None else load_file_from_url(url=self.model_url, model_dir=self.model_download_path, file_name="model.ckpt", progress=True)
49 |
50 | yaml = local_yaml_path if local_yaml_path is not None else load_file_from_url(url=self.yaml_url, model_dir=self.model_download_path, file_name="project.yaml", progress=True)
51 |
52 | try:
53 | return LDSR(model, yaml)
54 |
55 | except Exception:
56 | print("Error importing LDSR:", file=sys.stderr)
57 | print(traceback.format_exc(), file=sys.stderr)
58 | return None
59 |
60 | def do_upscale(self, img, path):
61 | ldsr = self.load_model(path)
62 | if ldsr is None:
63 | print("NO LDSR!")
64 | return img
65 | ddim_steps = shared.opts.ldsr_steps
66 | return ldsr.super_resolution(img, ddim_steps, self.scale)
67 |
68 |
69 | def on_ui_settings():
70 | import gradio as gr
71 |
72 | shared.opts.add_option("ldsr_steps", shared.OptionInfo(100, "LDSR processing steps. Lower = faster", gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}, section=('upscaling', "Upscaling")))
73 | shared.opts.add_option("ldsr_cached", shared.OptionInfo(False, "Cache LDSR model in memory", gr.Checkbox, {"interactive": True}, section=('upscaling', "Upscaling")))
74 |
75 |
76 | script_callbacks.on_ui_settings(on_ui_settings)
77 |
--------------------------------------------------------------------------------
/javascript/dragdrop.js:
--------------------------------------------------------------------------------
1 | // allows drag-dropping files into gradio image elements, and also pasting images from clipboard
2 |
3 | function isValidImageList(files) {
4 | return files && files?.length === 1 && ['image/png', 'image/gif', 'image/jpeg'].includes(files[0].type);
5 | }
6 |
7 | function dropReplaceImage(imgWrap, files) {
8 | if (!isValidImageList(files)) {
9 | return;
10 | }
11 |
12 | const tmpFile = files[0];
13 |
14 | imgWrap.querySelector('.modify-upload button + button, .touch-none + div button + button')?.click();
15 | const callback = () => {
16 | const fileInput = imgWrap.querySelector('input[type="file"]');
17 | if (fileInput) {
18 | if (files.length === 0) {
19 | files = new DataTransfer();
20 | files.items.add(tmpFile);
21 | fileInput.files = files.files;
22 | } else {
23 | fileInput.files = files;
24 | }
25 | fileInput.dispatchEvent(new Event('change'));
26 | }
27 | };
28 |
29 | if (imgWrap.closest('#pnginfo_image')) {
30 | // special treatment for PNG Info tab, wait for fetch request to finish
31 | const oldFetch = window.fetch;
32 | window.fetch = async(input, options) => {
33 | const response = await oldFetch(input, options);
34 | if ('api/predict/' === input) {
35 | const content = await response.text();
36 | window.fetch = oldFetch;
37 | window.requestAnimationFrame(() => callback());
38 | return new Response(content, {
39 | status: response.status,
40 | statusText: response.statusText,
41 | headers: response.headers
42 | });
43 | }
44 | return response;
45 | };
46 | } else {
47 | window.requestAnimationFrame(() => callback());
48 | }
49 | }
50 |
51 | window.document.addEventListener('dragover', e => {
52 | const target = e.composedPath()[0];
53 | const imgWrap = target.closest('[data-testid="image"]');
54 | if (!imgWrap && target.placeholder && target.placeholder.indexOf("Prompt") == -1) {
55 | return;
56 | }
57 | e.stopPropagation();
58 | e.preventDefault();
59 | e.dataTransfer.dropEffect = 'copy';
60 | });
61 |
62 | window.document.addEventListener('drop', e => {
63 | const target = e.composedPath()[0];
64 | if (target.placeholder.indexOf("Prompt") == -1) {
65 | return;
66 | }
67 | const imgWrap = target.closest('[data-testid="image"]');
68 | if (!imgWrap) {
69 | return;
70 | }
71 | e.stopPropagation();
72 | e.preventDefault();
73 | const files = e.dataTransfer.files;
74 | dropReplaceImage(imgWrap, files);
75 | });
76 |
77 | window.addEventListener('paste', e => {
78 | const files = e.clipboardData.files;
79 | if (!isValidImageList(files)) {
80 | return;
81 | }
82 |
83 | const visibleImageFields = [...gradioApp().querySelectorAll('[data-testid="image"]')]
84 | .filter(el => uiElementIsVisible(el))
85 | .sort((a, b) => uiElementInSight(b) - uiElementInSight(a));
86 |
87 |
88 | if (!visibleImageFields.length) {
89 | return;
90 | }
91 |
92 | const firstFreeImageField = visibleImageFields
93 | .filter(el => el.querySelector('input[type=file]'))?.[0];
94 |
95 | dropReplaceImage(
96 | firstFreeImageField ?
97 | firstFreeImageField :
98 | visibleImageFields[visibleImageFields.length - 1]
99 | , files
100 | );
101 | });
102 |
--------------------------------------------------------------------------------
/modules/sd_samplers_common.py:
--------------------------------------------------------------------------------
1 | from collections import namedtuple
2 | import numpy as np
3 | import torch
4 | from PIL import Image
5 | from modules import devices, processing, images, sd_vae_approx, sd_samplers, sd_vae_taesd
6 |
7 | from modules.shared import opts, state
8 | import modules.shared as shared
9 |
10 | SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])
11 |
12 |
13 | def setup_img2img_steps(p, steps=None):
14 | if opts.img2img_fix_steps or steps is not None:
15 | requested_steps = (steps or p.steps)
16 | steps = int(requested_steps / min(p.denoising_strength, 0.999)) if p.denoising_strength > 0 else 0
17 | t_enc = requested_steps - 1
18 | else:
19 | steps = p.steps
20 | t_enc = int(min(p.denoising_strength, 0.999) * steps)
21 |
22 | return steps, t_enc
23 |
24 |
25 | approximation_indexes = {"Full": 0, "Approx NN": 1, "Approx cheap": 2, "TAESD": 3}
26 |
27 |
28 | def single_sample_to_image(sample, approximation=None):
29 | if approximation is None:
30 | approximation = approximation_indexes.get(opts.show_progress_type, 0)
31 |
32 | if approximation == 2:
33 | x_sample = sd_vae_approx.cheap_approximation(sample) * 0.5 + 0.5
34 | elif approximation == 1:
35 | x_sample = sd_vae_approx.model()(sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach() * 0.5 + 0.5
36 | elif approximation == 3:
37 | x_sample = sample * 1.5
38 | x_sample = sd_vae_taesd.model()(x_sample.to(devices.device, devices.dtype).unsqueeze(0))[0].detach()
39 | else:
40 | x_sample = processing.decode_first_stage(shared.sd_model, sample.unsqueeze(0))[0] * 0.5 + 0.5
41 |
42 | x_sample = torch.clamp(x_sample, min=0.0, max=1.0)
43 | x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
44 | x_sample = x_sample.astype(np.uint8)
45 |
46 | return Image.fromarray(x_sample)
47 |
48 |
49 | def sample_to_image(samples, index=0, approximation=None):
50 | return single_sample_to_image(samples[index], approximation)
51 |
52 |
53 | def samples_to_image_grid(samples, approximation=None):
54 | return images.image_grid([single_sample_to_image(sample, approximation) for sample in samples])
55 |
56 |
57 | def store_latent(decoded):
58 | state.current_latent = decoded
59 |
60 | if opts.live_previews_enable and opts.show_progress_every_n_steps > 0 and shared.state.sampling_step % opts.show_progress_every_n_steps == 0:
61 | if not shared.parallel_processing_allowed:
62 | shared.state.assign_current_image(sample_to_image(decoded))
63 |
64 |
65 | def is_sampler_using_eta_noise_seed_delta(p):
66 | """returns whether sampler from config will use eta noise seed delta for image creation"""
67 |
68 | sampler_config = sd_samplers.find_sampler_config(p.sampler_name)
69 |
70 | eta = p.eta
71 |
72 | if eta is None and p.sampler is not None:
73 | eta = p.sampler.eta
74 |
75 | if eta is None and sampler_config is not None:
76 | eta = 0 if sampler_config.options.get("default_eta_is_0", False) else 1.0
77 |
78 | if eta == 0:
79 | return False
80 |
81 | return sampler_config.options.get("uses_ensd", False)
82 |
83 |
84 | class InterruptedException(BaseException):
85 | pass
86 |
87 |
88 | if opts.randn_source == "CPU":
89 | import torchsde._brownian.brownian_interval
90 |
91 | def torchsde_randn(size, dtype, device, seed):
92 | generator = torch.Generator(devices.cpu).manual_seed(int(seed))
93 | return torch.randn(size, dtype=dtype, device=devices.cpu, generator=generator).to(device)
94 |
95 | torchsde._brownian.brownian_interval._randn = torchsde_randn
96 |
--------------------------------------------------------------------------------
/modules/masking.py:
--------------------------------------------------------------------------------
1 | from PIL import Image, ImageFilter, ImageOps
2 |
3 |
4 | def get_crop_region(mask, pad=0):
5 | """finds a rectangular region that contains all masked ares in an image. Returns (x1, y1, x2, y2) coordinates of the rectangle.
6 | For example, if a user has painted the top-right part of a 512x512 image", the result may be (256, 0, 512, 256)"""
7 |
8 | h, w = mask.shape
9 |
10 | crop_left = 0
11 | for i in range(w):
12 | if not (mask[:, i] == 0).all():
13 | break
14 | crop_left += 1
15 |
16 | crop_right = 0
17 | for i in reversed(range(w)):
18 | if not (mask[:, i] == 0).all():
19 | break
20 | crop_right += 1
21 |
22 | crop_top = 0
23 | for i in range(h):
24 | if not (mask[i] == 0).all():
25 | break
26 | crop_top += 1
27 |
28 | crop_bottom = 0
29 | for i in reversed(range(h)):
30 | if not (mask[i] == 0).all():
31 | break
32 | crop_bottom += 1
33 |
34 | return (
35 | int(max(crop_left-pad, 0)),
36 | int(max(crop_top-pad, 0)),
37 | int(min(w - crop_right + pad, w)),
38 | int(min(h - crop_bottom + pad, h))
39 | )
40 |
41 |
42 | def expand_crop_region(crop_region, processing_width, processing_height, image_width, image_height):
43 | """expands crop region get_crop_region() to match the ratio of the image the region will processed in; returns expanded region
44 | for example, if user drew mask in a 128x32 region, and the dimensions for processing are 512x512, the region will be expanded to 128x128."""
45 |
46 | x1, y1, x2, y2 = crop_region
47 |
48 | ratio_crop_region = (x2 - x1) / (y2 - y1)
49 | ratio_processing = processing_width / processing_height
50 |
51 | if ratio_crop_region > ratio_processing:
52 | desired_height = (x2 - x1) / ratio_processing
53 | desired_height_diff = int(desired_height - (y2-y1))
54 | y1 -= desired_height_diff//2
55 | y2 += desired_height_diff - desired_height_diff//2
56 | if y2 >= image_height:
57 | diff = y2 - image_height
58 | y2 -= diff
59 | y1 -= diff
60 | if y1 < 0:
61 | y2 -= y1
62 | y1 -= y1
63 | if y2 >= image_height:
64 | y2 = image_height
65 | else:
66 | desired_width = (y2 - y1) * ratio_processing
67 | desired_width_diff = int(desired_width - (x2-x1))
68 | x1 -= desired_width_diff//2
69 | x2 += desired_width_diff - desired_width_diff//2
70 | if x2 >= image_width:
71 | diff = x2 - image_width
72 | x2 -= diff
73 | x1 -= diff
74 | if x1 < 0:
75 | x2 -= x1
76 | x1 -= x1
77 | if x2 >= image_width:
78 | x2 = image_width
79 |
80 | return x1, y1, x2, y2
81 |
82 |
83 | def fill(image, mask):
84 | """fills masked regions with colors from image using blur. Not extremely effective."""
85 |
86 | image_mod = Image.new('RGBA', (image.width, image.height))
87 |
88 | image_masked = Image.new('RGBa', (image.width, image.height))
89 | image_masked.paste(image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(mask.convert('L')))
90 |
91 | image_masked = image_masked.convert('RGBa')
92 |
93 | for radius, repeats in [(256, 1), (64, 1), (16, 2), (4, 4), (2, 2), (0, 1)]:
94 | blurred = image_masked.filter(ImageFilter.GaussianBlur(radius)).convert('RGBA')
95 | for _ in range(repeats):
96 | image_mod.alpha_composite(blurred)
97 |
98 | return image_mod.convert("RGB")
99 |
100 |
--------------------------------------------------------------------------------
/modules/sd_hijack_clip_old.py:
--------------------------------------------------------------------------------
1 | from modules import sd_hijack_clip
2 | from modules import shared
3 |
4 |
5 | def process_text_old(self: sd_hijack_clip.FrozenCLIPEmbedderWithCustomWordsBase, texts):
6 | id_start = self.id_start
7 | id_end = self.id_end
8 | maxlen = self.wrapped.max_length # you get to stay at 77
9 | used_custom_terms = []
10 | remade_batch_tokens = []
11 | hijack_comments = []
12 | hijack_fixes = []
13 | token_count = 0
14 |
15 | cache = {}
16 | batch_tokens = self.tokenize(texts)
17 | batch_multipliers = []
18 | for tokens in batch_tokens:
19 | tuple_tokens = tuple(tokens)
20 |
21 | if tuple_tokens in cache:
22 | remade_tokens, fixes, multipliers = cache[tuple_tokens]
23 | else:
24 | fixes = []
25 | remade_tokens = []
26 | multipliers = []
27 | mult = 1.0
28 |
29 | i = 0
30 | while i < len(tokens):
31 | token = tokens[i]
32 |
33 | embedding, embedding_length_in_tokens = self.hijack.embedding_db.find_embedding_at_position(tokens, i)
34 |
35 | mult_change = self.token_mults.get(token) if shared.opts.enable_emphasis else None
36 | if mult_change is not None:
37 | mult *= mult_change
38 | i += 1
39 | elif embedding is None:
40 | remade_tokens.append(token)
41 | multipliers.append(mult)
42 | i += 1
43 | else:
44 | emb_len = int(embedding.vec.shape[0])
45 | fixes.append((len(remade_tokens), embedding))
46 | remade_tokens += [0] * emb_len
47 | multipliers += [mult] * emb_len
48 | used_custom_terms.append((embedding.name, embedding.checksum()))
49 | i += embedding_length_in_tokens
50 |
51 | if len(remade_tokens) > maxlen - 2:
52 | vocab = {v: k for k, v in self.wrapped.tokenizer.get_vocab().items()}
53 | ovf = remade_tokens[maxlen - 2:]
54 | overflowing_words = [vocab.get(int(x), "") for x in ovf]
55 | overflowing_text = self.wrapped.tokenizer.convert_tokens_to_string(''.join(overflowing_words))
56 | hijack_comments.append(f"Warning: too many input tokens; some ({len(overflowing_words)}) have been truncated:\n{overflowing_text}\n")
57 |
58 | token_count = len(remade_tokens)
59 | remade_tokens = remade_tokens + [id_end] * (maxlen - 2 - len(remade_tokens))
60 | remade_tokens = [id_start] + remade_tokens[0:maxlen - 2] + [id_end]
61 | cache[tuple_tokens] = (remade_tokens, fixes, multipliers)
62 |
63 | multipliers = multipliers + [1.0] * (maxlen - 2 - len(multipliers))
64 | multipliers = [1.0] + multipliers[0:maxlen - 2] + [1.0]
65 |
66 | remade_batch_tokens.append(remade_tokens)
67 | hijack_fixes.append(fixes)
68 | batch_multipliers.append(multipliers)
69 | return batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count
70 |
71 |
72 | def forward_old(self: sd_hijack_clip.FrozenCLIPEmbedderWithCustomWordsBase, texts):
73 | batch_multipliers, remade_batch_tokens, used_custom_terms, hijack_comments, hijack_fixes, token_count = process_text_old(self, texts)
74 |
75 | self.hijack.comments += hijack_comments
76 |
77 | if len(used_custom_terms) > 0:
78 | embedding_names = ", ".join(f"{word} [{checksum}]" for word, checksum in used_custom_terms)
79 | self.hijack.comments.append(f"Used embeddings: {embedding_names}")
80 |
81 | self.hijack.fixes = hijack_fixes
82 | return self.process_tokens(remade_batch_tokens, batch_multipliers)
83 |
--------------------------------------------------------------------------------
/modules/call_queue.py:
--------------------------------------------------------------------------------
1 | import html
2 | import sys
3 | import threading
4 | import traceback
5 | import time
6 |
7 | from modules import shared, progress
8 |
9 | queue_lock = threading.Lock()
10 |
11 |
12 | def wrap_queued_call(func):
13 | def f(*args, **kwargs):
14 | with queue_lock:
15 | res = func(*args, **kwargs)
16 |
17 | return res
18 |
19 | return f
20 |
21 |
22 | def wrap_gradio_gpu_call(func, extra_outputs=None):
23 | def f(*args, **kwargs):
24 |
25 | # if the first argument is a string that says "task(...)", it is treated as a job id
26 | if len(args) > 0 and type(args[0]) == str and args[0][0:5] == "task(" and args[0][-1] == ")":
27 | id_task = args[0]
28 | progress.add_task_to_queue(id_task)
29 | else:
30 | id_task = None
31 |
32 | with queue_lock:
33 | shared.state.begin()
34 | progress.start_task(id_task)
35 |
36 | try:
37 | res = func(*args, **kwargs)
38 | progress.record_results(id_task, res)
39 | finally:
40 | progress.finish_task(id_task)
41 |
42 | shared.state.end()
43 |
44 | return res
45 |
46 | return wrap_gradio_call(f, extra_outputs=extra_outputs, add_stats=True)
47 |
48 |
49 | def wrap_gradio_call(func, extra_outputs=None, add_stats=False):
50 | def f(*args, extra_outputs_array=extra_outputs, **kwargs):
51 | run_memmon = shared.opts.memmon_poll_rate > 0 and not shared.mem_mon.disabled and add_stats
52 | if run_memmon:
53 | shared.mem_mon.monitor()
54 | t = time.perf_counter()
55 |
56 | try:
57 | res = list(func(*args, **kwargs))
58 | except Exception as e:
59 | # When printing out our debug argument list, do not print out more than a MB of text
60 | max_debug_str_len = 131072 # (1024*1024)/8
61 |
62 | print("Error completing request", file=sys.stderr)
63 | argStr = f"Arguments: {args} {kwargs}"
64 | print(argStr[:max_debug_str_len], file=sys.stderr)
65 | if len(argStr) > max_debug_str_len:
66 | print(f"(Argument list truncated at {max_debug_str_len}/{len(argStr)} characters)", file=sys.stderr)
67 |
68 | print(traceback.format_exc(), file=sys.stderr)
69 |
70 | shared.state.job = ""
71 | shared.state.job_count = 0
72 |
73 | if extra_outputs_array is None:
74 | extra_outputs_array = [None, '']
75 |
76 | error_message = f'{type(e).__name__}: {e}'
77 | res = extra_outputs_array + [f"{html.escape(error_message)}
"]
78 |
79 | shared.state.skipped = False
80 | shared.state.interrupted = False
81 | shared.state.job_count = 0
82 |
83 | if not add_stats:
84 | return tuple(res)
85 |
86 | elapsed = time.perf_counter() - t
87 | elapsed_m = int(elapsed // 60)
88 | elapsed_s = elapsed % 60
89 | elapsed_text = f"{elapsed_s:.2f}s"
90 | if elapsed_m > 0:
91 | elapsed_text = f"{elapsed_m}m "+elapsed_text
92 |
93 | if run_memmon:
94 | mem_stats = {k: -(v//-(1024*1024)) for k, v in shared.mem_mon.stop().items()}
95 | active_peak = mem_stats['active_peak']
96 | reserved_peak = mem_stats['reserved_peak']
97 | sys_peak = mem_stats['system_peak']
98 | sys_total = mem_stats['total']
99 | sys_pct = round(sys_peak/max(sys_total, 1) * 100, 2)
100 |
101 | vram_html = f"Torch active/reserved: {active_peak}/{reserved_peak} MiB, Sys VRAM: {sys_peak}/{sys_total} MiB ({sys_pct}%)
"
102 | else:
103 | vram_html = ''
104 |
105 | # last item is always HTML
106 | res[-1] += f""
107 |
108 | return tuple(res)
109 |
110 | return f
111 |
112 |
--------------------------------------------------------------------------------
/modules/sd_hijack_unet.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from packaging import version
3 |
4 | from modules import devices
5 | from modules.sd_hijack_utils import CondFunc
6 |
7 |
8 | class TorchHijackForUnet:
9 | """
10 | This is torch, but with cat that resizes tensors to appropriate dimensions if they do not match;
11 | this makes it possible to create pictures with dimensions that are multiples of 8 rather than 64
12 | """
13 |
14 | def __getattr__(self, item):
15 | if item == 'cat':
16 | return self.cat
17 |
18 | if hasattr(torch, item):
19 | return getattr(torch, item)
20 |
21 | raise AttributeError(f"'{type(self).__name__}' object has no attribute '{item}'")
22 |
23 | def cat(self, tensors, *args, **kwargs):
24 | if len(tensors) == 2:
25 | a, b = tensors
26 | if a.shape[-2:] != b.shape[-2:]:
27 | a = torch.nn.functional.interpolate(a, b.shape[-2:], mode="nearest")
28 |
29 | tensors = (a, b)
30 |
31 | return torch.cat(tensors, *args, **kwargs)
32 |
33 |
34 | th = TorchHijackForUnet()
35 |
36 |
37 | # Below are monkey patches to enable upcasting a float16 UNet for float32 sampling
38 | def apply_model(orig_func, self, x_noisy, t, cond, **kwargs):
39 |
40 | if isinstance(cond, dict):
41 | for y in cond.keys():
42 | cond[y] = [x.to(devices.dtype_unet) if isinstance(x, torch.Tensor) else x for x in cond[y]]
43 |
44 | with devices.autocast():
45 | return orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs).float()
46 |
47 |
48 | class GELUHijack(torch.nn.GELU, torch.nn.Module):
49 | def __init__(self, *args, **kwargs):
50 | torch.nn.GELU.__init__(self, *args, **kwargs)
51 | def forward(self, x):
52 | if devices.unet_needs_upcast:
53 | return torch.nn.GELU.forward(self.float(), x.float()).to(devices.dtype_unet)
54 | else:
55 | return torch.nn.GELU.forward(self, x)
56 |
57 |
58 | ddpm_edit_hijack = None
59 | def hijack_ddpm_edit():
60 | global ddpm_edit_hijack
61 | if not ddpm_edit_hijack:
62 | CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.decode_first_stage', first_stage_sub, first_stage_cond)
63 | CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.encode_first_stage', first_stage_sub, first_stage_cond)
64 | ddpm_edit_hijack = CondFunc('modules.models.diffusion.ddpm_edit.LatentDiffusion.apply_model', apply_model, unet_needs_upcast)
65 |
66 |
67 | unet_needs_upcast = lambda *args, **kwargs: devices.unet_needs_upcast
68 | CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model, unet_needs_upcast)
69 | CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast)
70 | if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available():
71 | CondFunc('ldm.modules.diffusionmodules.util.GroupNorm32.forward', lambda orig_func, self, *args, **kwargs: orig_func(self.float(), *args, **kwargs), unet_needs_upcast)
72 | CondFunc('ldm.modules.attention.GEGLU.forward', lambda orig_func, self, x: orig_func(self.float(), x.float()).to(devices.dtype_unet), unet_needs_upcast)
73 | CondFunc('open_clip.transformer.ResidualAttentionBlock.__init__', lambda orig_func, *args, **kwargs: kwargs.update({'act_layer': GELUHijack}) and False or orig_func(*args, **kwargs), lambda _, *args, **kwargs: kwargs.get('act_layer') is None or kwargs['act_layer'] == torch.nn.GELU)
74 |
75 | first_stage_cond = lambda _, self, *args, **kwargs: devices.unet_needs_upcast and self.model.diffusion_model.dtype == torch.float16
76 | first_stage_sub = lambda orig_func, self, x, **kwargs: orig_func(self, x.to(devices.dtype_vae), **kwargs)
77 | CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.decode_first_stage', first_stage_sub, first_stage_cond)
78 | CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.encode_first_stage', first_stage_sub, first_stage_cond)
79 | CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.get_first_stage_encoding', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).float(), first_stage_cond)
80 |
--------------------------------------------------------------------------------
/modules/models/diffusion/uni_pc/sampler.py:
--------------------------------------------------------------------------------
1 | """SAMPLING ONLY."""
2 |
3 | import torch
4 |
5 | from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC
6 | from modules import shared, devices
7 |
8 |
9 | class UniPCSampler(object):
10 | def __init__(self, model, **kwargs):
11 | super().__init__()
12 | self.model = model
13 | to_torch = lambda x: x.clone().detach().to(torch.float32).to(model.device)
14 | self.before_sample = None
15 | self.after_sample = None
16 | self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod))
17 |
18 | def register_buffer(self, name, attr):
19 | if type(attr) == torch.Tensor:
20 | if attr.device != devices.device:
21 | attr = attr.to(devices.device)
22 | setattr(self, name, attr)
23 |
24 | def set_hooks(self, before_sample, after_sample, after_update):
25 | self.before_sample = before_sample
26 | self.after_sample = after_sample
27 | self.after_update = after_update
28 |
29 | @torch.no_grad()
30 | def sample(self,
31 | S,
32 | batch_size,
33 | shape,
34 | conditioning=None,
35 | callback=None,
36 | normals_sequence=None,
37 | img_callback=None,
38 | quantize_x0=False,
39 | eta=0.,
40 | mask=None,
41 | x0=None,
42 | temperature=1.,
43 | noise_dropout=0.,
44 | score_corrector=None,
45 | corrector_kwargs=None,
46 | verbose=True,
47 | x_T=None,
48 | log_every_t=100,
49 | unconditional_guidance_scale=1.,
50 | unconditional_conditioning=None,
51 | # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
52 | **kwargs
53 | ):
54 | if conditioning is not None:
55 | if isinstance(conditioning, dict):
56 | ctmp = conditioning[list(conditioning.keys())[0]]
57 | while isinstance(ctmp, list):
58 | ctmp = ctmp[0]
59 | cbs = ctmp.shape[0]
60 | if cbs != batch_size:
61 | print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
62 |
63 | elif isinstance(conditioning, list):
64 | for ctmp in conditioning:
65 | if ctmp.shape[0] != batch_size:
66 | print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
67 |
68 | else:
69 | if conditioning.shape[0] != batch_size:
70 | print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
71 |
72 | # sampling
73 | C, H, W = shape
74 | size = (batch_size, C, H, W)
75 | # print(f'Data shape for UniPC sampling is {size}')
76 |
77 | device = self.model.betas.device
78 | if x_T is None:
79 | img = torch.randn(size, device=device)
80 | else:
81 | img = x_T
82 |
83 | ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)
84 |
85 | # SD 1.X is "noise", SD 2.X is "v"
86 | model_type = "v" if self.model.parameterization == "v" else "noise"
87 |
88 | model_fn = model_wrapper(
89 | lambda x, t, c: self.model.apply_model(x, t, c),
90 | ns,
91 | model_type=model_type,
92 | guidance_type="classifier-free",
93 | #condition=conditioning,
94 | #unconditional_condition=unconditional_conditioning,
95 | guidance_scale=unconditional_guidance_scale,
96 | )
97 |
98 | uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, variant=shared.opts.uni_pc_variant, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample, after_update=self.after_update)
99 | x = uni_pc.sample(img, steps=S, skip_type=shared.opts.uni_pc_skip_type, method="multistep", order=shared.opts.uni_pc_order, lower_order_final=shared.opts.uni_pc_lower_order_final)
100 |
101 | return x.to(device), None
102 |
--------------------------------------------------------------------------------
/modules/gfpgan_model.py:
--------------------------------------------------------------------------------
1 | import os
2 | import sys
3 | import traceback
4 |
5 | import facexlib
6 | import gfpgan
7 |
8 | import modules.face_restoration
9 | from modules import paths, shared, devices, modelloader
10 |
11 | model_dir = "GFPGAN"
12 | user_path = None
13 | model_path = os.path.join(paths.models_path, model_dir)
14 | model_url = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth"
15 | have_gfpgan = False
16 | loaded_gfpgan_model = None
17 |
18 |
19 | def gfpgann():
20 | global loaded_gfpgan_model
21 | global model_path
22 | if loaded_gfpgan_model is not None:
23 | loaded_gfpgan_model.gfpgan.to(devices.device_gfpgan)
24 | return loaded_gfpgan_model
25 |
26 | if gfpgan_constructor is None:
27 | return None
28 |
29 | models = modelloader.load_models(model_path, model_url, user_path, ext_filter="GFPGAN")
30 | if len(models) == 1 and "http" in models[0]:
31 | model_file = models[0]
32 | elif len(models) != 0:
33 | latest_file = max(models, key=os.path.getctime)
34 | model_file = latest_file
35 | else:
36 | print("Unable to load gfpgan model!")
37 | return None
38 | if hasattr(facexlib.detection.retinaface, 'device'):
39 | facexlib.detection.retinaface.device = devices.device_gfpgan
40 | model = gfpgan_constructor(model_path=model_file, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=devices.device_gfpgan)
41 | loaded_gfpgan_model = model
42 |
43 | return model
44 |
45 |
46 | def send_model_to(model, device):
47 | model.gfpgan.to(device)
48 | model.face_helper.face_det.to(device)
49 | model.face_helper.face_parse.to(device)
50 |
51 |
52 | def gfpgan_fix_faces(np_image):
53 | model = gfpgann()
54 | if model is None:
55 | return np_image
56 |
57 | send_model_to(model, devices.device_gfpgan)
58 |
59 | np_image_bgr = np_image[:, :, ::-1]
60 | cropped_faces, restored_faces, gfpgan_output_bgr = model.enhance(np_image_bgr, has_aligned=False, only_center_face=False, paste_back=True)
61 | np_image = gfpgan_output_bgr[:, :, ::-1]
62 |
63 | model.face_helper.clean_all()
64 |
65 | if shared.opts.face_restoration_unload:
66 | send_model_to(model, devices.cpu)
67 |
68 | return np_image
69 |
70 |
71 | gfpgan_constructor = None
72 |
73 |
74 | def setup_model(dirname):
75 | global model_path
76 | if not os.path.exists(model_path):
77 | os.makedirs(model_path)
78 |
79 | try:
80 | from gfpgan import GFPGANer
81 | from facexlib import detection, parsing # noqa: F401
82 | global user_path
83 | global have_gfpgan
84 | global gfpgan_constructor
85 |
86 | load_file_from_url_orig = gfpgan.utils.load_file_from_url
87 | facex_load_file_from_url_orig = facexlib.detection.load_file_from_url
88 | facex_load_file_from_url_orig2 = facexlib.parsing.load_file_from_url
89 |
90 | def my_load_file_from_url(**kwargs):
91 | return load_file_from_url_orig(**dict(kwargs, model_dir=model_path))
92 |
93 | def facex_load_file_from_url(**kwargs):
94 | return facex_load_file_from_url_orig(**dict(kwargs, save_dir=model_path, model_dir=None))
95 |
96 | def facex_load_file_from_url2(**kwargs):
97 | return facex_load_file_from_url_orig2(**dict(kwargs, save_dir=model_path, model_dir=None))
98 |
99 | gfpgan.utils.load_file_from_url = my_load_file_from_url
100 | facexlib.detection.load_file_from_url = facex_load_file_from_url
101 | facexlib.parsing.load_file_from_url = facex_load_file_from_url2
102 | user_path = dirname
103 | have_gfpgan = True
104 | gfpgan_constructor = GFPGANer
105 |
106 | class FaceRestorerGFPGAN(modules.face_restoration.FaceRestoration):
107 | def name(self):
108 | return "GFPGAN"
109 |
110 | def restore(self, np_image):
111 | return gfpgan_fix_faces(np_image)
112 |
113 | shared.face_restorers.append(FaceRestorerGFPGAN())
114 | except Exception:
115 | print("Error setting up GFPGAN:", file=sys.stderr)
116 | print(traceback.format_exc(), file=sys.stderr)
117 |
--------------------------------------------------------------------------------
/scripts/sd_upscale.py:
--------------------------------------------------------------------------------
1 | import math
2 |
3 | import modules.scripts as scripts
4 | import gradio as gr
5 | from PIL import Image
6 |
7 | from modules import processing, shared, images, devices
8 | from modules.processing import Processed
9 | from modules.shared import opts, state
10 |
11 |
12 | class Script(scripts.Script):
13 | def title(self):
14 | return "SD upscale"
15 |
16 | def show(self, is_img2img):
17 | return is_img2img
18 |
19 | def ui(self, is_img2img):
20 | info = gr.HTML("Will upscale the image by the selected scale factor; use width and height sliders to set tile size
")
21 | overlap = gr.Slider(minimum=0, maximum=256, step=16, label='Tile overlap', value=64, elem_id=self.elem_id("overlap"))
22 | 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"))
23 | 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"))
24 |
25 | return [info, overlap, upscaler_index, scale_factor]
26 |
27 | def run(self, p, _, overlap, upscaler_index, scale_factor):
28 | if isinstance(upscaler_index, str):
29 | upscaler_index = [x.name.lower() for x in shared.sd_upscalers].index(upscaler_index.lower())
30 | processing.fix_seed(p)
31 | upscaler = shared.sd_upscalers[upscaler_index]
32 |
33 | p.extra_generation_params["SD upscale overlap"] = overlap
34 | p.extra_generation_params["SD upscale upscaler"] = upscaler.name
35 |
36 | initial_info = None
37 | seed = p.seed
38 |
39 | init_img = p.init_images[0]
40 | init_img = images.flatten(init_img, opts.img2img_background_color)
41 |
42 | if upscaler.name != "None":
43 | img = upscaler.scaler.upscale(init_img, scale_factor, upscaler.data_path)
44 | else:
45 | img = init_img
46 |
47 | devices.torch_gc()
48 |
49 | grid = images.split_grid(img, tile_w=p.width, tile_h=p.height, overlap=overlap)
50 |
51 | batch_size = p.batch_size
52 | upscale_count = p.n_iter
53 | p.n_iter = 1
54 | p.do_not_save_grid = True
55 | p.do_not_save_samples = True
56 |
57 | work = []
58 |
59 | for _y, _h, row in grid.tiles:
60 | for tiledata in row:
61 | work.append(tiledata[2])
62 |
63 | batch_count = math.ceil(len(work) / batch_size)
64 | state.job_count = batch_count * upscale_count
65 |
66 | print(f"SD upscaling will process a total of {len(work)} images tiled as {len(grid.tiles[0][2])}x{len(grid.tiles)} per upscale in a total of {state.job_count} batches.")
67 |
68 | result_images = []
69 | for n in range(upscale_count):
70 | start_seed = seed + n
71 | p.seed = start_seed
72 |
73 | work_results = []
74 | for i in range(batch_count):
75 | p.batch_size = batch_size
76 | p.init_images = work[i * batch_size:(i + 1) * batch_size]
77 |
78 | state.job = f"Batch {i + 1 + n * batch_count} out of {state.job_count}"
79 | processed = processing.process_images(p)
80 |
81 | if initial_info is None:
82 | initial_info = processed.info
83 |
84 | p.seed = processed.seed + 1
85 | work_results += processed.images
86 |
87 | image_index = 0
88 | for _y, _h, row in grid.tiles:
89 | for tiledata in row:
90 | tiledata[2] = work_results[image_index] if image_index < len(work_results) else Image.new("RGB", (p.width, p.height))
91 | image_index += 1
92 |
93 | combined_image = images.combine_grid(grid)
94 | result_images.append(combined_image)
95 |
96 | if opts.samples_save:
97 | images.save_image(combined_image, p.outpath_samples, "", start_seed, p.prompt, opts.samples_format, info=initial_info, p=p)
98 |
99 | processed = Processed(p, result_images, seed, initial_info)
100 |
101 | return processed
102 |
--------------------------------------------------------------------------------
/javascript/aspectRatioOverlay.js:
--------------------------------------------------------------------------------
1 |
2 | let currentWidth = null;
3 | let currentHeight = null;
4 | let arFrameTimeout = setTimeout(function() {}, 0);
5 |
6 | function dimensionChange(e, is_width, is_height) {
7 |
8 | if (is_width) {
9 | currentWidth = e.target.value * 1.0;
10 | }
11 | if (is_height) {
12 | currentHeight = e.target.value * 1.0;
13 | }
14 |
15 | var inImg2img = gradioApp().querySelector("#tab_img2img").style.display == "block";
16 |
17 | if (!inImg2img) {
18 | return;
19 | }
20 |
21 | var targetElement = null;
22 |
23 | var tabIndex = get_tab_index('mode_img2img');
24 | if (tabIndex == 0) { // img2img
25 | targetElement = gradioApp().querySelector('#img2img_image div[data-testid=image] img');
26 | } else if (tabIndex == 1) { //Sketch
27 | targetElement = gradioApp().querySelector('#img2img_sketch div[data-testid=image] img');
28 | } else if (tabIndex == 2) { // Inpaint
29 | targetElement = gradioApp().querySelector('#img2maskimg div[data-testid=image] img');
30 | } else if (tabIndex == 3) { // Inpaint sketch
31 | targetElement = gradioApp().querySelector('#inpaint_sketch div[data-testid=image] img');
32 | }
33 |
34 |
35 | if (targetElement) {
36 |
37 | var arPreviewRect = gradioApp().querySelector('#imageARPreview');
38 | if (!arPreviewRect) {
39 | arPreviewRect = document.createElement('div');
40 | arPreviewRect.id = "imageARPreview";
41 | gradioApp().appendChild(arPreviewRect);
42 | }
43 |
44 |
45 |
46 | var viewportOffset = targetElement.getBoundingClientRect();
47 |
48 | var viewportscale = Math.min(targetElement.clientWidth / targetElement.naturalWidth, targetElement.clientHeight / targetElement.naturalHeight);
49 |
50 | var scaledx = targetElement.naturalWidth * viewportscale;
51 | var scaledy = targetElement.naturalHeight * viewportscale;
52 |
53 | var cleintRectTop = (viewportOffset.top + window.scrollY);
54 | var cleintRectLeft = (viewportOffset.left + window.scrollX);
55 | var cleintRectCentreY = cleintRectTop + (targetElement.clientHeight / 2);
56 | var cleintRectCentreX = cleintRectLeft + (targetElement.clientWidth / 2);
57 |
58 | var arscale = Math.min(scaledx / currentWidth, scaledy / currentHeight);
59 | var arscaledx = currentWidth * arscale;
60 | var arscaledy = currentHeight * arscale;
61 |
62 | var arRectTop = cleintRectCentreY - (arscaledy / 2);
63 | var arRectLeft = cleintRectCentreX - (arscaledx / 2);
64 | var arRectWidth = arscaledx;
65 | var arRectHeight = arscaledy;
66 |
67 | arPreviewRect.style.top = arRectTop + 'px';
68 | arPreviewRect.style.left = arRectLeft + 'px';
69 | arPreviewRect.style.width = arRectWidth + 'px';
70 | arPreviewRect.style.height = arRectHeight + 'px';
71 |
72 | clearTimeout(arFrameTimeout);
73 | arFrameTimeout = setTimeout(function() {
74 | arPreviewRect.style.display = 'none';
75 | }, 2000);
76 |
77 | arPreviewRect.style.display = 'block';
78 |
79 | }
80 |
81 | }
82 |
83 |
84 | onUiUpdate(function() {
85 | var arPreviewRect = gradioApp().querySelector('#imageARPreview');
86 | if (arPreviewRect) {
87 | arPreviewRect.style.display = 'none';
88 | }
89 | var tabImg2img = gradioApp().querySelector("#tab_img2img");
90 | if (tabImg2img) {
91 | var inImg2img = tabImg2img.style.display == "block";
92 | if (inImg2img) {
93 | let inputs = gradioApp().querySelectorAll('input');
94 | inputs.forEach(function(e) {
95 | var is_width = e.parentElement.id == "img2img_width";
96 | var is_height = e.parentElement.id == "img2img_height";
97 |
98 | if ((is_width || is_height) && !e.classList.contains('scrollwatch')) {
99 | e.addEventListener('input', function(e) {
100 | dimensionChange(e, is_width, is_height);
101 | });
102 | e.classList.add('scrollwatch');
103 | }
104 | if (is_width) {
105 | currentWidth = e.value * 1.0;
106 | }
107 | if (is_height) {
108 | currentHeight = e.value * 1.0;
109 | }
110 | });
111 | }
112 | }
113 | });
114 |
--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/bug_report.yml:
--------------------------------------------------------------------------------
1 | name: Bug Report
2 | description: You think somethings is broken in the UI
3 | title: "[Bug]: "
4 | labels: ["bug-report"]
5 |
6 | body:
7 | - type: checkboxes
8 | attributes:
9 | label: Is there an existing issue for this?
10 | description: Please search to see if an issue already exists for the bug you encountered, and that it hasn't been fixed in a recent build/commit.
11 | options:
12 | - label: I have searched the existing issues and checked the recent builds/commits
13 | required: true
14 | - type: markdown
15 | attributes:
16 | value: |
17 | *Please fill this form with as much information as possible, don't forget to fill "What OS..." and "What browsers" and *provide screenshots if possible**
18 | - type: textarea
19 | id: what-did
20 | attributes:
21 | label: What happened?
22 | description: Tell us what happened in a very clear and simple way
23 | validations:
24 | required: true
25 | - type: textarea
26 | id: steps
27 | attributes:
28 | label: Steps to reproduce the problem
29 | description: Please provide us with precise step by step information on how to reproduce the bug
30 | value: |
31 | 1. Go to ....
32 | 2. Press ....
33 | 3. ...
34 | validations:
35 | required: true
36 | - type: textarea
37 | id: what-should
38 | attributes:
39 | label: What should have happened?
40 | description: Tell what you think the normal behavior should be
41 | validations:
42 | required: true
43 | - type: input
44 | id: commit
45 | attributes:
46 | label: Commit where the problem happens
47 | description: Which commit are you running ? (Do not write *Latest version/repo/commit*, as this means nothing and will have changed by the time we read your issue. Rather, copy the **Commit** link at the bottom of the UI, or from the cmd/terminal if you can't launch it.)
48 | validations:
49 | required: true
50 | - type: dropdown
51 | id: py-version
52 | attributes:
53 | label: What Python version are you running on ?
54 | multiple: false
55 | options:
56 | - Python 3.10.x
57 | - Python 3.11.x (above, no supported yet)
58 | - Python 3.9.x (below, no recommended)
59 | - type: dropdown
60 | id: platforms
61 | attributes:
62 | label: What platforms do you use to access the UI ?
63 | multiple: true
64 | options:
65 | - Windows
66 | - Linux
67 | - MacOS
68 | - iOS
69 | - Android
70 | - Other/Cloud
71 | - type: dropdown
72 | id: device
73 | attributes:
74 | label: What device are you running WebUI on?
75 | multiple: true
76 | options:
77 | - Nvidia GPUs (RTX 20 above)
78 | - Nvidia GPUs (GTX 16 below)
79 | - AMD GPUs (RX 6000 above)
80 | - AMD GPUs (RX 5000 below)
81 | - CPU
82 | - Other GPUs
83 | - type: dropdown
84 | id: browsers
85 | attributes:
86 | label: What browsers do you use to access the UI ?
87 | multiple: true
88 | options:
89 | - Mozilla Firefox
90 | - Google Chrome
91 | - Brave
92 | - Apple Safari
93 | - Microsoft Edge
94 | - type: textarea
95 | id: cmdargs
96 | attributes:
97 | label: Command Line Arguments
98 | description: Are you using any launching parameters/command line arguments (modified webui-user .bat/.sh) ? If yes, please write them below. Write "No" otherwise.
99 | render: Shell
100 | validations:
101 | required: true
102 | - type: textarea
103 | id: extensions
104 | attributes:
105 | label: List of extensions
106 | description: Are you using any extensions other than built-ins? If yes, provide a list, you can copy it at "Extensions" tab. Write "No" otherwise.
107 | validations:
108 | required: true
109 | - type: textarea
110 | id: logs
111 | attributes:
112 | label: Console logs
113 | description: Please provide **full** cmd/terminal logs from the moment you started UI to the end of it, after your bug happened. If it's very long, provide a link to pastebin or similar service.
114 | render: Shell
115 | validations:
116 | required: true
117 | - type: textarea
118 | id: misc
119 | attributes:
120 | label: Additional information
121 | description: Please provide us with any relevant additional info or context.
122 |
--------------------------------------------------------------------------------
/modules/postprocessing.py:
--------------------------------------------------------------------------------
1 | import os
2 |
3 | from PIL import Image
4 |
5 | from modules import shared, images, devices, scripts, scripts_postprocessing, ui_common, generation_parameters_copypaste
6 | from modules.shared import opts
7 |
8 |
9 | def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, show_extras_results, *args, save_output: bool = True):
10 | devices.torch_gc()
11 |
12 | shared.state.begin()
13 | shared.state.job = 'extras'
14 |
15 | image_data = []
16 | image_names = []
17 | outputs = []
18 |
19 | if extras_mode == 1:
20 | for img in image_folder:
21 | if isinstance(img, Image.Image):
22 | image = img
23 | fn = ''
24 | else:
25 | image = Image.open(os.path.abspath(img.name))
26 | fn = os.path.splitext(img.orig_name)[0]
27 | image_data.append(image)
28 | image_names.append(fn)
29 | elif extras_mode == 2:
30 | assert not shared.cmd_opts.hide_ui_dir_config, '--hide-ui-dir-config option must be disabled'
31 | assert input_dir, 'input directory not selected'
32 |
33 | image_list = shared.listfiles(input_dir)
34 | for filename in image_list:
35 | try:
36 | image = Image.open(filename)
37 | except Exception:
38 | continue
39 | image_data.append(image)
40 | image_names.append(filename)
41 | else:
42 | assert image, 'image not selected'
43 |
44 | image_data.append(image)
45 | image_names.append(None)
46 |
47 | if extras_mode == 2 and output_dir != '':
48 | outpath = output_dir
49 | else:
50 | outpath = opts.outdir_samples or opts.outdir_extras_samples
51 |
52 | infotext = ''
53 |
54 | for image, name in zip(image_data, image_names):
55 | shared.state.textinfo = name
56 |
57 | existing_pnginfo = image.info or {}
58 |
59 | pp = scripts_postprocessing.PostprocessedImage(image.convert("RGB"))
60 |
61 | scripts.scripts_postproc.run(pp, args)
62 |
63 | if opts.use_original_name_batch and name is not None:
64 | basename = os.path.splitext(os.path.basename(name))[0]
65 | else:
66 | basename = ''
67 |
68 | infotext = ", ".join([k if k == v else f'{k}: {generation_parameters_copypaste.quote(v)}' for k, v in pp.info.items() if v is not None])
69 |
70 | if opts.enable_pnginfo:
71 | pp.image.info = existing_pnginfo
72 | pp.image.info["postprocessing"] = infotext
73 |
74 | if save_output:
75 | images.save_image(pp.image, path=outpath, basename=basename, seed=None, prompt=None, extension=opts.samples_format, info=infotext, short_filename=True, no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=existing_pnginfo, forced_filename=None)
76 |
77 | if extras_mode != 2 or show_extras_results:
78 | outputs.append(pp.image)
79 |
80 | devices.torch_gc()
81 |
82 | return outputs, ui_common.plaintext_to_html(infotext), ''
83 |
84 |
85 | def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility, upscale_first: bool, save_output: bool = True):
86 | """old handler for API"""
87 |
88 | args = scripts.scripts_postproc.create_args_for_run({
89 | "Upscale": {
90 | "upscale_mode": resize_mode,
91 | "upscale_by": upscaling_resize,
92 | "upscale_to_width": upscaling_resize_w,
93 | "upscale_to_height": upscaling_resize_h,
94 | "upscale_crop": upscaling_crop,
95 | "upscaler_1_name": extras_upscaler_1,
96 | "upscaler_2_name": extras_upscaler_2,
97 | "upscaler_2_visibility": extras_upscaler_2_visibility,
98 | },
99 | "GFPGAN": {
100 | "gfpgan_visibility": gfpgan_visibility,
101 | },
102 | "CodeFormer": {
103 | "codeformer_visibility": codeformer_visibility,
104 | "codeformer_weight": codeformer_weight,
105 | },
106 | })
107 |
108 | return run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, show_extras_results, *args, save_output=save_output)
109 |
--------------------------------------------------------------------------------
/modules/upscaler.py:
--------------------------------------------------------------------------------
1 | import os
2 | from abc import abstractmethod
3 |
4 | import PIL
5 | from PIL import Image
6 |
7 | import modules.shared
8 | from modules import modelloader, shared
9 |
10 | LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
11 | NEAREST = (Image.Resampling.NEAREST if hasattr(Image, 'Resampling') else Image.NEAREST)
12 |
13 |
14 | class Upscaler:
15 | name = None
16 | model_path = None
17 | model_name = None
18 | model_url = None
19 | enable = True
20 | filter = None
21 | model = None
22 | user_path = None
23 | scalers: []
24 | tile = True
25 |
26 | def __init__(self, create_dirs=False):
27 | self.mod_pad_h = None
28 | self.tile_size = modules.shared.opts.ESRGAN_tile
29 | self.tile_pad = modules.shared.opts.ESRGAN_tile_overlap
30 | self.device = modules.shared.device
31 | self.img = None
32 | self.output = None
33 | self.scale = 1
34 | self.half = not modules.shared.cmd_opts.no_half
35 | self.pre_pad = 0
36 | self.mod_scale = None
37 | self.model_download_path = None
38 |
39 | if self.model_path is None and self.name:
40 | self.model_path = os.path.join(shared.models_path, self.name)
41 | if self.model_path and create_dirs:
42 | os.makedirs(self.model_path, exist_ok=True)
43 |
44 | try:
45 | import cv2 # noqa: F401
46 | self.can_tile = True
47 | except Exception:
48 | pass
49 |
50 | @abstractmethod
51 | def do_upscale(self, img: PIL.Image, selected_model: str):
52 | return img
53 |
54 | def upscale(self, img: PIL.Image, scale, selected_model: str = None):
55 | self.scale = scale
56 | dest_w = int(img.width * scale)
57 | dest_h = int(img.height * scale)
58 |
59 | for _ in range(3):
60 | shape = (img.width, img.height)
61 |
62 | img = self.do_upscale(img, selected_model)
63 |
64 | if shape == (img.width, img.height):
65 | break
66 |
67 | if img.width >= dest_w and img.height >= dest_h:
68 | break
69 |
70 | if img.width != dest_w or img.height != dest_h:
71 | img = img.resize((int(dest_w), int(dest_h)), resample=LANCZOS)
72 |
73 | return img
74 |
75 | @abstractmethod
76 | def load_model(self, path: str):
77 | pass
78 |
79 | def find_models(self, ext_filter=None) -> list:
80 | return modelloader.load_models(model_path=self.model_path, model_url=self.model_url, command_path=self.user_path)
81 |
82 | def update_status(self, prompt):
83 | print(f"\nextras: {prompt}", file=shared.progress_print_out)
84 |
85 |
86 | class UpscalerData:
87 | name = None
88 | data_path = None
89 | scale: int = 4
90 | scaler: Upscaler = None
91 | model: None
92 |
93 | def __init__(self, name: str, path: str, upscaler: Upscaler = None, scale: int = 4, model=None):
94 | self.name = name
95 | self.data_path = path
96 | self.local_data_path = path
97 | self.scaler = upscaler
98 | self.scale = scale
99 | self.model = model
100 |
101 |
102 | class UpscalerNone(Upscaler):
103 | name = "None"
104 | scalers = []
105 |
106 | def load_model(self, path):
107 | pass
108 |
109 | def do_upscale(self, img, selected_model=None):
110 | return img
111 |
112 | def __init__(self, dirname=None):
113 | super().__init__(False)
114 | self.scalers = [UpscalerData("None", None, self)]
115 |
116 |
117 | class UpscalerLanczos(Upscaler):
118 | scalers = []
119 |
120 | def do_upscale(self, img, selected_model=None):
121 | return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=LANCZOS)
122 |
123 | def load_model(self, _):
124 | pass
125 |
126 | def __init__(self, dirname=None):
127 | super().__init__(False)
128 | self.name = "Lanczos"
129 | self.scalers = [UpscalerData("Lanczos", None, self)]
130 |
131 |
132 | class UpscalerNearest(Upscaler):
133 | scalers = []
134 |
135 | def do_upscale(self, img, selected_model=None):
136 | return img.resize((int(img.width * self.scale), int(img.height * self.scale)), resample=NEAREST)
137 |
138 | def load_model(self, _):
139 | pass
140 |
141 | def __init__(self, dirname=None):
142 | super().__init__(False)
143 | self.name = "Nearest"
144 | self.scalers = [UpscalerData("Nearest", None, self)]
145 |
--------------------------------------------------------------------------------
/modules/sd_hijack_inpainting.py:
--------------------------------------------------------------------------------
1 | import torch
2 |
3 | import ldm.models.diffusion.ddpm
4 | import ldm.models.diffusion.ddim
5 | import ldm.models.diffusion.plms
6 |
7 | from ldm.models.diffusion.ddim import noise_like
8 | from ldm.models.diffusion.sampling_util import norm_thresholding
9 |
10 |
11 | @torch.no_grad()
12 | def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
13 | temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
14 | unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None, dynamic_threshold=None):
15 | b, *_, device = *x.shape, x.device
16 |
17 | def get_model_output(x, t):
18 | if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
19 | e_t = self.model.apply_model(x, t, c)
20 | else:
21 | x_in = torch.cat([x] * 2)
22 | t_in = torch.cat([t] * 2)
23 |
24 | if isinstance(c, dict):
25 | assert isinstance(unconditional_conditioning, dict)
26 | c_in = {}
27 | for k in c:
28 | if isinstance(c[k], list):
29 | c_in[k] = [
30 | torch.cat([unconditional_conditioning[k][i], c[k][i]])
31 | for i in range(len(c[k]))
32 | ]
33 | else:
34 | c_in[k] = torch.cat([unconditional_conditioning[k], c[k]])
35 | else:
36 | c_in = torch.cat([unconditional_conditioning, c])
37 |
38 | e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
39 | e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
40 |
41 | if score_corrector is not None:
42 | assert self.model.parameterization == "eps"
43 | e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
44 |
45 | return e_t
46 |
47 | alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
48 | alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
49 | sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
50 | sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
51 |
52 | def get_x_prev_and_pred_x0(e_t, index):
53 | # select parameters corresponding to the currently considered timestep
54 | a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
55 | a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
56 | sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
57 | sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
58 |
59 | # current prediction for x_0
60 | pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
61 | if quantize_denoised:
62 | pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
63 | if dynamic_threshold is not None:
64 | pred_x0 = norm_thresholding(pred_x0, dynamic_threshold)
65 | # direction pointing to x_t
66 | dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
67 | noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
68 | if noise_dropout > 0.:
69 | noise = torch.nn.functional.dropout(noise, p=noise_dropout)
70 | x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
71 | return x_prev, pred_x0
72 |
73 | e_t = get_model_output(x, t)
74 | if len(old_eps) == 0:
75 | # Pseudo Improved Euler (2nd order)
76 | x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
77 | e_t_next = get_model_output(x_prev, t_next)
78 | e_t_prime = (e_t + e_t_next) / 2
79 | elif len(old_eps) == 1:
80 | # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
81 | e_t_prime = (3 * e_t - old_eps[-1]) / 2
82 | elif len(old_eps) == 2:
83 | # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
84 | e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
85 | elif len(old_eps) >= 3:
86 | # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
87 | e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
88 |
89 | x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
90 |
91 | return x_prev, pred_x0, e_t
92 |
93 |
94 | def do_inpainting_hijack():
95 | # p_sample_plms is needed because PLMS can't work with dicts as conditionings
96 |
97 | ldm.models.diffusion.plms.PLMSSampler.p_sample_plms = p_sample_plms
98 |
--------------------------------------------------------------------------------
/modules/mac_specific.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import platform
3 | from modules.sd_hijack_utils import CondFunc
4 | from packaging import version
5 |
6 |
7 | # has_mps is only available in nightly pytorch (for now) and macOS 12.3+.
8 | # check `getattr` and try it for compatibility
9 | def check_for_mps() -> bool:
10 | if not getattr(torch, 'has_mps', False):
11 | return False
12 | try:
13 | torch.zeros(1).to(torch.device("mps"))
14 | return True
15 | except Exception:
16 | return False
17 | has_mps = check_for_mps()
18 |
19 |
20 | # MPS workaround for https://github.com/pytorch/pytorch/issues/89784
21 | def cumsum_fix(input, cumsum_func, *args, **kwargs):
22 | if input.device.type == 'mps':
23 | output_dtype = kwargs.get('dtype', input.dtype)
24 | if output_dtype == torch.int64:
25 | return cumsum_func(input.cpu(), *args, **kwargs).to(input.device)
26 | elif output_dtype == torch.bool or cumsum_needs_int_fix and (output_dtype == torch.int8 or output_dtype == torch.int16):
27 | return cumsum_func(input.to(torch.int32), *args, **kwargs).to(torch.int64)
28 | return cumsum_func(input, *args, **kwargs)
29 |
30 |
31 | if has_mps:
32 | # MPS fix for randn in torchsde
33 | CondFunc('torchsde._brownian.brownian_interval._randn', lambda _, size, dtype, device, seed: torch.randn(size, dtype=dtype, device=torch.device("cpu"), generator=torch.Generator(torch.device("cpu")).manual_seed(int(seed))).to(device), lambda _, size, dtype, device, seed: device.type == 'mps')
34 |
35 | if platform.mac_ver()[0].startswith("13.2."):
36 | # MPS workaround for https://github.com/pytorch/pytorch/issues/95188, thanks to danieldk (https://github.com/explosion/curated-transformers/pull/124)
37 | CondFunc('torch.nn.functional.linear', lambda _, input, weight, bias: (torch.matmul(input, weight.t()) + bias) if bias is not None else torch.matmul(input, weight.t()), lambda _, input, weight, bias: input.numel() > 10485760)
38 |
39 | if version.parse(torch.__version__) < version.parse("1.13"):
40 | # PyTorch 1.13 doesn't need these fixes but unfortunately is slower and has regressions that prevent training from working
41 |
42 | # MPS workaround for https://github.com/pytorch/pytorch/issues/79383
43 | CondFunc('torch.Tensor.to', lambda orig_func, self, *args, **kwargs: orig_func(self.contiguous(), *args, **kwargs),
44 | lambda _, self, *args, **kwargs: self.device.type != 'mps' and (args and isinstance(args[0], torch.device) and args[0].type == 'mps' or isinstance(kwargs.get('device'), torch.device) and kwargs['device'].type == 'mps'))
45 | # MPS workaround for https://github.com/pytorch/pytorch/issues/80800
46 | CondFunc('torch.nn.functional.layer_norm', lambda orig_func, *args, **kwargs: orig_func(*([args[0].contiguous()] + list(args[1:])), **kwargs),
47 | lambda _, *args, **kwargs: args and isinstance(args[0], torch.Tensor) and args[0].device.type == 'mps')
48 | # MPS workaround for https://github.com/pytorch/pytorch/issues/90532
49 | CondFunc('torch.Tensor.numpy', lambda orig_func, self, *args, **kwargs: orig_func(self.detach(), *args, **kwargs), lambda _, self, *args, **kwargs: self.requires_grad)
50 | elif version.parse(torch.__version__) > version.parse("1.13.1"):
51 | cumsum_needs_int_fix = not torch.Tensor([1,2]).to(torch.device("mps")).equal(torch.ShortTensor([1,1]).to(torch.device("mps")).cumsum(0))
52 | cumsum_fix_func = lambda orig_func, input, *args, **kwargs: cumsum_fix(input, orig_func, *args, **kwargs)
53 | CondFunc('torch.cumsum', cumsum_fix_func, None)
54 | CondFunc('torch.Tensor.cumsum', cumsum_fix_func, None)
55 | CondFunc('torch.narrow', lambda orig_func, *args, **kwargs: orig_func(*args, **kwargs).clone(), None)
56 |
57 | # MPS workaround for https://github.com/pytorch/pytorch/issues/96113
58 | CondFunc('torch.nn.functional.layer_norm', lambda orig_func, x, normalized_shape, weight, bias, eps, **kwargs: orig_func(x.float(), normalized_shape, weight.float() if weight is not None else None, bias.float() if bias is not None else bias, eps).to(x.dtype), lambda _, input, *args, **kwargs: len(args) == 4 and input.device.type == 'mps')
59 |
60 | # MPS workaround for https://github.com/pytorch/pytorch/issues/92311
61 | if platform.processor() == 'i386':
62 | for funcName in ['torch.argmax', 'torch.Tensor.argmax']:
63 | CondFunc(funcName, lambda _, input, *args, **kwargs: torch.max(input.float() if input.dtype == torch.int64 else input, *args, **kwargs)[1], lambda _, input, *args, **kwargs: input.device.type == 'mps')
64 |
--------------------------------------------------------------------------------
/modules/lowvram.py:
--------------------------------------------------------------------------------
1 | import torch
2 | from modules import devices
3 |
4 | module_in_gpu = None
5 | cpu = torch.device("cpu")
6 |
7 |
8 | def send_everything_to_cpu():
9 | global module_in_gpu
10 |
11 | if module_in_gpu is not None:
12 | module_in_gpu.to(cpu)
13 |
14 | module_in_gpu = None
15 |
16 |
17 | def setup_for_low_vram(sd_model, use_medvram):
18 | parents = {}
19 |
20 | def send_me_to_gpu(module, _):
21 | """send this module to GPU; send whatever tracked module was previous in GPU to CPU;
22 | we add this as forward_pre_hook to a lot of modules and this way all but one of them will
23 | be in CPU
24 | """
25 | global module_in_gpu
26 |
27 | module = parents.get(module, module)
28 |
29 | if module_in_gpu == module:
30 | return
31 |
32 | if module_in_gpu is not None:
33 | module_in_gpu.to(cpu)
34 |
35 | module.to(devices.device)
36 | module_in_gpu = module
37 |
38 | # see below for register_forward_pre_hook;
39 | # first_stage_model does not use forward(), it uses encode/decode, so register_forward_pre_hook is
40 | # useless here, and we just replace those methods
41 |
42 | first_stage_model = sd_model.first_stage_model
43 | first_stage_model_encode = sd_model.first_stage_model.encode
44 | first_stage_model_decode = sd_model.first_stage_model.decode
45 |
46 | def first_stage_model_encode_wrap(x):
47 | send_me_to_gpu(first_stage_model, None)
48 | return first_stage_model_encode(x)
49 |
50 | def first_stage_model_decode_wrap(z):
51 | send_me_to_gpu(first_stage_model, None)
52 | return first_stage_model_decode(z)
53 |
54 | # for SD1, cond_stage_model is CLIP and its NN is in the tranformer frield, but for SD2, it's open clip, and it's in model field
55 | if hasattr(sd_model.cond_stage_model, 'model'):
56 | sd_model.cond_stage_model.transformer = sd_model.cond_stage_model.model
57 |
58 | # remove several big modules: cond, first_stage, depth/embedder (if applicable), and unet from the model and then
59 | # send the model to GPU. Then put modules back. the modules will be in CPU.
60 | stored = sd_model.cond_stage_model.transformer, sd_model.first_stage_model, getattr(sd_model, 'depth_model', None), getattr(sd_model, 'embedder', None), sd_model.model
61 | sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.embedder, sd_model.model = None, None, None, None, None
62 | sd_model.to(devices.device)
63 | sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.embedder, sd_model.model = stored
64 |
65 | # register hooks for those the first three models
66 | sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu)
67 | sd_model.first_stage_model.register_forward_pre_hook(send_me_to_gpu)
68 | sd_model.first_stage_model.encode = first_stage_model_encode_wrap
69 | sd_model.first_stage_model.decode = first_stage_model_decode_wrap
70 | if sd_model.depth_model:
71 | sd_model.depth_model.register_forward_pre_hook(send_me_to_gpu)
72 | if sd_model.embedder:
73 | sd_model.embedder.register_forward_pre_hook(send_me_to_gpu)
74 | parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model
75 |
76 | if hasattr(sd_model.cond_stage_model, 'model'):
77 | sd_model.cond_stage_model.model = sd_model.cond_stage_model.transformer
78 | del sd_model.cond_stage_model.transformer
79 |
80 | if use_medvram:
81 | sd_model.model.register_forward_pre_hook(send_me_to_gpu)
82 | else:
83 | diff_model = sd_model.model.diffusion_model
84 |
85 | # the third remaining model is still too big for 4 GB, so we also do the same for its submodules
86 | # so that only one of them is in GPU at a time
87 | stored = diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed
88 | diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = None, None, None, None
89 | sd_model.model.to(devices.device)
90 | diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = stored
91 |
92 | # install hooks for bits of third model
93 | diff_model.time_embed.register_forward_pre_hook(send_me_to_gpu)
94 | for block in diff_model.input_blocks:
95 | block.register_forward_pre_hook(send_me_to_gpu)
96 | diff_model.middle_block.register_forward_pre_hook(send_me_to_gpu)
97 | for block in diff_model.output_blocks:
98 | block.register_forward_pre_hook(send_me_to_gpu)
99 |
--------------------------------------------------------------------------------