generate.py prompt2png does not save metadata

[robotdad]

Describe your environment

• GPU: [cuda]
• VRAM: [10GB]
• CPU arch: [x86]
• OS: [Windows]
• Python: [miniconda]
• Branch: [main]
• Commit: [b0b82efffe2069d8041cf406008cd0dd43b1218a]

Describe the bug When using generate.py I would expect it to save metadata to the images, especially as I'm using this to bypass the UI to generate many images with variations. The call to pngwriter.save_image_and_prompt_to_png in this function does not pass metadata.

Expected behavior I would expect metadata to be saved. I also would expect for prompt2png to handle processing the provided options into an Args class for creating that metadata rather than making that the caller's responsibility. It would be simpler to use that way.

[lstein]

That's a good point. I'll fix it. Also a good opportunity to clean up the metadata handling code.

[robotdad]

I spent some time on this and decided to try the approach used in invoke.py of hooking the image_callback. This works in that I now have metadata in the image, but when I import them into the web ui the metadata isn't used. Is that not implemented to read metadata from an imported image?

Below are the functions I defined based on what is in invoke.py.

def prepare_image_metadata(
        opt,
        prefix,
        seed
):
    wildcards = dict(opt.__dict__)
    wildcards['prefix'] = prefix
    wildcards['seed'] = seed
    try:
        filename = opt.fnformat.format(**wildcards)
    except KeyError as e:
        print(f'** The filename format contains an unknown key \'{e.args[0]}\'. Will use \'{{prefix}}.{{seed}}.png\' instead')
        filename = f'{prefix}.{seed}.png'
    except IndexError as e:
        print(f'** The filename format is broken or complete. Will use \'{{prefix}}.{{seed}}.png\' instead')
        filename = f'{prefix}.{seed}.png'

    formatted_dream_prompt = opt.dream_prompt_str(seed=seed)
    return filename, formatted_dream_prompt

def image_writer(image, seed, first_seed=None, use_prefix=None):
    filename, dreamprompt = prepare_image_metadata(opt, prefix, seed)
    file_writer.save_image_and_prompt_to_png(
        image = image, 
        dream_prompt = dreamprompt, 
        metadata = metadata_dumps(
            opt, 
            seeds = [seed],
            model_hash = gen.model_hash), 
        name=filename)

This is the metadata that gets added to my files.

 {
    "model": "stable diffusion",
    "model_id": "stable-diffusion-1.4",
    "model_hash": "fe4efff1e174c627256e44ec2991ba279b3816e364b49f9be2abc0b3ff3f8556",
    "app_id": "invoke-ai/InvokeAI",
    "app_version": "v2.1.2",
    "image": {
        "width": 512,
        "facetool_strength": 0.0,
        "upscale": null,
        "cfg_scale": 9,
        "steps": 50,
        "perlin": 0.0,
        "facetool": "gfpgan",
        "seed": 1034856165,
        "init_mask": null,
        "prompt": [
            {
                "prompt": "an astronaut riding a horse",
                "weight": 1.0
            }
        ],
        "threshold": 0.0,
        "height": 512,
        "postprocessing": null,
        "sampler": "k_lms",
        "variations": [],
        "type": "txt2img"
    }
}

[robotdad]

!fetch in the cli does pull the info out of the files I'm generating so I'm happy with that.

[psychedelicious]

@robotdad , you need to save the metadata under a 'sd-metadata' key - doesn't look like you have got that in there?

Also, I am sorry to say taht the web UI metadata schema is slightly different than the CLI. The CLI attempts to adhere to a spec defined some time ago, but that spec wasn't able to provide the web UI the core functionality in some cases.

[robotdad]

I setup a GitHub repo here with minimal python scripts using InvokeAI based on the examples in ldm/generate.py. I based my metadata creation on script/invoke.py, if there is a better way that works with the CLI and WebUI that would be great, and I'd be happy to use it. Is there another spot I should look at?

I hope this project does care about usage through other Python scripts and not just the webui. I love the webui, but driving things from Python really opens up a lot more possibilities.

[robotdad]

BTW, I used sd-metadata.py to dump the example out above, that strips the sd-metadata key

[psychedelicious]

The UI metadata generation code is here: https://github.com/invoke-ai/InvokeAI/blob/main/backend/invoke_ai_web_server.py#L1144-L1248

For now, you may need to use this code to generate metadata (parameters is the same kwargs Generate.prompt2image() takes), or inspect the metadata from a UI-generated image and CLI-generated image and transform one to the other after generation.