{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Codec Engine with FP16 Support\n",
    "\n",
    "This notebook demonstrates using CodecEngine with FP16 precision for faster inference.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np\n",
    "\n",
    "from suno_utils.tasks.codec_engine import CodecEngine, Request, Audio\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import load_model\n",
    "from suno_utils.utils.s3 import read_from_s3\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Model and Create Engine (FP32)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_fp32 = load_model(\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\", use_fp16=False)\n",
    "\n",
    "engine_fp32 = CodecEngine(\n",
    "    model_fp32,\n",
    "    is_vae=True,\n",
    "    compile=True,\n",
    "    enable_profiler=False,\n",
    "    use_fp16=False,\n",
    ")\n",
    "\n",
    "print(f\"FP32 Engine created with model dtype: {next(model_fp32.parameters()).dtype}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Model and Create Engine (FP16)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_fp16 = load_model(\n",
    "    \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\", force_reload=True, use_fp16=True\n",
    ")\n",
    "\n",
    "engine_fp16 = CodecEngine(\n",
    "    model_fp16,\n",
    "    is_vae=True,\n",
    "    compile=True,\n",
    "    enable_profiler=False,\n",
    "    use_fp16=True,\n",
    ")\n",
    "\n",
    "print(f\"FP16 Engine created with model dtype: {next(model_fp16.parameters()).dtype}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Test Data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_id = \"76e47a5c-0f30-414c-93ff-251df117e130\"\n",
    "vae_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}_vae.npz\"\n",
    "\n",
    "print(f\"Loading VAE data from {vae_filepath}\")\n",
    "vae_data = read_from_s3(vae_filepath, read_f=np.load)\n",
    "vae_latents = vae_data[\"vae_latents\"]\n",
    "print(f\"VAE latents shape: {vae_latents.shape}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run FP32 Engine\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "req_fp32 = Request(\n",
    "    gen_id + \"_fp32\",\n",
    "    vae_latents,\n",
    "    input_tokens_finished=True,\n",
    ")\n",
    "\n",
    "job_fp32 = engine_fp32.run_request(req_fp32)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_fp32 = Audio.concatenate(job_fp32.generated_audios)\n",
    "audio_fp32.play()\n",
    "audio_fp32.write_opus(\"fp32.opus\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run FP16 Engine\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "req_fp16 = Request(\n",
    "    gen_id + \"_fp16\",\n",
    "    vae_latents,\n",
    "    input_tokens_finished=True,\n",
    ")\n",
    "\n",
    "job_fp16 = engine_fp16.run_request(req_fp16)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_fp16 = Audio.concatenate(job_fp16.generated_audios)\n",
    "audio_fp16.play()\n",
    "audio_fp16.write_opus(\"fp16.opus\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from suno_utils.diffusion import generation as diffusion_gen\n",
    "from suno_utils.tasks.codec_engine import CodecEngine, Request, Job, Audio\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import load_model, decode_stream_to_full_audio\n",
    "\n",
    "use_fp16 = True\n",
    "model = load_model(\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\", use_fp16=use_fp16)\n",
    "\n",
    "\n",
    "# Configure profiler output\n",
    "s3_profile_bucket = \"suno-data\"\n",
    "s3_profile_prefix = \"traces/codec_engine\" + (\"_fp16\" if use_fp16 else \"_fp32\")\n",
    "local_profile_cache = Path(\"./profile_logs/codec_engine\")\n",
    "local_profile_cache.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "engine = CodecEngine(\n",
    "    model,\n",
    "    is_vae=True,\n",
    "    compile=True,\n",
    "    enable_profiler=True,\n",
    "    use_fp16=use_fp16,\n",
    "    profiler_config={\n",
    "        \"wait_steps\": 0,\n",
    "        \"warmup_steps\": 1,\n",
    "        \"active_steps\": 3,\n",
    "        \"export_backend\": \"s3\",\n",
    "        \"s3_bucket\": s3_profile_bucket,\n",
    "        \"s3_prefix\": s3_profile_prefix,\n",
    "        \"local_directory\": str(local_profile_cache.parent),\n",
    "        \"project_name\": local_profile_cache.name,\n",
    "        \"record_shapes\": True,\n",
    "        \"profile_memory\": True,\n",
    "        \"with_stack\": True,\n",
    "    },\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n",
    "# gen_id = \"a5e2198a-f352-4abb-9a24-7f81b143ded3\"  # stone\n",
    "gen_id = \"c934f808-894b-40b7-87c0-defa814e6d44\"  # v5 jpop\n",
    "\n",
    "\n",
    "s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n",
    "mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "vae_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}_vae.npz\"\n",
    "print(s3_filepath)\n",
    "vae_data = read_from_s3(vae_filepath, read_f=np.load)\n",
    "\n",
    "print(list(vae_data.keys()))\n",
    "vae_data = vae_data[\"vae_latents\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}