{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d179a842",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2ac61060",
   "metadata": {},
   "outputs": [],
   "source": [
    "DEMO_PATH = \"/Users/glelan/projects/melodyflow/checkpoints/t18_32k_mono\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3b7480c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from audiocraft.models.melodyflow import MelodyFlow\n",
    "from audiocraft.data.audio import audio_read\n",
    "import IPython.display as ipd\n",
    "import random\n",
    "import numpy as np\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e04e72f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "melodyflow = MelodyFlow.get_pretrained(DEMO_PATH, device='mps')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "22e73988",
   "metadata": {},
   "outputs": [],
   "source": [
    "waveform, sr = audio_read('/Users/glelan/Music/californication.wav')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4b567c1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "duration = 10.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4fd89d39",
   "metadata": {},
   "outputs": [],
   "source": [
    "waveform = waveform[:, :int(duration * sr)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5db79484",
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_sequence = melodyflow.encode_audio(waveform[None, :].to('mps'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "26824720",
   "metadata": {},
   "outputs": [],
   "source": [
    "original_sample = melodyflow.generate_audio(latent_sequence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ec9a001",
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.Audio(original_sample[0].cpu().numpy(), rate=sr, autoplay=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bef8c82f",
   "metadata": {},
   "outputs": [],
   "source": [
    "melodyflow.lm.cfg_coef = 4.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae47b32c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e59d4a9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio, latents = melodyflow.generate_unconditional(2, steps=16, progress=True, return_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7c71fd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.Audio(audio[0].cpu().numpy(), rate=sr, autoplay=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f1f22e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.Audio(audio[1].cpu().numpy(), rate=sr, autoplay=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d9ccf9ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "melodyflow.lm.cfg_coef = 4.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4ef6a98",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio, latents = melodyflow.generate([\"A song about friendship.\"], solver=\"euler\", steps=32, progress=True, return_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a267aa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.Audio(audio[0].cpu().numpy(), rate=sr, autoplay=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "4a002f29",
   "metadata": {},
   "outputs": [],
   "source": [
    "melodyflow.lm.cfg_coef = 4.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d21a763",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio, final_latents = melodyflow.edit(latent_sequence,\n",
    "                        [\"Piano Intro.\"],\n",
    "                        inversion=\"renoise\",\n",
    "                        target_latent=0.0,\n",
    "                        steps=64,\n",
    "                        return_tokens=True,\n",
    "                        src_descriptions=[\"Bass Intro.\"],\n",
    "                        progress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbdd81a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ipd.Audio(audio[0].cpu().numpy(), rate=sr, autoplay=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4947271",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 5
}