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.

1 files changed, 14 insertions, 0 deletions

diff --git a/modules/images.py b/modules/images.py
index a41965ab..3abaf412 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<prompt1|default><prompt2>..]

         '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'