Support safetensors files for replicate/dreambooth

[Conaire]

Currently only ckpt files are supported for dreambooth using the ckpt_base field, would be great for support of safetensors to be added too!

Current output when attempting to use a safetensors file:

[Conaire]

Looking into this further, don't think this was specifically a safetensors issue, got the same error as above with a 'ckpt' file, the 'model_index.json' can be missing for both types.

[Pwntus]

Fixed in https://github.com/replicate/dreambooth/pull/33

[spuliz]

Hi @Pwntus this fix is not working. I am providing a .safetensors file path in order to fine tune my own SD model and this is the error I am getting:

Traceback (most recent call last):
File "/src/convert_original_stable_diffusion_to_diffusers.py", line 150, in <module>
pipe = download_from_original_stable_diffusion_ckpt(
File "/root/.pyenv/versions/3.10.10/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py", line 1056, in download_from_original_stable_diffusion_ckpt
checkpoint = torch.load(checkpoint_path, map_location=device)
File "/root/.pyenv/versions/3.10.10/lib/python3.10/site-packages/torch/serialization.py", line 795, in load
return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
File "/root/.pyenv/versions/3.10.10/lib/python3.10/site-packages/torch/serialization.py", line 1002, in _legacy_load
magic_number = pickle_module.load(f, **pickle_load_args)
_pickle.UnpicklingError: invalid load key, '`'.
Traceback (most recent call last):
File "/root/.pyenv/versions/3.10.10/lib/python3.10/site-packages/diffusers/configuration_utils.py", line 355, in load_config
config_file = hf_hub_download(
File "/root/.pyenv/versions/3.10.10/lib/python3.10/site-packages/huggingface_hub/utils/_validators.py", line 120, in _inner_fn
return fn(*args, **kwargs)
File "/root/.pyenv/versions/3.10.10/lib/python3.10/site-packages/huggingface_hub/file_download.py", line 1253, in hf_hub_download
raise LocalEntryNotFoundError(
huggingface_hub.utils._errors.LocalEntryNotFoundError: Cannot find the requested files in the disk cache and outgoing traffic has been disabled. To enable hf.co look-ups and downloads online, set 'local_files_only' to False.
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/root/.pyenv/versions/3.10.10/lib/python3.10/site-packages/cog/server/worker.py", line 217, in _predict
result = predict(**payload)
File "/src/predictor.py", line 309, in predict
main(args)
File "/src/dreambooth.py", line 554, in main
vae=AutoencoderKL.from_pretrained(
File "/root/.pyenv/versions/3.10.10/lib/python3.10/site-packages/diffusers/models/modeling_utils.py", line 491, in from_pretrained
config, unused_kwargs, commit_hash = cls.load_config(
File "/root/.pyenv/versions/3.10.10/lib/python3.10/site-packages/diffusers/configuration_utils.py", line 382, in load_config
raise EnvironmentError(
OSError: cog_custom_base_data/vae does not appear to have a file named config.json.

The error message suggests:

The Python script convert_original_stable_diffusion_to_diffusers.py is having trouble loading a checkpoint. Specifically, the torch.load() function is trying to unpickle a file and encountering an "invalid load key". This is often due to the file not being a valid PyTorch model or not being pickle-serialized.

To fix this: if path.endswith("ckpt"): sd = torch.load(path, map_location="cpu")["state_dict"] elif path.endswith("safetensors"): sd = load_safetensors(path)

Source Fix: https://github.com/Stability-AI/generative-models/blob/e5dc9669eddacb6606ac08f7e11dcf452e636bfe/sgm/models/diffusion.py#L89-L92

[j1-dev]

I still haven't been able to open any .safetensors file. Need fix asap, now that sdxl 1.0 will only be released as a .safetensor. either that or add sdxl 1.0 as a new trainer version ASAP. We need this

[spuliz]

Any update? 🙏

[spuliz]

Hi @Pwntus @j1-dev is there any way to run dreambooth training on .safetensors files on replicate?

[chen-rn]

Bump!

Having support for this would mean that we can simply get a download link of a .safetensor file from something like civitai and directly apply Dreambooth directly on top of that! Even though Replicate now has support for fine tunine SDXL, the resulting images are not nearly as good as a finetuned SD1.5 model.

Otherwise, we need to download the model, convert to ckpt, and host it somewhere accessible via http.

[DanielKwakye]

Any update on this?

[dongle132]

Any update on this OR any other way to train Dreambooth with external SAFETENSOR files ???