{ "cells": [ { "cell_type": "code", "execution_count": 2, "id": "acab479f", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "===================================BUG REPORT===================================\n", "Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n", "================================================================================\n" ] } ], "source": [ "import argparse\n", "import gc\n", "import hashlib\n", "import itertools\n", "import logging\n", "import math\n", "import os\n", "import threading\n", "import warnings\n", "from pathlib import Path\n", "from typing import Optional\n", "import psutil\n", "import json\n", "\n", "import torch\n", "import torch.nn.functional as F\n", "import torch.utils.checkpoint\n", "from torch.utils.data import Dataset\n", "\n", "import datasets\n", "import diffusers\n", "import transformers\n", "from accelerate import Accelerator\n", "from accelerate.logging import get_logger\n", "from accelerate.utils import set_seed\n", "from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel\n", "from diffusers import DDPMScheduler, PNDMScheduler, StableDiffusionPipeline\n", "from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker\n", "from diffusers.optimization import get_scheduler\n", "from diffusers.utils import check_min_version\n", "from diffusers.utils.import_utils import is_xformers_available\n", "from huggingface_hub import HfFolder, Repository, whoami\n", "from PIL import Image\n", "from torchvision import transforms\n", "from tqdm.auto import tqdm\n", "from transformers import AutoTokenizer, PretrainedConfig, CLIPFeatureExtractor\n", "from peft import PeftModel, LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict\n", "\n", "# Will error if the minimal version of diffusers is not installed. Remove at your own risks.\n", "check_min_version(\"0.10.0.dev0\")\n", "\n", "logger = get_logger(__name__)\n", "\n", "\n", "MODEL_NAME = \"CompVis/stable-diffusion-v1-4\" # \"stabilityai/stable-diffusion-2-1-base\"\n", "INSTANCE_PROMPT = \"a photo of sks dog\"\n", "base_path = \"/home/sourab/temp/\"" ] }, { "cell_type": "code", "execution_count": 3, "id": "06cfd506", "metadata": {}, "outputs": [], "source": [ "def get_lora_sd_pipeline(\n", " ckpt_dir, base_model_name_or_path=None, dtype=torch.float16, device=\"cuda\", adapter_name=\"default\"\n", "):\n", " unet_sub_dir = os.path.join(ckpt_dir, \"unet\")\n", " text_encoder_sub_dir = os.path.join(ckpt_dir, \"text_encoder\")\n", " if os.path.exists(text_encoder_sub_dir) and base_model_name_or_path is None:\n", " config = LoraConfig.from_pretrained(text_encoder_sub_dir)\n", " base_model_name_or_path = config.base_model_name_or_path\n", "\n", " if base_model_name_or_path is None:\n", " raise ValueError(\"Please specify the base model name or path\")\n", "\n", " pipe = StableDiffusionPipeline.from_pretrained(\n", " base_model_name_or_path, torch_dtype=dtype, requires_safety_checker=False\n", " ).to(device)\n", " pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)\n", "\n", " if os.path.exists(text_encoder_sub_dir):\n", " pipe.text_encoder = PeftModel.from_pretrained(\n", " pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name\n", " )\n", "\n", " if dtype in (torch.float16, torch.bfloat16):\n", " pipe.unet.half()\n", " pipe.text_encoder.half()\n", "\n", " pipe.to(device)\n", " return pipe\n", "\n", "\n", "def load_adapter(pipe, ckpt_dir, adapter_name):\n", " unet_sub_dir = os.path.join(ckpt_dir, \"unet\")\n", " text_encoder_sub_dir = os.path.join(ckpt_dir, \"text_encoder\")\n", " pipe.unet.load_adapter(unet_sub_dir, adapter_name=adapter_name)\n", " if os.path.exists(text_encoder_sub_dir):\n", " pipe.text_encoder.load_adapter(text_encoder_sub_dir, adapter_name=adapter_name)\n", "\n", "\n", "def set_adapter(pipe, adapter_name):\n", " pipe.unet.set_adapter(adapter_name)\n", " if isinstance(pipe.text_encoder, PeftModel):\n", " pipe.text_encoder.set_adapter(adapter_name)\n", "\n", "\n", "def merging_lora_with_base(pipe, ckpt_dir, adapter_name=\"default\"):\n", " unet_sub_dir = os.path.join(ckpt_dir, \"unet\")\n", " text_encoder_sub_dir = os.path.join(ckpt_dir, \"text_encoder\")\n", " if isinstance(pipe.unet, PeftModel):\n", " pipe.unet.set_adapter(adapter_name)\n", " else:\n", " pipe.unet = PeftModel.from_pretrained(pipe.unet, unet_sub_dir, adapter_name=adapter_name)\n", " pipe.unet = pipe.unet.merge_and_unload()\n", "\n", " if os.path.exists(text_encoder_sub_dir):\n", " if isinstance(pipe.text_encoder, PeftModel):\n", " pipe.text_encoder.set_adapter(adapter_name)\n", " else:\n", " pipe.text_encoder = PeftModel.from_pretrained(\n", " pipe.text_encoder, text_encoder_sub_dir, adapter_name=adapter_name\n", " )\n", " pipe.text_encoder = pipe.text_encoder.merge_and_unload()\n", "\n", " return pipe\n", "\n", "\n", "def create_weighted_lora_adapter(pipe, adapters, weights, adapter_name=\"default\"):\n", " pipe.unet.add_weighted_adapter(adapters, weights, adapter_name)\n", " if isinstance(pipe.text_encoder, PeftModel):\n", " pipe.text_encoder.add_weighted_adapter(adapters, weights, adapter_name)\n", "\n", " return pipe" ] }, { "cell_type": "code", "execution_count": 4, "id": "d4e888d2", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "9f12b2cca0784cba9dc14ec48de929d5", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Fetching 19 files: 0%| | 0/19 [00:00" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "prompt = \"sks dog playing fetch in the park\"\n", "negative_prompt = \"low quality, blurry, unfinished\"\n", "image = pipe(prompt, num_inference_steps=50, guidance_scale=7, negative_prompt=negative_prompt).images[0]\n", "image" ] }, { "cell_type": "code", "execution_count": 9, "id": "1e1d1f30", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "CPU times: user 7.74 ms, sys: 0 ns, total: 7.74 ms\n", "Wall time: 7.31 ms\n" ] } ], "source": [ "%%time\n", "set_adapter(pipe, adapter_name=\"toy\")" ] }, { "cell_type": "code", "execution_count": 13, "id": "0c50c03d", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "4f978437709b44b391744cf972415027", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/50 [00:00" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "prompt = \"narendra modi rendered in the style of <1>\"\n", "negative_prompt = \"low quality, blurry, unfinished\"\n", "image = pipe(prompt, num_inference_steps=50, guidance_scale=7, negative_prompt=negative_prompt).images[0]\n", "image" ] }, { "cell_type": "code", "execution_count": 24, "id": "e3b9a681", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "8871731127904802ba6123128ba22ecf", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/50 [00:00" ] }, "execution_count": 24, "metadata": {}, "output_type": "execute_result" } ], "source": [ "set_adapter(pipe, adapter_name=\"dog\")\n", "prompt = \"sks dog in a big red bucket\"\n", "negative_prompt = \"low quality, blurry, unfinished\"\n", "image = pipe(prompt, num_inference_steps=50, guidance_scale=7, negative_prompt=negative_prompt).images[0]\n", "image" ] }, { "cell_type": "code", "execution_count": 15, "id": "a659ca6e", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "b701336557e4417d8d39d500699c298b", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/50 [00:00" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "set_adapter(pipe, adapter_name=\"toy\")\n", "prompt = \"superman rendered in the style of <1>, close up potrait\"\n", "negative_prompt = \"low quality, blurry, unfinished\"\n", "image = pipe(prompt, num_inference_steps=50, guidance_scale=7, negative_prompt=negative_prompt).images[0]\n", "image" ] }, { "cell_type": "code", "execution_count": 34, "id": "1f0ecb40", "metadata": {}, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "b47385e953b14c768d82bea776129904", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/50 [00:00" ] }, "execution_count": 34, "metadata": {}, "output_type": "execute_result" } ], "source": [ "set_adapter(pipe, adapter_name=\"toy_dog\")\n", "prompt = \"sks dog rendered in the style of <1>, close up potrait, 4K HD\"\n", "negative_prompt = \"low quality, blurry, unfinished\"\n", "image = pipe(prompt, num_inference_steps=50, guidance_scale=7, negative_prompt=negative_prompt).images[0]\n", "image" ] }, { "cell_type": "code", "execution_count": null, "id": "29720cdb", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.4" }, "vscode": { "interpreter": { "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49" } } }, "nbformat": 4, "nbformat_minor": 5 }