Include the model name (or the SHA256 of the file) in the metadata #271

This commit is contained in:
AUTOMATIC 2022-09-12 20:47:46 +03:00
parent 35a4649c9e
commit 3de44fc580
4 changed files with 16 additions and 3 deletions

View File

@ -9,7 +9,7 @@ from fonts.ttf import Roboto
import string
import modules.shared
from modules import sd_samplers
from modules import sd_samplers, shared
from modules.shared import opts
LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
@ -278,6 +278,8 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
file_decoration = file_decoration.replace("[height]", str(p.height))
file_decoration = file_decoration.replace("[sampler]", sd_samplers.samplers[p.sampler_index].name)
file_decoration = file_decoration.replace("[model_hash]", shared.sd_model_hash)
if extension == 'png' and opts.enable_pnginfo and info is not None:
pnginfo = PngImagePlugin.PngInfo()

View File

@ -188,6 +188,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
"Seed": all_seeds[index],
"Face restoration": (opts.face_restoration_model if p.restore_faces else None),
"Size": f"{p.width}x{p.height}",
"Model hash": (None if not opts.add_model_hash_to_info or not shared.sd_model_hash else shared.sd_model_hash),
"Batch size": (None if p.batch_size < 2 else p.batch_size),
"Batch pos": (None if p.batch_size < 2 else position_in_batch),
"Variation seed": (None if p.subseed_strength == 0 else all_subseeds[index]),

View File

@ -97,7 +97,7 @@ class Options:
data = None
hide_dirs = {"visible": False} if cmd_opts.hide_ui_dir_config else None
data_labels = {
"samples_filename_format": OptionInfo("", "Samples filename format using following tags: [steps],[cfg],[prompt],[prompt_spaces],[width],[height],[sampler],[seed]. Leave blank for default."),
"samples_filename_format": OptionInfo("", "Samples filename format using following tags: [steps], [cfg], [prompt], [prompt_spaces], [width], [height], [sampler], [seed], [model_hash]. Leave blank for default."),
"outdir_samples": OptionInfo("", "Output directory for images; if empty, defaults to two directories below", component_args=hide_dirs),
"outdir_txt2img_samples": OptionInfo("outputs/txt2img-images", 'Output directory for txt2img images', component_args=hide_dirs),
"outdir_img2img_samples": OptionInfo("outputs/img2img-images", 'Output directory for img2img images', component_args=hide_dirs),
@ -120,6 +120,7 @@ class Options:
"jpeg_quality": OptionInfo(80, "Quality for saved jpeg images", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}),
"export_for_4chan": OptionInfo(True, "If PNG image is larger than 4MB or any dimension is larger than 4000, downscale and save copy as JPG"),
"enable_pnginfo": OptionInfo(True, "Save text information about generation parameters as chunks to png files"),
"add_model_hash_to_info": OptionInfo(False, "Add model hash to generation information"),
"font": OptionInfo("", "Font for image grids that have text"),
"enable_emphasis": OptionInfo(True, "Use (text) to make model pay more attention to text text and [text] to make it pay less attention"),
"save_txt": OptionInfo(False, "Create a text file next to every image with generation parameters."),
@ -178,6 +179,7 @@ if os.path.exists(config_filename):
sd_upscalers = []
sd_model = None
sd_model_hash = ''
progress_print_out = sys.stdout

View File

@ -35,7 +35,7 @@ realesrgan.setup_realesrgan()
def load_model_from_config(config, ckpt, verbose=False):
print(f"Loading model from {ckpt}")
print(f"Loading model [{shared.sd_model_hash}] from {ckpt}")
pl_sd = torch.load(ckpt, map_location="cpu")
if "global_step" in pl_sd:
print(f"Global Step: {pl_sd['global_step']}")
@ -89,6 +89,14 @@ try:
except Exception:
pass
with open(cmd_opts.ckpt, "rb") as file:
import hashlib
m = hashlib.sha256()
file.seek(0x100000)
m.update(file.read(0x10000))
shared.sd_model_hash = m.hexdigest()[0:8]
sd_config = OmegaConf.load(cmd_opts.config)
shared.sd_model = load_model_from_config(sd_config, cmd_opts.ckpt)
shared.sd_model = (shared.sd_model if cmd_opts.no_half else shared.sd_model.half())