mirror of
https://github.com/nod-ai/SHARK-Studio.git
synced 2026-01-08 05:24:00 -05:00
Small tweaks to ckpt processing, add tool to prefix params keys
This commit is contained in:
@@ -262,6 +262,8 @@ class StableDiffusion:
|
||||
control_mode,
|
||||
hints,
|
||||
):
|
||||
if seed == -1:
|
||||
seed = randint(0, sys.maxsize)
|
||||
img = self.sd_pipe.generate_images(
|
||||
prompt,
|
||||
negative_prompt,
|
||||
|
||||
@@ -71,7 +71,12 @@ def save_irpa(weights_path, prepend_str):
|
||||
new_key = prepend_str + key
|
||||
archive.add_tensor(new_key, weights[key])
|
||||
|
||||
irpa_file = weights_path.replace(".safetensors", ".irpa")
|
||||
if "safetensors" in weights_path:
|
||||
irpa_file = weights_path.replace(".safetensors", ".irpa")
|
||||
elif "irpa" in weights_path:
|
||||
irpa_file = weights_path
|
||||
else:
|
||||
return Exception("Invalid file format. Please provide a .safetensors or .irpa file.")
|
||||
archive.save(irpa_file)
|
||||
return irpa_file
|
||||
|
||||
|
||||
20
apps/shark_studio/tools/params_prefixer.py
Normal file
20
apps/shark_studio/tools/params_prefixer.py
Normal file
@@ -0,0 +1,20 @@
|
||||
from apps.shark_studio.modules.ckpt_processing import save_irpa
|
||||
import argparse
|
||||
import safetensors
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--input",
|
||||
type=str,
|
||||
default="",
|
||||
help="input safetensors/irpa",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prefix",
|
||||
type=str,
|
||||
default="",
|
||||
help="prefix to add to all the keys in the irpa",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
output_file = save_irpa(args.input, args.prefix)
|
||||
print("saved irpa to", output_file, "with prefix", args.prefix)
|
||||
Reference in New Issue
Block a user