mystic_ai_sdxl.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/AinzOwl/mysticai-colab/blob/main/mystic_ai_sdxl.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Mystic.Ai SDXL v1 with refiner model deployment script**\n",
        "\n",
        "Hello, I'm [Ainzoal](https://ainz.sbs) you can find me on mystic.ai discord server, I made this colab after suffering a lot to deploy my first model on v2 to point i quited but now v3 is out and with help of **Paul** i managed to deploy my first model, so i decided to give back and help other people who want to deploy on this super amazing platform.\n",
        "\n",
        "this colab is really easy, and if you face any errors just ask, anyway i tried to document it and make it as clear as possible so here you go.\n",
        "\n",
        "For V1.5 [click here](https://github.com/AinzOwl/mysticai-colab/blob/main/mystic_ai_sd1_5.ipynb)"
      ],
      "metadata": {
        "id": "m4mGNd2iVkqm"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Downloading packages**\n",
        "\n",
        "This was actually the error i faced i had missing packages, but worry not this will fix it just hit the download, also you dont need gpu for this so if it's enabled at top right u can just disable it to save on google resources."
      ],
      "metadata": {
        "id": "lYhFANqrV0cS"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0gyxUz4UpYYC"
      },
      "outputs": [],
      "source": [
        "#@title Download Packages\n",
        "\n",
        "!pip install pipeline-ai torch transformers diffusers accelerate xformers invisible_watermark safetensors"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **setup api key**\n",
        "\n",
        "When you run this it will ask you for your api key if you don't know where to find it it's [Here](https://mystic.ai/api-tokens) just type it and hit enter"
      ],
      "metadata": {
        "id": "LA4AROr2WENK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Setup Mystic.Ai Api Key\n",
        "\n",
        "api_key = input('Enter mystic.ai api key: ')\n",
        "!pipeline cluster login catalyst-api {api_key} -u https://www.mystic.ai -a"
      ],
      "metadata": {
        "id": "5_a5ckfDprQW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **import packages**\n",
        "Yeah just run this one and that's it, nothing to change"
      ],
      "metadata": {
        "id": "qWGjMVmnWNUu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Import packages\n",
        "\n",
        "from pathlib import Path\n",
        "from typing import List\n",
        "\n",
        "import torch\n",
        "from diffusers import DiffusionPipeline\n",
        "\n",
        "from pipeline import Pipeline, Variable, entity, pipe\n",
        "from pipeline.cloud import compute_requirements, environments, pipelines\n",
        "from pipeline.objects import File\n",
        "from pipeline.objects.graph import InputField, InputSchema\n"
      ],
      "metadata": {
        "id": "esTbvnwjps-c"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **inputs (please read)**\n",
        "\n",
        "this is the configuration for the model inputs, you might want to set up this one, the default values are set a bit high for the model set in this demo, but its easy to change them just search for name of value you want to change and change its default, for example you want to set default number of inference steps to 30 then go to num_inference_steps and change 50 with 30\n",
        "\n",
        "for negative prompt you could leave it empty by doing \"\" but as of making this there's a bug that make it required so i suggest typing a generic negative prompt, the one used is one suggested in animagine-xl huggingface."
      ],
      "metadata": {
        "id": "sWZFbGIxWYBx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Setup inputs\n",
        "\n",
        "class ModelKwargs(InputSchema):\n",
        "    num_inference_steps: int | None = InputField(\n",
        "        title=\"Number of inference steps\",\n",
        "        default=50,\n",
        "    )\n",
        "\n",
        "    guidance_scale: float | None = InputField(\n",
        "        default=12,\n",
        "        title=\"Guidance scale\",\n",
        "        ge=7.5,\n",
        "        le=12.5,\n",
        "    )\n",
        "\n",
        "    denoising_end: float | None = InputField(\n",
        "        default=0.8,\n",
        "        title=\"Denoising end\",\n",
        "        ge=0.0,\n",
        "        le=1.0,\n",
        "    )\n",
        "\n",
        "    negative_prompt: str | None = InputField(\n",
        "        default='lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry',\n",
        "        title=\"Negative Prompt\",\n",
        "    )\n",
        "\n",
        "    height: int | None = InputField(\n",
        "        default=1024,\n",
        "        title=\"Height\",\n",
        "        ge=64,\n",
        "        le=1024,\n",
        "        multiple_of=8,\n",
        "    )\n",
        "\n",
        "    width: int | None = InputField(\n",
        "        default=1024,\n",
        "        title=\"Width\",\n",
        "        ge=64,\n",
        "        le=1024,\n",
        "        multiple_of=8,\n",
        "    )\n",
        "\n"
      ],
      "metadata": {
        "id": "XlnIR3jvpvqa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Setting up the model**\n",
        "\n",
        "here you select your model and refiner\n",
        "just note this uses diffusers so the model must be in huggingface.co.\n",
        "\n",
        "if it doesn't work then the moddel could be missing something this colab was tried with Linaqruf/animagine-xl and it worked perfectly.\n",
        "\n",
        "To change the model change ``Linaqruf/animagine-xl`` in the code.\n",
        "\n",
        "To change the refiner change ``stabilityai/stable-diffusion-xl-refiner-1.0`` in the code\n"
      ],
      "metadata": {
        "id": "pCSf0eDrXXkk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Setup Model\n",
        "\n",
        "@entity\n",
        "class StableDiffusionModel:\n",
        "    @pipe(on_startup=True, run_once=True)\n",
        "    def load(self):\n",
        "        self.base = DiffusionPipeline.from_pretrained(\n",
        "            \"Linaqruf/animagine-xl\",\n",
        "            torch_dtype=torch.float16,\n",
        "            use_safetensors=True,\n",
        "        )\n",
        "        self.base.to(\"cuda\")\n",
        "        self.refiner = DiffusionPipeline.from_pretrained(\n",
        "            \"stabilityai/stable-diffusion-xl-refiner-1.0\",\n",
        "            text_encoder_2=self.base.text_encoder_2,\n",
        "            vae=self.base.vae,\n",
        "            torch_dtype=torch.float16,\n",
        "            use_safetensors=True,\n",
        "            variant=\"fp16\",\n",
        "        )\n",
        "        self.refiner.to(\"cuda\")\n",
        "\n",
        "    @pipe\n",
        "    def predict(self, prompt: str, kwargs: ModelKwargs) -> List[File]:\n",
        "        image = self.base(\n",
        "            prompt=prompt,\n",
        "            negative_prompt=kwargs.negative_prompt,\n",
        "            height=kwargs.height,\n",
        "            width=kwargs.width,\n",
        "            num_inference_steps=kwargs.num_inference_steps,\n",
        "            denoising_end=kwargs.denoising_end,\n",
        "            guidance_scale=kwargs.guidance_scale,\n",
        "            output_type=\"latent\",\n",
        "        ).images\n",
        "        images = self.refiner(\n",
        "            prompt=prompt,\n",
        "            negative_prompt=kwargs.negative_prompt,\n",
        "            num_inference_steps=kwargs.num_inference_steps,\n",
        "            guidance_scale=kwargs.guidance_scale,\n",
        "            denoising_start=kwargs.denoising_end,\n",
        "            image=image,\n",
        "        ).images\n",
        "\n",
        "        output_images = []\n",
        "        for i, image in enumerate(images):\n",
        "            path = Path(f\"/tmp/sd/image-{i}.jpg\")\n",
        "            path.parent.mkdir(parents=True, exist_ok=True)\n",
        "            image.save(str(path))\n",
        "            output_images.append(File(path=path, allow_out_of_context_creation=True))\n",
        "\n",
        "        return output_images"
      ],
      "metadata": {
        "id": "uJy6qOtZpwVn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Environment**\n",
        "\n",
        "Env is the one holding the packages that your model uses, the ones selected are the ones needed to get this diffusers model to run and same env can be used for different models so you kinda really just need to run this one ones and you're good"
      ],
      "metadata": {
        "id": "2xqKGSdgYBWp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title setup mystic environment\n",
        "\n",
        "with Pipeline() as builder:\n",
        "    prompt = Variable(\n",
        "        str,\n",
        "        title=\"Prompt\",\n",
        "    )\n",
        "    kwargs = Variable(\n",
        "        ModelKwargs,\n",
        "        title=\"Model kwargs\",\n",
        "    )\n",
        "\n",
        "    model = StableDiffusionModel()\n",
        "\n",
        "    model.load()\n",
        "\n",
        "    output = model.predict(prompt, kwargs)\n",
        "\n",
        "    builder.output(output)\n",
        "\n",
        "my_pl = builder.get_pipeline()\n",
        "\n",
        "try:\n",
        "    environments.create_environment(\n",
        "        \"stable-diffusion-xl-refiner-1.0\",\n",
        "        python_requirements=[\n",
        "            \"torch==2.0.1\",\n",
        "            \"transformers==4.30.2\",\n",
        "            \"diffusers==0.19.3\",\n",
        "            \"accelerate==0.21.0\",\n",
        "            \"xformers==0.0.21\",\n",
        "            \"invisible_watermark==0.2.0\",\n",
        "            \"safetensors==0.3.3\",\n",
        "        ],\n",
        "    )\n",
        "except Exception:\n",
        "    pass"
      ],
      "metadata": {
        "id": "XTDGbti8p0xd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Pushing the model**\n",
        "\n",
        "pushing is easy you just replace ``animagine-xl`` with what you named the model and thats it, if you used a different name for environement you can change it.\n",
        "\n",
        "For vram it depends on model how much it need but dont overkill it just try to find the balance.\n",
        "\n",
        "For accelerators check this page [Mystic Ai Docs](https://docs.mystic.ai/docs/gpus-and-accelerators) you'll find the id of each and pick one you want depending on price and vram"
      ],
      "metadata": {
        "id": "ciY9gI-0YKef"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Push the model (open this to change required gpu type)\n",
        "\n",
        "remote_pipeline = pipelines.upload_pipeline(\n",
        "    my_pl,\n",
        "    \"animagine-xl:latest\",\n",
        "    environment_id_or_name=\"stable-diffusion-xl-refiner-1.0\",\n",
        "    required_gpu_vram_mb=20_000,\n",
        "    accelerators=[\n",
        "        compute_requirements.Accelerator.nvidia_a100,\n",
        "    ],\n",
        ")\n",
        "\n",
        "print(remote_pipeline.id)"
      ],
      "metadata": {
        "id": "nJxYFmYYp6LQ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}