Use one counter across full save path
This commit is contained in:
+5
-5
@@ -1,7 +1,7 @@
|
|||||||
from .dumas_image_nodes import (
|
from .dumas_image_nodes import (
|
||||||
NODE_CLASS_MAPPINGS as IMAGE_NODE_CLASS_MAPPINGS,
|
NODE_CLASS_MAPPINGS as IMAGE_NODE_CLASS_MAPPINGS,
|
||||||
NODE_DISPLAY_NAME_MAPPINGS as IMAGE_NODE_DISPLAY_NAME_MAPPINGS,
|
NODE_DISPLAY_NAME_MAPPINGS as IMAGE_NODE_DISPLAY_NAME_MAPPINGS,
|
||||||
_next_counter,
|
_next_counter_for_relative_path,
|
||||||
resolve_serve_token,
|
resolve_serve_token,
|
||||||
)
|
)
|
||||||
from .dumas_json_nodes import (
|
from .dumas_json_nodes import (
|
||||||
@@ -102,10 +102,10 @@ try:
|
|||||||
parts = [part for part in name.replace("\\", "/").split("/") if part]
|
parts = [part for part in name.replace("\\", "/").split("/") if part]
|
||||||
if not parts:
|
if not parts:
|
||||||
return 1, ""
|
return 1, ""
|
||||||
frame_dir = os.path.join(base, *parts[:-1])
|
relative_template = "/".join(parts)
|
||||||
counter = _next_counter(frame_dir, parts[-1])
|
counter = _next_counter_for_relative_path(base, relative_template)
|
||||||
resolved_name = parts[-1].replace("%counter%", str(counter).zfill(digits))
|
resolved_relative = relative_template.replace("%counter%", str(counter).zfill(digits))
|
||||||
return counter, "/".join(parts[:-1] + [resolved_name])
|
return counter, resolved_relative
|
||||||
|
|
||||||
try:
|
try:
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
|
|||||||
+20
-4
@@ -136,6 +136,20 @@ def _next_counter(directory, filename_template):
|
|||||||
return highest + 1
|
return highest + 1
|
||||||
|
|
||||||
|
|
||||||
|
def _next_counter_for_relative_path(base_directory, relative_template):
|
||||||
|
os.makedirs(base_directory, exist_ok=True)
|
||||||
|
if "%counter%" not in relative_template:
|
||||||
|
return 1
|
||||||
|
|
||||||
|
counter = 1
|
||||||
|
while True:
|
||||||
|
candidate = relative_template.replace("%counter%", str(counter).zfill(3))
|
||||||
|
full_path = os.path.join(base_directory, *[part for part in candidate.split("/") if part])
|
||||||
|
if not os.path.exists(full_path):
|
||||||
|
return counter
|
||||||
|
counter += 1
|
||||||
|
|
||||||
|
|
||||||
def _build_pnginfo(prompt=None, extra_pnginfo=None):
|
def _build_pnginfo(prompt=None, extra_pnginfo=None):
|
||||||
try:
|
try:
|
||||||
pnginfo = Image.PngImagePlugin.PngInfo()
|
pnginfo = Image.PngImagePlugin.PngInfo()
|
||||||
@@ -362,11 +376,13 @@ class DumasSaveImageNode:
|
|||||||
for batch_index in range(images.shape[0]):
|
for batch_index in range(images.shape[0]):
|
||||||
frame_pattern = resolved_pattern.replace("%batch_num%", str(batch_index))
|
frame_pattern = resolved_pattern.replace("%batch_num%", str(batch_index))
|
||||||
frame_parts = [part for part in frame_pattern.split("/") if part]
|
frame_parts = [part for part in frame_pattern.split("/") if part]
|
||||||
sub_dirs = frame_parts[:-1]
|
relative_template = "/".join(frame_parts[:-1] + [((frame_parts[-1] if frame_parts else "image_%counter%") + extension)])
|
||||||
filename_template = (frame_parts[-1] if frame_parts else "image_%counter%") + extension
|
counter = _next_counter_for_relative_path(target_dir, relative_template)
|
||||||
|
resolved_relative = relative_template.replace("%counter%", str(counter).zfill(3))
|
||||||
|
resolved_parts = [part for part in resolved_relative.split("/") if part]
|
||||||
|
sub_dirs = resolved_parts[:-1]
|
||||||
|
filename = resolved_parts[-1] if resolved_parts else f"image_{str(counter).zfill(3)}{extension}"
|
||||||
frame_dir = os.path.join(target_dir, *sub_dirs)
|
frame_dir = os.path.join(target_dir, *sub_dirs)
|
||||||
counter = _next_counter(frame_dir, filename_template)
|
|
||||||
filename = filename_template.replace("%counter%", str(counter).zfill(3))
|
|
||||||
image = _tensor_image_to_pil_image(images[batch_index : batch_index + 1])
|
image = _tensor_image_to_pil_image(images[batch_index : batch_index + 1])
|
||||||
full_path = os.path.join(frame_dir, filename)
|
full_path = os.path.join(frame_dir, filename)
|
||||||
|
|
||||||
|
|||||||
@@ -204,6 +204,23 @@ class DumasImageNodeTests(unittest.TestCase):
|
|||||||
saved_names = [os.path.basename(path) for path, _args, _kwargs in FakePILImage.saved_paths]
|
saved_names = [os.path.basename(path) for path, _args, _kwargs in FakePILImage.saved_paths]
|
||||||
self.assertEqual(saved_names, ["counter_001.png", "counter_002.png"])
|
self.assertEqual(saved_names, ["counter_001.png", "counter_002.png"])
|
||||||
|
|
||||||
|
def test_save_image_uses_same_counter_for_folder_and_filename(self):
|
||||||
|
node = self.image_nodes.DumasSaveImageNode()
|
||||||
|
image = FakeTensorBatch()
|
||||||
|
|
||||||
|
node.save_images(
|
||||||
|
images=image,
|
||||||
|
folder=self.temp_dir,
|
||||||
|
pattern="Char%counter%/Char%counter%",
|
||||||
|
format="png",
|
||||||
|
quality=100,
|
||||||
|
embed_workflow=False,
|
||||||
|
save_on_run=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
saved_path = FakePILImage.saved_paths[0][0].replace("\\", "/")
|
||||||
|
self.assertTrue(saved_path.endswith("/Char001/Char001.png"))
|
||||||
|
|
||||||
def test_save_image_returns_ui_entries_for_output_folder(self):
|
def test_save_image_returns_ui_entries_for_output_folder(self):
|
||||||
node = self.image_nodes.DumasSaveImageNode()
|
node = self.image_nodes.DumasSaveImageNode()
|
||||||
image = FakeTensorBatch()
|
image = FakeTensorBatch()
|
||||||
|
|||||||
Reference in New Issue
Block a user