Skip to content

Commit

Permalink
feat: add jpg and webp support, add exif data handling for metadata (#…
Browse files Browse the repository at this point in the history
…1863)

* feature: added flag, config and ui update for image extension change #1789

* moved function to config module

* moved image extension to webui via async worker. Passing as parameter to log and get_current_html_path functions per feedback

* check flag before displaying image extension radio button

* disabled if image log flag is passed in

* fix: add missing image_extension parameter to log call

* refactor: change label

* feat: add webp to image_extensions

supported image extemsions: see https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html

* feat: use consistent file name in gradio

returns and uses filepaths instead of numpy image by saving to temp dir
uses double the temp dir file storage on disk as it saves to temp dir and gradio temp dir when displaying the image, but reuses logged output image

* feat: delete temp images after yielding to gradio

* feat: use args temp path if given

* chore: code cleanup, remove redundant if statement

* feat: always show image_extension element

this is now possible due to image extension support in gradio via #1932

* refactor: rename image_extension to image_file_extension

* feat: use optimized jpg parameters when saving the image

quality=95
optimize=True
progressive=True

* refactor: rename image_file_extension to output_format

* feat: add exif handling

* refactor: code cleanup, remove items from metadata output

---------

Co-authored-by: Manuel Schmid <[email protected]>
Co-authored-by: Manuel Schmid <[email protected]>
Co-authored-by: Manuel Schmid <[email protected]>
Co-authored by: eddyizm <[email protected]>
  • Loading branch information
mashb1t committed Feb 26, 2024
1 parent ba9eadb commit b6d2367
Show file tree
Hide file tree
Showing 6 changed files with 104 additions and 28 deletions.
7 changes: 4 additions & 3 deletions modules/async_worker.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import threading
import os
from modules.patch import PatchSettings, patch_settings, patch_all

patch_all()
Expand Down Expand Up @@ -142,6 +141,7 @@ def handler(async_task):
performance_selection = Performance(args.pop())
aspect_ratios_selection = args.pop()
image_number = args.pop()
output_format = args.pop()
image_seed = args.pop()
sharpness = args.pop()
guidance_scale = args.pop()
Expand Down Expand Up @@ -414,6 +414,7 @@ def handler(async_task):

progressbar(async_task, 3, 'Processing prompts ...')
tasks = []

for i in range(image_number):
if disable_seed_increment:
task_seed = seed
Expand Down Expand Up @@ -553,7 +554,7 @@ def handler(async_task):

if direct_return:
d = [('Upscale (Fast)', '2x')]
uov_input_image_path = log(uov_input_image, d)
uov_input_image_path = log(uov_input_image, d, output_format)
yield_result(async_task, uov_input_image_path, do_not_show_finished_images=True)
return

Expand Down Expand Up @@ -863,7 +864,7 @@ def callback(step, x0, x, total_steps, y):
d.append((f'LoRA {li + 1}', f'lora_combined_{li + 1}', f'{n} : {w}'))

d.append(('Version', 'version', 'Fooocus v' + fooocus_version.version))
img_paths.append(log(x, d, metadata_parser))
img_paths.append(log(x, d, metadata_parser, output_format))

yield_result(async_task, img_paths, do_not_show_finished_images=len(tasks) == 1 or disable_intermediate_results)
except ldm_patched.modules.model_management.InterruptProcessingException as e:
Expand Down
5 changes: 5 additions & 0 deletions modules/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,11 @@ def get_config_item_or_set_default(key, default_value, validator, disable_empty_
default_value=32,
validator=lambda x: isinstance(x, int) and x >= 1
)
default_output_format = get_config_item_or_set_default(
key='default_output_format',
default_value='png',
validator=lambda x: x in modules.flags.output_formats
)
default_image_number = get_config_item_or_set_default(
key='default_image_number',
default_value=2,
Expand Down
2 changes: 2 additions & 0 deletions modules/flags.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,8 @@
cn_ip: (0.5, 0.6), cn_ip_face: (0.9, 0.75), cn_canny: (0.5, 1.0), cn_cpds: (0.5, 1.0)
} # stop, weight

output_formats = ['png', 'jpg', 'webp']

inpaint_engine_versions = ['None', 'v1', 'v2.5', 'v2.6']
inpaint_option_default = 'Inpaint or Outpaint (default)'
inpaint_option_detail = 'Improve Detail (face, hand, eyes, etc.)'
Expand Down
66 changes: 61 additions & 5 deletions modules/meta_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import gradio as gr
from PIL import Image

import fooocus_version
import modules.config
import modules.sdxl_styles
from modules.flags import MetadataScheme, Performance, Steps
Expand Down Expand Up @@ -181,13 +182,43 @@ def get_lora(key: str, fallback: str | None, source_dict: dict, results: list):

def get_sha256(filepath):
global hash_cache

if filepath not in hash_cache:
hash_cache[filepath] = calculate_sha256(filepath)

