klimaleksus / stable-diffusion-webui-embedding-merge

Extension for AUTOMATIC1111/stable-diffusion-webui for creating and merging Textual Inversion embeddings at runtime from string literals.
The Unlicense
106 stars 10 forks source link

Safetensor vs. unsafe pickle support #14

Open maxxrox opened 3 months ago

maxxrox commented 3 months ago

First up: amazing work. I use this extension every time I'm prompting and it's a thing of beauty.

The issue: Automatic1111 doesn't enable unsafe unpickle by default, so unless the below-mentioned flag is passed to webui at startup (not a great idea, security-wise), creating embeddings via the Embedding Merge extension fails with the following error in the console:

*** Error verifying pickled file from F:\Utility\Automatic1111\stable-diffusion-webui\embeddings\_EmbeddingMerge_temp.pt
*** The file may be malicious, so the program is not going to read it.
*** You can skip this check with --disable-safe-unpickle commandline argument.
***
    Traceback (most recent call last):
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\safe.py", line 137, in load_with_extra
        check_pt(filename, extra_handler)
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\safe.py", line 84, in check_pt
        check_zip_filenames(filename, z.namelist())
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\safe.py", line 76, in check_zip_filenames
        raise Exception(f"bad file inside {filename}: {name}")
    Exception: bad file inside F:\Utility\Automatic1111\stable-diffusion-webui\embeddings\_EmbeddingMerge_temp.pt: _EmbeddingMerge_temp/byteorder

---
Traceback (most recent call last):
  File "F:\Utility\Automatic1111\stable-diffusion-webui\extensions\stable-diffusion-webui-embedding-merge\scripts\embedding_merge.py", line 1151, in need_save_embed
    token = list(pt['string_to_param'].keys())[0]
TypeError: 'NoneType' object is not subscriptable

I started playing around with the EM scripts but it looks like there's a call to Automatic1111's textual_inversion.py (specifically, create_embedding). I tried pulling just that bit out separately, but kept getting errors; I have no idea what I'm missing (not coded in Python before).

Would it be possible to support saving (both the interim _EmbeddingMerge_temp file and the final output) in safetensors format instead of .pt?

Please let me know if I can provide additional information.

aleksusklim commented 3 months ago

I think I know how to save in .safetensors, I did this in other scripts.

