{html.escape(x)}
\n" for x in text.split('\n')]) return text def image_grid(imgs, batch_size=1, rows=None): if rows is None: if opts.n_rows > 0: rows = opts.n_rows elif opts.n_rows == 0: rows = batch_size else: rows = math.sqrt(len(imgs)) rows = round(rows) cols = math.ceil(len(imgs) / rows) w, h = imgs[0].size grid = Image.new('RGB', size=(cols * w, rows * h), color='black') for i, img in enumerate(imgs): grid.paste(img, box=(i % cols * w, i // cols * h)) return grid Grid = namedtuple("Grid", ["tiles", "tile_w", "tile_h", "image_w", "image_h", "overlap"]) def split_grid(image, tile_w=512, tile_h=512, overlap=64): w = image.width h = image.height now = tile_w - overlap # non-overlap width noh = tile_h - overlap cols = math.ceil((w - overlap) / now) rows = math.ceil((h - overlap) / noh) grid = Grid([], tile_w, tile_h, w, h, overlap) for row in range(rows): row_images = [] y = row * noh if y + tile_h >= h: y = h - tile_h for col in range(cols): x = col * now if x+tile_w >= w: x = w - tile_w tile = image.crop((x, y, x + tile_w, y + tile_h)) row_images.append([x, tile_w, tile]) grid.tiles.append([y, tile_h, row_images]) return grid def combine_grid(grid): def make_mask_image(r): r = r * 255 / grid.overlap r = r.astype(np.uint8) return Image.fromarray(r, 'L') mask_w = make_mask_image(np.arange(grid.overlap, dtype=np.float32).reshape((1, grid.overlap)).repeat(grid.tile_h, axis=0)) mask_h = make_mask_image(np.arange(grid.overlap, dtype=np.float32).reshape((grid.overlap, 1)).repeat(grid.image_w, axis=1)) combined_image = Image.new("RGB", (grid.image_w, grid.image_h)) for y, h, row in grid.tiles: combined_row = Image.new("RGB", (grid.image_w, h)) for x, w, tile in row: if x == 0: combined_row.paste(tile, (0, 0)) continue combined_row.paste(tile.crop((0, 0, grid.overlap, h)), (x, 0), mask=mask_w) combined_row.paste(tile.crop((grid.overlap, 0, w, h)), (x + grid.overlap, 0)) if y == 0: combined_image.paste(combined_row, (0, 0)) continue combined_image.paste(combined_row.crop((0, 0, combined_row.width, grid.overlap)), (0, y), mask=mask_h) combined_image.paste(combined_row.crop((0, grid.overlap, combined_row.width, h)), (0, y + grid.overlap)) return combined_image class GridAnnotation: def __init__(self, text='', is_active=True): self.text = text self.is_active = is_active self.size = None def draw_grid_annotations(im, width, height, hor_texts, ver_texts): def wrap(drawing, text, font, line_length): lines = [''] for word in text.split(): line = f'{lines[-1]} {word}'.strip() if drawing.textlength(line, font=font) <= line_length: lines[-1] = line else: lines.append(word) return lines def draw_texts(drawing, draw_x, draw_y, lines): for i, line in enumerate(lines): drawing.multiline_text((draw_x, draw_y + line.size[1] / 2), line.text, font=fnt, fill=color_active if line.is_active else color_inactive, anchor="mm", align="center") if not line.is_active: drawing.line((draw_x - line.size[0]//2, draw_y + line.size[1]//2, draw_x + line.size[0]//2, draw_y + line.size[1]//2), fill=color_inactive, width=4) draw_y += line.size[1] + line_spacing fontsize = (width + height) // 25 line_spacing = fontsize // 2 fnt = ImageFont.truetype(opts.font, fontsize) color_active = (0, 0, 0) color_inactive = (153, 153, 153) pad_left = width * 3 // 4 if len(ver_texts) > 0 else 0 cols = im.width // width rows = im.height // height assert cols == len(hor_texts), f'bad number of horizontal texts: {len(hor_texts)}; must be {cols}' assert rows == len(ver_texts), f'bad number of vertical texts: {len(ver_texts)}; must be {rows}' calc_img = Image.new("RGB", (1, 1), "white") calc_d = ImageDraw.Draw(calc_img) for texts, allowed_width in zip(hor_texts + ver_texts, [width] * len(hor_texts) + [pad_left] * len(ver_texts)): items = [] + texts texts.clear() for line in items: wrapped = wrap(calc_d, line.text, fnt, allowed_width) texts += [GridAnnotation(x, line.is_active) for x in wrapped] for line in texts: bbox = calc_d.multiline_textbbox((0, 0), line.text, font=fnt) line.size = (bbox[2] - bbox[0], bbox[3] - bbox[1]) hor_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing for lines in hor_texts] ver_text_heights = [sum([line.size[1] + line_spacing for line in lines]) - line_spacing * len(lines) for lines in ver_texts] pad_top = max(hor_text_heights) + line_spacing * 2 result = Image.new("RGB", (im.width + pad_left, im.height + pad_top), "white") result.paste(im, (pad_left, pad_top)) d = ImageDraw.Draw(result) for col in range(cols): x = pad_left + width * col + width / 2 y = pad_top / 2 - hor_text_heights[col] / 2 draw_texts(d, x, y, hor_texts[col]) for row in range(rows): x = pad_left / 2 y = pad_top + height * row + height / 2 - ver_text_heights[row] / 2 draw_texts(d, x, y, ver_texts[row]) return result def draw_prompt_matrix(im, width, height, all_prompts): prompts = all_prompts[1:] boundary = math.ceil(len(prompts) / 2) prompts_horiz = prompts[:boundary] prompts_vert = prompts[boundary:] hor_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_horiz)] for pos in range(1 << len(prompts_horiz))] ver_texts = [[GridAnnotation(x, is_active=pos & (1 << i) != 0) for i, x in enumerate(prompts_vert)] for pos in range(1 << len(prompts_vert))] return draw_grid_annotations(im, width, height, hor_texts, ver_texts) def draw_xy_grid(xs, ys, x_label, y_label, cell): res = [] ver_texts = [[GridAnnotation(y_label(y))] for y in ys] hor_texts = [[GridAnnotation(x_label(x))] for x in xs] for y in ys: for x in xs: state.job = f"{x + y * len(xs)} out of {len(xs) * len(ys)}" res.append(cell(x, y)) grid = image_grid(res, rows=len(ys)) grid = draw_grid_annotations(grid, res[0].width, res[0].height, hor_texts, ver_texts) return grid def resize_image(resize_mode, im, width, height): if resize_mode == 0: res = im.resize((width, height), resample=LANCZOS) elif resize_mode == 1: ratio = width / height src_ratio = im.width / im.height src_w = width if ratio > src_ratio else im.width * height // im.height src_h = height if ratio <= src_ratio else im.height * width // im.width resized = im.resize((src_w, src_h), resample=LANCZOS) res = Image.new("RGB", (width, height)) res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2)) else: ratio = width / height src_ratio = im.width / im.height src_w = width if ratio < src_ratio else im.width * height // im.height src_h = height if ratio >= src_ratio else im.height * width // im.width resized = im.resize((src_w, src_h), resample=LANCZOS) res = Image.new("RGB", (width, height)) res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2)) if ratio < src_ratio: fill_height = height // 2 - src_h // 2 res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0)) res.paste(resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)), box=(0, fill_height + src_h)) elif ratio > src_ratio: fill_width = width // 2 - src_w // 2 res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0)) res.paste(resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)), box=(fill_width + src_w, 0)) return res def wrap_gradio_gpu_call(func): def f(*args, **kwargs): with queue_lock: res = func(*args, **kwargs) return res return wrap_gradio_call(f) def wrap_gradio_call(func): def f(*args, **kwargs): t = time.perf_counter() try: res = list(func(*args, **kwargs)) except Exception as e: print("Error completing request", file=sys.stderr) print("Arguments:", args, kwargs, file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) res = [None, '', f"Time taken: {elapsed:.2f}s
" state.interrupted = False return tuple(res) return f class StableDiffusionModelHijack: ids_lookup = {} word_embeddings = {} word_embeddings_checksums = {} fixes = None comments = [] dir_mtime = None def load_textual_inversion_embeddings(self, dirname, model): mt = os.path.getmtime(dirname) if self.dir_mtime is not None and mt <= self.dir_mtime: return self.dir_mtime = mt self.ids_lookup.clear() self.word_embeddings.clear() tokenizer = model.cond_stage_model.tokenizer def const_hash(a): r = 0 for v in a: r = (r * 281 ^ int(v) * 997) & 0xFFFFFFFF return r def process_file(path, filename): name = os.path.splitext(filename)[0] data = torch.load(path) param_dict = data['string_to_param'] if hasattr(param_dict, '_parameters'): param_dict = getattr(param_dict, '_parameters') # fix for torch 1.12.1 loading saved file from torch 1.11 assert len(param_dict) == 1, 'embedding file has multiple terms in it' emb = next(iter(param_dict.items()))[1].reshape(768) self.word_embeddings[name] = emb self.word_embeddings_checksums[name] = f'{const_hash(emb)&0xffff:04x}' ids = tokenizer([name], add_special_tokens=False)['input_ids'][0] first_id = ids[0] if first_id not in self.ids_lookup: self.ids_lookup[first_id] = [] self.ids_lookup[first_id].append((ids, name)) for fn in os.listdir(dirname): try: process_file(os.path.join(dirname, fn), fn) except Exception: print(f"Error loading emedding {fn}:", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) continue print(f"Loaded a total of {len(self.word_embeddings)} text inversion embeddings.") def hijack(self, m): model_embeddings = m.cond_stage_model.transformer.text_model.embeddings model_embeddings.token_embedding = EmbeddingsWithFixes(model_embeddings.token_embedding, self) m.cond_stage_model = FrozenCLIPEmbedderWithCustomWords(m.cond_stage_model, self) class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module): def __init__(self, wrapped, hijack): super().__init__() self.wrapped = wrapped self.hijack = hijack self.tokenizer = wrapped.tokenizer self.max_length = wrapped.max_length self.token_mults = {} tokens_with_parens = [(k, v) for k, v in self.tokenizer.get_vocab().items() if '(' in k or ')' in k or '[' in k or ']' in k] for text, ident in tokens_with_parens: mult = 1.0 for c in text: if c == '[': mult /= 1.1 if c == ']': mult *= 1.1 if c == '(': mult *= 1.1 if c == ')': mult /= 1.1 if mult != 1.0: self.token_mults[ident] = mult def forward(self, text): self.hijack.fixes = [] self.hijack.comments = [] remade_batch_tokens = [] id_start = self.wrapped.tokenizer.bos_token_id id_end = self.wrapped.tokenizer.eos_token_id maxlen = self.wrapped.max_length - 2 used_custom_terms = [] cache = {} batch_tokens = self.wrapped.tokenizer(text, truncation=False, add_special_tokens=False)["input_ids"] batch_multipliers = [] for tokens in batch_tokens: tuple_tokens = tuple(tokens) if tuple_tokens in cache: remade_tokens, fixes, multipliers = cache[tuple_tokens] else: fixes = [] remade_tokens = [] multipliers = [] mult = 1.0 i = 0 while i < len(tokens): token = tokens[i] possible_matches = self.hijack.ids_lookup.get(token, None) mult_change = self.token_mults.get(token) if opts.enable_emphasis else None if mult_change is not None: mult *= mult_change elif possible_matches is None: remade_tokens.append(token) multipliers.append(mult) else: found = False for ids, word in possible_matches: if tokens[i:i+len(ids)] == ids: fixes.append((len(remade_tokens), word)) remade_tokens.append(777) multipliers.append(mult) i += len(ids) - 1 found = True used_custom_terms.append((word, self.hijack.word_embeddings_checksums[word])) break if not found: remade_tokens.append(token) multipliers.append(mult) i += 1 if len(remade_tokens) > maxlen - 2: vocab = {v: k for k, v in self.wrapped.tokenizer.get_vocab().items()} ovf = remade_tokens[maxlen - 2:] overflowing_words = [vocab.get(int(x), "") for x in ovf] overflowing_text = self.wrapped.tokenizer.convert_tokens_to_string(''.join(overflowing_words)) self.hijack.comments.append(f"Warning: too many input tokens; some ({len(overflowing_words)}) have been truncated:\n{overflowing_text}\n") remade_tokens = remade_tokens + [id_end] * (maxlen - 2 - len(remade_tokens)) remade_tokens = [id_start] + remade_tokens[0:maxlen-2] + [id_end] cache[tuple_tokens] = (remade_tokens, fixes, multipliers) multipliers = multipliers + [1.0] * (maxlen - 2 - len(multipliers)) multipliers = [1.0] + multipliers[0:maxlen - 2] + [1.0] remade_batch_tokens.append(remade_tokens) self.hijack.fixes.append(fixes) batch_multipliers.append(multipliers) if len(used_custom_terms) > 0: self.hijack.comments.append("Used custom terms: " + ", ".join([f'{word} [{checksum}]' for word, checksum in used_custom_terms])) tokens = torch.asarray(remade_batch_tokens).to(device) outputs = self.wrapped.transformer(input_ids=tokens) z = outputs.last_hidden_state # restoring original mean is likely not correct, but it seems to work well to prevent artifacts that happen otherwise batch_multipliers = torch.asarray(np.array(batch_multipliers)).to(device) original_mean = z.mean() z *= batch_multipliers.reshape(batch_multipliers.shape + (1,)).expand(z.shape) new_mean = z.mean() z *= original_mean / new_mean return z class EmbeddingsWithFixes(nn.Module): def __init__(self, wrapped, embeddings): super().__init__() self.wrapped = wrapped self.embeddings = embeddings def forward(self, input_ids): batch_fixes = self.embeddings.fixes self.embeddings.fixes = None inputs_embeds = self.wrapped(input_ids) if batch_fixes is not None: for fixes, tensor in zip(batch_fixes, inputs_embeds): for offset, word in fixes: tensor[offset] = self.embeddings.word_embeddings[word] return inputs_embeds class StableDiffusionProcessing: def __init__(self, outpath_samples=None, outpath_grids=None, prompt="", seed=-1, sampler_index=0, batch_size=1, n_iter=1, steps=50, cfg_scale=7.0, width=512, height=512, prompt_matrix=False, use_GFPGAN=False, do_not_save_samples=False, do_not_save_grid=False, extra_generation_params=None, overlay_images=None, negative_prompt=None): self.outpath_samples: str = outpath_samples self.outpath_grids: str = outpath_grids self.prompt: str = prompt self.negative_prompt: str = (negative_prompt or "") self.seed: int = seed self.sampler_index: int = sampler_index self.batch_size: int = batch_size self.n_iter: int = n_iter self.steps: int = steps self.cfg_scale: float = cfg_scale self.width: int = width self.height: int = height self.prompt_matrix: bool = prompt_matrix self.use_GFPGAN: bool = use_GFPGAN self.do_not_save_samples: bool = do_not_save_samples self.do_not_save_grid: bool = do_not_save_grid self.extra_generation_params: dict = extra_generation_params self.overlay_images = overlay_images self.paste_to = None def init(self): pass def sample(self, x, conditioning, unconditional_conditioning): raise NotImplementedError() def p_sample_ddim_hook(sampler_wrapper, x_dec, cond, ts, *args, **kwargs): if sampler_wrapper.mask is not None: img_orig = sampler_wrapper.sampler.model.q_sample(sampler_wrapper.init_latent, ts) x_dec = img_orig * sampler_wrapper.mask + sampler_wrapper.nmask * x_dec return sampler_wrapper.orig_p_sample_ddim(x_dec, cond, ts, *args, **kwargs) class VanillaStableDiffusionSampler: def __init__(self, constructor): self.sampler = constructor(sd_model) self.orig_p_sample_ddim = self.sampler.p_sample_ddim if hasattr(self.sampler, 'p_sample_ddim') else None self.mask = None self.nmask = None self.init_latent = None def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning): t_enc = int(min(p.denoising_strength, 0.999) * p.steps) # existing code fails with cetin step counts, like 9 try: self.sampler.make_schedule(ddim_num_steps=p.steps, verbose=False) except Exception: self.sampler.make_schedule(ddim_num_steps=p.steps+1, verbose=False) x1 = self.sampler.stochastic_encode(x, torch.tensor([t_enc] * int(x.shape[0])).to(device), noise=noise) self.sampler.p_sample_ddim = lambda x_dec, cond, ts, *args, **kwargs: p_sample_ddim_hook(self, x_dec, cond, ts, *args, **kwargs) self.mask = p.mask self.nmask = p.nmask self.init_latent = p.init_latent samples = self.sampler.decode(x1, conditioning, t_enc, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning) return samples def sample(self, p: StableDiffusionProcessing, x, conditioning, unconditional_conditioning): samples_ddim, _ = self.sampler.sample(S=p.steps, conditioning=conditioning, batch_size=int(x.shape[0]), shape=x[0].shape, verbose=False, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning, x_T=x) return samples_ddim class CFGDenoiser(nn.Module): def __init__(self, model): super().__init__() self.inner_model = model self.mask = None self.nmask = None self.init_latent = None def forward(self, x, sigma, uncond, cond, cond_scale): if batch_cond_uncond: x_in = torch.cat([x] * 2) sigma_in = torch.cat([sigma] * 2) cond_in = torch.cat([uncond, cond]) uncond, cond = self.inner_model(x_in, sigma_in, cond=cond_in).chunk(2) denoised = uncond + (cond - uncond) * cond_scale else: uncond = self.inner_model(x, sigma, cond=uncond) cond = self.inner_model(x, sigma, cond=cond) denoised = uncond + (cond - uncond) * cond_scale if self.mask is not None: denoised = self.init_latent * self.mask + self.nmask * denoised return denoised def extended_trange(*args, **kwargs): for x in tqdm.trange(*args, desc=state.job, **kwargs): if state.interrupted: break yield x class KDiffusionSampler: def __init__(self, funcname): self.model_wrap = k_diffusion.external.CompVisDenoiser(sd_model) self.funcname = funcname self.func = getattr(k_diffusion.sampling, self.funcname) self.model_wrap_cfg = CFGDenoiser(self.model_wrap) def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning): t_enc = int(min(p.denoising_strength, 0.999) * p.steps) sigmas = self.model_wrap.get_sigmas(p.steps) noise = noise * sigmas[p.steps - t_enc - 1] xi = x + noise sigma_sched = sigmas[p.steps - t_enc - 1:] self.model_wrap_cfg.mask = p.mask self.model_wrap_cfg.nmask = p.nmask self.model_wrap_cfg.init_latent = p.init_latent if hasattr(k_diffusion.sampling, 'trange'): k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs) return self.func(self.model_wrap_cfg, xi, sigma_sched, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False) def sample(self, p: StableDiffusionProcessing, x, conditioning, unconditional_conditioning): sigmas = self.model_wrap.get_sigmas(p.steps) x = x * sigmas[0] if hasattr(k_diffusion.sampling, 'trange'): k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs) def cb(d): n = d['i'] img = d['denoised'] x_samples_ddim = sd_model.decode_first_stage(img) x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) for i, x_sample in enumerate(x_samples_ddim): x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2) x_sample = x_sample.astype(np.uint8) image = Image.fromarray(x_sample) image.save(f'a/{n}.png') samples_ddim = self.func(self.model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False) return samples_ddim class Processed: def __init__(self, p: StableDiffusionProcessing, images, seed, info): self.images = images self.prompt = p.prompt self.seed = seed self.info = info self.width = p.width self.height = p.height self.sampler = samplers[p.sampler_index].name self.cfg_scale = p.cfg_scale self.steps = p.steps def js(self): obj = { "prompt": self.prompt, "seed": int(self.seed), "width": self.width, "height": self.height, "sampler": self.sampler, "cfg_scale": self.cfg_scale, "steps": self.steps, } return json.dumps(obj) def process_images(p: StableDiffusionProcessing) -> Processed: """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch""" prompt = p.prompt model = sd_model assert p.prompt is not None torch_gc() seed = int(random.randrange(4294967294) if p.seed == -1 else p.seed) os.makedirs(p.outpath_samples, exist_ok=True) os.makedirs(p.outpath_grids, exist_ok=True) comments = [] prompt_matrix_parts = [] if p.prompt_matrix: all_prompts = [] prompt_matrix_parts = prompt.split("|") combination_count = 2 ** (len(prompt_matrix_parts) - 1) for combination_num in range(combination_count): selected_prompts = [text.strip().strip(',') for n, text in enumerate(prompt_matrix_parts[1:]) if combination_num & (1 << n)] if opts.prompt_matrix_add_to_start: selected_prompts = selected_prompts + [prompt_matrix_parts[0]] else: selected_prompts = [prompt_matrix_parts[0]] + selected_prompts all_prompts.append(", ".join(selected_prompts)) p.n_iter = math.ceil(len(all_prompts) / p.batch_size) all_seeds = len(all_prompts) * [seed] print(f"Prompt matrix will create {len(all_prompts)} images using a total of {p.n_iter} batches.") else: all_prompts = p.batch_size * p.n_iter * [prompt] all_seeds = [seed + x for x in range(len(all_prompts))] def infotext(iteration=0, position_in_batch=0): generation_params = { "Steps": p.steps, "Sampler": samplers[p.sampler_index].name, "CFG scale": p.cfg_scale, "Seed": all_seeds[position_in_batch + iteration * p.batch_size], "GFPGAN": ("GFPGAN" if p.use_GFPGAN else None) } if p.extra_generation_params is not None: generation_params.update(p.extra_generation_params) generation_params_text = ", ".join([k if k == v else f'{k}: {v}' for k, v in generation_params.items() if v is not None]) return f"{prompt}\n{generation_params_text}".strip() + "".join(["\n\n" + x for x in comments]) if os.path.exists(cmd_opts.embeddings_dir): model_hijack.load_textual_inversion_embeddings(cmd_opts.embeddings_dir, model) output_images = [] precision_scope = autocast if cmd_opts.precision == "autocast" else nullcontext ema_scope = (nullcontext if cmd_opts.lowvram else model.ema_scope) with torch.no_grad(), precision_scope("cuda"), ema_scope(): p.init() for n in range(p.n_iter): if state.interrupted: break prompts = all_prompts[n * p.batch_size:(n + 1) * p.batch_size] seeds = all_seeds[n * p.batch_size:(n + 1) * p.batch_size] uc = model.get_learned_conditioning(len(prompts) * [p.negative_prompt]) c = model.get_learned_conditioning(prompts) if len(model_hijack.comments) > 0: comments += model_hijack.comments # we manually generate all input noises because each one should have a specific seed x = create_random_tensors([opt_C, p.height // opt_f, p.width // opt_f], seeds=seeds) if p.n_iter > 1: state.job = f"Batch {n+1} out of {p.n_iter}" samples_ddim = p.sample(x=x, conditioning=c, unconditional_conditioning=uc) x_samples_ddim = model.decode_first_stage(samples_ddim) x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) for i, x_sample in enumerate(x_samples_ddim): x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2) x_sample = x_sample.astype(np.uint8) if p.use_GFPGAN: torch_gc() gfpgan_model = gfpgan() x_sample = gfpgan_fix_faces(gfpgan_model, x_sample) image = Image.fromarray(x_sample) if p.overlay_images is not None and i < len(p.overlay_images): overlay = p.overlay_images[i] if p.paste_to is not None: x, y, w, h = p.paste_to base_image = Image.new('RGBA', (overlay.width, overlay.height)) image = resize_image(1, image, w, h) base_image.paste(image, (x, y)) image = base_image image = image.convert('RGBA') image.alpha_composite(overlay) image = image.convert('RGB') if opts.samples_save and not p.do_not_save_samples: save_image(image, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i)) output_images.append(image) unwanted_grid_because_of_img_count = len(output_images) < 2 and opts.grid_only_if_multiple if not p.do_not_save_grid and not unwanted_grid_because_of_img_count: return_grid = opts.return_grid if p.prompt_matrix: grid = image_grid(output_images, p.batch_size, rows=1 << ((len(prompt_matrix_parts)-1)//2)) try: grid = draw_prompt_matrix(grid, p.width, p.height, prompt_matrix_parts) except Exception: import traceback print("Error creating prompt_matrix text:", file=sys.stderr) print(traceback.format_exc(), file=sys.stderr) return_grid = True else: grid = image_grid(output_images, p.batch_size) if return_grid: output_images.insert(0, grid) if opts.grid_save: save_image(grid, p.outpath_grids, "grid", seed, prompt, opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename) torch_gc() return Processed(p, output_images, seed, infotext()) class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): sampler = None def init(self): self.sampler = samplers[self.sampler_index].constructor() def sample(self, x, conditioning, unconditional_conditioning): samples_ddim = self.sampler.sample(self, x, conditioning, unconditional_conditioning) return samples_ddim def txt2img(prompt: str, negative_prompt: str, steps: int, sampler_index: int, use_GFPGAN: bool, prompt_matrix: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, height: int, width: int, code: str): p = StableDiffusionProcessingTxt2Img( outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples, outpath_grids=opts.outdir_grids or opts.outdir_txt2img_grids, prompt=prompt, negative_prompt=negative_prompt, seed=seed, sampler_index=sampler_index, batch_size=batch_size, n_iter=n_iter, steps=steps, cfg_scale=cfg_scale, width=width, height=height, prompt_matrix=prompt_matrix, use_GFPGAN=use_GFPGAN ) if code != '' and cmd_opts.allow_code: p.do_not_save_grid = True p.do_not_save_samples = True display_result_data = [[], -1, ""] def display(imgs, s=display_result_data[1], i=display_result_data[2]): display_result_data[0] = imgs display_result_data[1] = s display_result_data[2] = i from types import ModuleType compiled = compile(code, '', 'exec') module = ModuleType("testmodule") module.__dict__.update(globals()) module.p = p module.display = display exec(compiled, module.__dict__) processed = Processed(p, *display_result_data) else: processed = process_images(p) return processed.images, processed.js(), plaintext_to_html(processed.info) def image_from_url_text(filedata): if filedata.startswith("data:image/png;base64,"): filedata = filedata[len("data:image/png;base64,"):] filedata = base64.decodebytes(filedata.encode('utf-8')) image = Image.open(io.BytesIO(filedata)) return image def send_gradio_gallery_to_image(x): if len(x) == 0: return None return image_from_url_text(x[0]) def save_files(js_data, images): import csv os.makedirs(opts.outdir_save, exist_ok=True) filenames = [] data = json.loads(js_data) with open("log/log.csv", "a", encoding="utf8", newline='') as file: at_start = file.tell() == 0 writer = csv.writer(file) if at_start: writer.writerow(["prompt", "seed", "width", "height", "sampler", "cfgs", "steps", "filename"]) filename_base = str(int(time.time() * 1000)) for i, filedata in enumerate(images): filename = filename_base + ("" if len(images) == 1 else "-" + str(i + 1)) + ".png" filepath = os.path.join(opts.outdir_save, filename) if filedata.startswith("data:image/png;base64,"): filedata = filedata[len("data:image/png;base64,"):] with open(filepath, "wb") as imgfile: imgfile.write(base64.decodebytes(filedata.encode('utf-8'))) filenames.append(filename) writer.writerow([data["prompt"], data["seed"], data["width"], data["height"], data["sampler"], data["cfg_scale"], data["steps"], filenames[0]]) return '', '', plaintext_to_html(f"Saved: {filenames[0]}") with gr.Blocks(analytics_enabled=False) as txt2img_interface: with gr.Row(): prompt = gr.Textbox(label="Prompt", elem_id="txt2img_prompt", show_label=False, placeholder="Prompt", lines=1) negative_prompt = gr.Textbox(label="Negative prompt", elem_id="txt2img_negative_prompt", show_label=False, placeholder="Negative prompt", lines=1, visible=False) submit = gr.Button('Generate', elem_id="txt2img_generate", variant='primary') with gr.Row().style(equal_height=False): with gr.Column(variant='panel'): steps = gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20) sampler_index = gr.Radio(label='Sampling method', elem_id="txt2img_sampling", choices=[x.name for x in samplers], value=samplers[0].name, type="index") with gr.Row(): use_GFPGAN = gr.Checkbox(label='GFPGAN', value=False, visible=have_gfpgan) prompt_matrix = gr.Checkbox(label='Prompt matrix', value=False) with gr.Row(): batch_count = gr.Slider(minimum=1, maximum=cmd_opts.max_batch_count, step=1, label='Batch count', value=1) batch_size = gr.Slider(minimum=1, maximum=8, step=1, label='Batch size', value=1) cfg_scale = gr.Slider(minimum=1.0, maximum=15.0, step=0.5, label='CFG Scale', value=7.0) with gr.Group(): height = gr.Slider(minimum=64, maximum=2048, step=64, label="Height", value=512) width = gr.Slider(minimum=64, maximum=2048, step=64, label="Width", value=512) seed = gr.Number(label='Seed', value=-1) code = gr.Textbox(label="Python script", visible=cmd_opts.allow_code, lines=1) with gr.Column(variant='panel'): with gr.Group(): txt2img_gallery = gr.Gallery(label='Output') with gr.Group(): with gr.Row(): save = gr.Button('Save') send_to_img2img = gr.Button('Send to img2img') send_to_inpaint = gr.Button('Send to inpaint') send_to_extras = gr.Button('Send to extras') interrupt = gr.Button('Interrupt') with gr.Group(): html_info = gr.HTML() generation_info = gr.Textbox(visible=False) txt2img_args = dict( fn=wrap_gradio_gpu_call(txt2img), inputs=[ prompt, negative_prompt, steps, sampler_index, use_GFPGAN, prompt_matrix, batch_count, batch_size, cfg_scale, seed, height, width, code ], outputs=[ txt2img_gallery, generation_info, html_info ] ) prompt.submit(**txt2img_args) submit.click(**txt2img_args) interrupt.click( fn=lambda: state.interrupt(), inputs=[], outputs=[], ) save.click( fn=wrap_gradio_call(save_files), inputs=[ generation_info, txt2img_gallery, ], outputs=[ html_info, html_info, html_info, ] ) def get_crop_region(mask, pad=0): h, w = mask.shape crop_left = 0 for i in range(w): if not (mask[:,i] == 0).all(): break crop_left += 1 crop_right = 0 for i in reversed(range(w)): if not (mask[:,i] == 0).all(): break crop_right += 1 crop_top = 0 for i in range(h): if not (mask[i] == 0).all(): break crop_top += 1 crop_bottom = 0 for i in reversed(range(h)): if not (mask[i] == 0).all(): break crop_bottom += 1 return ( int(max(crop_left-pad, 0)), int(max(crop_top-pad, 0)), int(min(w - crop_right + pad, w)), int(min(h - crop_bottom + pad, h)) ) def fill(image, mask): image_mod = Image.new('RGBA', (image.width, image.height)) image_masked = Image.new('RGBa', (image.width, image.height)) image_masked.paste(image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(mask.convert('L'))) image_masked = image_masked.convert('RGBa') for radius, repeats in [(64, 1), (16, 2), (4, 4), (2, 2), (0, 1)]: blurred = image_masked.filter(ImageFilter.GaussianBlur(radius)).convert('RGBA') for _ in range(repeats): image_mod.alpha_composite(blurred) return image_mod.convert("RGB") class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): sampler = None def __init__(self, init_images=None, resize_mode=0, denoising_strength=0.75, mask=None, mask_blur=4, inpainting_fill=0, inpaint_full_res=True, **kwargs): super().__init__(**kwargs) self.init_images = init_images self.resize_mode: int = resize_mode self.denoising_strength: float = denoising_strength self.init_latent = None self.image_mask = mask self.mask_for_overlay = None self.mask_blur = mask_blur self.inpainting_fill = inpainting_fill self.inpaint_full_res = inpaint_full_res self.mask = None self.nmask = None def init(self): self.sampler = samplers_for_img2img[self.sampler_index].constructor() crop_region = None if self.image_mask is not None: if self.mask_blur > 0: self.image_mask = self.image_mask.filter(ImageFilter.GaussianBlur(self.mask_blur)).convert('L') if self.inpaint_full_res: self.mask_for_overlay = self.image_mask mask = self.image_mask.convert('L') crop_region = get_crop_region(np.array(mask), 64) x1, y1, x2, y2 = crop_region mask = mask.crop(crop_region) self.image_mask = resize_image(2, mask, self.width, self.height) self.paste_to = (x1, y1, x2-x1, y2-y1) else: self.image_mask = resize_image(self.resize_mode, self.image_mask, self.width, self.height) self.mask_for_overlay = self.image_mask self.overlay_images = [] imgs = [] for img in self.init_images: image = img.convert("RGB") if crop_region is None: image = resize_image(self.resize_mode, image, self.width, self.height) if self.image_mask is not None: if self.inpainting_fill != 1: image = fill(image, self.mask_for_overlay) image_masked = Image.new('RGBa', (image.width, image.height)) image_masked.paste(image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(self.mask_for_overlay.convert('L'))) self.overlay_images.append(image_masked.convert('RGBA')) if crop_region is not None: image = image.crop(crop_region) image = resize_image(2, image, self.width, self.height) image = np.array(image).astype(np.float32) / 255.0 image = np.moveaxis(image, 2, 0) imgs.append(image) if len(imgs) == 1: batch_images = np.expand_dims(imgs[0], axis=0).repeat(self.batch_size, axis=0) if self.overlay_images is not None: self.overlay_images = self.overlay_images * self.batch_size elif len(imgs) <= self.batch_size: self.batch_size = len(imgs) batch_images = np.array(imgs) else: raise RuntimeError(f"bad number of images passed: {len(imgs)}; expecting {self.batch_size} or less") image = torch.from_numpy(batch_images) image = 2. * image - 1. image = image.to(device) self.init_latent = sd_model.get_first_stage_encoding(sd_model.encode_first_stage(image)) if self.image_mask is not None: latmask = self.image_mask.convert('RGB').resize((self.init_latent.shape[3], self.init_latent.shape[2])) latmask = np.moveaxis(np.array(latmask, dtype=np.float64), 2, 0) / 255 latmask = latmask[0] latmask = np.tile(latmask[None], (4, 1, 1)) self.mask = torch.asarray(1.0 - latmask).to(device).type(sd_model.dtype) self.nmask = torch.asarray(latmask).to(device).type(sd_model.dtype) if self.inpainting_fill == 2: self.init_latent = self.init_latent * self.mask + create_random_tensors(self.init_latent.shape[1:], [self.seed + x + 1 for x in range(self.init_latent.shape[0])]) * self.nmask elif self.inpainting_fill == 3: self.init_latent = self.init_latent * self.mask def sample(self, x, conditioning, unconditional_conditioning): samples = self.sampler.sample_img2img(self, self.init_latent, x, conditioning, unconditional_conditioning) if self.mask is not None: samples = samples * self.nmask + self.init_latent * self.mask return samples def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index: int, mask_blur: int, inpainting_fill: int, use_GFPGAN: bool, prompt_matrix, mode: int, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, height: int, width: int, resize_mode: int, upscaler_name: str, upscale_overlap: int, inpaint_full_res: bool): is_classic = mode == 0 is_inpaint = mode == 1 is_loopback = mode == 2 is_upscale = mode == 3 if is_inpaint: image = init_img_with_mask['image'] mask = init_img_with_mask['mask'] else: image = init_img mask = None assert 0. <= denoising_strength <= 1., 'can only work with strength in [0.0, 1.0]' p = StableDiffusionProcessingImg2Img( outpath_samples=opts.outdir_samples or opts.outdir_img2img_samples, outpath_grids=opts.outdir_grids or opts.outdir_img2img_grids, prompt=prompt, seed=seed, sampler_index=sampler_index, batch_size=batch_size, n_iter=n_iter, steps=steps, cfg_scale=cfg_scale, width=width, height=height, prompt_matrix=prompt_matrix, use_GFPGAN=use_GFPGAN, init_images=[image], mask=mask, mask_blur=mask_blur, inpainting_fill=inpainting_fill, resize_mode=resize_mode, denoising_strength=denoising_strength, inpaint_full_res=inpaint_full_res, extra_generation_params={"Denoising Strength": denoising_strength} ) if is_loopback: output_images, info = None, None history = [] initial_seed = None initial_info = None for i in range(n_iter): p.n_iter = 1 p.batch_size = 1 p.do_not_save_grid = True state.job = f"Batch {i + 1} out of {n_iter}" processed = process_images(p) if initial_seed is None: initial_seed = processed.seed initial_info = processed.info p.init_images = [processed.images[0]] p.seed = processed.seed + 1 p.denoising_strength = max(p.denoising_strength * 0.95, 0.1) history.append(processed.images[0]) grid = image_grid(history, batch_size, rows=1) save_image(grid, p.outpath_grids, "grid", initial_seed, prompt, opts.grid_format, info=info, short_filename=not opts.grid_extended_filename) processed = Processed(p, history, initial_seed, initial_info) elif is_upscale: initial_seed = None initial_info = None upscaler = sd_upscalers.get(upscaler_name, next(iter(sd_upscalers.values()))) img = upscaler(init_img) torch_gc() grid = split_grid(img, tile_w=width, tile_h=height, overlap=upscale_overlap) p.n_iter = 1 p.do_not_save_grid = True p.do_not_save_samples = True work = [] work_results = [] for y, h, row in grid.tiles: for tiledata in row: work.append(tiledata[2]) batch_count = math.ceil(len(work) / p.batch_size) print(f"SD upscaling will process a total of {len(work)} images tiled as {len(grid.tiles[0][2])}x{len(grid.tiles)} in a total of {batch_count} batches.") for i in range(batch_count): p.init_images = work[i*p.batch_size:(i+1)*p.batch_size] state.job = f"Batch {i + 1} out of {batch_count}" processed = process_images(p) if initial_seed is None: initial_seed = processed.seed initial_info = processed.info p.seed = processed.seed + 1 work_results += processed.images image_index = 0 for y, h, row in grid.tiles: for tiledata in row: tiledata[2] = work_results[image_index] if image_index{plaintext_to_html(str(key))}
{plaintext_to_html(str(text))}
{message}