return hash_cache[filepath]


def parse_meta_from_preset(preset_content):
assert isinstance(preset_content, dict)
preset_prepared = {}
items = preset_content

for settings_key, meta_key in modules.config.possible_preset_keys.items():
if settings_key == "default_loras":
loras = getattr(modules.config, settings_key)
if settings_key in items:
loras = items[settings_key]
for index, lora in enumerate(loras[:5]):
preset_prepared[f'lora_combined_{index + 1}'] = ' : '.join(map(str, lora))
elif settings_key == "default_aspect_ratio":
if settings_key in items and items[settings_key] is not None:
default_aspect_ratio = items[settings_key]
width, height = default_aspect_ratio.split('*')
else:
default_aspect_ratio = getattr(modules.config, settings_key)
width, height = default_aspect_ratio.split('×')
height = height[:height.index(" ")]
preset_prepared[meta_key] = (width, height)
else:
preset_prepared[meta_key] = items[settings_key] if settings_key in items and items[
settings_key] is not None else getattr(modules.config, settings_key)

if settings_key == "default_styles" or settings_key == "default_aspect_ratio":
preset_prepared[meta_key] = str(preset_prepared[meta_key])

return preset_prepared


class MetadataParser(ABC):
def __init__(self):
self.raw_prompt: str = ''
Expand All @@ -213,7 +244,8 @@ def parse_json(self, metadata: dict | str) -> dict:
def parse_string(self, metadata: dict) -> str:
raise NotImplementedError

def set_data(self, raw_prompt, full_prompt, raw_negative_prompt, full_negative_prompt, steps, base_model_name, refiner_model_name, loras):
def set_data(self, raw_prompt, full_prompt, raw_negative_prompt, full_negative_prompt, steps, base_model_name,
refiner_model_name, loras):
self.raw_prompt = raw_prompt
self.full_prompt = full_prompt
self.raw_negative_prompt = raw_negative_prompt
Expand Down Expand Up @@ -492,16 +524,28 @@ def get_metadata_parser(metadata_scheme: MetadataScheme) -> MetadataParser:
raise NotImplementedError


def read_info_from_image(filepath) -> tuple[str | None, dict, MetadataScheme | None]:
def read_info_from_image(filepath) -> tuple[str | None, MetadataScheme | None]:
with Image.open(filepath) as image:
items = (image.info or {}).copy()

parameters = items.pop('parameters', None)
metadata_scheme = items.pop('fooocus_scheme', None)
exif = items.pop('exif', None)

if parameters is not None and is_json(parameters):
parameters = json.loads(parameters)
elif exif is not None:
exif = image.getexif()
# 0x9286 = UserComment
parameters = exif.get(0x9286, None)
# 0x927C = MakerNote
metadata_scheme = exif.get(0x927C, None)

if is_json(parameters):
parameters = json.loads(parameters)

try:
metadata_scheme = MetadataScheme(items.pop('fooocus_scheme', None))
metadata_scheme = MetadataScheme(metadata_scheme)
except ValueError:
metadata_scheme = None

Expand All @@ -512,4 +556,16 @@ def read_info_from_image(filepath) -> tuple[str | None, dict, MetadataScheme | N
if isinstance(parameters, str):
metadata_scheme = MetadataScheme.A1111

return parameters, items, metadata_scheme
return parameters, metadata_scheme


def get_exif(metadata: str | None, metadata_scheme: str):
exif = Image.Exif()
# tags see see https://github.com/python-pillow/Pillow/blob/9.2.x/src/PIL/ExifTags.py
# 0x9286 = UserComment
exif[0x9286] = metadata
# 0x0131 = Software
exif[0x0131] = 'Fooocus v' + fooocus_version.version
# 0x927C = MakerNote
exif[0x927C] = metadata_scheme
return exif
32 changes: 20 additions & 12 deletions modules/private_logger.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,34 +7,42 @@
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from modules.util import generate_temp_filename
from modules.meta_parser import MetadataParser
from tempfile import gettempdir
from modules.meta_parser import MetadataParser, get_exif

log_cache = {}


def get_current_html_path():
def get_current_html_path(output_format=None):
output_format = output_format if output_format else modules.config.default_output_format
date_string, local_temp_filename, only_name = generate_temp_filename(folder=modules.config.path_outputs,
extension='png')
extension=output_format)
html_name = os.path.join(os.path.dirname(local_temp_filename), 'log.html')
return html_name


def log(img, metadata, metadata_parser: MetadataParser | None = None) -> str:
def log(img, metadata, metadata_parser: MetadataParser | None = None, output_format=None) -> str:
path_outputs = args_manager.args.temp_path if args_manager.args.disable_image_log else modules.config.path_outputs
date_string, local_temp_filename, only_name = generate_temp_filename(folder=path_outputs, extension='png')
output_format = output_format if output_format else modules.config.default_output_format
date_string, local_temp_filename, only_name = generate_temp_filename(folder=path_outputs, extension=output_format)
os.makedirs(os.path.dirname(local_temp_filename), exist_ok=True)