Are any .pt considered "unsafe"? What about other TI embeddings out there (for example https://sliders.baulab.info/weights/xl_sliders/), are they all effectively deprecated?

Why I didn't get this error while I'm sure that I tested at default Forge installation, and also in A1111 WebUI without any fancy command-line arguments?

I mean, is the problem really just in ".pt" but not in something else too?

maxxrox commented 3 months ago

I'm not actually quite certain why some .pt are marked unsafe and others aren't. Perhaps 1 in 10 or 20 from civitai are considered unsafe and require conversion prior to use in Automatic1111.

I'd be fine converting the outputs from Embedding Merge, but in this case, without unsafe pickles enabled, the final embed never gets created.

maxxrox commented 3 months ago

I pulled down cartoon_style.pt from your link and reloaded the embedding list, but getting a separate error in Automatic1111; unsure if this is related to the unsafe pickling issue, though; normally I see this if a file's in an unexpected format for the loader (like having a LoRa in an embedding folder):

*** Error loading embedding cartoon_style.pt
    Traceback (most recent call last):
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\textual_inversion\textual_inversion.py", line 203, in load_from_dir
        self.load_from_file(fullfn, fn)
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\textual_inversion\textual_inversion.py", line 184, in load_from_file
        embedding = create_embedding_from_data(data, name, filename=filename, filepath=path)
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\textual_inversion\textual_inversion.py", line 306, in create_embedding_from_data
        raise Exception(f"Couldn't identify {filename} as neither textual inversion embedding nor diffuser concept.")
    Exception: Couldn't identify cartoon_style.pt as neither textual inversion embedding nor diffuser concept.

---
aleksusklim commented 3 months ago

Those were embeddings for SDXL model; for SD1 you would need https://sliders.baulab.info/weights/sd14_sliders/

aleksusklim commented 3 months ago

Oh wait, they all are LoRAs, I'm dumb.

aleksusklim commented 3 months ago

I use negative embeddings more often than any positive one, so examples are: https://huggingface.co/klnaD/negative_embeddings/tree/main https://huggingface.co/gemasai/realisticvision-negative-embedding/tree/main https://huggingface.co/datasets/Nerfgun3/bad_prompt/tree/main

maxxrox commented 3 months ago

If they're LoRas, definitely makes sense to get that error, hah!

I downloaded every single one of those .pt negative embeddings you linked from HuggingFace, no issues (naturally).

Here's one from Civitai that required conversion before Automatic1111 would load it: https://civitai.com/models/353809/mandip-gill-british-actress.

Conversion script I've been using is attached, in case you're curious. It's in .psm1 format as I load it as a module, but it can simply be renamed to a .ps1 file and called directly (or contents pasted into a PowerShell prompt).

Convert-EmEmbeddings.psm1.txt

maxxrox commented 3 months ago

Forgot to add the python bits called via the PowerShell script; whoops. The attached was adapted from DiffusionDalmation's code here: https://github.com/DiffusionDalmation/pt_to_safetensors_converter_notebook

convertPtToSafetensors.py.txt

aleksusklim commented 3 months ago

So, with those TI .pt you don't have errors, but only in EM saving you have the issue?

maxxrox commented 3 months ago

No errors with the PT you linked on HuggingFace, no.

I receive errors in the Automatic1111 console when attempting to load A) the interim temporary embedding created by EM, B) the final .pt file created by EM (prior to conversion), and C) some .pt files (but not all) from Civitai prior to conversion to .safetensors.

Without the unsafe unpickle flag set, I never receive a final EM embedding file; only the temp .pt file is created, but when the code attempts to read the temp file and generate the final .pt, Automatic1111 will refuse to read it.

If I set the unsafe unpickle flag, I can create the embeddings with EM (because Automatic1111 can read the temp .pt), but the final .pt will still throw an error if I subsequently remove the unsafe unpickle flag (default, safer config). Conversion to .safetensors is the only way I've found around that.

maxxrox commented 3 months ago

Did a bit more reading; if it's helpful, I suspect this may be the change in Automatic1111 that's foundationally causing issues with .pt files: https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/14261 - byteorder referenced here: https://github.com/pytorch/pytorch/issues/101688

If that is indeed the issue, there could be a few fixes:

aleksusklim commented 3 months ago

The only place where I do torch.load is the explicit saving of a new SD1 embedding. What if I'll just pass weights_only=True when doing so?

For me, there was no error in the first place. And now nothing breaks, everything still works. Also I've added a fallback to older call in case this one throws.

Can you check now? Is the error still there?

maxxrox commented 3 months ago

Same issue, sadly; only the line number changed due to the commits:

*** Error verifying pickled file from F:\Utility\Automatic1111\stable-diffusion-webui\embeddings\_EmbeddingMerge_temp.pt
*** The file may be malicious, so the program is not going to read it.
*** You can skip this check with --disable-safe-unpickle commandline argument.
***
    Traceback (most recent call last):
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\safe.py", line 137, in load_with_extra
        check_pt(filename, extra_handler)
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\safe.py", line 84, in check_pt
        check_zip_filenames(filename, z.namelist())
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\safe.py", line 76, in check_zip_filenames
        raise Exception(f"bad file inside {filename}: {name}")
    Exception: bad file inside F:\Utility\Automatic1111\stable-diffusion-webui\embeddings\_EmbeddingMerge_temp.pt: _EmbeddingMerge_temp/byteorder

---
Traceback (most recent call last):
  File "F:\Utility\Automatic1111\stable-diffusion-webui\extensions\stable-diffusion-webui-embedding-merge\scripts\embedding_merge.py", line 1155, in need_save_embed
    token = list(pt['string_to_param'].keys())[0]
