Repository: nod-ai/SHARK

Issue:
RuntimeError: duplicate registrations
`python apps/shark_studio/web/index.py`
```
/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.
  _torch_pytree._register_pytree_node(
Traceback (most recent call last):
  File "/home/rch/Downloads/SHARK/apps/shark_studio/web/index.py", line 5, in <module>
    from ui.chat import chat_element
  File "/home/rch/Downloads/SHARK/apps/shark_studio/web/ui/chat.py", line 11, in <module>
    from apps.shark_studio.api.llm import (
  File "/home/rch/Downloads/SHARK/apps/shark_studio/api/llm.py", line 1, in <module>
    from turbine_models.custom_models import stateless_llama
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/turbine_models/custom_models/stateless_llama.py", line 9, in <module>
    from shark_turbine.aot import *
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/shark_turbine/aot/__init__.py", line 7, in <module>
    from .compiled_module import CompiledModule
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/shark_turbine/aot/compiled_module.py", line 18, in <module>
    from . import builtins
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/shark_turbine/aot/builtins/__init__.py", line 8, in <module>
    from .jittable import jittable
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/shark_turbine/aot/builtins/jittable.py", line 25, in <module>
    from ...dynamo.passes import (
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/shark_turbine/dynamo/passes.py", line 4, in <module>
    from shark_turbine.dynamo import utils
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/shark_turbine/dynamo/utils.py", line 13, in <module>
    @register_decomposition(torch.ops.aten._scaled_dot_product_flash_attention.default)
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/torch/_decomp/__init__.py", line 185, in decomposition_decorator
    pytree.tree_map_(register, aten_op)
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/torch/utils/_pytree.py", line 798, in tree_map_
    deque(map(func, *flat_args), maxlen=0)  # consume and exhaust the iterable
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/torch/_decomp/__init__.py", line 182, in register
    _add_op_to_registry(registry, op, fn)
  File "/home/rch/Downloads/SHARK/shark.venv/lib/python3.11/site-packages/torch/_decomp/__init__.py", line 55, in _add_op_to_registry
    raise RuntimeError(f"duplicate registrations for {op_overload}")
RuntimeError: duplicate registrations for aten._scaled_dot_product_flash_attention.default
```

What could be the problem?


Discussion:
That line was introduced in this PR. 

https://github.com/nod-ai/SHARK-Turbine/pull/294/files/1ba10b7a049f13c29288ca75d1a32085f73036c2#r1437280202


Temporary fix is to pip install torch==2.1.0

@gpetters94 @aviator19941 any updates on a better solution to SDPFA decomposition? If not, we need to make sure the workaround decomposition isn't added if user's torch install has it already 
It seems like an if statement should be possible here.

Resolving patch:
diff --git a/apps/shark_studio/api/sd.py b/apps/shark_studio/api/sd.py
--- a/apps/shark_studio/api/sd.py
+++ b/apps/shark_studio/api/sd.py
@@ -96,6 +96,7 @@ def __init__(
         num_loras: int = 0,
         import_ir: bool = True,
         is_controlled: bool = False,
+        hf_auth_token=None,
     ):
         self.model_max_length = 77
         self.batch_size = batch_size
