init fix + don't reinit + buttons

app.py

@@ -52,12 +52,15 @@ class Convert:
 def init():
   global model
 
-  model = JukeboxVQVAE.from_pretrained(
-    model_id,
-    device_map = "auto",
-    torch_dtype = t.float16,
-    cache_dir = f"{cache_path}/jukebox/models"
-  )
+  try:
+    model
+    print("Model already initialized")
+  except NameError:
+    model = JukeboxVQVAE.from_pretrained(
+      model_id,
+      torch_dtype = t.float16,
+      cache_dir = f"{cache_path}/jukebox/models"
+    )
 
 def validate_tokens_list(tokens_list):
   # Make sure that:
@@ -85,8 +88,8 @@ with gr.Blocks() as ui:
   audio = gr.Audio(label='audio')
   
   # Buttons to convert from music tokens to audio (primary) and vice versa (secondary)
-  gr.Button(label="Convert tokens to audio", primary=True).click(Convert.TokensFile.to_audio, tokens, audio)
-  gr.Button(label="Convert audio to tokens", primary=False).click(Convert.Audio.to_tokens_file, audio, tokens)
+  gr.Button("Convert tokens to audio", variant='primary').click(Convert.TokensFile.to_audio, tokens, audio)
+  gr.Button("Convert audio to tokens", variant='secondary').click(Convert.Audio.to_tokens_file, audio, tokens)
   
 if __name__ == '__main__':
   init()

jukwi-vqvae.ipynb

@@ -81,12 +81,15 @@
         "def init():\n",
         "  global model\n",
         "\n",
-        "  model = JukeboxVQVAE.from_pretrained(\n",
-        "    model_id,\n",
-        "    device_map = \"auto\",\n",
-        "    torch_dtype = t.float16,\n",
-        "    cache_dir = f\"{cache_path}/jukebox/models\"\n",
-        "  )\n",
+        "  try:\n",
+        "    model\n",
+        "    print(\"Model already initialized\")\n",
+        "  except NameError:\n",
+        "    model = JukeboxVQVAE.from_pretrained(\n",
+        "      model_id,\n",
+        "      torch_dtype = t.float16,\n",
+        "      cache_dir = f\"{cache_path}/jukebox/models\"\n",
+        "    )\n",
         "\n",
         "def validate_tokens_list(tokens_list):\n",
         "  # Make sure that:\n",
@@ -114,8 +117,8 @@
         "  audio = gr.Audio(label='audio')\n",
         "\n",
         "  # Buttons to convert from music tokens to audio (primary) and vice versa (secondary)\n",
-        "  gr.Button(label=\"Convert tokens to audio\", primary=True).click(Convert.TokensFile.to_audio, tokens, audio)\n",
-        "  gr.Button(label=\"Convert audio to tokens\", primary=False).click(Convert.Audio.to_tokens_file, audio, tokens)\n",
+        "  gr.Button(\"Convert tokens to audio\", variant='primary').click(Convert.TokensFile.to_audio, tokens, audio)\n",
+        "  gr.Button(\"Convert audio to tokens\", variant='secondary').click(Convert.Audio.to_tokens_file, audio, tokens)\n",
         "\n",
         "if __name__ == '__main__':\n",
         "  init()\n",