TypeError: 'NoneType' object is not subscriptable
aleksusklim commented 3 months ago

What is your system? Is it an ordinary Windows? Or a some rare kind of Linux? What is your PC ? Is it, like, a normal computer/laptop or some kind of weird device?

Try this:

  1. Run WebUI without any unsafe-pickle flags
  2. Load any SD1 model as normally
  3. Got to Train tab, with Create embedding default subtab there
  4. Type test in Name, and also test in Initialization text
  5. Click Create embedding (leaving Number of vectors per token = 1)
  6. Generate an image with prompt test and observe "TI hashes" in metadata

Do you get any error here?

maxxrox commented 3 months ago

Windows 10; desktop PC, Intel i7-9700K processor, 32GB RAM, nVidia 2080ti. Pretty standard setup.

Using the native "Train" tab in Automatic1111 results in the same sort of error. The empty embedding does get created, but doesn't look like the system can load it, either:

*** Error verifying pickled file from F:\Utility\Automatic1111\stable-diffusion-webui\embeddings\test123.pt
*** The file may be malicious, so the program is not going to read it.
*** You can skip this check with --disable-safe-unpickle commandline argument.
***
    Traceback (most recent call last):
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\safe.py", line 137, in load_with_extra
        check_pt(filename, extra_handler)
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\safe.py", line 84, in check_pt
        check_zip_filenames(filename, z.namelist())
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\safe.py", line 76, in check_zip_filenames
        raise Exception(f"bad file inside {filename}: {name}")
    Exception: bad file inside F:\Utility\Automatic1111\stable-diffusion-webui\embeddings\test123.pt: test123/byteorder

---
*** Error loading embedding test123.pt
    Traceback (most recent call last):
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\textual_inversion\textual_inversion.py", line 202, in load_from_dir
        self.load_from_file(fullfn, fn)
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\textual_inversion\textual_inversion.py", line 183, in load_from_file
        embedding = create_embedding_from_data(data, name, filename=filename, filepath=path)
      File "F:\Utility\Automatic1111\stable-diffusion-webui\modules\textual_inversion\textual_inversion.py", line 283, in create_embedding_from_data
        if 'string_to_param' in data:  # textual inversion embeddings
    TypeError: argument of type 'NoneType' is not iterable

---
aleksusklim commented 3 months ago

I found a proper Issue where to report this: https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/15214#issuecomment-2018976262

Probably I will work around this in the same way I save SDXL embeddings: by directly writing them manually instead of asking WebUI to create a temporary one.

Pretty standard setup.

True. Very strange! Thank you.

maxxrox commented 3 months ago

As I read more, I would suspect the issue is that changes to pyTorch haven't been addressed by the Automatic1111 team, inferring from the thread here: https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/10179 and the mention of not handling metadata in the safety checker here: https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/14261

Interestingly, the .pt file can be edited and the generated bytorder/data folders manually removed to prevent the safety checker from bailing, per the comment here: https://github.com/ai-dock/stable-diffusion-webui/issues/10#issuecomment-1958504201 . I was able to manually confirm this by creating a blank embedding with EM (it threw an error, as expected), then opening the _EmbeddingMerge_Temp.pt file with 7zip and deleting the .data folder and the byteorder file. Refreshing the embedding list shows this in the console:

Calculating sha256 for F:\Utility\Automatic1111\stable-diffusion-webui\embeddings\_EmbeddingMerge_temp.pt: 33bb5e2b0bb9e1111f4dbbaad382f0d0cc59a7d9647cf086db92afc5f9db6ac5

You did mention you couldn't replicate the issue - are you on a Torch version lower than v2.1 (I'm on 2.1.1)?

aleksusklim commented 3 months ago

(Leaving this Issue open until the upstream pickle bug is fixed too; but my extension should work now!)