@@ -112,7 +113,7 @@ def __init__(
             "unet": {
                 "hf_model_name": base_model_id,
                 "unet_model": unet.UnetModel(
-                    hf_model_name=base_model_id, hf_auth_token=None
+                    hf_model_name=base_model_id, hf_auth_token=hf_auth_token
                 ),
                 "batch_size": batch_size,
                 # "is_controlled": is_controlled,
@@ -125,8 +126,8 @@ def __init__(
             "vae_encode": {
                 "hf_model_name": base_model_id,
                 "vae_model": vae.VaeModel(
-                    hf_model_name=base_model_id,
-                    custom_vae=custom_vae,
+                    hf_model_name=custom_vae if custom_vae else base_model_id,
+                    hf_auth_token=hf_auth_token,
                 ),
                 "batch_size": batch_size,
                 "height": height,
@@ -136,8 +137,8 @@ def __init__(
             "vae_decode": {
                 "hf_model_name": base_model_id,
                 "vae_model": vae.VaeModel(
-                    hf_model_name=base_model_id,
-                    custom_vae=custom_vae,
+                    hf_model_name=custom_vae if custom_vae else base_model_id,
+                    hf_auth_token=hf_auth_token,
                 ),
                 "batch_size": batch_size,
                 "height": height,
@@ -184,7 +185,7 @@ def prepare_pipe(self, custom_weights, adapters, embeddings, is_img2img):
                 custom_weights_params, _ = process_custom_pipe_weights(custom_weights)
                 if submodel not in ["clip", "clip2"]:
                     self.static_kwargs[submodel][
-                        "external_weight_file"
+                        "external_weights"
                     ] = custom_weights_params
                 else:
                     self.static_kwargs[submodel]["external_weight_path"] = os.path.join(
diff --git a/apps/shark_studio/web/ui/sd.py b/apps/shark_studio/web/ui/sd.py
--- a/apps/shark_studio/web/ui/sd.py
+++ b/apps/shark_studio/web/ui/sd.py
@@ -31,6 +31,11 @@
     resize_stencil,
 )
 from apps.shark_studio.modules.shared_cmd_opts import cmd_opts
+from apps.shark_studio.web.ui.utils import (
+    nodlogo_loc,
+    none_to_str_none,
+    str_none_to_none,
+)
 from apps.shark_studio.web.utils.state import (
     status_label,
 )
@@ -119,7 +124,7 @@ def pull_sd_configs(
     controlnets,
     embeddings,
 ):
-    sd_args = locals()
+    sd_args = str_none_to_none(locals())
     sd_cfg = {}
     for arg in sd_args:
         if arg in [
@@ -135,11 +140,12 @@ def pull_sd_configs(
                 sd_cfg[arg] = {}
         else:
             sd_cfg[arg] = sd_args[arg]
-    return sd_cfg
+
+    return json.dumps(sd_cfg)
 
 
 def load_sd_cfg(sd_json: dict, load_sd_config: str):
-    new_sd_config = json.loads(view_json_file(load_sd_config))
+    new_sd_config = none_to_str_none(json.loads(view_json_file(load_sd_config)))
     if sd_json:
         for key in new_sd_config:
             sd_json[key] = new_sd_config[key]
@@ -696,6 +702,63 @@ def base_model_changed(base_model_id):
                                     inputs=[sd_json, sd_config_name],
                                     outputs=[sd_config_name],
                                 )
+                            )
+                        )
+                    with gr.Column(scale=1):
+                        clear_sd_config = gr.ClearButton(
+                            value="Clear Config", size="sm", components=sd_json
+                        )
+                        with gr.Row():
+                            save_sd_config = gr.Button(value="Save Config", size="sm")
+                            sd_config_name = gr.Textbox(
+                                value="Config Name",
+                                info="Name of the file this config will be saved to.",
+                                interactive=True,
+                            )
+                        load_sd_config = gr.FileExplorer(
+                            label="Load Config",
+                            file_count="single",
+                            root=(
+                                cmd_opts.configs_path
+                                if cmd_opts.configs_path
+                                else get_configs_path()
+                            ),
+                            height=75,
+                        )
+                        load_sd_config.change(
+                            fn=load_sd_cfg,
+                            inputs=[sd_json, load_sd_config],
+                            outputs=[
+                                prompt,
+                                negative_prompt,
+                                sd_init_image,
+                                height,
+                                width,
+                                steps,
+                                strength,
+                                guidance_scale,
+                                seed,
+                                batch_count,
+                                batch_size,
+                                scheduler,
+                                base_model_id,
+                                custom_weights,
+                                custom_vae,
+                                precision,
+                                device,
+                                ondemand,
+                                repeatable_seeds,
+                                resample_type,
+                                cnet_config,
+                                embeddings_config,
+                                sd_json,
+                            ],
+                        )
+                        save_sd_config.click(
+                            fn=save_sd_cfg,
+                            inputs=[sd_json, sd_config_name],
+                            outputs=[sd_config_name],
+                        )
 
     pull_kwargs = dict(
         fn=pull_sd_configs,
diff --git a/apps/shark_studio/web/ui/utils.py b/apps/shark_studio/web/ui/utils.py
--- a/apps/shark_studio/web/ui/utils.py
+++ b/apps/shark_studio/web/ui/utils.py
@@ -29,3 +29,15 @@ def hsl_color(alpha: float, start, end):
 
     # Return a CSS HSL string
     return f"hsl({math.floor(result)}, 80%, 35%)"
+
+
+def none_to_str_none(props: dict):
+    for key in props:
+        props[key] = "None" if props[key] == None else props[key]
+    return props
+
+
+def str_none_to_none(props: dict):
+    for key in props:
+        props[key] = None if props[key] == "None" else props[key]
+    return props


Test patch: