Small tweaks to ckpt processing, add tool to prefix params keys

apps/shark_studio/api/sd.py

@@ -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,

apps/shark_studio/modules/ckpt_processing.py

@@ -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
 

apps/shark_studio/tools/params_prefixer.py

@@ -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)

requirements.txt

@@ -8,8 +8,8 @@ wheel
 
 torch==2.3.0
 shark-turbine @ git+https://github.com/iree-org/iree-turbine.git@main
-turbine-models @ git+https://github.com/nod-ai/SHARK-Turbine.git@deprecated-constraints#subdirectory=models
-diffusers @ git+https://github.com/nod-ai/diffusers@0.29.0.dev0-shark
+turbine-models @ git+https://github.com/nod-ai/SHARK-Turbine.git@ean-unify-sd#subdirectory=models
+diffusers @ git+https://github.com/nod-ai/diffusers@v0.24.0-release
 brevitas @ git+https://github.com/Xilinx/brevitas.git@6695e8df7f6a2c7715b9ed69c4b78157376bb60b
 
 # SHARK Runner