All SD1/SD2 embeddings that EM creates were containing: "name": "_EmbeddingMerge_temp" – Which is very bad on its own, because it reveals that this file was made by EmbeddingMerge (while if that would be an intended behavior, I would have been doing this for SDXL embeddings too, but I don't; maybe it is worth including the "formula" of the embedding instead?)

This was a side-effect for asking WebUI to create a temporary one under this name – and that name was in fact inserted as metadata that I hadn't stripped afterwards!

With moving to manually created structure of

{
    "string_to_token": {
        "*": 265
    },
    "string_to_param": {
        "*": "torch.float32 = [<VECTORS>, <768|1024>]"
    },
    "name": "<NAME>",
    "step": 0,
    "sd_checkpoint": null,
    "sd_checkpoint_name": null
}

– I am not exposing an explicit relation to EM anymore. (Though, the embedding itself with step=0 and nulled checkpoints looks forensically suspicious anyway; if anybody would want to disguise it as a trained one, he should probably grab a real embedding from Kohya's trainer and replace the raw tensor inside)

Here is my little script to convert .pt to .json to read internals:

# model2json.py

import sys

if len(sys.argv)<2:
    print('Usage: python model2json.py "file or folder" [...]')
    print('Dumps shape of ckpt/safetensors/pt to .json')
    exit()

import torch
import json
from safetensors import safe_open
from safetensors.torch import save_file

def load_st_ckpt(path,only_st=False):
    try:
        tensors = {}
        with safe_open(path,framework='pt',device='cpu') as model:
            for k in model.keys():
                tensors[k] = model.get_tensor(k)
        return tensors
    except:
        pass
    if only_st:
        return None
    try:
        tensors = {}
        model = torch.load(path,map_location='cpu')
        for k in model.keys():
            tensors[k] = model[k]
        return tensors
    except:
        return None

def save_st_ckpt(path,data,is_safetensors=True):
    if is_safetensors:
        save_file(data,path)
    else:
        torch.save(data,path)

def model2json(tensors,size_as_string=True):
    if torch.is_tensor(tensors):
        if size_as_string:
            return str(tensors.dtype)+' = '+str(list(tensors.size()))
        else:
            return tensors.size();
    if isinstance(tensors,list) or isinstance(tensors,tuple):
        return [model2json(elem) for elem in tensors]
    if isinstance(tensors,dict):
        res = {}
        for key,value in tensors.items():
            res[key] = model2json(value)
        return res
    return tensors

def save_json(path,data):
    if data is None:
        raise Exception('null')
    with open(path,'w',encoding='utf8') as file:
        file.write(json.dumps(data,indent=4))

def work_one_file(file):
    try:
        save_json(file+'.json',model2json(load_st_ckpt(file)))
        print('OK for "'+file+'"')
    except:
        print('FAIL for "'+file+'"')

from os import listdir
from os.path import isfile, join, splitext

good = {
  '.pt': True,
  '.ckpt': True,
  '.safetensors': True,
}
argv = sys.argv
for i in range(1,len(argv)):
    path = argv[i]
    files = None
    try:
        files = [name for name in listdir(path) if isfile(join(path,name))]
    except:
        pass
    if files is None:
        work_one_file(path)
    else:
        for name in files:
            full = join(path,name)
            ext = splitext(full)[1].lower()
            if ext in good:
                work_one_file(full)

#EOF

To use it, either just drag-and-drop .pt onto .py, or (if you don't have torch and safetensors in the default Python environment) – open CMD with activated VENV from WebUI there, and run as any other normal python script.

maxxrox commented 3 months ago

Hooray!

EM now outputs a .pt file (A1111 still complains when it tries to read it, but it writes!), which I can easily convert to .safetensors and run as expected.

Rock on - I have a backlog of custom merges to complete, now :)

aleksusklim commented 3 months ago

A1111 still complains when it tries to read it

Really? Hm-m, isn't it torch.save a real culprit here…

aleksusklim commented 3 months ago

What if I'll try to torch.load right away after saving, and if it throws – would re-save as .safetensors ?

aleksusklim commented 3 months ago

I don't know whether it would still print an error to the console or not, though.

maxxrox commented 3 months ago

I'll update the extension and test with the new commits once the GPU is free :)

maxxrox commented 3 months ago

It works!

The first attempt to read the resultant .pt file throws the usual error in the console, but your new code then creates a final output as a .safetensors file that is readable by Automatic1111.

monk91 commented 2 months ago

Hi everyone, I think you are talking about a similar problem to mine, but I haven't figured out how to solve it. Using FORGE, EM generates .pt files for me, but A1111 doesn't recognize them, how can I use forge and transform the files into .safetensors?

aleksusklim commented 2 months ago

You can use this script, it produces the same files as EM would. Save it to em_pt2safetensors.py and if you have python+pytorch+safetensors in global environment, you may just drag-and-drop a bunch of files or the whole folder directly onto it in Explorer! Otherwise, call it in command-line from WebUI activated venv as any other python script:

# em_pt2safetensors.py

import sys

if len(sys.argv)<2:
    print('Usage: python em_pt2safetensors.py "embedding.pt" | "/folder/" [...]')
    print('Converts SD1/SD2/SDXL embeddings to .safetensors')
    exit()

import torch
from safetensors import safe_open
from safetensors.torch import save_file

def load_pt(path):
    tensors = {}
    model = torch.load(path,map_location='cpu')
    for k in model.keys():
        tensors[k] = model[k]
    return tensors

def save_st(path,data):
    if data is None:
        raise Exception('fail')
    save_file(data,path)

def pt2st(tensors):
    res = None
    try:
        if ('string_to_param' in tensors) and ('*' in tensors['string_to_param']):
            res = {
              'emb_params': tensors['string_to_param']['*'],
            }
    except:
        pass
    if res is None and ('emb_params' in tensors):
        res = {
          'emb_params': tensors['emb_params'],
        }
    if res is None and ('clip_l' in tensors) and ('clip_g' in tensors):
        res = {
          'clip_g': tensors['clip_g'],
          'clip_l': tensors['clip_l'],
        }
    return res

def work_one_file(file):
    try:
        save_st(file[:-3]+'.safetensors',pt2st(load_pt(file)))
        print('OK for "'+file+'"')
    except:
        print('FAIL for "'+file+'"')

from os import listdir
from os.path import isfile, join, splitext

good = {
  '.pt': True,
}
argv = sys.argv
for i in range(1,len(argv)):
    path = argv[i]
    files = None
    try:
        files = [name for name in listdir(path) if isfile(join(path,name))]
    except:
        pass
    if files is None:
        work_one_file(path)
    else:
        for name in files:
            full = join(path,name)
            ext = splitext(full)[1].lower()
            if ext in good:
                work_one_file(full)

#EOF
monk91 commented 2 months ago

I created the file as you told me, but I wasn't able to use it, how do I activate it from the webui? I'm not very expert

aleksusklim commented 2 months ago

Go to stable-diffusion-webui folder. Open cmd.exe from there (you can type "cmd" into the address bar of Explorer, or Shift+RightClick inside the folder and choose CMD/PowerShell; if you do get a blue powershell window – type cmd + Enter there!)

Inside the console, type venv\Scripts\activate (you can press Tab to autocomplete this by parts of the path; just don't run accelerate instead of activate accidentally) and press Enter. You should see (venv) at the beginning of the command-line prompt.

Make sure em_pt2safetensors.py is saved to this folder (WebUI's root) and run: python em_pt2safetensors.py embeddings (where "embeddings" is the standard name of the folder with them; also you might need embeddings\embedding_merge specifically since my script is non-recursive); again, paths can be autocompleted by Tab.

Then you probably would need to get rid of old .pt files manually.

If you see .py files as executable in Explorer, to be able to drag-and-drop embeddings to this script directly without opening venv – you can try this: open the console in any folder (do not activate venv!) and type there pip install torch safetensors I think this is enough to get my script working stand-alone.

monk91 commented 2 months ago

fantastic, thanks