Add .csv and .json param info to output gallery (#1495)

apps/stable_diffusion/web/ui/outputgallery_ui.py

@@ -12,10 +12,7 @@ from apps.stable_diffusion.web.ui.utils import nodlogo_loc
 from apps.stable_diffusion.web.utils.gradio_configs import (
     gradio_tmp_galleries_folder,
 )
-from apps.stable_diffusion.web.utils.png_metadata import (
-    parse_generation_parameters,
-)
-from apps.stable_diffusion.web.utils.exif_metadata import parse_exif
+from apps.stable_diffusion.web.utils.metadata import displayable_metadata
 
 # -- Functions for file, directory and image info querying
 
@@ -35,37 +32,6 @@ def outputgallery_filenames(subdir) -> list[str]:
         return []
 
 
-def parameters_for_display(image_filename) -> tuple[str, list[list[str]]]:
-    pil_image = Image.open(image_filename)
-
-    # we have PNG generation parameters
-    if "parameters" in pil_image.info:
-        params = parse_generation_parameters(pil_image.info["parameters"])
-
-        # make showing the sizes more compact by using only one line each
-        if params.keys() & {"Size-1", "Size-2"}:
-            params["Size"] = f"{params.pop('Size-1')}x{params.pop('Size-2')}"
-
-        if params.keys() & {"Hires resize-1", "Hires resize-1"}:
-            hires_x = params.pop("Hires resize-1")
-            hires_y = params.pop("Hires resize-2")
-
-            if hires_x == 0 and hires_y == 0:
-                params["Hires resize"] = "None"
-            else:
-                params["Hires resize"] = f"{hires_x}x{hires_y}"
-
-        return "params", list(map(list, params.items()))
-
-    # we have EXIF data, but no generation parameters we know how to read
-    elif pil_image.getexif():
-        return "exif", list(map(list, parse_exif(pil_image).items()))
-
-    # couldn't find anything
-    else:
-        return None, None
-
-
 def output_subdirs() -> list[str]:
     # Gets a list of subdirectories of output_dir and below, as relative paths.
     relative_paths = [
@@ -80,8 +46,9 @@ def output_subdirs() -> list[str]:
     if get_generated_imgs_todays_subdir() not in relative_paths:
         relative_paths.append(get_generated_imgs_todays_subdir())
 
-    # sort subdirectories so that that the date named ones we probably created in this or previous sessions
-    # come first, sorted with the most recent first. Other subdirs are listed after.
+    # sort subdirectories so that that the date named ones we probably created in this or
+    # previous sessions come first, sorted with the most recent first. Other subdirs are listed
+    # after.
     generated_paths = sorted(
         [path for path in relative_paths if path.isnumeric()], reverse=True
     )
@@ -288,20 +255,18 @@ with gr.Blocks() as outputgallery_web:
     def on_select_image(images: list[str], evt: gr.SelectData) -> list:
         # evt.index is an index into the full list of filenames for the current subdirectory
         filename = images[evt.index]
+        params = displayable_metadata(filename)
 
-        # this gets the parameters in the form our dataframe is expecting (list of lists)
-        params_type, params = parameters_for_display(filename)
+        if params:
+            return [
+                filename,
+                list(map(list, params["parameters"].items())),
+            ]
 
-        if params_type == "params":
-            new_parameters = params
-        elif params_type == "exif":
-            new_parameters = [
-                ["Status", "No PNG parameters found, showing EXIF metadata"]
-            ] + params
-        else:
-            new_parameters = [["Status", "No parameters found"]]
-
-        return [filename, new_parameters]
+        return [
+            filename,
+            [["Status", "No parameters found"]],
+        ]
 
     def on_outputgallery_filename_change(filename: str) -> list:
         exists = filename != "None" and os.path.exists(filename)

apps/stable_diffusion/web/ui/txt2img_ui.py

@@ -16,7 +16,7 @@ from apps.stable_diffusion.web.ui.utils import (
     predefined_models,
     cancel_sd,
 )
-from apps.stable_diffusion.web.utils.png_metadata import import_png_metadata
+from apps.stable_diffusion.web.utils.metadata import import_png_metadata
 from apps.stable_diffusion.web.utils.common_label_calc import status_label
 from apps.stable_diffusion.src import (
     args,

apps/stable_diffusion/web/utils/metadata/__init__.py

@@ -0,0 +1,6 @@
+from .png_metadata import (
+    import_png_metadata,
+)
+from .display import (
+    displayable_metadata,
+)

apps/stable_diffusion/web/utils/metadata/csv_metadata.py

@@ -0,0 +1,31 @@
+import csv
+import os
+from .format import humanize, humanizable
+
+
+def csv_path(image_filename: str):
+    return os.path.join(os.path.dirname(image_filename), "imgs_details.csv")
+
+
+def has_csv(image_filename: str) -> bool:
+    return os.path.exists(csv_path(image_filename))
+
+
+def parse_csv(image_filename: str):
+    # We use a reader instead of a DictReader here for images_details.csv files due to the lack of
+    # headers, and then match up the return list for each row with our guess at which column format
+    # the file is using.
+
+    # we assume the final column of the csv has the original filename with full path and match that
+    # against the image_filename. We then exclude the filename from the output, hence the -1's.
+    csv_filename = csv_path(image_filename)
+
+    matches = [
+        humanize(row)
+        for row in csv.reader(open(csv_filename, "r", newline=""))
+        if row
+        and humanizable(row)
+        and os.path.basename(image_filename) in row[-1]
+    ]
+
+    return matches[0] if matches else {}

apps/stable_diffusion/web/utils/metadata/display.py

@@ -0,0 +1,50 @@
+import json
+import os
+from PIL import Image
+from .png_metadata import parse_generation_parameters
+from .exif_metadata import has_exif, parse_exif
+from .csv_metadata import has_csv, parse_csv
+from .format import compact, humanize
+
+
+def displayable_metadata(image_filename: str) -> dict:
+    pil_image = Image.open(image_filename)
+
+    # we have PNG generation parameters (preferred, as it's what the txt2img dropzone reads,
+    # and we go via that for SendTo, and is directly tied to the image)
+    if "parameters" in pil_image.info:
+        return {
+            "source": "png",
+            "parameters": compact(
+                parse_generation_parameters(pil_image.info["parameters"])
+            ),
+        }
+
+    # we have a matching json file (next most likely to be accurate when it's there)
+    json_path = os.path.splitext(image_filename)[0] + ".json"
+    if os.path.isfile(json_path):
+        with open(json_path) as params_file:
+            return {
+                "source": "json",
+                "parameters": compact(
+                    humanize(json.load(params_file), includes_filename=False)
+                ),
+            }
+
+    # we have a CSV file so try that (can be different shapes, and it usually has no
+    # headers/param names so of the things we we *know* have parameters, it's the
+    # last resort)
+    if has_csv(image_filename):
+        params = parse_csv(image_filename)
+        if params:  # we might not have found the filename in the csv
+            return {
+                "source": "csv",
+                "parameters": compact(params),  # already humanized
+            }
+
+    # EXIF data, probably a .jpeg, may well not include parameters, but at least it's *something*
+    if has_exif(image_filename):
+        return {"source": "exif", "parameters": parse_exif(pil_image)}
+
+    # we've got nothing
+    return None

apps/stable_diffusion/web/utils/metadata/exif_metadata.py

@@ -2,6 +2,10 @@ from PIL import Image
 from PIL.ExifTags import Base as EXIFKeys, TAGS, IFD, GPSTAGS
 
 
+def has_exif(image_filename: str) -> bool:
+    return True if Image.open(image_filename).getexif() else False
+
+
 def parse_exif(pil_image: Image) -> dict:
     img_exif = pil_image.getexif()
 

apps/stable_diffusion/web/utils/metadata/format.py

@@ -0,0 +1,115 @@
+# As SHARK has evolved more columns have been added to images_details.csv. However, since
+# no version of the CSV has any headers (yet) we don't actually have anything within the
+# file that tells us which parameter each column is for. So this is a list of known patterns
+# indexed by length which is what we're going to have to use to guess which columns are the
+# right ones for the file we're looking at.
+
+# The same ordering is used for JSON, but these do have key names, however they are not very
+# human friendly, nor do they match up with the what is written to the .png headers
+
+# So these are functions to try and get something consistent out the raw input from all
+# these sources
+
+PARAMS_FORMATS = {
+    9: {
+        "VARIANT": "Model",
+        "SCHEDULER": "Sampler",
+        "PROMPT": "Prompt",
+        "NEG_PROMPT": "Negative prompt",
+        "SEED": "Seed",
+        "CFG_SCALE": "CFG scale",
+        "PRECISION": "Precision",
+        "STEPS": "Steps",
+        "OUTPUT": "Filename",
+    },
+    10: {
+        "MODEL": "Model",
+        "VARIANT": "Variant",
+        "SCHEDULER": "Sampler",
+        "PROMPT": "Prompt",
+        "NEG_PROMPT": "Negative prompt",
+        "SEED": "Seed",
+        "CFG_SCALE": "CFG scale",
+        "PRECISION": "Precision",
+        "STEPS": "Steps",
+        "OUTPUT": "Filename",
+    },
+    12: {
+        "VARIANT": "Model",
+        "SCHEDULER": "Sampler",
+        "PROMPT": "Prompt",
+        "NEG_PROMPT": "Negative prompt",
+        "SEED": "Seed",
+        "CFG_SCALE": "CFG scale",
+        "PRECISION": "Precision",
+        "STEPS": "Steps",
+        "HEIGHT": "Height",
+        "WIDTH": "Width",
+        "MAX_LENGTH": "Max Length",
+        "OUTPUT": "Filename",
+    },
+}
+
+PARAMS_FORMAT_LONGEST = PARAMS_FORMATS[max(PARAMS_FORMATS.keys())]
+
+
+def compact(metadata: dict) -> dict:
+    # we don't want to alter the original dictionary
+    result = dict(metadata)
+
+    # discard the filename because we should already have it
+    if result.keys() & {"Filename"}:
+        result.pop("Filename")
+
+    # make showing the sizes more compact by using only one line each
+    if result.keys() & {"Size-1", "Size-2"}:
+        result["Size"] = f"{result.pop('Size-1')}x{result.pop('Size-2')}"
+    elif result.keys() & {"Height", "Width"}:
+        result["Size"] = f"{result.pop('Height')}x{result.pop('Width')}"
+
+    if result.keys() & {"Hires resize-1", "Hires resize-1"}:
+        hires_y = result.pop("Hires resize-1")
+        hires_x = result.pop("Hires resize-2")
+
+        if hires_x == 0 and hires_y == 0:
+            result["Hires resize"] = "None"
+        else:
+            result["Hires resize"] = f"{hires_y}x{hires_x}"
+
+    return result
+
+
+def humanizable(metadata: dict | list[str], includes_filename=True) -> dict:
+    lookup_key = len(metadata) + (0 if includes_filename else 1)
+    return lookup_key in PARAMS_FORMATS.keys()
+
+
+def humanize(metadata: dict | list[str], includes_filename=True) -> dict:
+    lookup_key = len(metadata) + (0 if includes_filename else 1)
+
+    # For lists we can only work based on the length, we have no other information
+    if isinstance(metadata, list):
+        if humanizable(metadata, includes_filename):
+            return dict(zip(PARAMS_FORMATS[lookup_key].values(), metadata))
+        else:
+            raise KeyError(
+                f"Humanize could not find the format for a parameter list of length {len(metadata)}"
+            )
+
+    # For dictionaries we try to use the matching length parameter format if
+    # available, otherwise we use the longest. Then we swap keys in the
+    # metadata that match keys in the format for the friendlier name that we
+    # have set in the format value
+    if isinstance(metadata, dict):
+        if humanizable(metadata, includes_filename):
+            format = PARAMS_FORMATS[lookup_key]
+        else:
+            format = PARAMS_FORMAT_LONGEST
+
+        return {
+            format[key]: value
+            for (key, value) in metadata.items()
+            if key in format.keys()
+        }
+
+    raise TypeError("Can only humanize parameter lists or dictionaries")