V 2.0.0 - Universal api #63 - file saver
This commit is contained in:
+36
-114
@@ -18,6 +18,7 @@ import comfy_extras.nodes_flux as nodes_flux
|
||||
import torch
|
||||
from ..components import utility
|
||||
from ..components import primeresamplers
|
||||
from ..components import file_output
|
||||
from server import PromptServer
|
||||
from ..components.tree import PRIMERE_ROOT
|
||||
from comfy.cli_args import args
|
||||
@@ -28,7 +29,7 @@ from transformers import pipeline
|
||||
from torchvision.transforms import functional as TF
|
||||
import comfy_extras.nodes_model_advanced as nodes_model_advanced
|
||||
|
||||
ALLOWED_EXT = ('.jpeg', '.jpg', '.png', '.tiff', '.gif', '.bmp', '.webp')
|
||||
ALLOWED_EXT = file_output.ALLOWED_EXT
|
||||
|
||||
class PrimereMetaSave:
|
||||
RETURN_TYPES = ("STRING",)
|
||||
@@ -105,10 +106,10 @@ class PrimereMetaSave:
|
||||
|
||||
delimiter = filename_delimiter
|
||||
number_padding = filename_number_padding
|
||||
tokens = TextTokens()
|
||||
tokens = file_output.TextTokens()
|
||||
|
||||
original_output = self.output_dir
|
||||
filename_prefix = tokens.parseTokens(filename_prefix)
|
||||
filename_prefix = file_output.sanitize_path_part(tokens.parseTokens(filename_prefix))
|
||||
nowdate = datetime.datetime.now()
|
||||
if image_metadata is None:
|
||||
image_metadata = {}
|
||||
@@ -137,39 +138,28 @@ class PrimereMetaSave:
|
||||
if 'width' in image_metadata:
|
||||
filename_prefix = filename_prefix + '_' + str(image_metadata['saved_image_width']) + 'x' + str(image_metadata['saved_image_heigth'])
|
||||
|
||||
if output_path in [None, '', "none", "."]:
|
||||
output_path = self.output_dir
|
||||
else:
|
||||
output_path = tokens.parseTokens(output_path)
|
||||
if not os.path.isabs(output_path):
|
||||
output_path = os.path.join(self.output_dir, output_path)
|
||||
output_path = file_output.parse_output_path_base(output_path, self.output_dir)
|
||||
base_output = os.path.basename(output_path)
|
||||
if output_path.endswith("ComfyUI/output") or output_path.endswith("ComfyUI\output"):
|
||||
base_output = ""
|
||||
|
||||
if add_concept_to_path == True and 'model_concept' in image_metadata:
|
||||
path = Path(output_path)
|
||||
ConceptStartPath = output_path.replace(path.stem, '')
|
||||
ConceptPath = image_metadata['model_concept']
|
||||
if image_metadata['model_concept'] == 'Auto':
|
||||
if 'model_version' in image_metadata:
|
||||
match image_metadata['model_version']:
|
||||
case 'SDXL_2048':
|
||||
ConceptPath = 'SDXL'
|
||||
case 'BaseModel_768':
|
||||
ConceptPath = 'SD1'
|
||||
case 'SD3_1024':
|
||||
ConceptPath = 'SD3'
|
||||
case 'Stable_Zero123_768':
|
||||
ConceptPath = 'Stable_Zero'
|
||||
ConceptPath = Path(ConceptPath)
|
||||
subdirs = []
|
||||
|
||||
output_path = ConceptStartPath + ConceptPath.stem.upper() + os.sep + path.stem
|
||||
if add_concept_to_path == True and 'model_concept' in image_metadata:
|
||||
concept_name = image_metadata['model_concept']
|
||||
if concept_name == 'Auto' and 'model_version' in image_metadata:
|
||||
match image_metadata['model_version']:
|
||||
case 'SDXL_2048':
|
||||
concept_name = 'SDXL'
|
||||
case 'BaseModel_768':
|
||||
concept_name = 'SD1'
|
||||
case 'SD3_1024':
|
||||
concept_name = 'SD3'
|
||||
case 'Stable_Zero123_768':
|
||||
concept_name = 'Stable_Zero'
|
||||
subdirs.append(file_output.sanitize_path_part(Path(concept_name).stem.upper()))
|
||||
|
||||
if add_modelname_to_path == True and 'model' in image_metadata:
|
||||
path = Path(output_path)
|
||||
ModelStartPath = output_path.replace(path.stem, '')
|
||||
|
||||
if 'model_concept' in image_metadata and 'model_version' in image_metadata:
|
||||
original_model_concept_selector = 'Auto'
|
||||
if extra_pnginfo is not None:
|
||||
@@ -184,55 +174,27 @@ class PrimereMetaSave:
|
||||
image_metadata['model'] = image_metadata['concept_data']['flux_diffusion']
|
||||
case 'StableCascade':
|
||||
image_metadata['model'] = image_metadata['concept_data']['cascade_stage_c']
|
||||
subdirs.append(file_output.sanitize_path_part(Path(image_metadata['model']).stem.upper()))
|
||||
|
||||
|
||||
ModelPath = Path(image_metadata['model'])
|
||||
|
||||
if subpath_priority == True and 'preferred' in image_metadata and type(image_metadata['preferred']).__name__ == 'dict' and len(image_metadata['preferred']) > 0 and 'subpath' in image_metadata['preferred']:
|
||||
if image_metadata['preferred']['subpath'] is not None and len(image_metadata['preferred']['subpath'].strip()) > 0:
|
||||
subpath = image_metadata['preferred']['subpath']
|
||||
output_path = ModelStartPath + ModelPath.stem.upper() + os.sep + subpath + os.sep + path.stem
|
||||
else:
|
||||
if subpath_priority == False and subpath is not None and subpath != 'None' and len(subpath.strip()) > 0:
|
||||
output_path = ModelStartPath + ModelPath.stem.upper() + os.sep + subpath + os.sep + path.stem
|
||||
else:
|
||||
output_path = ModelStartPath + ModelPath.stem.upper() + os.sep + path.stem
|
||||
else:
|
||||
if subpath_priority == True and 'preferred' in image_metadata and type(image_metadata['preferred']).__name__ == 'dict' and len(image_metadata['preferred']) > 0 and 'subpath' in image_metadata['preferred'] and image_metadata['preferred']['subpath'] is not None and len(image_metadata['preferred']['subpath'].strip()) > 0:
|
||||
path = Path(output_path)
|
||||
ModelStartPath = output_path.replace(path.stem, '')
|
||||
if subpath_priority == True and 'preferred' in image_metadata and type(image_metadata['preferred']).__name__ == 'dict' and len(image_metadata['preferred']) > 0 and 'subpath' in image_metadata['preferred']:
|
||||
if image_metadata['preferred']['subpath'] is not None and len(image_metadata['preferred']['subpath'].strip()) > 0:
|
||||
subpath = image_metadata['preferred']['subpath']
|
||||
output_path = ModelStartPath + os.sep + subpath + os.sep + path.stem
|
||||
else:
|
||||
path = Path(output_path)
|
||||
ModelStartPath = output_path.replace(path.stem, '')
|
||||
if subpath is not None and subpath != 'None' and len(subpath.strip()) > 0:
|
||||
output_path = ModelStartPath + os.sep + subpath + os.sep + path.stem
|
||||
else:
|
||||
output_path = ModelStartPath + os.sep + path.stem
|
||||
subdirs.append(file_output.sanitize_path_part(subpath))
|
||||
elif subpath_priority == False and subpath is not None and subpath != 'None' and len(subpath.strip()) > 0:
|
||||
subdirs.append(file_output.sanitize_path_part(subpath))
|
||||
|
||||
if output_path.strip() != '':
|
||||
if not os.path.isabs(output_path):
|
||||
output_path = os.path.join(folder_paths.output_directory, output_path)
|
||||
if not os.path.exists(output_path.strip()):
|
||||
print(f'The path `{output_path.strip()}` specified doesn\'t exist! Creating directory.')
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
output_path = file_output.append_subdirs_before_stem(output_path, subdirs)
|
||||
output_path = file_output.ensure_output_dir(output_path)
|
||||
|
||||
if filename_number_start == 'true':
|
||||
pattern = f"(\\d{{{filename_number_padding}}}){re.escape(delimiter)}{re.escape(filename_prefix)}"
|
||||
else:
|
||||
pattern = f"{re.escape(filename_prefix)}{re.escape(delimiter)}(\\d{{{filename_number_padding}}})"
|
||||
existing_counters = [int(re.search(pattern, filename).group(1)) for filename in os.listdir(output_path) if re.match(pattern, os.path.basename(filename))]
|
||||
existing_counters.sort(reverse=True)
|
||||
|
||||
if existing_counters:
|
||||
counter = existing_counters[0] + 1
|
||||
else:
|
||||
counter = 1
|
||||
|
||||
file_extension = '.' + extension
|
||||
if file_extension not in ALLOWED_EXT:
|
||||
file_extension = "jpg"
|
||||
file, counter = file_output.build_filename_and_counter(
|
||||
output_path=output_path,
|
||||
prefix=filename_prefix,
|
||||
delimiter=delimiter,
|
||||
number_padding=number_padding,
|
||||
number_start=filename_number_start,
|
||||
extension=extension,
|
||||
overwrite_mode=overwrite_mode,
|
||||
)
|
||||
|
||||
results = list()
|
||||
# for image in images:
|
||||
@@ -250,16 +212,6 @@ class PrimereMetaSave:
|
||||
for x in extra_pnginfo:
|
||||
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
||||
|
||||
if overwrite_mode == 'prefix_as_filename':
|
||||
file = f"{filename_prefix}{file_extension}"
|
||||
else:
|
||||
if filename_number_start == 'true':
|
||||
file = f"{counter:0{number_padding}}{delimiter}{filename_prefix}{file_extension}"
|
||||
else:
|
||||
file = f"{filename_prefix}{delimiter}{counter:0{number_padding}}{file_extension}"
|
||||
if os.path.exists(os.path.join(output_path, file)):
|
||||
counter += 1
|
||||
|
||||
exif_metadata_A11 = None
|
||||
try:
|
||||
if 'positive' in image_metadata and 'negative' in image_metadata:
|
||||
@@ -384,36 +336,6 @@ Steps: {str(image_metadata['steps'])}, Sampler: {a11samplername}, CFG scale: {st
|
||||
subfolder_path = os.sep.join(subfolder_parts[:-1])
|
||||
return subfolder_path
|
||||
|
||||
class TextTokens:
|
||||
def __init__(self):
|
||||
self.tokens = {
|
||||
'[time]': str(time.time()).replace('.', '_')
|
||||
}
|
||||
if '.' in self.tokens['[time]']: self.tokens['[time]'] = self.tokens['[time]'].split('.')[0]
|
||||
|
||||
def format_time(self, format_code):
|
||||
return time.strftime(format_code, time.localtime(time.time()))
|
||||
|
||||
def parseTokens(self, text):
|
||||
tokens = self.tokens.copy()
|
||||
|
||||
# Update time
|
||||
tokens['[time]'] = str(time.time())
|
||||
if '.' in tokens['[time]']:
|
||||
tokens['[time]'] = tokens['[time]'].split('.')[0]
|
||||
|
||||
for token, value in tokens.items():
|
||||
if token.startswith('[time('):
|
||||
continue
|
||||
text = text.replace(token, value)
|
||||
|
||||
def replace_custom_time(match):
|
||||
format_code = match.group(1)
|
||||
return self.format_time(format_code)
|
||||
|
||||
text = re.sub(r'\[time\((.*?)\)\]', replace_custom_time, text)
|
||||
return text
|
||||
|
||||
class AnyType(str):
|
||||
def __ne__(self, __value: object) -> bool:
|
||||
return False
|
||||
|
||||
+86
-2
@@ -11,16 +11,19 @@ import random
|
||||
import argparse
|
||||
import json
|
||||
import copy
|
||||
import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
import requests
|
||||
import sys
|
||||
from PIL import Image
|
||||
from io import BytesIO
|
||||
import numpy as np
|
||||
|
||||
from ..components.API import api_json_to_requestbody
|
||||
from ..components.API import external_api_backend
|
||||
from ..components.API import api_schema_registry
|
||||
from ..components import file_output
|
||||
|
||||
class PrimereApiProcessor:
|
||||
CATEGORY = TREE_API
|
||||
@@ -40,7 +43,23 @@ class PrimereApiProcessor:
|
||||
"api_provider": (external_api_backend.provider_list(cls),),
|
||||
"api_service": (external_api_backend.service_list(cls),),
|
||||
"prompt": ("STRING", {"forceInput": True}),
|
||||
"batch": ("INT", {"default": 1, "max": 10, "min": 1, "step": 1})
|
||||
"batch": ("INT", {"default": 1, "max": 10, "min": 1, "step": 1}),
|
||||
"auto_save_result": ("BOOLEAN", {"default": False, "label_on": "Save result", "label_off": "Don't save result"}),
|
||||
"output_path": ("STRING", {"default": '[time(%Y-%m-%d)]', "multiline": False}),
|
||||
"subpath": (["None", "Dev", "Test", "Serie", "Production", "Preview", "NewModel", "Project", "Portfolio", "Civitai", "Behance", "Facebook", "Instagram", "Character", "Style", "Product", "Fun", "SFW", "NSFW"], {"default": "Project"}),
|
||||
"add_provider_to_path": ("BOOLEAN", {"default": False}),
|
||||
"add_service_to_path": ("BOOLEAN", {"default": False}),
|
||||
"add_model_to_path": ("BOOLEAN", {"default": False}),
|
||||
"filename_prefix": ("STRING", {"default": "API"}),
|
||||
"filename_delimiter": ("STRING", {"default": "_"}),
|
||||
"add_date_to_filename": ("BOOLEAN", {"default": True}),
|
||||
"add_time_to_filename": ("BOOLEAN", {"default": True}),
|
||||
"filename_number_padding": ("INT", {"default": 2, "min": 1, "max": 9, "step": 1}),
|
||||
"filename_number_start": ("BOOLEAN", {"default": False}),
|
||||
"image_extension": ([ext.lstrip('.') for ext in file_output.ALLOWED_EXT], {"default": "jpg"}),
|
||||
"image_quality": ("INT", {"default": 95, "min": 1, "max": 100, "step": 1}),
|
||||
"save_data_to_json": ("BOOLEAN", {"default": False}),
|
||||
"save_data_to_txt": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
|
||||
cls.optional_inputs = {
|
||||
@@ -205,6 +224,7 @@ class PrimereApiProcessor:
|
||||
api_result = None
|
||||
api_error = None
|
||||
result_image = None
|
||||
save_bytes = None
|
||||
batch = max(1, int(batch))
|
||||
sdk_context = {}
|
||||
response_url = None
|
||||
@@ -276,6 +296,70 @@ class PrimereApiProcessor:
|
||||
|
||||
if api_error is None:
|
||||
response_context = {"response_url": response_url, "call_url": response_url, "loaded_client": loaded_client, "client": client, "sdk_context": sdk_context}
|
||||
result_image = external_api_backend.apply_response_handler(schema, api_result, provider=api_provider, service=(selected_service or api_service), response_context=response_context)
|
||||
handler_result = external_api_backend.apply_response_handler(schema, api_result, provider=api_provider, service=(selected_service or api_service), response_context=response_context)
|
||||
if isinstance(handler_result, list) and len(handler_result) == 2:
|
||||
result_image, save_bytes = handler_result
|
||||
else:
|
||||
result_image = handler_result
|
||||
|
||||
# --- File naming and output path resolution ---
|
||||
auto_save_result = kwargs.get('auto_save_result', False)
|
||||
if auto_save_result and result_image is not None:
|
||||
output_path_input = kwargs.get('output_path', '[time(%Y-%m-%d)]')
|
||||
subpath = kwargs.get('subpath', 'None')
|
||||
add_provider_to_path = kwargs.get('add_provider_to_path', False)
|
||||
add_service_to_path = kwargs.get('add_service_to_path', False)
|
||||
filename_prefix = kwargs.get('filename_prefix', 'API')
|
||||
filename_delimiter = kwargs.get('filename_delimiter', '_')
|
||||
add_date_to_filename = kwargs.get('add_date_to_filename', True)
|
||||
add_time_to_filename = kwargs.get('add_time_to_filename', True)
|
||||
filename_number_padding = kwargs.get('filename_number_padding', 2)
|
||||
filename_number_start = kwargs.get('filename_number_start', False)
|
||||
image_extension = kwargs.get('image_extension', 'jpg')
|
||||
image_quality = kwargs.get('image_quality', 95)
|
||||
save_data_to_json = kwargs.get('save_data_to_json', False)
|
||||
save_data_to_txt = kwargs.get('save_data_to_txt', False)
|
||||
add_model_to_path = kwargs.get('add_model_to_path', False)
|
||||
|
||||
model_subdir = None
|
||||
if add_model_to_path:
|
||||
model_subdir = next((custom_values[k] for k in ('model', 'model_name', 'version') if custom_values.get(k)), None)
|
||||
if model_subdir:
|
||||
model_subdir = file_output.sanitize_path_part(model_subdir)
|
||||
|
||||
if not os.path.isabs(output_path_input):
|
||||
output_path_input = file_output.sanitize_path_part(output_path_input)
|
||||
filename_prefix = file_output.sanitize_path_part(filename_prefix)
|
||||
|
||||
subdirs = []
|
||||
if add_provider_to_path and api_provider:
|
||||
subdirs.append(file_output.sanitize_path_part(api_provider))
|
||||
if add_service_to_path and (selected_service or api_service):
|
||||
subdirs.append(file_output.sanitize_path_part(selected_service or api_service))
|
||||
if model_subdir:
|
||||
subdirs.append(model_subdir)
|
||||
if subpath and subpath != 'None' and subpath.strip():
|
||||
subdirs.append(file_output.sanitize_path_part(subpath))
|
||||
|
||||
output_file, json_file, txt_file = file_output.resolve_output_file(
|
||||
output_path_input, folder_paths.output_directory, subdirs,
|
||||
filename_prefix, filename_delimiter,
|
||||
add_date_to_filename, add_time_to_filename,
|
||||
filename_number_padding, filename_number_start, image_extension,
|
||||
)
|
||||
|
||||
Path(folder_paths.temp_directory).mkdir(parents=True, exist_ok=True)
|
||||
file_output.save_bytes_to_file(save_bytes, output_file, image_extension, image_quality, folder_paths.temp_directory)
|
||||
|
||||
save_data = {
|
||||
"provider": api_provider,
|
||||
"service": selected_service or api_service,
|
||||
"selected_parameters": selected_parameters_output,
|
||||
"used_values": used_values_output,
|
||||
# "rendered": rendered_payload,
|
||||
"raw_payload": raw_payload,
|
||||
# "api_result": api_result_debug,
|
||||
}
|
||||
file_output.save_metadata(save_data, json_file, txt_file, save_data_to_json, save_data_to_txt, used_values_output)
|
||||
|
||||
return (result_image, client, api_provider, schema, rendered_payload, raw_payload, used_values_output, api_schemas, api_result_debug)
|
||||
File diff suppressed because one or more lines are too long
@@ -25,6 +25,7 @@
|
||||
- 5.2 [`reference_images_handler`](#52-reference_images_handler)
|
||||
6. [Runtime rules and validation](#6-runtime-rules-and-validation)
|
||||
7. [Debug outputs — understanding and using them](#7-debug-outputs--understanding-and-using-them)
|
||||
8. [File save settings](#8-file-save-settings)
|
||||
|
||||
---
|
||||
|
||||
@@ -727,3 +728,66 @@ When writing a new service schema from scratch or adapting an existing one, the
|
||||
4. **Use `debug_mode = ON`** during the design phase — the API call is never made, so there is no cost and no rate limit risk. Iterate on the schema until all three payload outputs look correct, then switch to production mode.
|
||||
|
||||
5. **Read `API_SCHEMAS`** after a failed production call. The `api_error` field contains the provider error message. Combined with `rendered` and `api_result` in the same object, you can diagnose whether the error is a structural problem (wrong key, wrong nesting, wrong type) or a credential/quota problem.
|
||||
|
||||
---
|
||||
|
||||
## 8) File save settings
|
||||
|
||||
File saving only runs when `auto_save_result` is ON and the API returned a valid result. If the response is `None` (no API error but no result), nothing is written.
|
||||
|
||||
### Output path
|
||||
|
||||
| Input | Default | Description |
|
||||
|---|---|---|
|
||||
| `output_path` | `[time(%Y-%m-%d)]` | Base output directory. Supports `[time(...)]` tokens. Relative paths are anchored to the ComfyUI output directory. Absolute paths are used as-is. |
|
||||
| `subpath` | `Project` | Fixed subdirectory appended after provider/service/model dirs. Select `None` to skip. |
|
||||
| `add_provider_to_path` | OFF | Adds the API provider name as a subdirectory (e.g. `Gemini`). |
|
||||
| `add_service_to_path` | OFF | Adds the selected service name as a subdirectory (e.g. `Imagen`). |
|
||||
| `add_model_to_path` | OFF | Adds the model identifier as a subdirectory. Reads `model_name` first, then `model`, then `version` from service parameters. |
|
||||
|
||||
Directory structure example with all path options enabled:
|
||||
|
||||
```
|
||||
<output_path> / <provider> / <service> / <model> / <subpath> / <filename>
|
||||
```
|
||||
|
||||
User-supplied strings (provider, service, model, subpath, output_path) are automatically sanitized before use as path components: spaces and special characters (` / \ . , ; - `) are replaced with `_`, consecutive underscores collapsed to one.
|
||||
|
||||
### Filename
|
||||
|
||||
| Input | Default | Description |
|
||||
|---|---|---|
|
||||
| `filename_prefix` | `API` | Base name for the saved file. |
|
||||
| `filename_delimiter` | `_` | Separator between prefix, date, time, and counter parts. |
|
||||
| `add_date_to_filename` | ON | Appends current date (`YYYY-MM-DD`) to filename. |
|
||||
| `add_time_to_filename` | ON | Appends current time (`HHMMSS`) to filename. |
|
||||
| `filename_number_padding` | `2` | Zero-padding width for the auto-increment counter. |
|
||||
| `filename_number_start` | OFF | If ON, counter is placed before the prefix instead of after. |
|
||||
|
||||
### Image format
|
||||
|
||||
| Input | Default | Description |
|
||||
|---|---|---|
|
||||
| `image_extension` | `jpg` | Target image format. Options: `jpeg jpg png tiff gif bmp webp`. |
|
||||
| `image_quality` | `95` | Compression quality for JPEG and WEBP. PNG, TIFF, GIF ignore this. |
|
||||
|
||||
### Non-image results
|
||||
|
||||
The actual file type is detected from the API response bytes (MIME detection), not assumed from `image_extension`. If the API returns audio, video, or text, the correct extension is used automatically:
|
||||
|
||||
| MIME type | Saved extension |
|
||||
|---|---|
|
||||
| `image/*` | uses `image_extension` input |
|
||||
| `audio/*` | `.mp3` |
|
||||
| `video/*` | `.mp4` |
|
||||
| `text/*` | `.txt` (UTF-8) |
|
||||
| other / unknown | extension from `image_extension` |
|
||||
|
||||
### Metadata files
|
||||
|
||||
| Input | Default | Description |
|
||||
|---|---|---|
|
||||
| `save_data_to_json` | OFF | Saves a `.json` file alongside the result containing provider, service, selected parameters, used values, and raw payload. |
|
||||
| `save_data_to_txt` | OFF | Saves a `.txt` file alongside the result containing provider, service, and all used parameter values (flattened key: value lines). |
|
||||
|
||||
Both files share the same base path and filename as the saved result, only the extension differs.
|
||||
|
||||
@@ -54,4 +54,5 @@ def handle_response(api_result, schema=None, loaded_client=None, response_url=No
|
||||
if not image_url:
|
||||
raise RuntimeError("No result image...")
|
||||
|
||||
return response_helper.url_to_tensor(image_url)
|
||||
image_bytes = response_helper.fetch_url_bytes(image_url)
|
||||
return [response_helper.bytes_to_tensor(image_bytes), image_bytes]
|
||||
@@ -15,13 +15,17 @@ def handle_response(api_result, schema=None, loaded_client=None, response_url=No
|
||||
|
||||
remote_images = json_object.get("images", [])
|
||||
image_tensors = []
|
||||
first_bytes = None
|
||||
for remote_image in remote_images:
|
||||
image_url = remote_image.get("url") if isinstance(remote_image, dict) else None
|
||||
if not image_url:
|
||||
continue
|
||||
|
||||
tensor = response_helper.url_to_tensor(image_url)
|
||||
image_bytes = response_helper.fetch_url_bytes(image_url)
|
||||
tensor = response_helper.bytes_to_tensor(image_bytes)
|
||||
if tensor is not None:
|
||||
image_tensors.append(tensor)
|
||||
if first_bytes is None:
|
||||
first_bytes = image_bytes
|
||||
|
||||
return response_helper.stack_image_tensors(image_tensors)
|
||||
return [response_helper.stack_image_tensors(image_tensors), first_bytes]
|
||||
|
||||
@@ -5,8 +5,8 @@ from . import response_helper
|
||||
def handle_response(api_result, schema=None, loaded_client=None, response_url=None):
|
||||
generated_images = getattr(api_result, "generated_images", None)
|
||||
if not generated_images:
|
||||
return None
|
||||
return [None, None]
|
||||
|
||||
image = getattr(generated_images[0], "image", None)
|
||||
image_bytes = getattr(image, "image_bytes", None)
|
||||
return response_helper.bytes_to_tensor(image_bytes)
|
||||
return [response_helper.bytes_to_tensor(image_bytes), image_bytes]
|
||||
|
||||
@@ -5,19 +5,22 @@ from . import response_helper
|
||||
def handle_response(api_result, schema=None, loaded_client=None, response_url=None):
|
||||
candidates = getattr(api_result, "candidates", None)
|
||||
if not candidates:
|
||||
return None
|
||||
return [None, None]
|
||||
|
||||
content = getattr(candidates[0], "content", None)
|
||||
parts = getattr(content, "parts", None)
|
||||
if not parts:
|
||||
return None
|
||||
return [None, None]
|
||||
|
||||
image_tensors = []
|
||||
first_bytes = None
|
||||
for part in parts:
|
||||
inline_data = getattr(part, "inline_data", None)
|
||||
image_bytes = getattr(inline_data, "data", None)
|
||||
tensor = response_helper.bytes_to_tensor(image_bytes)
|
||||
if tensor is not None:
|
||||
image_tensors.append(tensor)
|
||||
if first_bytes is None:
|
||||
first_bytes = image_bytes
|
||||
|
||||
return response_helper.stack_image_tensors(image_tensors)
|
||||
return [response_helper.stack_image_tensors(image_tensors), first_bytes]
|
||||
|
||||
@@ -6,9 +6,10 @@ from . import response_helper
|
||||
def handle_response(api_result, schema=None, loaded_client=None, response_url=None):
|
||||
data_items = getattr(api_result, "data", None)
|
||||
if not data_items:
|
||||
return None
|
||||
return [None, None]
|
||||
|
||||
image_tensors = []
|
||||
first_bytes = None
|
||||
for item in data_items:
|
||||
image_base64 = getattr(item, "b64_json", None)
|
||||
if not image_base64:
|
||||
@@ -18,5 +19,7 @@ def handle_response(api_result, schema=None, loaded_client=None, response_url=No
|
||||
tensor = response_helper.bytes_to_tensor(image_bytes)
|
||||
if tensor is not None:
|
||||
image_tensors.append(tensor)
|
||||
if first_bytes is None:
|
||||
first_bytes = image_bytes
|
||||
|
||||
return response_helper.stack_image_tensors(image_tensors)
|
||||
return [response_helper.stack_image_tensors(image_tensors), first_bytes]
|
||||
|
||||
@@ -0,0 +1,195 @@
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import datetime
|
||||
import random
|
||||
import json
|
||||
from io import BytesIO
|
||||
from mimetypes import MimeTypes
|
||||
from pathlib import Path
|
||||
import folder_paths
|
||||
import magic
|
||||
from PIL import Image
|
||||
|
||||
ALLOWED_EXT = ('.jpeg', '.jpg', '.png', '.tiff', '.gif', '.bmp', '.webp')
|
||||
|
||||
class TextTokens:
|
||||
def __init__(self):
|
||||
self.tokens = {
|
||||
'[time]': str(time.time()).replace('.', '_')
|
||||
}
|
||||
if '.' in self.tokens['[time]']:
|
||||
self.tokens['[time]'] = self.tokens['[time]'].split('.')[0]
|
||||
|
||||
def format_time(self, format_code):
|
||||
return time.strftime(format_code, time.localtime(time.time()))
|
||||
|
||||
def parseTokens(self, text):
|
||||
tokens = self.tokens.copy()
|
||||
tokens['[time]'] = str(time.time())
|
||||
if '.' in tokens['[time]']:
|
||||
tokens['[time]'] = tokens['[time]'].split('.')[0]
|
||||
|
||||
for token, value in tokens.items():
|
||||
if token.startswith('[time('):
|
||||
continue
|
||||
text = text.replace(token, value)
|
||||
|
||||
def replace_custom_time(match):
|
||||
format_code = match.group(1)
|
||||
return self.format_time(format_code)
|
||||
|
||||
text = re.sub(r'\[time\((.*?)\)\]', replace_custom_time, text)
|
||||
return text
|
||||
|
||||
def parse_output_path_base(output_path, base_dir):
|
||||
tokens = TextTokens()
|
||||
if output_path in [None, '', "none", "."]:
|
||||
return base_dir
|
||||
output_path = tokens.parseTokens(output_path)
|
||||
if not os.path.isabs(output_path):
|
||||
output_path = os.path.join(base_dir, output_path)
|
||||
return output_path
|
||||
|
||||
|
||||
def append_subdirs_before_stem(output_path, subdirs):
|
||||
path = Path(output_path)
|
||||
valid = [str(s) for s in subdirs if s and str(s) != 'None' and str(s).strip()]
|
||||
if valid:
|
||||
return os.path.join(str(path.parent), *valid, path.stem)
|
||||
return os.path.join(str(path.parent), path.stem)
|
||||
|
||||
def ensure_output_dir(output_path):
|
||||
if output_path.strip():
|
||||
if not os.path.isabs(output_path):
|
||||
output_path = os.path.join(folder_paths.output_directory, output_path)
|
||||
if not os.path.exists(output_path.strip()):
|
||||
print(f"The path `{output_path.strip()}` specified doesn't exist! Creating directory.")
|
||||
os.makedirs(output_path, exist_ok=True)
|
||||
return output_path
|
||||
|
||||
def build_filename_and_counter(output_path, prefix, delimiter, number_padding, number_start, extension, overwrite_mode='false'):
|
||||
number_start_bool = number_start is True or str(number_start).lower() == 'true'
|
||||
file_extension = ('.' + extension) if not extension.startswith('.') else extension
|
||||
if file_extension not in ALLOWED_EXT:
|
||||
file_extension = '.jpg'
|
||||
if number_start_bool:
|
||||
pattern = f"(\\d{{{number_padding}}}){re.escape(delimiter)}{re.escape(prefix)}"
|
||||
else:
|
||||
pattern = f"{re.escape(prefix)}{re.escape(delimiter)}(\\d{{{number_padding}}})"
|
||||
existing_counters = [
|
||||
int(re.search(pattern, fn).group(1))
|
||||
for fn in os.listdir(output_path)
|
||||
if re.match(pattern, os.path.basename(fn))
|
||||
]
|
||||
existing_counters.sort(reverse=True)
|
||||
counter = existing_counters[0] + 1 if existing_counters else 1
|
||||
if overwrite_mode == 'prefix_as_filename':
|
||||
file = f"{prefix}{file_extension}"
|
||||
else:
|
||||
if number_start_bool:
|
||||
file = f"{counter:0{number_padding}}{delimiter}{prefix}{file_extension}"
|
||||
else:
|
||||
file = f"{prefix}{delimiter}{counter:0{number_padding}}{file_extension}"
|
||||
if os.path.exists(os.path.join(output_path, file)):
|
||||
counter += 1
|
||||
if number_start_bool:
|
||||
file = f"{counter:0{number_padding}}{delimiter}{prefix}{file_extension}"
|
||||
else:
|
||||
file = f"{prefix}{delimiter}{counter:0{number_padding}}{file_extension}"
|
||||
|
||||
return file, counter
|
||||
|
||||
def resolve_output_file(output_path_input, base_dir, subdirs, prefix, delimiter, add_date, add_time, number_padding, number_start, extension):
|
||||
resolved = parse_output_path_base(output_path_input, base_dir)
|
||||
resolved = append_subdirs_before_stem(resolved, subdirs)
|
||||
resolved = ensure_output_dir(resolved)
|
||||
tokens = TextTokens()
|
||||
prefix_parsed = tokens.parseTokens(prefix)
|
||||
nowdate = datetime.datetime.now()
|
||||
if add_date:
|
||||
prefix_parsed = prefix_parsed + delimiter + nowdate.strftime("%Y-%m-%d")
|
||||
if add_time:
|
||||
prefix_parsed = prefix_parsed + delimiter + nowdate.strftime("%H%M%S")
|
||||
|
||||
filename, counter = build_filename_and_counter(
|
||||
output_path=resolved,
|
||||
prefix=prefix_parsed,
|
||||
delimiter=delimiter,
|
||||
number_padding=number_padding,
|
||||
number_start=number_start,
|
||||
extension=extension,
|
||||
)
|
||||
|
||||
output_file = os.path.abspath(os.path.join(resolved, filename))
|
||||
json_file = os.path.splitext(output_file)[0] + '.json'
|
||||
txt_file = os.path.splitext(output_file)[0] + '.txt'
|
||||
|
||||
return output_file, json_file, txt_file
|
||||
|
||||
|
||||
def detect_mime(save_bytes, temp_directory):
|
||||
tmp_path = os.path.join(temp_directory, f"api_mime_{random.randint(1000, 9999)}")
|
||||
with open(tmp_path, 'wb') as f:
|
||||
f.write(save_bytes)
|
||||
try:
|
||||
mime = magic.from_file(tmp_path, mime=True)
|
||||
except Exception:
|
||||
mime, _ = MimeTypes().guess_type(tmp_path)
|
||||
return mime or 'application/octet-stream'
|
||||
|
||||
def save_bytes_to_file(save_bytes, output_file, image_extension, image_quality, temp_directory):
|
||||
if save_bytes is None:
|
||||
return output_file
|
||||
mime = detect_mime(save_bytes, temp_directory)
|
||||
stem = os.path.splitext(output_file)[0]
|
||||
if mime.startswith('image/'):
|
||||
fmt_map = {'jpg': 'JPEG', 'jpeg': 'JPEG', 'webp': 'WEBP', 'png': 'PNG', 'tiff': 'TIFF', 'gif': 'GIF'}
|
||||
fmt = fmt_map.get(image_extension.lower(), image_extension.upper())
|
||||
img = Image.open(BytesIO(save_bytes))
|
||||
if fmt in ('PNG', 'TIFF', 'GIF'):
|
||||
img.save(output_file, format=fmt)
|
||||
elif fmt == 'JPEG':
|
||||
img.save(output_file, format='JPEG', quality=image_quality, optimize=True)
|
||||
else:
|
||||
img.save(output_file, format=fmt, quality=image_quality)
|
||||
elif mime.startswith('audio/'):
|
||||
output_file = stem + '.mp3'
|
||||
with open(output_file, 'wb') as f:
|
||||
f.write(save_bytes)
|
||||
elif mime.startswith('video/'):
|
||||
output_file = stem + '.mp4'
|
||||
with open(output_file, 'wb') as f:
|
||||
f.write(save_bytes)
|
||||
elif mime.startswith('text/'):
|
||||
output_file = stem + '.txt'
|
||||
with open(output_file, 'w', encoding='utf-8') as f:
|
||||
f.write(save_bytes.decode('utf-8', errors='replace'))
|
||||
else:
|
||||
with open(output_file, 'wb') as f:
|
||||
f.write(save_bytes)
|
||||
return output_file
|
||||
|
||||
def save_metadata(save_data, json_file, txt_file, save_data_to_json, save_data_to_txt, used_values):
|
||||
if save_data_to_json:
|
||||
with open(json_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(save_data, f, ensure_ascii=False, indent=4)
|
||||
if save_data_to_txt:
|
||||
with open(txt_file, 'w', encoding='utf-8') as f:
|
||||
f.write(f"provider: {save_data.get('provider', '')}\n")
|
||||
f.write(f"service: {save_data.get('service', '')}\n")
|
||||
if isinstance(used_values, dict):
|
||||
for k, v in flatten_dict(used_values):
|
||||
f.write(f"{k}: {v}\n")
|
||||
|
||||
def sanitize_path_part(value):
|
||||
result = re.sub(r'[ /\\.,;`\-]+', '_', str(value))
|
||||
return re.sub(r'_+', '_', result).strip('_')
|
||||
|
||||
def flatten_dict(d, prefix=''):
|
||||
for k, v in d.items():
|
||||
key = f"{prefix}.{k}" if prefix else str(k)
|
||||
if isinstance(v, dict):
|
||||
yield from flatten_dict(v, key)
|
||||
else:
|
||||
yield key, v
|
||||
+2
-1
@@ -51,4 +51,5 @@ anthropic
|
||||
openai
|
||||
google-genai>=1.65.0
|
||||
dotenv
|
||||
elevenlabs
|
||||
elevenlabs
|
||||
python-magic-bin
|
||||
Reference in New Issue
Block a user