sdxl_v1.0_autotrain_train_dreambooth_lora_colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/sdxl-colab/blob/main/sdxl_v1.0_autotrain_train_dreambooth_lora_colab.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!pip install -q autotrain-advanced -U\n",
        "!autotrain setup\n",
        "\n",
        "# !mkdir /content/images\n",
        "# !pip install datasets\n",
        "# from datasets import load_dataset\n",
        "# dataset = load_dataset(\"camenduru/plushies\", split=\"train\")\n",
        "# for idx, image in enumerate(dataset['image']):\n",
        "#   image.save(f\"/content/images/{idx}.jpg\")\n",
        "\n",
        "# !wget https://huggingface.co/test/resolve/main/test.zip -O /content/test.zip\n",
        "# !mkdir /content/images\n",
        "# !unzip /content/test.zip\n",
        "# !mv /content/*.jpg /content/images"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SaAJk33ppFw1"
      },
      "outputs": [],
      "source": [
        "!autotrain dreambooth \\\n",
        "--model='stabilityai/stable-diffusion-xl-base-1.0' \\\n",
        "--output='lora' \\\n",
        "--image-path='/content/images' \\\n",
        "--prompt='a photo of a plushie toy' \\\n",
        "--resolution=1024 \\\n",
        "--batch-size=2 \\\n",
        "--num-steps=500 \\\n",
        "--gradient-accumulation=4 \\\n",
        "--lr=1e-4 \\\n",
        "--fp16 \\\n",
        "--xformers \\\n",
        "--use-8bit-adam \\\n",
        "--gradient-checkpointing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "from diffusers import DiffusionPipeline, AutoencoderKL\n",
        "\n",
        "vae = AutoencoderKL.from_pretrained(\"madebyollin/sdxl-vae-fp16-fix\", torch_dtype=torch.float16)\n",
        "pipe = DiffusionPipeline.from_pretrained(\"stabilityai/stable-diffusion-xl-base-1.0\", vae=vae, torch_dtype=torch.float16, variant=\"fp16\", use_safetensors=True).to(\"cuda\")\n",
        "pipe.load_lora_weights(\"/content/lora\", weight_name=\"pytorch_lora_weights.bin\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "while True:\n",
        "  prompt = \"a photo of a plushie toy\"\n",
        "  negative_prompt = \"\"\n",
        "  import random\n",
        "  random = random.randint(0, 9223372036854775807)\n",
        "  generator = torch.Generator(\"cuda\").manual_seed(random)\n",
        "  image = pipe(prompt=prompt, negative_prompt=negative_prompt, generator=generator, num_inference_steps=20).images[0]\n",
        "  image.save(f\"/content/output/{random}.jpg\")\n",
        "  # from IPython.display import Image\n",
        "  # Image(f\"/content/output/{random}.jpg\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "del pipe\n",
        "refiner = DiffusionPipeline.from_pretrained(\"stabilityai/stable-diffusion-xl-refiner-1.0\", variant=\"fp16\", torch_dtype=torch.float16, use_safetensors=True).to(\"cuda\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "rimage = refiner(prompt=prompt, negative_prompt=negative_prompt, generator=generator, num_inference_steps=20, image=image).images[0]\n",
        "rimage"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "private_outputs": true,
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}