From da06e685c1bbf8417695897efa5418716be26402 Mon Sep 17 00:00:00 2001 From: veni-vivi-code Date: Fri, 9 Sep 2022 14:32:58 +0200 Subject: [PATCH 1/2] Adds Exif and PNG Info to all scripts --- optimizedSD/img2img_gradio.py | 12 ++++++++---- optimizedSD/inpaint_gradio.py | 11 +++++++---- optimizedSD/optimUtils.py | 31 +++++++++++++++++++++++++++++-- optimizedSD/optimized_img2img.py | 13 +++++++++---- optimizedSD/optimized_txt2img.py | 13 +++++++++---- optimizedSD/txt2img_gradio.py | 15 ++++++++------- 6 files changed, 70 insertions(+), 25 deletions(-) diff --git a/optimizedSD/img2img_gradio.py b/optimizedSD/img2img_gradio.py index 65d844d3b..de1f2ee58 100644 --- a/optimizedSD/img2img_gradio.py +++ b/optimizedSD/img2img_gradio.py @@ -21,7 +21,8 @@ from ldm.util import instantiate_from_config from transformers import logging import pandas as pd -from optimUtils import split_weighted_subprompts, logger +from optimUtils import split_weighted_subprompts, logger, create_exif_info + logging.set_verbosity_error() import mimetypes mimetypes.init() @@ -221,9 +222,12 @@ def generate( x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) all_samples.append(x_sample.to("cpu")) x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), "c h w -> h w c") - Image.fromarray(x_sample.astype(np.uint8)).save( - os.path.join(sample_path, "seed_" + str(seed) + "_" + f"{base_count:05}.{img_format}") - ) + img = Image.fromarray(x_sample.astype(np.uint8)) + target_path = os.path.join(sample_path, "seed_" + str(seed) + "_" + + f"{base_count:05}.{img_format}") + info, exif = create_exif_info(img, seed, prompt, ddim_steps, sampler, scale, + full_precision, batch_size, i, "text2img_gradio.py") + img.save(target_path, pnginfo=info, exif=exif) seeds += str(seed) + "," seed += 1 base_count += 1 diff --git a/optimizedSD/inpaint_gradio.py b/optimizedSD/inpaint_gradio.py index d42930244..4a13e1943 100644 --- a/optimizedSD/inpaint_gradio.py +++ b/optimizedSD/inpaint_gradio.py @@ -19,7 +19,7 @@ from transformers import logging from ldm.util import instantiate_from_config -from optimUtils import split_weighted_subprompts, logger +from optimUtils import split_weighted_subprompts, logger, create_exif_info logging.set_verbosity_error() import mimetypes @@ -220,9 +220,12 @@ def generate( x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) all_samples.append(x_sample.to("cpu")) x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), "c h w -> h w c") - Image.fromarray(x_sample.astype(np.uint8)).save( - os.path.join(sample_path, "seed_" + str(seed) + "_" + f"{base_count:05}.{img_format}") - ) + img = Image.fromarray(x_sample.astype(np.uint8)) + target_path = os.path.join(sample_path, "seed_" + str(seed) + "_" + + f"{base_count:05}.{img_format}") + info, exif = create_exif_info(img, seed, prompt, ddim_steps, sampler, scale, + full_precision, batch_size, i, "text2img_gradio.py") + img.save(target_path, pnginfo=info, exif=exif) seeds += str(seed) + "," seed += 1 base_count += 1 diff --git a/optimizedSD/optimUtils.py b/optimizedSD/optimUtils.py index 18b996792..1e49e93c4 100644 --- a/optimizedSD/optimUtils.py +++ b/optimizedSD/optimUtils.py @@ -1,5 +1,7 @@ import os import pandas as pd +from PIL import Image +from PIL.PngImagePlugin import PngInfo def split_weighted_subprompts(text): @@ -51,7 +53,7 @@ def logger(params, log_csv): os.makedirs('logs', exist_ok=True) cols = [arg for arg, _ in params.items()] if not os.path.exists(log_csv): - df = pd.DataFrame(columns=cols) + df = pd.DataFrame(columns=cols) df.to_csv(log_csv, index=False) df = pd.read_csv(log_csv) @@ -70,4 +72,29 @@ def logger(params, log_csv): li[col] = '' df = pd.DataFrame(li,index = [0]) - df.to_csv(log_csv,index=False, mode='a', header=False) \ No newline at end of file + df.to_csv(log_csv,index=False, mode='a', header=False) + + +def create_exif_info(img, seed, prompt, ddim_steps, sampler, scale, full_precision, batch_size, + batch_index, scriptname): + info = PngInfo() + info.add_itxt("Software", f"Stable Diffusion (github.com/basujindal/stable-diffusion), " + f"{scriptname}", zip=True) + info.add_itxt("Seed", str(seed)) + info.add_itxt("Prompt", prompt, zip=True) + info.add_itxt("ddim_steps", str(ddim_steps)) + info.add_itxt("sampler", str(sampler)) + info.add_itxt("Scale", str(scale)) + info.add_itxt("FullPrecision", str(full_precision)) + info.add_itxt("Batch Size", str(batch_size)) + info.add_itxt("Batch Index", str(batch_index)) + exif = img.getexif() + exif[0xA401] = 1 # Yes, this has had special processing on it + exif[0xA40B] = f"Seed: {seed}, ddim_steps: {ddim_steps}, sampler: {sampler}, scale: {scale}, " \ + f"precision: {str(full_precision)}, batchsize: {batch_size}, batchindex: {batch_index}, " \ + f"prompt: {prompt}" + exif[0x927C] = exif[0xA40B] + exif[0x9286] = exif[0xA40B] + exif[0x010E] = prompt + exif[0x0131] = f"Stable Diffusion (github.com/basujindal/stable-diffusion) {scriptname}" + return info, exif diff --git a/optimizedSD/optimized_img2img.py b/optimizedSD/optimized_img2img.py index 76e6419fd..a936e6878 100644 --- a/optimizedSD/optimized_img2img.py +++ b/optimizedSD/optimized_img2img.py @@ -14,7 +14,7 @@ from contextlib import contextmanager, nullcontext from einops import rearrange, repeat from ldm.util import instantiate_from_config -from optimUtils import split_weighted_subprompts, logger +from optimUtils import split_weighted_subprompts, logger, create_exif_info from transformers import logging import pandas as pd logging.set_verbosity_error() @@ -243,7 +243,8 @@ def load_img(path, h0, w0): else: print(f"reading prompts from {opt.from_file}") with open(opt.from_file, "r") as f: - data = f.read().splitlines() + prompt = f.read() + data = prompt.splitlines() data = batch_size * list(data) data = list(chunk(sorted(data), batch_size)) @@ -332,8 +333,12 @@ def load_img(path, h0, w0): x_samples_ddim = modelFS.decode_first_stage(samples_ddim[i].unsqueeze(0)) x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), "c h w -> h w c") - Image.fromarray(x_sample.astype(np.uint8)).save( - os.path.join(sample_path, "seed_" + str(opt.seed) + "_" + f"{base_count:05}.{opt.format}") + img = Image.fromarray(x_sample.astype(np.uint8)) + info, exif = create_exif_info(img, opt.seed, prompt, opt.ddim_steps, opt.sampler, opt.scale, + opt.precision, batch_size, i, "optimized_img2img.py") + img.save( + os.path.join(sample_path, "seed_" + str(opt.seed) + "_" + f"{base_count:05}.{opt.format}"), + pnginfo=info ) seeds += str(opt.seed) + "," opt.seed += 1 diff --git a/optimizedSD/optimized_txt2img.py b/optimizedSD/optimized_txt2img.py index 61d11316d..d71e4f94b 100644 --- a/optimizedSD/optimized_txt2img.py +++ b/optimizedSD/optimized_txt2img.py @@ -13,7 +13,7 @@ from torch import autocast from contextlib import contextmanager, nullcontext from ldm.util import instantiate_from_config -from optimUtils import split_weighted_subprompts, logger +from optimUtils import split_weighted_subprompts, logger, create_exif_info from transformers import logging # from samplers import CompVisDenoiser logging.set_verbosity_error() @@ -237,7 +237,8 @@ def load_model_from_config(ckpt, verbose=False): else: print(f"reading prompts from {opt.from_file}") with open(opt.from_file, "r") as f: - data = f.read().splitlines() + prompt = f.read() + data = prompt.splitlines() data = batch_size * list(data) data = list(chunk(sorted(data), batch_size)) @@ -309,8 +310,12 @@ def load_model_from_config(ckpt, verbose=False): x_samples_ddim = modelFS.decode_first_stage(samples_ddim[i].unsqueeze(0)) x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), "c h w -> h w c") - Image.fromarray(x_sample.astype(np.uint8)).save( - os.path.join(sample_path, "seed_" + str(opt.seed) + "_" + f"{base_count:05}.{opt.format}") + img = Image.fromarray(x_sample.astype(np.uint8)) + info, exif = create_exif_info(img, opt.seed, prompt, opt.ddim_steps, opt.sampler, opt.scale, + opt.precision, batch_size, i, "optimized_txt2img.py") + img.save( + os.path.join(sample_path, "seed_" + str(opt.seed) + "_" + f"{base_count:05}.{opt.format}"), + pnginfo=info ) seeds += str(opt.seed) + "," opt.seed += 1 diff --git a/optimizedSD/txt2img_gradio.py b/optimizedSD/txt2img_gradio.py index d9775735a..e8f1b3ef0 100644 --- a/optimizedSD/txt2img_gradio.py +++ b/optimizedSD/txt2img_gradio.py @@ -20,7 +20,7 @@ from torch import autocast from contextlib import nullcontext from ldm.util import instantiate_from_config -from optimUtils import split_weighted_subprompts, logger +from optimUtils import split_weighted_subprompts, logger, create_exif_info from transformers import logging logging.set_verbosity_error() import mimetypes @@ -95,7 +95,6 @@ def generate( full_precision, sampler, ): - C = 4 f = 8 start_code = None @@ -178,20 +177,22 @@ def generate( unconditional_conditioning=uc, eta=ddim_eta, x_T=start_code, - sampler = sampler, + sampler=sampler, ) modelFS.to(device) print("saving images") for i in range(batch_size): - x_samples_ddim = modelFS.decode_first_stage(samples_ddim[i].unsqueeze(0)) x_sample = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0) all_samples.append(x_sample.to("cpu")) x_sample = 255.0 * rearrange(x_sample[0].cpu().numpy(), "c h w -> h w c") - Image.fromarray(x_sample.astype(np.uint8)).save( - os.path.join(sample_path, "seed_" + str(seed) + "_" + f"{base_count:05}.{img_format}") - ) + img = Image.fromarray(x_sample.astype(np.uint8)) + target_path = os.path.join(sample_path, "seed_" + str(seed) + "_" + + f"{base_count:05}.{img_format}") + info, exif = create_exif_info(img, seed, prompt, ddim_steps, sampler, scale, + full_precision, batch_size, i, "text2img_gradio.py") + img.save(target_path, pnginfo=info, exif=exif) seeds += str(seed) + "," seed += 1 base_count += 1 From d148b1d06d75a0437e1e5b6106eb6a0af239e66d Mon Sep 17 00:00:00 2001 From: veni-vivi-code Date: Fri, 9 Sep 2022 14:49:12 +0200 Subject: [PATCH 2/2] fix: Wrong script titles in exifs --- optimizedSD/img2img_gradio.py | 2 +- optimizedSD/inpaint_gradio.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/optimizedSD/img2img_gradio.py b/optimizedSD/img2img_gradio.py index de1f2ee58..217ea9f80 100644 --- a/optimizedSD/img2img_gradio.py +++ b/optimizedSD/img2img_gradio.py @@ -226,7 +226,7 @@ def generate( target_path = os.path.join(sample_path, "seed_" + str(seed) + "_" + f"{base_count:05}.{img_format}") info, exif = create_exif_info(img, seed, prompt, ddim_steps, sampler, scale, - full_precision, batch_size, i, "text2img_gradio.py") + full_precision, batch_size, i, "img2img_gradio.py") img.save(target_path, pnginfo=info, exif=exif) seeds += str(seed) + "," seed += 1 diff --git a/optimizedSD/inpaint_gradio.py b/optimizedSD/inpaint_gradio.py index 4a13e1943..fd0d532a0 100644 --- a/optimizedSD/inpaint_gradio.py +++ b/optimizedSD/inpaint_gradio.py @@ -224,7 +224,7 @@ def generate( target_path = os.path.join(sample_path, "seed_" + str(seed) + "_" + f"{base_count:05}.{img_format}") info, exif = create_exif_info(img, seed, prompt, ddim_steps, sampler, scale, - full_precision, batch_size, i, "text2img_gradio.py") + full_precision, batch_size, i, "inpaint_gradio.py") img.save(target_path, pnginfo=info, exif=exif) seeds += str(seed) + "," seed += 1