parsed_parameters = metadata_parser.parse_string(metadata) if metadata_parser is not None else ''
image = Image.fromarray(img)

if parsed_parameters != '':
pnginfo = PngInfo()
pnginfo.add_text('parameters', parsed_parameters)
pnginfo.add_text('fooocus_scheme', metadata_parser.get_scheme().value)
if output_format == 'png':
if parsed_parameters != '':
pnginfo = PngInfo()
pnginfo.add_text('parameters', parsed_parameters)
pnginfo.add_text('fooocus_scheme', metadata_parser.get_scheme().value)
else:
pnginfo = None
image.save(local_temp_filename, pnginfo=pnginfo)
elif output_format == 'jpg':
image.save(local_temp_filename, quality=95, optimize=True, progressive=True, exif=get_exif(parsed_parameters, metadata_parser.get_scheme().value) if metadata_parser else Image.Exif())
elif output_format == 'webp':
image.save(local_temp_filename, quality=95, lossless=False, exif=get_exif(parsed_parameters, metadata_parser.get_scheme().value) if metadata_parser else Image.Exif())
else:
pnginfo = None
image.save(local_temp_filename, pnginfo=pnginfo)
image.save(local_temp_filename)

if args_manager.args.disable_image_log:
return local_temp_filename
Expand Down
20 changes: 12 additions & 8 deletions webui.py
Original file line number Diff line number Diff line change
Expand Up @@ -224,15 +224,12 @@ def ip_advance_checked(x):
metadata_import_button = gr.Button(value='Apply Metadata')

def trigger_metadata_preview(filepath):
parameters, items, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
parameters, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)

results = {}
if parameters is not None:
results['parameters'] = parameters

if items:
results['items'] = items

if isinstance(metadata_scheme, flags.MetadataScheme):
results['metadata_scheme'] = metadata_scheme.value

Expand Down Expand Up @@ -263,6 +260,11 @@ def trigger_metadata_preview(filepath):
value=modules.config.default_aspect_ratio, info='width × height',
elem_classes='aspect_ratios')
image_number = gr.Slider(label='Image Number', minimum=1, maximum=modules.config.default_max_image_number, step=1, value=modules.config.default_image_number)

output_format = gr.Radio(label='Output Format',
choices=modules.flags.output_formats,
value=modules.config.default_output_format)

negative_prompt = gr.Textbox(label='Negative Prompt', show_label=True, placeholder="Type prompt here.",
info='Describing what you do not want to see.', lines=2,
elem_id='negative_prompt',
Expand Down Expand Up @@ -292,7 +294,7 @@ def update_history_link():
if args_manager.args.disable_image_log:
return gr.update(value='')

return gr.update(value=f'<a href="file={get_current_html_path()}" target="_blank">\U0001F4DA History Log</a>')
return gr.update(value=f'<a href="file={get_current_html_path(output_format)}" target="_blank">\U0001F4DA History Log</a>')

history_link = gr.HTML()
shared.gradio_root.load(update_history_link, outputs=history_link, queue=False, show_progress=False)
Expand Down Expand Up @@ -532,7 +534,9 @@ def model_refresh_clicked():
adm_scaler_negative, refiner_switch, refiner_model, sampler_name,
scheduler_name, adaptive_cfg, refiner_swap_method, negative_prompt, disable_intermediate_results
], queue=False, show_progress=False)


output_format.input(lambda x: gr.update(output_format=x), inputs=output_format)

advanced_checkbox.change(lambda x: gr.update(visible=x), advanced_checkbox, advanced_column,
queue=False, show_progress=False) \
.then(fn=lambda: None, _js='refresh_grid_delayed', queue=False, show_progress=False)
Expand Down Expand Up @@ -573,7 +577,7 @@ def inpaint_mode_change(mode):
ctrls = [currentTask, generate_image_grid]
ctrls += [
prompt, negative_prompt, style_selections,
performance_selection, aspect_ratios_selection, image_number, image_seed, sharpness, guidance_scale
performance_selection, aspect_ratios_selection, image_number, output_format, image_seed, sharpness, guidance_scale
]

ctrls += [base_model, refiner_model, refiner_switch] + lora_ctrls
Expand Down Expand Up @@ -622,7 +626,7 @@ def parse_meta(raw_prompt_txt, is_generating):
load_parameter_button.click(modules.meta_parser.load_parameter_button_click, inputs=[prompt, state_is_generating], outputs=load_data_outputs, queue=False, show_progress=False)

def trigger_metadata_import(filepath, state_is_generating):
parameters, items, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
parameters, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
if parameters is None:
print('Could not find metadata in the image!')
parsed_parameters = {}
Expand Down

0 comments on commit b6d2367

Please sign in to comment.