From 1db7d212836e0fd8a4eff6922c13a54a372175b2 Mon Sep 17 00:00:00 2001 From: "fumitaka.yano" Date: Tue, 23 May 2023 15:56:08 +0900 Subject: [PATCH] Subject:. Improvements to handle VAE filenames in generated image filenames Body:. 1) Added new line 24 to import sd_vae module. 2) Added new method get_vae_filename at lines 340-349 to obtain the VAE filename to be used for image generation and further process it to extract only the filename by splitting it with a dot symbol. 3) Added a new lambda function 'vae_filename' at line 373 to handle VAE filenames. Reason:. A function was needed to get the VAE filename and handle it in the program. Test:. We tested whether we could use this new functionality to get the expected file names. The correct behaviour was confirmed for the following commonly distributed VAE files. vae-ft-mse-840000-ema-pruned.safetensors -> vae-ft-mse-840000-ema-pruned anything-v4.0.vae.pt -> anything-v4.0 ruff response:. There were no problems with the code I added. There was a minor configuration error in a line I did not modify, but I did not modify it as it was not relevant to this modification. Logged. images.py:426:56: F841 [*] Local variable `_` is assigned to but never used images.py:432:43: F841 [*] Local variable `_` is assigned to but never used Impact:. This change makes it easier to retrieve the VAE filename used for image generation and use it in the programme. --- modules/images.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/modules/images.py b/modules/images.py index a41965ab6..3abaf4121 100644 --- a/modules/images.py +++ b/modules/images.py @@ -21,6 +21,8 @@ import hashlib from modules import sd_samplers, shared, script_callbacks, errors from modules.shared import opts, cmd_opts +import modules.sd_vae as sd_vae + LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS) @@ -335,6 +337,16 @@ def sanitize_filename_part(text, replace_spaces=True): class FilenameGenerator: + def get_vae_filename(self): #get the name of the VAE file. + if sd_vae.loaded_vae_file is None: + return "NoneType" + file_name = os.path.basename(sd_vae.loaded_vae_file) + split_file_name = file_name.split('.') + if len(split_file_name) > 1 and split_file_name[0] == '': + return split_file_name[1] # if the first character of the filename is "." then [1] is obtained. + else: + return split_file_name[0] + replacements = { 'seed': lambda self: self.seed if self.seed is not None else '', 'steps': lambda self: self.p and self.p.steps, @@ -358,6 +370,8 @@ class FilenameGenerator: 'hasprompt': lambda self, *args: self.hasprompt(*args), # accepts formats:[hasprompt..] 'clip_skip': lambda self: opts.data["CLIP_stop_at_last_layers"], 'denoising': lambda self: self.p.denoising_strength if self.p and self.p.denoising_strength else NOTHING_AND_SKIP_PREVIOUS_TEXT, + 'vae_filename': lambda self: self.get_vae_filename(), + } default_time_format = '%Y%m%d%H%M%S'