svd_xt_demo_v2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "be6d1db7-70ed-4652-9e4d-1d9124966c67",
   "metadata": {},
   "source": [
    "# Stable Video Diffusion XT 1.1 on Amazon SageMaker\n",
    "\n",
    "Stability AI's [Stable Video Diffusion XT (SVT-XT) 1.1](https://medium.com/r/?url=https%3A%2F%2Fstability.ai%2Fstable-video) foundation model is a diffusion model that takes in a still image as a conditioning frame and generates a short 4 second video. The notebook walks through configuring, creating, and invoking an [Asynchronous Inference Endpoint](https://docs.aws.amazon.com/sagemaker/latest/dg/async-inference.html) backed by the SVT-XT foundation model on Amazon SageMaker.\n",
    "\n",
    "Version 2 of the Notebook passes an encoded image in the request payload used to invoke the model. Use corresponding custom inference script. Use the corresponding custom inference script, [inference_v2/inference.py](inference_v2/inference.py) when preparing the model archive. \n",
    "\n",
    "**Author:** Gary A. Stafford  \n",
    "**Date:** 2024-05-08\n",
    "\n",
    "![Architecture V2](architecture/async_inference_v2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6477bffd-639d-4480-9f18-ba4bb3a1993a",
   "metadata": {},
   "source": [
    "## 1: Install Required Packages\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1521e122-4cae-4e6f-95d3-90ef474de547",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%sh\n",
    "\n",
    "# optional: update OS packages in Amazon SageMaker Studio Ubuntu environment\n",
    "sudo apt-get update -qq -y && sudo apt-get upgrade -qq -y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "614d104c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%sh\n",
    "\n",
    "sudo apt-get install git libgl1 ffmpeg git-lfs wget -y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40524589-c9dd-4d48-bf12-c7d2d8e87c83",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install sagemaker boto3 botocore ffmpeg-python ipython diffusers pywget -Uq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b4aba5e-1d59-446c-a121-19715ae2f5e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# restart kernel 1x when installing new packages\n",
    "\n",
    "import os\n",
    "\n",
    "os._exit(00)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2267437f-9900-42a6-b086-8c00611f10c7",
   "metadata": {},
   "source": [
    "## 2: Prepare the SVD-XT Model for Inference\n",
    "\n",
    "Steps to prepare the model for inference: 1/ Download the model artifacts from Hugging Face, 2/ add the custom inference script, 3/ create an archive file from the model artifacts, and 4/ upload the archive file to Amazon S3 for deployment.\n",
    "\n",
    "Alternately, for steps 2.2-2.4, below, if the model archive is already available from Amazon S3, see '2.2-2.4: Alternate Method if Model Already Exists in S3', below.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e4b52d1",
   "metadata": {},
   "source": [
    "### 2.1: Import Packages and Set SageMaker Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3328a4e-f3f3-4a38-9eec-9383b8890ebc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "import boto3\n",
    "from botocore.exceptions import ClientError\n",
    "import os\n",
    "from sagemaker.huggingface.model import HuggingFaceModel\n",
    "from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig\n",
    "from sagemaker.s3 import s3_path_join"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d091f7a5-a217-43dc-8782-6c00b07517a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_session_bucket = None\n",
    "\n",
    "sm_session = sagemaker.Session()\n",
    "\n",
    "if sm_session_bucket is None and sm_session is not None:\n",
    "    # set to default bucket if a bucket name is not given\n",
    "    sm_session_bucket = sm_session.default_bucket()\n",
    "try:\n",
    "    sm_role = sagemaker.get_execution_role()\n",
    "except ValueError:\n",
    "    iam_client = boto3.client(\"iam\")\n",
    "    sm_role = iam_client.get_role(RoleName=\"sagemaker_execution_role\")[\"Role\"][\"Arn\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa2a2e63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# name of packaged model archive file\n",
    "MODEL_ARCHIVE = \"model_v2.tar.gz\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a604a64-bab3-4efa-a442-4e9e16a84b73",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"sagemaker role arn: {sm_role}\")\n",
    "print(f\"sagemaker bucket: {sm_session.default_bucket()}\")\n",
    "print(f\"sagemaker session region: {sm_session.boto_region_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a53689c-07b2-483d-b768-0d943efc6daa",
   "metadata": {},
   "source": [
    "### 2.2: Download the Model Artifacts from Hugging Face\n",
    "\n",
    "It will take 6-7 minutes to download model artifacts from Hugging Face. You will need a Hugging Face account to get your personal access token. Requires approximately 34 GB of space.\n",
    "\n",
    "Check the `/dev/nvme1n1` volume, mounted to `/home/sagemaker-user` to ensure it has enough space.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68fe03ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%sh\n",
    "\n",
    "df -h /home/sagemaker-user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69543da6-b78f-4064-bd10-a340189a8dd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%sh\n",
    "\n",
    "git lfs install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6931008-d698-4536-9ad1-94286ff8bf50",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "%%sh\n",
    "\n",
    "user_name=\"<YOUR_HUGGINGFACE_USERNAME>\"\n",
    "access_token=\"<YOUR_HUGGING_FACE_ACCESS_TOKEN>\"\n",
    "\n",
    "git lfs clone \"https://${user_name}:${access_token}@huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt-1-1.git\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e1f05b2",
   "metadata": {},
   "source": [
    "### 2.3: Add Custom Inference Script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bd55c45-194b-4d6c-bc1a-34210c50d423",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "\n",
    "destination = \"stable-video-diffusion-img2vid-xt-1-1\"\n",
    "\n",
    "shutil.copy(\"inference_v2/inference.py\", destination)\n",
    "shutil.copy(\"inference_v2/requirements.txt\", destination)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19f456f2-5d0f-48db-b8e1-68173c1f510f",
   "metadata": {},
   "source": [
    "### 2.4: TAR GZIP Model Artifacts\n",
    "\n",
    "Important: Final model archive file will be 14-15 GB and could take 35-40 minutes to package and compress.\n",
    "\n",
    "Continuously poll the size of the model archive file file every 15 seconds from your terminal:\n",
    "\n",
    "```sh\n",
    "while sleep 15; do ls -la model_v2.tar.gz; done\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47e7bcc3-b50c-4a27-aa22-37461ab5f668",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "import os\n",
    "import tarfile\n",
    "\n",
    "TAR_MODE = \"w:gz\"\n",
    "\n",
    "\n",
    "def create_tar_archive(folder_path, output_tar_file):\n",
    "    \"\"\"\n",
    "    Create a tar archive from a folder, excluding hidden files.\n",
    "\n",
    "    :param folder_path: The path to the folder to be archived.\n",
    "    :param output_tar_file: The path to the output tar file.\n",
    "    \"\"\"\n",
    "    with tarfile.open(output_tar_file, TAR_MODE) as tar:\n",
    "        for root, dirs, files in os.walk(folder_path):\n",
    "            files = [f for f in files if not f[0] == \".\"]\n",
    "            dirs[:] = [d for d in dirs if not d[0] == \".\"]\n",
    "            for file in files:\n",
    "                file_path = os.path.join(root, file)\n",
    "                tar.add(file_path, arcname=os.path.relpath(file_path, folder_path))\n",
    "                print(f\"Added {file_path} to the archive.\")\n",
    "\n",
    "\n",
    "create_tar_archive(destination, MODEL_ARCHIVE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dccb67e8",
   "metadata": {},
   "source": [
    "### 2.2-2.4: Alternate Method if Model Already Exists in S3\n",
    "\n",
    "If the model archive file already exists in S3, skip steps 1-3 above. Create an Amazon S3 presigned URL and use the URL to download the model package. This replaces the two steps above: downloading the model artifacts and TAR GZIP. This step takes 4-7 minutes in the same AWS Region.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eafaece3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "import os\n",
    "from pywget import wget\n",
    "\n",
    "presigned_s3_url = \"<YOUR_PRESIGNED_URL_GOES_HERE>\"\n",
    "\n",
    "wget.download(presigned_s3_url, MODEL_ARCHIVE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6dfccbd-c114-48c3-9e48-ad40a49ca715",
   "metadata": {},
   "source": [
    "### 2.5: Copy Model Artifacts to S3\n",
    "\n",
    "This step takes 2-3 minutes in the same AWS Region to copy model archive file to Amazon S3, which is approximately 14 GB.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9c4d5ea-10c3-4a11-80fc-cab527c4e623",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "s3_client = boto3.client(\"s3\")\n",
    "\n",
    "response = s3_client.upload_file(\n",
    "    MODEL_ARCHIVE,\n",
    "    sm_session_bucket,\n",
    "    f\"async_inference/model/{MODEL_ARCHIVE}\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "301b0cf9-f391-42c4-806f-46fc7b3da1ce",
   "metadata": {},
   "source": [
    "## 3: Deploy Model to Amazon SageMaker Endpoint\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "368c2f73",
   "metadata": {},
   "source": [
    "### 3.1: Deploy Model to Amazon SageMaker Endpoint\n",
    "\n",
    "Deploying the Amazon SageMaker Asynchronous Inference Endpoint takes 5-7 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "578ea117-2bd4-442b-af0d-2789b980412c",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = {\n",
    "    \"SAGEMAKER_MODEL_SERVER_TIMEOUT\": \"3600\",\n",
    "    \"TS_MAX_RESPONSE_SIZE\": \"1000000000\",\n",
    "    \"TS_MAX_REQUEST_SIZE\": \"1000000000\",\n",
    "    \"MMS_MAX_RESPONSE_SIZE\": \"1000000000\",\n",
    "    \"MMS_MAX_REQUEST_SIZE\": \"1000000000\",\n",
    "}\n",
    "\n",
    "huggingface_model = HuggingFaceModel(\n",
    "    model_data=s3_path_join(\n",
    "        \"s3://\", sm_session_bucket, f\"async_inference/model/{MODEL_ARCHIVE}\"\n",
    "    ),\n",
    "    transformers_version=\"4.37.0\",\n",
    "    pytorch_version=\"2.1.0\",\n",
    "    py_version=\"py310\",\n",
    "    env=env,\n",
    "    role=sm_role,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c38bf8bf-2681-4475-b49b-89db367a600c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# where the response payload or error will be stored\n",
    "\n",
    "async_config = AsyncInferenceConfig(\n",
    "    output_path=s3_path_join(\"s3://\", sm_session_bucket, \"async_inference/output\"),\n",
    "    failure_path=s3_path_join(\n",
    "        \"s3://\", sm_session_bucket, \"async_inference/output_errors\"\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "783b2a82-8c38-4312-a48c-b1b930acb4d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# also successfully tested with a ml.g5.2xlarge instance\n",
    "\n",
    "predictor = huggingface_model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=\"ml.g5.4xlarge\",\n",
    "    async_inference_config=async_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "961d575b-0990-4118-a503-74193829f4ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name = predictor.endpoint_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65b86978",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(endpoint_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39c0cada",
   "metadata": {},
   "source": [
    "### 3.2: Optional: Set Endpoint Name Manually\n",
    "\n",
    "If the model was previously deployed to an endpoint, then uncomment and set the `endpoint_name` variable manually.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e5db6f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# endpoint_name = \"<YOUR_MODEL_ENDPOINT_NAME>\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "082d6ce4-9f1f-40f4-9e02-d86ccaf404f8",
   "metadata": {},
   "source": [
    "## 4: Choosing Conditioning Images and Inference Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f765f2a0",
   "metadata": {},
   "source": [
    "### 4.1: Method for Encoding Conditioning Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b7f200c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "\n",
    "\n",
    "def encode_image(image_path):\n",
    "    with open(image_path, \"rb\") as image_file:\n",
    "        return base64.b64encode(image_file.read()).decode(\"utf-8\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f001dcc",
   "metadata": {},
   "source": [
    "### 4.2 Set Movie Name and Inference Parameters\n",
    "\n",
    "Select one of the sets of inference parameters below and run that cell, or create your own set of parameters. Each variation creates a different video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59dcbbc0-7ac4-4a41-aa6f-0ca712f61d6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"rocket_1.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/rocket.png\"),\n",
    "    \"width\": 1024,\n",
    "    \"height\": 576,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 25,\n",
    "    \"min_guidance_scale\": 1.0,\n",
    "    \"max_guidance_scale\": 3.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 127,\n",
    "    \"noise_aug_strength\": 0.02,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 42,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f0a76c2-b464-45e5-95ee-849f43c0a024",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"rocket_2.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/rocket.png\"),\n",
    "    \"width\": 1024,\n",
    "    \"height\": 576,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 25,\n",
    "    \"min_guidance_scale\": 1.0,\n",
    "    \"max_guidance_scale\": 3.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 180,\n",
    "    \"noise_aug_strength\": 0.02,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 42,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e3c2db7-5972-4542-97a8-435ca5a122a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"smoke.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/smoke.jpg\"),\n",
    "    \"width\": 576,\n",
    "    \"height\": 1024,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 50,\n",
    "    \"min_guidance_scale\": 0.5,\n",
    "    \"max_guidance_scale\": 1.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 25,\n",
    "    \"noise_aug_strength\": 0.8,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 111142,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a5a69fc-3d19-4444-a7b3-683e0e617c46",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"color_smoke.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/colored_smoke.jpg\"),\n",
    "    \"width\": 576,\n",
    "    \"height\": 1024,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 50,\n",
    "    \"min_guidance_scale\": 0.5,\n",
    "    \"max_guidance_scale\": 1.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 25,\n",
    "    \"noise_aug_strength\": 0.8,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 111142,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eb5b21e-360e-4a5d-a9c9-2b69bbd0b24e",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"beach_bike_1.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/beach_bike.jpg\"),\n",
    "    \"width\": 1024,\n",
    "    \"height\": 576,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 25,\n",
    "    \"min_guidance_scale\": 1.0,\n",
    "    \"max_guidance_scale\": 3.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 127,\n",
    "    \"noise_aug_strength\": 0.02,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 1234567890,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2ecc9fa-7370-4a53-84ce-927fa17c00f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"beach_bike_2.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/beach_bike.jpg\"),\n",
    "    \"width\": 1024,\n",
    "    \"height\": 576,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 25,\n",
    "    \"min_guidance_scale\": 1.0,\n",
    "    \"max_guidance_scale\": 3.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 127,\n",
    "    \"noise_aug_strength\": 0.02,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 123,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "882b9f42-e470-40ff-afee-d5befd93749b",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"waterfall.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/waterfall.jpg\"),\n",
    "    \"width\": 1024,\n",
    "    \"height\": 576,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 25,\n",
    "    \"min_guidance_scale\": 1.0,\n",
    "    \"max_guidance_scale\": 3.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 127,\n",
    "    \"noise_aug_strength\": 0.02,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 1234567890,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a425633f-4952-4ff9-ba88-bc462f9b0ff3",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"boat_ocean.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/boat_ocean.jpg\"),\n",
    "    \"width\": 1024,\n",
    "    \"height\": 576,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 25,\n",
    "    \"min_guidance_scale\": 1.0,\n",
    "    \"max_guidance_scale\": 3.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 127,\n",
    "    \"noise_aug_strength\": 0.02,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 42,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94e4ad19-8dfb-4d71-bcb7-3552258eb7f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"red_car.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/red_car.jpg\"),\n",
    "    \"width\": 1024,\n",
    "    \"height\": 576,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 25,\n",
    "    \"min_guidance_scale\": 1.0,\n",
    "    \"max_guidance_scale\": 3.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 127,\n",
    "    \"noise_aug_strength\": 0.02,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 42,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbbb9b52-9e87-41da-a666-692ae3b85915",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"coffee.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/coffee_stream.jpg\"),\n",
    "    \"width\": 576,\n",
    "    \"height\": 1024,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 25,\n",
    "    \"min_guidance_scale\": 1.0,\n",
    "    \"max_guidance_scale\": 3.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 127,\n",
    "    \"noise_aug_strength\": 0.02,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 42,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dd2ec84-48cf-4a4f-a2f3-fac763674b20",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"koi.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/koi.jpg\"),\n",
    "    \"width\": 1024,\n",
    "    \"height\": 576,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 25,\n",
    "    \"min_guidance_scale\": 1.0,\n",
    "    \"max_guidance_scale\": 3.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 127,\n",
    "    \"noise_aug_strength\": 0.02,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 9288258982,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2f4efc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_title = \"champagne.mp4\"\n",
    "\n",
    "data = {\n",
    "    \"image\": encode_image(\"images_scaled/champagne2.jpg\"),\n",
    "    \"width\": 576,\n",
    "    \"height\": 1024,\n",
    "    \"num_frames\": 25,\n",
    "    \"num_inference_steps\": 25,\n",
    "    \"min_guidance_scale\": 1.0,\n",
    "    \"max_guidance_scale\": 3.0,\n",
    "    \"fps\": 6,\n",
    "    \"motion_bucket_id\": 127,\n",
    "    \"noise_aug_strength\": 0.02,\n",
    "    \"decode_chunk_size\": 8,\n",
    "    \"seed\": 42,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bb617a0-2de8-416b-917c-44964d1d21ff",
   "metadata": {},
   "source": [
    "## 5: Upload Request Payload and Invoke Endpoint\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9df94df8",
   "metadata": {},
   "source": [
    "### 5.1: Upload Request Payload\n",
    "\n",
    "Upload the JSON request payload to Amazon S3 and invoke the endpoint for inference.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd3d92ba-8454-4b9c-8a2f-d8945dcc64b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def upload_file(input_location):\n",
    "    return sm_session.upload_data(\n",
    "        input_location,\n",
    "        bucket=sm_session.default_bucket(),\n",
    "        key_prefix=\"async_inference/input\",\n",
    "        extra_args={\"ContentType\": \"application/json\"},\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e49a137-605c-4812-8d7d-8bca6bdcb395",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "file_name = \"request_payloads/payload.json\"\n",
    "\n",
    "with open(file_name, \"w\", encoding=\"utf-8\") as f:\n",
    "    json.dump(data, f, ensure_ascii=False, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f2e736f-c7e8-4488-82c4-c9a0b7015518",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_s3_location = upload_file(file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cea1baf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Request payload location: {input_s3_location}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b6b1643",
   "metadata": {},
   "source": [
    "### 5.2: Invoke the Endpoint for Inference\n",
    "\n",
    "Invocation time for a video with 25 inference steps is about 2 minutes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "952ca07c-3edd-46dc-bfdb-8f8abb77c856",
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_runtime_client = boto3.client(\"sagemaker-runtime\")\n",
    "\n",
    "response = sm_runtime_client.invoke_endpoint_async(\n",
    "    EndpointName=endpoint_name,\n",
    "    InputLocation=input_s3_location,\n",
    "    InvocationTimeoutSeconds=3600,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "496f49a5-1b07-44fa-b14d-0ff71b0738f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Model response payload location: {response['OutputLocation']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a48fe1d9",
   "metadata": {},
   "source": [
    "### 5.3: Poll for Model Response\n",
    "\n",
    "Poll the Amazon S3 bucket for a response from the model invocation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f9a1ada-dbe6-4dee-8650-750d4246470a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import urllib\n",
    "import time\n",
    "\n",
    "\n",
    "# function reference: https://github.com/aws/amazon-sagemaker-examples/blob/main/async-inference/Async-Inference-Walkthrough-SageMaker-Python-SDK.ipynb\n",
    "def get_output(output_location):\n",
    "    output_url = urllib.parse.urlparse(output_location)\n",
    "    bucket = output_url.netloc\n",
    "    key = output_url.path[1:]\n",
    "    while True:\n",
    "        try:\n",
    "            return sm_session.read_s3_file(bucket=bucket, key_prefix=key)\n",
    "        except ClientError as e:\n",
    "            if e.response[\"Error\"][\"Code\"] == \"NoSuchKey\":\n",
    "                print(\"Waiting for model output...\")\n",
    "                time.sleep(15)\n",
    "                continue\n",
    "            raise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d40572-c442-4214-9549-97e24e23912d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "output = get_output(response[\"OutputLocation\"])\n",
    "print(f\"Sample of output: {output[:500]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d767ffce-faae-465b-a32a-2b913974b284",
   "metadata": {},
   "source": [
    "## 6: Frames to MP4 Video"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84703ecb",
   "metadata": {},
   "source": [
    "### 6.1: Frames to MP4 Video\n",
    "\n",
    "Convert binary objects in response to JPEGs of each frame, then combined into MP4 using Hugging Face's `diffusers.utils.export_to_video` method.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75f81507-cf99-4454-b9e8-796a00500358",
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "def load_video_frames(video_frames):\n",
    "    loaded_video_frames = []\n",
    "\n",
    "    for idx, video_frame in enumerate(video_frames):\n",
    "        frame = bytes(video_frame, \"raw_unicode_escape\")\n",
    "        frame_name = (\n",
    "            f\"frames_out/frame_0{idx+1}.jpg\"\n",
    "            if idx < 9\n",
    "            else f\"frames_out/frame_{idx+1}.jpg\"\n",
    "        )\n",
    "        with open(frame_name, \"wb\") as fh:\n",
    "            fh.write(base64.decodebytes(frame))\n",
    "\n",
    "        image = Image.open(frame_name, mode=\"r\")\n",
    "        loaded_video_frames.append(image)\n",
    "\n",
    "    return loaded_video_frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61a2f5c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers.utils import export_to_video\n",
    "\n",
    "output = get_output(response[\"OutputLocation\"])\n",
    "data = json.loads(output)\n",
    "loaded_video_frames = load_video_frames(data[\"frames\"])\n",
    "\n",
    "export_to_video(loaded_video_frames, f\"video_out/{movie_title}\", fps=6)\n",
    "print(f\"Video created: {movie_title}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11dd820f-e4ab-49de-be64-b3b33c658646",
   "metadata": {},
   "source": [
    "### 6.2: Display Frames as Grid\n",
    "\n",
    "Display the 25 frames as a 5x5 grid using Hugging Face's `diffusers.utils.make_image_grid` method.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b1eda93-f62c-4962-81c0-24895a3c8b0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers.utils import make_image_grid\n",
    "\n",
    "image = make_image_grid(loaded_video_frames, 5, 5)\n",
    "(width, height) = (image.width // 2, image.height // 2)\n",
    "im_resized = image.resize((width, height))\n",
    "display(im_resized)\n",
    "im_resized.save(\"frames.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de629f04-0580-4cb9-8eb6-1cf4147d90a9",
   "metadata": {},
   "source": [
    "### 6.3: Display Video\n",
    "\n",
    "Convert video CODEC to H.264 and display in notebook at 50% of actual size.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbd8391f-d738-4a8c-be3d-8bc64ade7ad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ffmpeg\n",
    "\n",
    "output_options = {\n",
    "    \"crf\": 20,\n",
    "    \"preset\": \"slower\",\n",
    "    \"movflags\": \"faststart\",\n",
    "    \"pix_fmt\": \"yuv420p\",\n",
    "    \"vcodec\": \"libx264\",\n",
    "}\n",
    "\n",
    "(\n",
    "    ffmpeg.input(f\"video_out/{movie_title}\")\n",
    "    .output(\"video_out/tmp.mp4\", **output_options)\n",
    "    .run(overwrite_output=True, quiet=True)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90fbcd7d-a53f-4b7f-894c-d9326cab7569",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Video\n",
    "\n",
    "Video(\n",
    "    url=\"video_out/tmp.mp4\",\n",
    "    width=(loaded_video_frames[0].width // 2),\n",
    "    html_attributes=\"controls muted autoplay loop\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "faa0a6c2",
   "metadata": {},
   "source": [
    "## 7: Generating of Multiple Video Variations\n",
    "\n",
    "Generating multiple videos variations by combining the above code in a loop. In this example we are creating five variations, changing the seed each time.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2632215a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import json\n",
    "from diffusers.utils import export_to_video\n",
    "\n",
    "sm_runtime_client = boto3.client(\"sagemaker-runtime\")\n",
    "\n",
    "for i in range(3):\n",
    "    seed = random.randrange(1, 9999999999)\n",
    "    data = {\n",
    "        \"image\": encode_image(\"images_scaled/red_car.jpg\"),\n",
    "        \"width\": 1024,\n",
    "        \"height\": 576,\n",
    "        \"num_frames\": 25,\n",
    "        \"num_inference_steps\": 25,\n",
    "        \"min_guidance_scale\": 1.0,\n",
    "        \"max_guidance_scale\": 3.0,\n",
    "        \"fps\": 6,\n",
    "        \"motion_bucket_id\": 127,\n",
    "        \"noise_aug_strength\": 0.02,\n",
    "        \"decode_chunk_size\": 8,\n",
    "        \"seed\": seed,\n",
    "    }\n",
    "    movie_title = f\"red_car_{seed}.mp4\"\n",
    "\n",
    "    file_name = f\"request_payloads/payload_{i}.json\"\n",
    "    with open(file_name, \"w\", encoding=\"utf-8\") as f:\n",
    "        json.dump(data, f, ensure_ascii=False, indent=4)\n",
    "    input_s3_location = upload_file(file_name)\n",
    "\n",
    "    response = sm_runtime_client.invoke_endpoint_async(\n",
    "        EndpointName=endpoint_name,\n",
    "        InputLocation=input_s3_location,\n",
    "        InvocationTimeoutSeconds=3600,\n",
    "    )\n",
    "\n",
    "    output = get_output(response[\"OutputLocation\"])\n",
    "    data = json.loads(output)\n",
    "    loaded_video_frames = load_video_frames(data[\"frames\"])\n",
    "\n",
    "    export_to_video(loaded_video_frames, f\"video_out/{movie_title}\", fps=6)\n",
    "    print(f\"Video created: {movie_title}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95a85c6b",
   "metadata": {},
   "source": [
    "## 8: Delete Amazon SageMaker Endpoint\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2006585f-5f17-424c-b17f-adbd1ce369db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# client_sm = boto3.client(\"sagemaker\")\n",
    "\n",
    "# client_sm.delete_endpoint(EndpointName=endpoint_name)"
   ]
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}