Small fixes to sd, pin mpmath

nod-ai · Mar 1, 2024 · f7d1af4 · f7d1af4
1 parent 60c013e
commit f7d1af4
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 24 deletions.
diff --git a/apps/shark_studio/api/initializers.py b/apps/shark_studio/api/initializers.py
@@ -8,11 +8,12 @@
 
 from apps.shark_studio.modules.timer import startup_timer
 
-# from apps.shark_studio.web.utils.tmp_configs import (
-#     config_tmp,
-#     clear_tmp_mlir,
-#     clear_tmp_imgs,
-# )
+from apps.shark_studio.web.utils.tmp_configs import (
+    config_tmp,
+    clear_tmp_mlir,
+    clear_tmp_imgs,
+    shark_tmp,
+)
 
 
 def imports():
@@ -47,9 +48,9 @@ def initialize():
     # existing temporary images there if they exist. Then we can import gradio.
     # It has to be in this order or gradio ignores what we've set up.
 
-    # config_tmp()
+    config_tmp()
     # clear_tmp_mlir()
-    # clear_tmp_imgs()
+    clear_tmp_imgs()
 
     from apps.shark_studio.web.utils.file_utils import (
         create_checkpoint_folders,
@@ -82,8 +83,8 @@ def dumpstacks():
             code.append(f"""File: "{filename}", line {lineno}, in {name}""")
             if line:
                 code.append("  " + line.strip())
-
-    print("\n".join(code))
+    with open(os.path.join(shark_tmp, "stack_dump.log"), "w") as f:
+        f.write("\n".join(code)) 
 
 
 def setup_middleware(app):

diff --git a/apps/shark_studio/api/sd.py b/apps/shark_studio/api/sd.py
@@ -42,34 +42,37 @@
             "--iree-flow-collapse-reduction-dims",
             "--iree-opt-const-expr-hoisting=False",
             "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807",
+            "--iree-flow-inline-constants-max-byte-length=0",
             "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))",
         ],
     },
-    "vae_encode": {
-        "initializer": vae.export_vae_model,
-        "ireec_flags": [
-            "--iree-flow-collapse-reduction-dims",
-            "--iree-opt-const-expr-hoisting=False",
-            "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807",
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))",
-        ],
-    },
+    # "vae_encode": {
+    #     "initializer": vae.export_vae_model,
+    #     "ireec_flags": [
+    #         "--iree-flow-collapse-reduction-dims",
+    #         "--iree-opt-const-expr-hoisting=False",
+    #         "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807",
+    #         "--iree-flow-inline-constants-max-byte-length=0",
+    #         "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul))",
+    #     ],
+    # },
     "unet": {
         "initializer": unet.export_unet_model,
         "ireec_flags": [
             "--iree-flow-collapse-reduction-dims",
             "--iree-opt-const-expr-hoisting=False",
             "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807",
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))",
+            "--iree-flow-inline-constants-max-byte-length=0",
+            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-pad-linalg-ops{pad-size=32}))",
         ],
     },
     "vae_decode": {
         "initializer": vae.export_vae_model,
         "ireec_flags": [
-            "--iree-flow-collapse-reduction-dims",
             "--iree-opt-const-expr-hoisting=False",
             "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807",
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))",
+            "--iree-flow-inline-constants-max-byte-length=0",
+            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-pad-linalg-ops{pad-size=32}))",
         ],
     },
 }
@@ -111,7 +114,7 @@ def __init__(
             "unet": {
                 "hf_model_name": base_model_id,
                 "unet_model": unet.UnetModel(
-                    hf_model_name=base_model_id, hf_auth_token=hf_auth_token
+                    hf_model_name=base_model_id
                 ),
                 "batch_size": batch_size,
                 # "is_controlled": is_controlled,
@@ -125,7 +128,6 @@ def __init__(
                 "hf_model_name": base_model_id,
                 "vae_model": vae.VaeModel(
                     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,7 +138,6 @@ def __init__(
                 "hf_model_name": base_model_id,
                 "vae_model": vae.VaeModel(
                     hf_model_name=custom_vae if custom_vae else base_model_id,
-                    hf_auth_token=hf_auth_token,
                 ),
                 "batch_size": batch_size,
                 "height": height,

diff --git a/requirements.txt b/requirements.txt
@@ -44,6 +44,7 @@ timm # for MiniGPT4
 langchain
 einops # for zoedepth
 pydantic==2.4.1 # pin until pyinstaller-hooks-contrib works with beta versions
+mpmath==1.3.0
 
 # Keep PyInstaller at the end. Sometimes Windows Defender flags it but most folks can continue even if it errors
 pefile