sdxl_v1.0_diffusers_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_diffusers_train_dreambooth_lora_colab.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SaAJk33ppFw1"
      },
      "outputs": [],
      "source": [
        "!pip install xformers bitsandbytes transformers accelerate -q\n",
        "!pip install git+https://github.com/huggingface/diffusers.git -q\n",
        "!wget https://raw.githubusercontent.com/huggingface/diffusers/main/examples/dreambooth/train_dreambooth_lora_sdxl.py -O /content/train_dreambooth_lora_sdxl.py\n",
        "!accelerate config default\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": {},
      "outputs": [],
      "source": [
        "!accelerate launch train_dreambooth_lora_sdxl.py \\\n",
        "  --pretrained_model_name_or_path=\"stabilityai/stable-diffusion-xl-base-1.0\" \\\n",
        "  --pretrained_vae_model_name_or_path=\"madebyollin/sdxl-vae-fp16-fix\" \\\n",
        "  --instance_data_dir=\"images\" \\\n",
        "  --output_dir=\"lora\" \\\n",
        "  --mixed_precision=\"fp16\" \\\n",
        "  --instance_prompt=\"a photo of a plushie toy\" \\\n",
        "  --resolution=1024 \\\n",
        "  --train_batch_size=2 \\\n",
        "  --gradient_accumulation_steps=2 \\\n",
        "  --gradient_checkpointing \\\n",
        "  --learning_rate=1e-4 \\\n",
        "  --lr_scheduler=\"constant\" \\\n",
        "  --lr_warmup_steps=0 \\\n",
        "  --enable_xformers_memory_efficient_attention \\\n",
        "  --mixed_precision=\"fp16\" \\\n",
        "  --use_8bit_adam \\\n",
        "  --max_train_steps=500 \\\n",
        "  --checkpointing_steps=717 \\\n",
        "  --seed=\"0\""
      ]
    },
    {
      "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
}