Kandinsky_V2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1PN13Iti7H9m"
      },
      "outputs": [],
      "source": [
        "!pip install git+https://github.com/ai-forever/diffusers.git\n",
        "!pip install transformers\n",
        "!pip install accelerate\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import sys\n",
        "from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline\n",
        "import torch\n",
        "import PIL\n",
        "import torch\n",
        "from diffusers.utils import load_image\n",
        "from torchvision import transforms\n",
        "from transformers import CLIPVisionModelWithProjection\n",
        "from diffusers.models import UNet2DConditionModel\n",
        "import numpy as np\n",
        "\n",
        "DEVICE = torch.device('cuda:0')\n",
        "\n",
        "image_encoder = CLIPVisionModelWithProjection.from_pretrained(\n",
        "    'kandinsky-community/kandinsky-2-2-prior',\n",
        "    subfolder='image_encoder'\n",
        ").half().to(DEVICE)\n",
        "\n",
        "unet = UNet2DConditionModel.from_pretrained(\n",
        "    'kandinsky-community/kandinsky-2-2-decoder',\n",
        "    subfolder='unet'\n",
        ").half().to(DEVICE)\n",
        "\n",
        "prior = KandinskyV22PriorPipeline.from_pretrained(\n",
        "    'kandinsky-community/kandinsky-2-2-prior',\n",
        "    image_encoder=image_encoder,\n",
        "    torch_dtype=torch.float16\n",
        ").to(DEVICE)\n",
        "\n",
        "decoder = KandinskyV22Pipeline.from_pretrained(\n",
        "    'kandinsky-community/kandinsky-2-2-decoder',\n",
        "    unet=unet,\n",
        "    torch_dtype=torch.float16\n",
        ").to(DEVICE)"
      ],
      "metadata": {
        "id": "ZDW2Ipyn9Xl7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(42)\n",
        "\n",
        "negative_prior_prompt ='lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature'\n",
        "img_emb = prior(\n",
        "    prompt = input(\"Enter a description for the image: \"),\n",
        "    num_inference_steps=25,\n",
        "    num_images_per_prompt=1\n",
        ")\n",
        "\n",
        "negative_emb = prior(\n",
        "    prompt=negative_prior_prompt,\n",
        "    num_inference_steps=25,\n",
        "    num_images_per_prompt=1\n",
        ")\n",
        "\n",
        "images = decoder(\n",
        "    image_embeds=img_emb.image_embeds,\n",
        "    negative_image_embeds=negative_emb.image_embeds,\n",
        "    num_inference_steps=75,\n",
        "    height=512,\n",
        "    width=512)"
      ],
      "metadata": {
        "id": "i0DfHLxN9zZM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "images.images[0]  # Run this to retrieve your generated image!"
      ],
      "metadata": {
        "id": "pksBjC-M_z_j"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}