CheckpointTesting.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fL7u2wgpHilQ"
      },
      "outputs": [],
      "source": [
        "# Install dependencies.\n",
        "!pip install bitsandbytes transformers accelerate -q\n",
        "!pip install git+https://github.com/huggingface/diffusers.git -q\n",
        "from huggingface_hub import login\n",
        "login(\"hf_...\")\n",
        "# Set the details for your model here:\n",
        "my_prompt_trigger = \"wny man\"\n",
        "hf_username = \"TimothyAlexisVass\"\n",
        "hf_modelname = \"lora-without-reg\"\n",
        "\n",
        "import torch\n",
        "import os\n",
        "\n",
        "torch.cuda.empty_cache()\n",
        "vae = None\n",
        "base = None\n",
        "\n",
        "from diffusers import DiffusionPipeline, AutoencoderKL\n",
        "\n",
        "vae = AutoencoderKL.from_pretrained(\"madebyollin/sdxl-vae-fp16-fix\", torch_dtype=torch.float16)\n",
        "base = DiffusionPipeline.from_pretrained(\"stabilityai/stable-diffusion-xl-base-1.0\",\n",
        "    vae=vae,\n",
        "    torch_dtype=torch.float16,\n",
        "    variant=\"fp16\",\n",
        "    use_safetensors=True,\n",
        ")\n",
        "\n",
        "output_dir = \"/content/model_testing\" # You could save to Google drive if you want\n",
        "os.makedirs(output_dir, exist_ok=True)\n",
        "\n",
        "for i in range(2, 16):\n",
        "  checkpoint = f\"sdxl-dreambooth-lora-{str(i).zfill(6)}.safetensors\"\n",
        "  base.load_lora_weights(f\"{hf_username}/{hf_modelname}\", weight_name=checkpoint)\n",
        "\n",
        "  _ = base.to(\"cuda\")\n",
        "\n",
        "  for e in range(10):\n",
        "    image = base(\n",
        "        my_prompt_trigger,\n",
        "        height=1024,\n",
        "        width=1024,\n",
        "        num_images_per_prompt=1,\n",
        "        num_inference_steps=30,\n",
        "        guidance_scale=8,\n",
        "        generator=torch.manual_seed(1337+e)\n",
        "    ).images[0]\n",
        "\n",
        "    img_path = os.path.join(output_dir, f\"checkpoint{i}_{e}.png\")\n",
        "    image.save(img_path)\n",
        "    # os.system(f'curl -T \"{img_path}\" ftp://T...:T...@ftpcluster.loopia.se') # I send my images to an FTP server\n",
        "  base.unload_lora_weights()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}