{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49a3d5bb-3d04-4d11-aba6-043fc5667abd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from kandinsky import get_T2V_pipeline\n",
    "from IPython.display import Video\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9bcf5ed-f813-47e1-ade6-3f655d35d0e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = get_T2V_pipeline(\n",
    "    device_map={\"dit\": \"cuda:0\", \"vae\": \"cuda:0\", \"text_embedder\": \"cuda:0\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f1563bb-0ebc-4f7c-97f4-6918a55d774f",
   "metadata": {},
   "source": [
    "# Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48bc4ef9-0492-41a1-8133-c7e10a894d88",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = pipe(\"a cat in a red hat\", time_length=2, width=768, height=512, save_path='./test.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3379b1a9-b987-443e-aad0-a7f8ac62b69d",
   "metadata": {},
   "outputs": [],
   "source": [
    "Video('./test.mp4')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4de2c5e-2784-4827-92b8-6997d3012e79",
   "metadata": {},
   "source": [
    "# Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9cbc0fc-173f-463e-896b-b29b727b0e65",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = pipe(\n",
    "    \"a dog in a red boots\", \n",
    "    time_length=0, width=768, height=512, \n",
    "    save_path='./image.png'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a2c7741-9d7e-4396-a434-d65ff9e224e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "image = Image.open(\"./image.png\")\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea82ab34-ea66-4559-ac86-ec63af0e0f7a",
   "metadata": {},
   "source": [
    "# Distilled model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d42c0c0-92a9-4f91-a263-b3b9e58dbf99",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = get_T2V_pipeline(\n",
    "    device_map={\"dit\": \"cuda:0\", \"vae\": \"cuda:0\", \"text_embedder\": \"cuda:0\"},\n",
    "    conf_path=\"./configs/config_5s_distil.yaml\"\n",
    ") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14e5d407-135b-451c-a092-840bae230524",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = pipe(\"cheburashka in a blue hat\", time_length=5, width=768, height=512, guidance_weight=1.0, num_steps=16, scheduler_scale=5, save_path='./test2.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53cda76a-a60e-4e62-9cb4-3b4fcf2f0b4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "Video('./test2.mp4')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "004cc1da-ceb8-4e11-9774-8aa231b3fece",
   "metadata": {},
   "source": [
    "# 10s video NABLA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea8fb5f8-ad6d-4e04-a4e3-0b381e3259db",
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe = get_T2V_pipeline(\n",
    "    device_map={\"dit\": \"cuda:0\", \"vae\": \"cuda:0\", \"text_embedder\": \"cuda:0\"},\n",
    "    conf_path=\"./configs/config_10s_sft.yaml\"\n",
    ") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ffa470b-5c29-4e05-88af-fe26c4d0ada1",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = pipe(\"Shiba Inu is driving a car\", time_length=10, width=768, height=512, save_path='./test3.mp4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4bbd345-9ae0-48a6-bdb9-24f4c63d48d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "Video('./test3.mp4')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "test",
   "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}