From 3c38b720398bfb587e0c42e8cd91f4c05293c82a Mon Sep 17 00:00:00 2001 From: Or-Tal <47923357+Or-Tal@users.noreply.github.com> Date: Wed, 15 Jan 2025 10:52:27 +0200 Subject: [PATCH] Jasco release jan12 (#527) * Updated Changelog and version for JASCO release * release completions --- README.md | 3 +- demos/jasco_demo.ipynb | 49 +---- docs/JASCO.md | 2 + jasco_demo.ipynb | 489 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 495 insertions(+), 48 deletions(-) create mode 100644 jasco_demo.ipynb diff --git a/README.md b/README.md index 7e4012e3..e06013ae 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,7 @@ At the moment, AudioCraft contains the training code and inference code for: * [MAGNeT](./docs/MAGNET.md): A state-of-the-art non-autoregressive model for text-to-music and text-to-sound. * [AudioSeal](./docs/WATERMARKING.md): A state-of-the-art audio watermarking. * [MusicGen Style](./docs/MUSICGEN_STYLE.md): A state-of-the-art text-and-style-to-music model. +* [JASCO](./docs/JASCO.md): A state-of-the-art text-and-style-to-music model. ## Training code @@ -60,7 +61,7 @@ We provide some [API documentation](https://facebookresearch.github.io/audiocraf #### Is the training code available? -Yes! We provide the training code for [EnCodec](./docs/ENCODEC.md), [MusicGen](./docs/MUSICGEN.md) and [Multi Band Diffusion](./docs/MBD.md). +Yes! We provide the training code for [EnCodec](./docs/ENCODEC.md), [MusicGen](./docs/MUSICGEN.md),[Multi Band Diffusion](./docs/MBD.md) and [JASCO](./docs/JASCO.md). #### Where are the models stored? diff --git a/demos/jasco_demo.ipynb b/demos/jasco_demo.ipynb index 6f0afbd3..3973118d 100644 --- a/demos/jasco_demo.ipynb +++ b/demos/jasco_demo.ipynb @@ -19,58 +19,13 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "ename": "KeyboardInterrupt", - "evalue": "", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mKeyboardInterrupt\u001b[0m Traceback (most recent call last)", - "Cell \u001b[0;32mIn[1], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mos\u001b[39;00m \n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01maudiocraft\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmodels\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m JASCO\n\u001b[1;32m 4\u001b[0m chords_mapping_path \u001b[38;5;241m=\u001b[39m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mabspath(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m../../assets/chord_to_index_mapping.pkl\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m 5\u001b[0m model \u001b[38;5;241m=\u001b[39m JASCO\u001b[38;5;241m.\u001b[39mget_pretrained(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mfacebook/jasco-chords-drums-melody-400M\u001b[39m\u001b[38;5;124m'\u001b[39m, chords_mapping_path\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m../assets/chord_to_index_mapping.pkl\u001b[39m\u001b[38;5;124m'\u001b[39m)\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/audiocraft/__init__.py:24\u001b[0m\n\u001b[1;32m 6\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 7\u001b[0m \u001b[38;5;124;03mAudioCraft is a general framework for training audio generative models.\u001b[39;00m\n\u001b[1;32m 8\u001b[0m \u001b[38;5;124;03mAt the moment we provide the training code for:\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 20\u001b[0m \u001b[38;5;124;03m improves the perceived quality and reduces the artifacts coming from adversarial decoders.\u001b[39;00m\n\u001b[1;32m 21\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 23\u001b[0m \u001b[38;5;66;03m# flake8: noqa\u001b[39;00m\n\u001b[0;32m---> 24\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m data, modules, models\n\u001b[1;32m 26\u001b[0m __version__ \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m1.3.0\u001b[39m\u001b[38;5;124m'\u001b[39m\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/audiocraft/data/__init__.py:10\u001b[0m\n\u001b[1;32m 6\u001b[0m \u001b[38;5;124;03m\"\"\"Audio loading and writing support. Datasets for raw audio\u001b[39;00m\n\u001b[1;32m 7\u001b[0m \u001b[38;5;124;03mor also including some metadata.\"\"\"\u001b[39;00m\n\u001b[1;32m 9\u001b[0m \u001b[38;5;66;03m# flake8: noqa\u001b[39;00m\n\u001b[0;32m---> 10\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m audio, audio_dataset, info_audio_dataset, music_dataset, sound_dataset\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/audiocraft/data/info_audio_dataset.py:19\u001b[0m\n\u001b[1;32m 17\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01maudio_dataset\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AudioDataset, AudioMeta\n\u001b[1;32m 18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01menvironment\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AudioCraftEnvironment\n\u001b[0;32m---> 19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmodules\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mconditioners\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m SegmentWithAttributes, ConditioningAttributes\n\u001b[1;32m 22\u001b[0m logger \u001b[38;5;241m=\u001b[39m logging\u001b[38;5;241m.\u001b[39mgetLogger(\u001b[38;5;18m__name__\u001b[39m)\n\u001b[1;32m 25\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_clusterify_meta\u001b[39m(meta: AudioMeta) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m AudioMeta:\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/audiocraft/modules/__init__.py:22\u001b[0m\n\u001b[1;32m 20\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlstm\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m StreamableLSTM\n\u001b[1;32m 21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mseanet\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m SEANetEncoder, SEANetDecoder\n\u001b[0;32m---> 22\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtransformer\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m StreamingTransformer\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/audiocraft/modules/transformer.py:23\u001b[0m\n\u001b[1;32m 21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m functional \u001b[38;5;28;01mas\u001b[39;00m F\n\u001b[1;32m 22\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcheckpoint\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m checkpoint \u001b[38;5;28;01mas\u001b[39;00m torch_checkpoint\n\u001b[0;32m---> 23\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mxformers\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ops\n\u001b[1;32m 25\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mrope\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m RotaryEmbedding\n\u001b[1;32m 26\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mstreaming\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m StreamingModule\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/xformers/__init__.py:12\u001b[0m\n\u001b[1;32m 9\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\n\u001b[1;32m 11\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _cpp_lib\n\u001b[0;32m---> 12\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcheckpoint\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ( \u001b[38;5;66;03m# noqa: E402, F401\u001b[39;00m\n\u001b[1;32m 13\u001b[0m checkpoint,\n\u001b[1;32m 14\u001b[0m get_optimal_checkpoint_policy,\n\u001b[1;32m 15\u001b[0m list_operators,\n\u001b[1;32m 16\u001b[0m selective_checkpoint_wrapper,\n\u001b[1;32m 17\u001b[0m )\n\u001b[1;32m 19\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 20\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mversion\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m __version__ \u001b[38;5;66;03m# noqa: F401\u001b[39;00m\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/xformers/checkpoint.py:475\u001b[0m\n\u001b[1;32m 471\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcounter \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[1;32m 472\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moptim_output[count] \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[0;32m--> 475\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mSelectiveCheckpointWrapper\u001b[39;00m(ActivationWrapper):\n\u001b[1;32m 476\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, mod, memory_budget\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, policy_fn\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[1;32m 477\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m__version__ \u001b[38;5;241m<\u001b[39m (\u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m1\u001b[39m):\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/xformers/checkpoint.py:496\u001b[0m, in \u001b[0;36mSelectiveCheckpointWrapper\u001b[0;34m()\u001b[0m\n\u001b[1;32m 492\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m:\n\u001b[1;32m 493\u001b[0m \u001b[38;5;28;01mpass\u001b[39;00m\n\u001b[1;32m 495\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;129;43m@torch\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcompiler\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdisable\u001b[49m\n\u001b[0;32m--> 496\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;28;43;01mdef\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;21;43m_get_policy_fn\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m 497\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mnot\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mis_grad_enabled\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[1;32m 498\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;66;43;03m# no need to compute a policy as it won't be used\u001b[39;49;00m\n\u001b[1;32m 499\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43;01mreturn\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43m[\u001b[49m\u001b[43m]\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/compiler/__init__.py:152\u001b[0m, in \u001b[0;36mdisable\u001b[0;34m(fn, recursive)\u001b[0m\n\u001b[1;32m 143\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdisable\u001b[39m(fn\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m, recursive\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m):\n\u001b[1;32m 144\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 145\u001b[0m \u001b[38;5;124;03m This function provides both a decorator and a context manager to disable compilation on a function\u001b[39;00m\n\u001b[1;32m 146\u001b[0m \u001b[38;5;124;03m It also provides the option of recursively disabling called functions\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 150\u001b[0m \u001b[38;5;124;03m recursive (optional): A boolean value indicating whether the disabling should be recursive.\u001b[39;00m\n\u001b[1;32m 151\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m--> 152\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_dynamo\u001b[39;00m\n\u001b[1;32m 154\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_dynamo\u001b[38;5;241m.\u001b[39mdisable(fn, recursive)\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/_dynamo/__init__.py:2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m convert_frame, eval_frame, resume_execution\n\u001b[1;32m 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackends\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mregistry\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m list_backends, lookup_backend, register_backend\n\u001b[1;32m 4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcallback\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m callback_handler, on_compile_end, on_compile_start\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py:48\u001b[0m\n\u001b[1;32m 45\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_python_dispatch\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _disable_current_modes\n\u001b[1;32m 46\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_traceback\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m format_traceback_short\n\u001b[0;32m---> 48\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m config, exc, trace_rules\n\u001b[1;32m 49\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbackends\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mregistry\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m CompilerFn\n\u001b[1;32m 50\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbytecode_analysis\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m remove_dead_code, remove_pointless_jumps\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/_dynamo/trace_rules.py:52\u001b[0m\n\u001b[1;32m 49\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mresume_execution\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m TORCH_DYNAMO_RESUME_IN_PREFIX\n\u001b[1;32m 50\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m getfile, hashable, NP_SUPPORTED_MODULES, unwrap_if_wrapper\n\u001b[0;32m---> 52\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mvariables\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 53\u001b[0m BuiltinVariable,\n\u001b[1;32m 54\u001b[0m FunctorchHigherOrderVariable,\n\u001b[1;32m 55\u001b[0m NestedUserFunctionVariable,\n\u001b[1;32m 56\u001b[0m SkipFunctionVariable,\n\u001b[1;32m 57\u001b[0m TorchInGraphFunctionVariable,\n\u001b[1;32m 58\u001b[0m UserFunctionVariable,\n\u001b[1;32m 59\u001b[0m UserMethodVariable,\n\u001b[1;32m 60\u001b[0m )\n\u001b[1;32m 63\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m typing\u001b[38;5;241m.\u001b[39mTYPE_CHECKING:\n\u001b[1;32m 64\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mvariables\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbase\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m VariableTracker\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/_dynamo/variables/__init__.py:38\u001b[0m\n\u001b[1;32m 30\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdistributed\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m BackwardHookVariable, DistributedVariable, PlacementVariable\n\u001b[1;32m 31\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfunctions\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 32\u001b[0m FunctoolsPartialVariable,\n\u001b[1;32m 33\u001b[0m NestedUserFunctionVariable,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 36\u001b[0m UserMethodVariable,\n\u001b[1;32m 37\u001b[0m )\n\u001b[0;32m---> 38\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mhigher_order_ops\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 39\u001b[0m FunctorchHigherOrderVariable,\n\u001b[1;32m 40\u001b[0m TorchHigherOrderOperatorVariable,\n\u001b[1;32m 41\u001b[0m )\n\u001b[1;32m 42\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01miter\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 43\u001b[0m CountIteratorVariable,\n\u001b[1;32m 44\u001b[0m CycleIteratorVariable,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 47\u001b[0m RepeatIteratorVariable,\n\u001b[1;32m 48\u001b[0m )\n\u001b[1;32m 49\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mlazy\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LazyVariableTracker\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/_dynamo/variables/higher_order_ops.py:14\u001b[0m\n\u001b[1;32m 12\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfx\u001b[39;00m\n\u001b[1;32m 13\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m\n\u001b[0;32m---> 14\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01moperators\u001b[39;00m\n\u001b[1;32m 15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_dynamo\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m get_fake_value\n\u001b[1;32m 16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_dynamo\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mvariables\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ConstantVariable\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/onnx/__init__.py:11\u001b[0m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_C\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _onnx \u001b[38;5;28;01mas\u001b[39;00m _C_onnx\n\u001b[1;32m 4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_C\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_onnx\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 5\u001b[0m _CAFFE2_ATEN_FALLBACK,\n\u001b[1;32m 6\u001b[0m OperatorExportTypes,\n\u001b[1;32m 7\u001b[0m TensorProtoDataType,\n\u001b[1;32m 8\u001b[0m TrainingMode,\n\u001b[1;32m 9\u001b[0m )\n\u001b[0;32m---> 11\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ( \u001b[38;5;66;03m# usort:skip. Keep the order instead of sorting lexicographically\u001b[39;00m\n\u001b[1;32m 12\u001b[0m _deprecation,\n\u001b[1;32m 13\u001b[0m errors,\n\u001b[1;32m 14\u001b[0m symbolic_caffe2,\n\u001b[1;32m 15\u001b[0m symbolic_helper,\n\u001b[1;32m 16\u001b[0m symbolic_opset7,\n\u001b[1;32m 17\u001b[0m symbolic_opset8,\n\u001b[1;32m 18\u001b[0m symbolic_opset9,\n\u001b[1;32m 19\u001b[0m symbolic_opset10,\n\u001b[1;32m 20\u001b[0m symbolic_opset11,\n\u001b[1;32m 21\u001b[0m symbolic_opset12,\n\u001b[1;32m 22\u001b[0m symbolic_opset13,\n\u001b[1;32m 23\u001b[0m symbolic_opset14,\n\u001b[1;32m 24\u001b[0m symbolic_opset15,\n\u001b[1;32m 25\u001b[0m symbolic_opset16,\n\u001b[1;32m 26\u001b[0m symbolic_opset17,\n\u001b[1;32m 27\u001b[0m symbolic_opset18,\n\u001b[1;32m 28\u001b[0m symbolic_opset19,\n\u001b[1;32m 29\u001b[0m symbolic_opset20,\n\u001b[1;32m 30\u001b[0m utils,\n\u001b[1;32m 31\u001b[0m )\n\u001b[1;32m 33\u001b[0m \u001b[38;5;66;03m# TODO(After 1.13 release): Remove the deprecated SymbolicContext\u001b[39;00m\n\u001b[1;32m 34\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_exporter_states\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m ExportTypes, SymbolicContext\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/onnx/errors.py:9\u001b[0m\n\u001b[1;32m 7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _C\n\u001b[1;32m 8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _constants\n\u001b[0;32m----> 9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m diagnostics\n\u001b[1;32m 11\u001b[0m __all__ \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m 12\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mOnnxExporterError\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m 13\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mOnnxExporterWarning\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 16\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mUnsupportedOperatorError\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m 17\u001b[0m ]\n\u001b[1;32m 20\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mOnnxExporterWarning\u001b[39;00m(\u001b[38;5;167;01mUserWarning\u001b[39;00m):\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/onnx/_internal/diagnostics/__init__.py:1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_diagnostic\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 2\u001b[0m create_export_diagnostic_context,\n\u001b[1;32m 3\u001b[0m diagnose,\n\u001b[1;32m 4\u001b[0m engine,\n\u001b[1;32m 5\u001b[0m export_context,\n\u001b[1;32m 6\u001b[0m ExportDiagnosticEngine,\n\u001b[1;32m 7\u001b[0m TorchScriptOnnxExportDiagnostic,\n\u001b[1;32m 8\u001b[0m )\n\u001b[1;32m 9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_rules\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m rules\n\u001b[1;32m 10\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m levels\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/onnx/_internal/diagnostics/_diagnostic.py:12\u001b[0m\n\u001b[1;32m 8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m List, Optional\n\u001b[1;32m 10\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\n\u001b[0;32m---> 12\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m infra\n\u001b[1;32m 13\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m formatter, sarif\n\u001b[1;32m 14\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msarif\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m version \u001b[38;5;28;01mas\u001b[39;00m sarif_version\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/onnx/_internal/diagnostics/infra/__init__.py:1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_infra\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 2\u001b[0m DiagnosticOptions,\n\u001b[1;32m 3\u001b[0m Graph,\n\u001b[1;32m 4\u001b[0m Invocation,\n\u001b[1;32m 5\u001b[0m Level,\n\u001b[1;32m 6\u001b[0m levels,\n\u001b[1;32m 7\u001b[0m Location,\n\u001b[1;32m 8\u001b[0m Rule,\n\u001b[1;32m 9\u001b[0m RuleCollection,\n\u001b[1;32m 10\u001b[0m Stack,\n\u001b[1;32m 11\u001b[0m StackFrame,\n\u001b[1;32m 12\u001b[0m Tag,\n\u001b[1;32m 13\u001b[0m ThreadFlowLocation,\n\u001b[1;32m 14\u001b[0m )\n\u001b[1;32m 15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcontext\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Diagnostic, DiagnosticContext, RuntimeErrorWithDiagnostic\n\u001b[1;32m 17\u001b[0m __all__ \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m 18\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDiagnostic\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m 19\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDiagnosticContext\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 32\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mThreadFlowLocation\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m 33\u001b[0m ]\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/onnx/_internal/diagnostics/infra/_infra.py:11\u001b[0m\n\u001b[1;32m 8\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mlogging\u001b[39;00m\n\u001b[1;32m 9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FrozenSet, List, Mapping, Optional, Sequence, Tuple\n\u001b[0;32m---> 11\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m formatter, sarif\n\u001b[1;32m 14\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mLevel\u001b[39;00m(enum\u001b[38;5;241m.\u001b[39mIntEnum):\n\u001b[1;32m 15\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"The level of a diagnostic.\u001b[39;00m\n\u001b[1;32m 16\u001b[0m \n\u001b[1;32m 17\u001b[0m \u001b[38;5;124;03m This class is used to represent the level of a diagnostic. The levels are defined\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 34\u001b[0m \u001b[38;5;124;03m Level.ERROR = logging.ERROR = 40\u001b[39;00m\n\u001b[1;32m 35\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/onnx/_internal/diagnostics/infra/formatter.py:11\u001b[0m\n\u001b[1;32m 9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_logging\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LazyString\n\u001b[1;32m 10\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _beartype\n\u001b[0;32m---> 11\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m sarif\n\u001b[1;32m 14\u001b[0m \u001b[38;5;66;03m# A list of types in the SARIF module to support pretty printing.\u001b[39;00m\n\u001b[1;32m 15\u001b[0m \u001b[38;5;66;03m# This is solely for type annotation for the functions below.\u001b[39;00m\n\u001b[1;32m 16\u001b[0m _SarifClass \u001b[38;5;241m=\u001b[39m Union[\n\u001b[1;32m 17\u001b[0m sarif\u001b[38;5;241m.\u001b[39mSarifLog,\n\u001b[1;32m 18\u001b[0m sarif\u001b[38;5;241m.\u001b[39mRun,\n\u001b[1;32m 19\u001b[0m sarif\u001b[38;5;241m.\u001b[39mReportingDescriptor,\n\u001b[1;32m 20\u001b[0m sarif\u001b[38;5;241m.\u001b[39mResult,\n\u001b[1;32m 21\u001b[0m ]\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/onnx/_internal/diagnostics/infra/sarif/__init__.py:71\u001b[0m\n\u001b[1;32m 67\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msarif\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_result\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Result\n\u001b[1;32m 68\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msarif\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_result_provenance\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 69\u001b[0m ResultProvenance,\n\u001b[1;32m 70\u001b[0m )\n\u001b[0;32m---> 71\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msarif\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_run\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Run\n\u001b[1;32m 72\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msarif\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_run_automation_details\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 73\u001b[0m RunAutomationDetails,\n\u001b[1;32m 74\u001b[0m )\n\u001b[1;32m 75\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msarif\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_sarif_log\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m SarifLog\n", - "File \u001b[0;32m~/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/torch/onnx/_internal/diagnostics/infra/sarif/_run.py:9\u001b[0m\n\u001b[1;32m 6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mdataclasses\u001b[39;00m\n\u001b[1;32m 7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Any, List, Literal, Optional\n\u001b[0;32m----> 9\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01monnx\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m_internal\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdiagnostics\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01minfra\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01msarif\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m 10\u001b[0m _address,\n\u001b[1;32m 11\u001b[0m _artifact,\n\u001b[1;32m 12\u001b[0m _conversion,\n\u001b[1;32m 13\u001b[0m _external_property_file_references,\n\u001b[1;32m 14\u001b[0m _graph,\n\u001b[1;32m 15\u001b[0m _invocation,\n\u001b[1;32m 16\u001b[0m _logical_location,\n\u001b[1;32m 17\u001b[0m _property_bag,\n\u001b[1;32m 18\u001b[0m _result,\n\u001b[1;32m 19\u001b[0m _run_automation_details,\n\u001b[1;32m 20\u001b[0m _special_locations,\n\u001b[1;32m 21\u001b[0m _thread_flow_location,\n\u001b[1;32m 22\u001b[0m _tool,\n\u001b[1;32m 23\u001b[0m _tool_component,\n\u001b[1;32m 24\u001b[0m _version_control_details,\n\u001b[1;32m 25\u001b[0m _web_request,\n\u001b[1;32m 26\u001b[0m _web_response,\n\u001b[1;32m 27\u001b[0m )\n\u001b[1;32m 30\u001b[0m \u001b[38;5;129m@dataclasses\u001b[39m\u001b[38;5;241m.\u001b[39mdataclass\n\u001b[1;32m 31\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mRun\u001b[39;00m(\u001b[38;5;28mobject\u001b[39m):\n\u001b[1;32m 32\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Describes a single run of an analysis tool, and contains the reported output of that run.\"\"\"\u001b[39;00m\n", - "File \u001b[0;32m:1007\u001b[0m, in \u001b[0;36m_find_and_load\u001b[0;34m(name, import_)\u001b[0m\n", - "File \u001b[0;32m:982\u001b[0m, in \u001b[0;36m_find_and_load_unlocked\u001b[0;34m(name, import_)\u001b[0m\n", - "File \u001b[0;32m:925\u001b[0m, in \u001b[0;36m_find_spec\u001b[0;34m(name, path, target)\u001b[0m\n", - "File \u001b[0;32m:1423\u001b[0m, in \u001b[0;36mfind_spec\u001b[0;34m(cls, fullname, path, target)\u001b[0m\n", - "File \u001b[0;32m:1395\u001b[0m, in \u001b[0;36m_get_spec\u001b[0;34m(cls, fullname, path, target)\u001b[0m\n", - "File \u001b[0;32m:1555\u001b[0m, in \u001b[0;36mfind_spec\u001b[0;34m(self, fullname, target)\u001b[0m\n", - "File \u001b[0;32m:156\u001b[0m, in \u001b[0;36m_path_isfile\u001b[0;34m(path)\u001b[0m\n", - "File \u001b[0;32m:148\u001b[0m, in \u001b[0;36m_path_is_mode_type\u001b[0;34m(path, mode)\u001b[0m\n", - "File \u001b[0;32m:142\u001b[0m, in \u001b[0;36m_path_stat\u001b[0;34m(path)\u001b[0m\n", - "\u001b[0;31mKeyboardInterrupt\u001b[0m: " - ] - } - ], + "outputs": [], "source": [ "import os \n", "from audiocraft.models import JASCO\n", "\n", - "chords_mapping_path = os.path.abspath('../../assets/chord_to_index_mapping.pkl')\n", "model = JASCO.get_pretrained('facebook/jasco-chords-drums-melody-400M', chords_mapping_path='../assets/chord_to_index_mapping.pkl')\n" ] }, diff --git a/docs/JASCO.md b/docs/JASCO.md index cf723939..3c7de25f 100644 --- a/docs/JASCO.md +++ b/docs/JASCO.md @@ -35,6 +35,8 @@ We currently offer two ways to interact with JASCO: We provide a simple API and pre-trained models: - `facebook/jasco-chords-drums-400M`: 400M model, text to music with chords and drums support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-400M) - `facebook/jasco-chords-drums-1B`: 1B model, text to music with chords and drums support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-1B) +- `facebook/jasco-chords-drums-melody-400M`: 400M model, text to music with chords, drums and melody support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-melody-400M) +- `facebook/jasco-chords-drums-melody-1B`: 1B model, text to music with chords, drums and melody support, generates 10-second samples - [🤗 Hub](https://huggingface.co/facebook/jasco-chords-drums-melody-1B) See after a quick example for using the API. diff --git a/jasco_demo.ipynb b/jasco_demo.ipynb new file mode 100644 index 00000000..f408eefb --- /dev/null +++ b/jasco_demo.ipynb @@ -0,0 +1,489 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# JASCO\n", + "Welcome to JASCO's demo jupyter notebook. \n", + "Here you will find a self-contained example of how to use JASCO for temporally controlled music generation.\n", + "\n", + "You can choose a model from the following selection:\n", + "1. facebook/jasco-chords-drums-400M - 10s music generation conditioned on text, chords and drums, 400M parameters\n", + "2. facebook/jasco-chords-drums-1B - 10s music generation conditioned on text, chords and drums, 1B parameters\n", + "\n", + "\n", + "First, we start by initializing the JASCO model:" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/private/home/ortal1/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\n", + " @torch.library.impl_abstract(\"xformers_flash::flash_fwd\")\n", + "/private/home/ortal1/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we will remove `torch.library.impl_abstract` in a future version of PyTorch.\n", + " @torch.library.impl_abstract(\"xformers_flash::flash_bwd\")\n", + "/private/home/ortal1/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n", + "/checkpoint/ortal1/Projects/jasco_release/audiocraft/models/loaders.py:71: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", + " return torch.load(file, map_location=device)\n", + "/private/home/ortal1/miniconda3/envs/jasco_dev/lib/python3.9/site-packages/transformers/models/encodec/modeling_encodec.py:124: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n", + " self.register_buffer(\"padding_total\", torch.tensor(kernel_size - stride, dtype=torch.int64), persistent=False)\n" + ] + } + ], + "source": [ + "import os \n", + "from audiocraft.models import JASCO\n", + "\n", + "chords_mapping_path = os.path.abspath('./assets/chord_to_index_mapping.pkl')\n", + "model = JASCO.get_pretrained('facebook/jasco-chords-drums-1B', chords_mapping_path='./assets/chord_to_index_mapping.pkl')\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next, let us configure the generation parameters. Specifically, you can control the following:\n", + "* `cfg_coef_all` (float, optional): Coefficient used for classifier free guidance - fully conditional term. \n", + " Defaults to 5.0.\n", + "* `cfg_coef_txt` (float, optional): Coefficient used for classifier free guidance - additional text conditional term. \n", + " Defaults to 0.0.\n", + "\n", + "When left unchanged, JASCO will revert to its default parameters." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [], + "source": [ + "model.set_generation_params(\n", + " cfg_coef_all=0.0,\n", + " cfg_coef_txt=5.0\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next, we can go ahead and start generating music given textual prompts." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Text-conditional Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from audiocraft.utils.notebook import display_audio\n", + "\n", + "# set textual prompt\n", + "text = \"Funky groove with electric piano playing blue chords rhythmically\"\n", + "\n", + "# run the model\n", + "print(\"Generating...\") \n", + "output = model.generate(descriptions=[text], progress=True)\n", + "\n", + "# display the result\n", + "print(f\"Text: {text}\\n\")\n", + "display_audio(output, sample_rate=model.compression_model.sample_rate)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we can start adding temporal controls! We begin with conditioning on chord progressions:" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Chords-conditional Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from audiocraft.utils.notebook import display_audio\n", + "\n", + "model.set_generation_params(\n", + " cfg_coef_all=1.5,\n", + " cfg_coef_txt=2.5\n", + ")\n", + "\n", + "# set textual prompt\n", + "text = \"Strings, woodwind, orchestral, symphony.\"\n", + "\n", + "# define chord progression\n", + "chords = [('C', 0.0), ('D', 2.0), ('F', 4.0), ('Ab', 6.0), ('Bb', 7.0), ('C', 8.0)]\n", + "\n", + "# run the model\n", + "print(\"Generating...\")\n", + "output = model.generate_music(descriptions=[text], chords=chords, progress=True)\n", + "\n", + "# display the result\n", + "print(f'Text: {text}')\n", + "print(f'Chord progression: {chords}')\n", + "display_audio(output, sample_rate=model.compression_model.sample_rate)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next, we can condition the generation on drum tracks:" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Drums-conditional Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torchaudio\n", + "from audiocraft.utils.notebook import display_audio\n", + "\n", + "\n", + "# load drum prompt\n", + "drums_waveform, sr = torchaudio.load(\"./assets/sep_drums_1.mp3\")\n", + "\n", + "# set textual prompt \n", + "text = \"distortion guitars, heavy rock, catchy beat\"\n", + "\n", + "# run the model\n", + "print(\"Generating...\")\n", + "output = model.generate_music(\n", + " descriptions=[text],\n", + " drums_wav=drums_waveform,\n", + " drums_sample_rate=sr,\n", + " progress=True\n", + ")\n", + "\n", + "# display the result\n", + "print('drum prompt:')\n", + "display_audio(drums_waveform, sample_rate=sr)\n", + "print(f'Text: {text}')\n", + "display_audio(output, sample_rate=model.compression_model.sample_rate)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also combine multiple temporal controls! Let's move on to generating with both chords and drums conditioning:" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Drums + Chords conditioning" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torchaudio\n", + "from audiocraft.utils.notebook import display_audio\n", + "\n", + "\n", + "# load drum prompt\n", + "drums_waveform, sr = torchaudio.load(\"./assets/sep_drums_1.mp3\")\n", + "\n", + "# set textual prompt \n", + "text = \"string quartet, orchestral, dramatic\"\n", + "\n", + "# define chord progression\n", + "chords = [('C', 0.0), ('D', 2.0), ('F', 4.0), ('Ab', 6.0), ('Bb', 7.0), ('C', 8.0)]\n", + "\n", + "# run the model\n", + "print(\"Generating...\")\n", + "output = model.generate_music(\n", + " descriptions=[text],\n", + " drums_wav=drums_waveform,\n", + " drums_sample_rate=sr,\n", + " chords=chords,\n", + " progress=True\n", + ")\n", + "\n", + "# display the result\n", + "print('drum prompt:')\n", + "display_audio(drums_waveform, sample_rate=sr)\n", + "print(f'Chord progression: {chords}')\n", + "print(f'Text: {text}')\n", + "display_audio(output, sample_rate=model.compression_model.sample_rate)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Melody + Drums + Chords conditioning - inference example" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Source melody:\n" + ] + }, + { + "data": { + "text/html": [ + "\n", + " \n", + " " + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAABj0AAADQCAYAAABcDaP2AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8hTgPZAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAYDElEQVR4nO3df2jc9f0H8NfFtKm2udTUNVlpioXJanGt2GobHPuhmdkoY8UOHJStlqIwErFmbKOwVTYGEQfqCv5iP3R/WCsddGKZSqkuMoxaUwrVzbKB0GCXRJEmNdC0Np/vH6P37fVnkia5yzuPBxz0Pvfu3evzufe97uzT9+eTy7IsCwAAAAAAgCmuotQFAAAAAAAAjAehBwAAAAAAkAShBwAAAAAAkAShBwAAAAAAkAShBwAAAAAAkAShBwAAAAAAkAShBwAAAAAAkAShBwAAAAAAkITKUhdwtuHh4Thy5EhUV1dHLpcrdTkAAAAAAEAJZVkWx44diwULFkRFxcXXcpRd6HHkyJFoaGgodRkAAAAAAEAZ6e7ujoULF150TNmFHtXV1RHxv+Lz+XyJqwEAAAAAAEppYGAgGhoaCvnBxZRd6HH6lFb5fF7oAQAAAAAARESM6JIYLmQOAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAk4bJCj4ceeihyuVxs3ry5sO348ePR0tIS8+bNizlz5sS6deuit7f3cusEAAAAAAC4qDGHHvv27Yunn346li1bVrT9gQceiJdeeil27twZHR0dceTIkbjzzjsvu1AAAAAAAICLGVPo8dlnn8X69evj97//fVx99dWF7f39/fHHP/4xHnnkkbjttttixYoV8cwzz8Sbb74Zb7311rgVDQAAAAAAcLYxhR4tLS2xZs2aaGpqKtre1dUVJ0+eLNq+ZMmSWLRoUXR2dl5epQAAAAAAABdROdq/sGPHjti/f3/s27fvnMd6enpi5syZMXfu3KLtdXV10dPTc97nGxoaiqGhocL9gYGB0ZYEAAAAAAAwupUe3d3dcf/998dzzz0Xs2bNGpcC2tvbo6ampnBraGgYl+cFAAAAAACml1GFHl1dXdHX1xc33XRTVFZWRmVlZXR0dMS2bduisrIy6urq4sSJE3H06NGiv9fb2xv19fXnfc4tW7ZEf39/4dbd3T3mnQEAAAAAAKavUZ3e6vbbb4+DBw8Wbdu4cWMsWbIkfv7zn0dDQ0PMmDEj9u7dG+vWrYuIiEOHDsXhw4ejsbHxvM9ZVVUVVVVVYywfAAAAAADgf0YVelRXV8cNN9xQtG327Nkxb968wvZNmzZFW1tb1NbWRj6fj/vuuy8aGxtj9erV41c1AAAAAADAWUZ9IfNLefTRR6OioiLWrVsXQ0ND0dzcHE888cR4vwwAAAAAAECRXJZlWamLONPAwEDU1NREf39/5PP5UpcDAAAAAACU0Ghyg1FdyBwAAAAAAKBcCT0AAAAAAIAkCD0AAAAAAIAkCD0AAAAAAIAkVJa6ACZXLpeb8NfIsuySr32hMUy+iZoTI5kH4/WcAACjNRm/iy+H3z2lNdnzw39DAQClUsrfxRP1+8ZKDwAAAAAAIAlCDwAAAAAAIAlObzXNlHJJtOXY5Wmy3xfzAAAoB36TcDHlMj/KpQ4AIF0p/t6w0gMAAAAAAEiC0AMAAAAAAEiC0AMAAAAAAEiC0AMAAAAAAEiC0AMAAAAAAEhCZakLSFkulxvV+CzLJqgSSmm082CipDi/JvrYpnjMUlcun7eRML+AlJWyH+uv5aUcv5un0xy5nOM/nY4TAIynifr9c+Z383i+Rorf+VZ6AAAAAAAASRB6AAAAAAAASXB6qwmU4tIgRs88mDiOLWczJwDKg37MaeZCaTn+ADD5JuP713f8xVnpAQAAAAAAJEHoAQAAAAAAJMHprc5yoSvfWzKUvgu99yNljqRnqs2Jy633fCZqHyai1gifQ4CRmqg+fJp+PPX4buZs4zknJmMeTKXfwgCUp/H6LvH9UXpWegAAAAAAAEkQegAAAAAAAElwequzWH40fXnvOdtUmxNTqd6pVCtAivRhzmZOcLapNiemWr0AlB/fJemw0gMAAAAAAEiC0AMAAAAAAEiC0AMAAAAAAEjClLqmRy6XG7fnco62NI12jpgHlMJ49rIzTcR8nkq1AqRIH2YsxjJvzAnKSbnM4YnqwRPlzGMwktrL/XN/sX0o99qB8THZfVhvSYeVHgAAAAAAQBKEHgAAAAAAQBKm1OmtLDHiUswRpoKpNE+nUq0AKdKHGQvzhqmuXOZwudQxFlO59tNS2Afg8ugDjJWVHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBKEHgAAAAAAQBIqS10AAAAAAMB0kMvlCn/OsqyElUC6rPQAAAAAAACSIPQAAAAAAACSIPQAAAAAAACSMKrQo729PW6++eaorq6O+fPnx9q1a+PQoUNFY44fPx4tLS0xb968mDNnTqxbty56e3vHtWgAAAAAgKkmy7LCDZgYowo9Ojo6oqWlJd56663Ys2dPnDx5Mu64444YHBwsjHnggQfipZdeip07d0ZHR0ccOXIk7rzzznEvHAAAAAAA4Ey57DJixY8//jjmz58fHR0d8bWvfS36+/vjC1/4Qmzfvj2+//3vR0TEBx98ENdff310dnbG6tWrL/mcAwMDUVNTE/39/ZHP58daGgAAAAAAkIDR5AaXdU2P/v7+iIiora2NiIiurq44efJkNDU1FcYsWbIkFi1aFJ2dnZfzUgAAAAAAABdVOda/ODw8HJs3b45bb701brjhhoiI6OnpiZkzZ8bcuXOLxtbV1UVPT895n2doaCiGhoYK9wcGBsZaEgAAAAAAMI2NeaVHS0tLvPfee7Fjx47LKqC9vT1qamoKt4aGhst6PgAAAAAAYHoaU+jR2toau3fvjtdffz0WLlxY2F5fXx8nTpyIo0ePFo3v7e2N+vr68z7Xli1bor+/v3Dr7u4eS0kAAAAAAMA0N6rQI8uyaG1tjV27dsVrr70WixcvLnp8xYoVMWPGjNi7d29h26FDh+Lw4cPR2Nh43uesqqqKfD5fdAMAAAAAABitUV3To6WlJbZv3x4vvvhiVFdXF67TUVNTE1deeWXU1NTEpk2boq2tLWprayOfz8d9990XjY2NsXr16gnZAQAAAAAAgIiIXJZl2YgH53Ln3f7MM8/E3XffHRERx48fj5/85Cfx/PPPx9DQUDQ3N8cTTzxxwdNbnW1gYCBqamqiv7/fqg8AAAAAAJjmRpMbjCr0mAxCDwAAAAAA4LTR5AZjupA5AAAAAABAuRnVNT0mU01NzTnbSrko5UKn9rqYyax3pPWV2cKeMbnYvo52/858rhSOTTm50PvkOMPkGcl3w3h+Ji+np47le3aijKT2yT62AOfjtyxQjkb7u24q9y99GKA8WekBAAAAAAAkQegBAAAAAAAkoWxPb1VuFzIv92WK5V7feBrPfZ1Ox22yObZQepP9Obyc15tqPWOq1QukSS8CytF06k3TaV8BphIrPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCRUlroAAAAARiaXyxX+nGXZJbczNZz5/p3pQu/xSJkLcK6J6pcj+RxTWiPpo1Pt/Upxn2A8WOkBAAAAAAAkQegBAAAAAAAkwemtAAAApogLnaLCqSumtpG8f95jGB8T9VnyGS1/Kb5HKe4TjAcrPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCQIPQAAAAAAgCRUlroAAADSlMvlRjQuy7IJrgQAAIDpwkoPAAAAAAAgCUIPAAAAAAAgCU5vBQDAhHDaKgAAACablR4AAAAAAEAShB4AAAAAAEASnN6KcZfL5c67PfVTXJT7fpd7fSlyzBmLC82bsZiouWZuAwAAAJdrJP8GMpZ/a7DSAwAAAAAASILQAwAAAAAASILTWzHupuvpTcp9v8u9vhQ55ozFVJg3U6FGAAAAoLxN1L8vWOkBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkQegBAAAAAAAkobLUBQAAAMBkyeVylxyTZdkkVDJ9jOSYn8nxh/Mb7WdpMvi8AuXISg8AAAAAACAJQg8AAAAAACAJExZ6PP7443HttdfGrFmzYtWqVfHOO+9M1EsBAADAiGRZdskb42skx9zxh0sb7WdpMm4A5WhCQo8XXngh2tra4sEHH4z9+/fH8uXLo7m5Ofr6+ibi5QAAAAAAACYm9HjkkUfinnvuiY0bN8bSpUvjqaeeiquuuir+9Kc/TcTLAQAAAAAAjH/oceLEiejq6oqmpqb/f5GKimhqaorOzs7xfjkAAAAAAICIiKgc7yf85JNP4tSpU1FXV1e0va6uLj744INzxg8NDcXQ0FDh/sDAwHiXBAAAAAAATAPjHnqMVnt7e/zqV786Z7vwAwAAAAAAOJ0XZFl2ybHjHnpcc801ccUVV0Rvb2/R9t7e3qivrz9n/JYtW6Ktra1w/6OPPoqlS5dGQ0PDeJcGAAAAAABMUceOHYuampqLjhn30GPmzJmxYsWK2Lt3b6xduzYiIoaHh2Pv3r3R2tp6zviqqqqoqqoq3J8zZ050d3dHlmWxaNGi6O7ujnw+P95lAky6gYGBaGho0NeAZOhrQGr0NSA1+hqQiizL4tixY7FgwYJLjp2Q01u1tbXFhg0bYuXKlXHLLbfEY489FoODg7Fx48ZL/t2KiopYuHBhYblKPp/XlIGk6GtAavQ1IDX6GpAafQ1IwaVWeJw2IaHHXXfdFR9//HFs3bo1enp64sYbb4xXXnnlnIubAwAAAAAAjJcJu5B5a2vreU9nBQAAAAAAMBEqSl3AhVRVVcWDDz5YdL0PgKlMXwNSo68BqdHXgNToa8B0lMuyLCt1EQAAAAAAAJerbFd6AAAAAAAAjIbQAwAAAAAASILQAwAAAAAASILQAwAAAAAASEJZhh6PP/54XHvttTFr1qxYtWpVvPPOO6UuCeC83njjjfjud78bCxYsiFwuF3/961+LHs+yLLZu3Rpf/OIX48orr4ympqb497//XTTm008/jfXr10c+n4+5c+fGpk2b4rPPPpvEvQD4f+3t7XHzzTdHdXV1zJ8/P9auXRuHDh0qGnP8+PFoaWmJefPmxZw5c2LdunXR29tbNObw4cOxZs2auOqqq2L+/Pnx05/+ND7//PPJ3BWAiIh48sknY9myZZHP5yOfz0djY2O8/PLLhcf1NGAqe+ihhyKXy8XmzZsL2/Q1YLoru9DjhRdeiLa2tnjwwQdj//79sXz58mhubo6+vr5SlwZwjsHBwVi+fHk8/vjj53384Ycfjm3btsVTTz0Vb7/9dsyePTuam5vj+PHjhTHr16+P999/P/bs2RO7d++ON954I+69997J2gWAIh0dHdHS0hJvvfVW7NmzJ06ePBl33HFHDA4OFsY88MAD8dJLL8XOnTujo6Mjjhw5EnfeeWfh8VOnTsWaNWvixIkT8eabb8af//znePbZZ2Pr1q2l2CVgmlu4cGE89NBD0dXVFe+++27cdttt8b3vfS/ef//9iNDTgKlr37598fTTT8eyZcuKtutrwLSXlZlbbrkla2lpKdw/depUtmDBgqy9vb2EVQFcWkRku3btKtwfHh7O6uvrs9/+9reFbUePHs2qqqqy559/PsuyLPvnP/+ZRUS2b9++wpiXX345y+Vy2UcffTRptQNcSF9fXxYRWUdHR5Zl/+tjM2bMyHbu3FkY869//SuLiKyzszPLsiz729/+llVUVGQ9PT2FMU8++WSWz+ezoaGhyd0BgPO4+uqrsz/84Q96GjBlHTt2LLvuuuuyPXv2ZF//+tez+++/P8syv9UAsizLymqlx4kTJ6KrqyuampoK2yoqKqKpqSk6OztLWBnA6H344YfR09NT1NNqampi1apVhZ7W2dkZc+fOjZUrVxbGNDU1RUVFRbz99tuTXjPA2fr7+yMiora2NiIiurq64uTJk0W9bcmSJbFo0aKi3vaVr3wl6urqCmOam5tjYGCg8H9WA5TCqVOnYseOHTE4OBiNjY16GjBltbS0xJo1a4r6V4TfagAREZWlLuBMn3zySZw6daqo6UZE1NXVxQcffFCiqgDGpqenJyLivD3t9GM9PT0xf/78oscrKyujtra2MAagVIaHh2Pz5s1x6623xg033BAR/+tbM2fOjLlz5xaNPbu3na/3nX4MYLIdPHgwGhsb4/jx4zFnzpzYtWtXLF26NA4cOKCnAVPOjh07Yv/+/bFv375zHvNbDaDMQg8AAMpHS0tLvPfee/GPf/yj1KUAXJYvf/nLceDAgejv74+//OUvsWHDhujo6Ch1WQCj1t3dHffff3/s2bMnZs2aVepyAMpSWZ3e6pprrokrrrgient7i7b39vZGfX19iaoCGJvTfetiPa2+vj76+vqKHv/888/j008/1feAkmptbY3du3fH66+/HgsXLixsr6+vjxMnTsTRo0eLxp/d287X+04/BjDZZs6cGV/60pdixYoV0d7eHsuXL4/f/e53ehow5XR1dUVfX1/cdNNNUVlZGZWVldHR0RHbtm2LysrKqKur09eAaa+sQo+ZM2fGihUrYu/evYVtw8PDsXfv3mhsbCxhZQCjt3jx4qivry/qaQMDA/H2228XelpjY2McPXo0urq6CmNee+21GB4ejlWrVk16zQBZlkVra2vs2rUrXnvttVi8eHHR4ytWrIgZM2YU9bZDhw7F4cOHi3rbwYMHi0LdPXv2RD6fj6VLl07OjgBcxPDwcAwNDelpwJRz++23x8GDB+PAgQOF28qVK2P9+vWFP+trwHRXdqe3amtriw0bNsTKlSvjlltuicceeywGBwdj48aNpS4N4ByfffZZ/Oc//ync//DDD+PAgQNRW1sbixYtis2bN8dvfvObuO6662Lx4sXxy1/+MhYsWBBr166NiIjrr78+vv3tb8c999wTTz31VJw8eTJaW1vjBz/4QSxYsKBEewVMZy0tLbF9+/Z48cUXo7q6unBe55qamrjyyiujpqYmNm3aFG1tbVFbWxv5fD7uu+++aGxsjNWrV0dExB133BFLly6NH/7wh/Hwww9HT09P/OIXv4iWlpaoqqoq5e4B09CWLVviO9/5TixatCiOHTsW27dvj7///e/x6quv6mnAlFNdXV241tpps2fPjnnz5hW262vAdFd2ocddd90VH3/8cWzdujV6enrixhtvjFdeeeWcCywBlIN33303vvnNbxbut7W1RUTEhg0b4tlnn42f/exnMTg4GPfee28cPXo0vvrVr8Yrr7xSdO7V5557LlpbW+P222+PioqKWLduXWzbtm3S9wUgIuLJJ5+MiIhvfOMbRdufeeaZuPvuuyMi4tFHHy30q6GhoWhubo4nnniiMPaKK66I3bt3x49//ONobGyM2bNnx4YNG+LXv/71ZO0GQEFfX1/86Ec/iv/+979RU1MTy5Yti1dffTW+9a1vRYSeBqRHXwOmu1yWZVmpiwAAAAAAALhcZXVNDwAAAAAAgLESegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEkQegAAAAAAAEn4P8oFZFHDCXiRAAAAAElFTkSuQmCC", + "text/plain": [ + "
" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Chords:\n", + "[('N', 0.0), ('C', 0.32), ('Dm7', 3.456), ('Am', 4.608), ('F', 8.32), ('C', 9.216)]\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Separated drums:\n" + ] + }, + { + "data": { + "text/html": [ + "\n", + " \n", + " " + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Generating...\n" + ] + }, + { + "data": { + "text/html": [ + "\n", + " \n", + " " + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "\n", + " \n", + " " + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "\n", + " \n", + " " + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "%matplotlib inline\n", + "import torchaudio \n", + "from audiocraft.models import JASCO\n", + "from demucs import pretrained\n", + "from demucs.apply import apply_model\n", + "from demucs.audio import convert_audio\n", + "import torch\n", + "from audiocraft.utils.notebook import display_audio\n", + "import matplotlib.pyplot as plt\n", + "\n", + "# --------------------------\n", + "# First, choose file to load\n", + "# --------------------------\n", + "fnames = ['salience_1', 'salience_2']\n", + "chords = [\n", + " [('N', 0.0), ('Eb7', 1.088000000), ('C#', 4.352000000), ('D', 4.864000000), ('Dm7', 6.720000000), ('G7', 8.256000000), ('Am7b5/G', 9.152000000)], # for salience 1\n", + " [('N', 0.0), ('C', 0.320000000), ('Dm7', 3.456000000), ('Am', 4.608000000), ('F', 8.320000000), ('C', 9.216000000)] # for salience 2\n", + "]\n", + "file_idx = 1 # either 0 or 1\n", + "\n", + "\n", + "# ------------------------------------\n", + "# display audio, melody map and chords\n", + "# ------------------------------------\n", + "def plot_chromagram(tensor):\n", + " # Check if tensor is a PyTorch tensor\n", + " if not torch.is_tensor(tensor):\n", + " raise ValueError('Input should be a PyTorch tensor')\n", + " tensor = tensor.numpy().T # C, T\n", + " plt.figure(figsize=(20, 20))\n", + " plt.imshow(tensor, cmap='binary', interpolation='nearest', origin='lower')\n", + " plt.show()\n", + "\n", + "# load salience and display the corresponding wav\n", + "melody_prompt_wav, melody_prompt_sr = torchaudio.load(f\"./assets/{fnames[file_idx]}.wav\")\n", + "print(\"Source melody:\")\n", + "display_audio(melody_prompt_wav, sample_rate=melody_prompt_sr)\n", + "melody = torch.load(f\"./assets/{fnames[file_idx]}.th\", weights_only=True)\n", + "plot_chromagram(melody)\n", + "print(\"Chords:\")\n", + "print(chords[file_idx])\n", + "\n", + "# --------------------------------------------------\n", + "# use demucs to seperate the drums stem from src mix\n", + "# --------------------------------------------------\n", + "def _get_drums_stem(wav: torch.Tensor, sample_rate: int) -> torch.Tensor:\n", + " \"\"\"Get parts of the wav that holds the drums, extracting the main stems from the wav.\"\"\"\n", + " demucs_model = pretrained.get_model('htdemucs').to('cuda')\n", + " wav = convert_audio(\n", + " wav, sample_rate, demucs_model.samplerate, demucs_model.audio_channels) # type: ignore\n", + " stems = apply_model(demucs_model, wav.cuda().unsqueeze(0), device='cuda').squeeze(0)\n", + " drum_stem = stems[demucs_model.sources.index('drums')] # extract relevant stems for drums conditioning\n", + " return convert_audio(drum_stem.cpu(), demucs_model.samplerate, sample_rate, 1) # type: ignore\n", + "drums_wav = _get_drums_stem(melody_prompt_wav, melody_prompt_sr)\n", + "print(\"Separated drums:\")\n", + "display_audio(drums_wav, sample_rate=melody_prompt_sr)\n", + "\n", + "# ----------------------------------\n", + "# Generate using the loaded controls\n", + "# ----------------------------------\n", + "# these are free-form texts written randomly\n", + "texts = [\n", + " '90s rock with heavy drums and hammond',\n", + " '80s pop with groovy synth bass and drum machine',\n", + " 'folk song with leading accordion',\n", + "]\n", + "\n", + "print(\"Generating...\")\n", + "# replacing dynammic solver with simple euler solver\n", + "model.set_generation_params(cfg_coef_all=1.5, cfg_coef_txt=2.5, euler=True, euler_steps=50) # manually set with euler solver\n", + "output = model.generate_music(\n", + " descriptions=texts,\n", + " chords=chords[file_idx],\n", + " drums_wav=drums_wav,\n", + " drums_sample_rate=melody_prompt_sr,\n", + " melody_salience_matrix=melody.permute(1, 0),\n", + " progress=True\n", + ")\n", + "display_audio(output, sample_rate=model.compression_model.sample_rate)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "jasco_dev", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.19" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +}