sdxl_1_0_lora.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyOQObnCMCD6Dbf1Jr7D7yOx",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/detektor777/colab_list_image/blob/main/sdxl_1_0_lora.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w3qRB3KfG62c"
      },
      "outputs": [],
      "source": [
        "#@title ##**Install** { display-mode: \"form\" }\n",
        "%%capture\n",
        "\n",
        "!pip install -q invisible_watermark transformers accelerate safetensors diffusers xformers mediapy\n",
        "!pip install gdown\n",
        "import torch\n",
        "from diffusers import StableDiffusionXLImg2ImgPipeline, DiffusionPipeline, KDPM2AncestralDiscreteScheduler, StableDiffusionXLPipeline, AutoencoderKL\n",
        "import gc\n",
        "from PIL import Image\n",
        "import requests\n",
        "from io import BytesIO\n",
        "from IPython.display import display\n",
        "import os\n",
        "import random\n",
        "import re\n",
        "import gdown\n",
        "from IPython.display import display, Image\n",
        "from IPython.display import clear_output\n",
        "import ipywidgets as widgets\n",
        "\n",
        "model_base = \"stabilityai/stable-diffusion-xl-base-1.0\"\n",
        "v_autoencoder = \"madebyollin/sdxl-vae-fp16-fix\" # fix vae for run in fp16 precision without generating NaNs\n",
        "\n",
        "vae = AutoencoderKL.from_pretrained(v_autoencoder, torch_dtype=torch.float16)\n",
        "\n",
        "pipe = StableDiffusionXLPipeline.from_pretrained(\n",
        "    model_base,\n",
        "    vae=vae,\n",
        "    torch_dtype=torch.float16,\n",
        "    use_safetensors=True,\n",
        "    variant=\"fp16\",\n",
        "    add_watermarker=False, # no watermarker\n",
        "    )\n",
        "\n",
        "pipe.safety_checker = None\n",
        "\n",
        "pipe.to(\"cuda\")\n",
        "model_refiner = \"stabilityai/stable-diffusion-xl-refiner-1.0\"\n",
        "\n",
        "pipe_refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(\n",
        "    model_refiner,\n",
        "    torch_dtype=torch.float16,\n",
        "    use_safetensors=True,\n",
        "    variant=\"fp16\",\n",
        "    add_watermarker=False, # no watermarker\n",
        "    )\n",
        "\n",
        "#pipe_refiner.to(\"cuda\")\n",
        "pipe_refiner.enable_model_cpu_offload()\n",
        "#(Optional) Change the scheduler\n",
        "\n",
        "#pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(\n",
        " # pipe.scheduler.config, use_karras_sigmas=False\n",
        "#)\n",
        "\n",
        "#generator = torch.Generator().manual_seed(42)\n",
        "\n",
        "from diffusers.schedulers.scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler\n",
        "from diffusers import StableDiffusionXLPipeline, StableDiffusionPipeline\n",
        "from diffusers import DPMSolverMultistepScheduler, EulerDiscreteScheduler\n",
        "from diffusers import KDPM2DiscreteScheduler, KDPM2AncestralDiscreteScheduler\n",
        "from diffusers import EulerAncestralDiscreteScheduler, HeunDiscreteScheduler\n",
        "from diffusers import LMSDiscreteScheduler\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Load lora model (Optional)**  { display-mode: \"form\" }\n",
        "#@markdown Load a model from a link or in this folder '/content'\n",
        "#%%capture\n",
        "\n",
        "lora_name = '' #@param {type:\"string\"}\n",
        "model_link = \"\" #@param {type:\"string\"}\n",
        "\n",
        "if not lora_name or not model_link:\n",
        "    print(\"lora_name or model_link is empty\")\n",
        "else:\n",
        "    model_name = f\"{lora_name}.safetensors\"\n",
        "    url = model_link\n",
        "\n",
        "    if \"drive.google.com\" in url:\n",
        "        file_id_match = re.search(r'file/d/(.*?)/', url)\n",
        "        file_id = file_id_match.group(1)\n",
        "        print(f\"file_id: {file_id}\")\n",
        "        url = f'https://drive.google.com/uc?id={file_id}'\n",
        "        destination_path = f'/content/{lora_name}.safetensors'\n",
        "        gdown.download(url, destination_path, quiet=False)\n",
        "    else:\n",
        "        print(f\"link\")\n",
        "        !wget \"{url}\" -O \"{model_name}\"\n",
        "\n",
        "    clear_output(wait=True)"
      ],
      "metadata": {
        "id": "SuqHi6GLHY4a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Select lora model (Optional)**  { display-mode: \"form\" }\n",
        "\n",
        "folder_path = '/content'\n",
        "\n",
        "files = [f for f in os.listdir(folder_path) if f.endswith('.safetensors')]\n",
        "\n",
        "if not files:\n",
        "    print(\"not lora models\")\n",
        "else:\n",
        "    files_str = ', '.join(files)\n",
        "    print(f\"Lora models: {files_str}\")\n",
        "\n",
        "    selected_files = widgets.Dropdown(\n",
        "        options=files,\n",
        "        value=files[0],\n",
        "        description='Select:'\n",
        "    )\n",
        "\n",
        "    selected_file_value = selected_files.value\n",
        "\n",
        "    def on_value_change(change):\n",
        "        global selected_file_value\n",
        "        selected_file_value = change['new']\n",
        "        print(f\"Selected: {selected_file_value}\")\n",
        "\n",
        "    selected_files.observe(on_value_change, names='value')\n",
        "\n",
        "    display(selected_files)\n"
      ],
      "metadata": {
        "id": "onNprha6HILh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Create and show** { display-mode: \"form\" }\n",
        "\n",
        "lora_status = \"Yes\" #@param [\"Yes\",\"No\"]\n",
        "lora_scale = \"1.0\" #@param [\"0.0\", \"0.1\", \"0.2\", \"0.3\", \"0.4\", \"0.5\", \"0.6\", \"0.7\", \"0.8\", \"0.9\", \"1.0\",\"1.5\",\"2.0\"]\n",
        "lora_scale = float(lora_scale)\n",
        "\n",
        "if 'selected_file_value' in locals() and selected_file_value and lora_status == \"Yes\":\n",
        "    pipe.load_lora_weights(selected_file_value)\n",
        "    pipe.fuse_lora(lora_scale=lora_scale)\n",
        "    print(f\"Selected: {selected_file_value}\")\n",
        "else:\n",
        "    pipe.unfuse_lora()\n",
        "    pipe.unload_lora_weights()\n",
        "    print(\"Lora is empty.\")\n",
        "\n",
        "prompt =  \"cyberpunk, cat, Ukrainian flag\" #@param {type:\"string\"}\n",
        "prompt2 = \"\"\n",
        "\n",
        "negative_prompt = ''#@param {type:\"string\"}\n",
        "negative_prompt2 = ''\n",
        "\n",
        "scale = 10 #@param {type:\"integer\"}\n",
        "steps = 50 #@param {type:\"integer\"}\n",
        "\n",
        "height = 1024 #@param {type:\"integer\"}\n",
        "width = 1024 #@param {type:\"integer\"}\n",
        "\n",
        "scheduler_name = \"LMS\" #@param [\"Euler_a\",\"Euler\", \"Euler_K\", \"LMS\", \"LMS_Karras\", \"Heun\", \"DPMPP_2M\", \"DPMPP_2M_Karras\", \"DPMPP_2M_Lu\", \"DPMPP_2M_Stable\", \"DPMPP_2M_SDE\", \"DPMPP_2M_SDE_Karras\", \"DPMPP_2M_SDE_Lu\", \"DPMPP_2M_SDE_Stable\", \"DPM++_2M_Karras\", \"DPM_Fast\"]\n",
        "\n",
        "schedulers = {\n",
        "    \"Euler_K\": (EulerDiscreteScheduler, {\"use_karras_sigmas\": True}),\n",
        "\n",
        "    \"DPMPP_2M\": (DPMSolverMultistepScheduler, {}),\n",
        "    \"DPMPP_2M_Karras\": (DPMSolverMultistepScheduler, {\"use_karras_sigmas\": True}),\n",
        "    \"DPMPP_2M_Lu\": (DPMSolverMultistepScheduler, {\"use_lu_lambdas\": True}),\n",
        "    \"DPMPP_2M_Stable\": (DPMSolverMultistepScheduler, {\"euler_at_final\": True}),\n",
        "\n",
        "    \"DPMPP_2M_SDE\": (DPMSolverMultistepScheduler, {\"algorithm_type\": \"sde-dpmsolver++\"}),\n",
        "    \"DPMPP_2M_SDE_Karras\": (DPMSolverMultistepScheduler, {\"use_karras_sigmas\": True, \"algorithm_type\": \"sde-dpmsolver++\"}),\n",
        "    \"DPMPP_2M_SDE_Lu\": (DPMSolverMultistepScheduler, {\"use_lu_lambdas\": True, \"algorithm_type\": \"sde-dpmsolver++\"}),\n",
        "    \"DPMPP_2M_SDE_Stable\": (DPMSolverMultistepScheduler, {\"algorithm_type\": \"sde-dpmsolver++\", \"euler_at_final\": True}),\n",
        "\n",
        "    \"DPM++_2M_Karras\": (DPMSolverMultistepScheduler, {\"use_karras_sigmas\": True, \"algorithm_type\": \"sde-dpmsolver++\", \"steps\": 70}),\n",
        "    \"DPM_Fast\": (DPMSolverMultistepScheduler, {\"algorithm_type\": \"sde-dpmsolver++\", \"steps\": 100}),\n",
        "\n",
        "    \"DPM2\": (KDPM2DiscreteScheduler, {}),\n",
        "    \"DPM2_Karras\": (KDPM2DiscreteScheduler, {\"use_karras_sigmas\": True}),\n",
        "    \"DPM2_a\": (KDPM2AncestralDiscreteScheduler, {}),\n",
        "    \"DPM2_a_Karras\": (KDPM2AncestralDiscreteScheduler, {\"use_karras_sigmas\": True}),\n",
        "\n",
        "    \"Euler\": (EulerDiscreteScheduler, {'beta_start': 0.0006, 'beta_end': 0.012, 'beta_schedule': \"scaled_linear\", 'steps_offset': 0}),\n",
        "    \"Euler_a\": (EulerAncestralDiscreteScheduler, {}),\n",
        "\n",
        "    \"Heun\": (HeunDiscreteScheduler, {}),\n",
        "\n",
        "    \"LMS\": (LMSDiscreteScheduler, {}),\n",
        "    \"LMS_Karras\": (LMSDiscreteScheduler, {\"use_karras_sigmas\": True}),\n",
        "}\n",
        "\n",
        "#'beta_start': 0.0006, 'beta_end': 0.012, 'beta_schedule': \"scaled_linear\", 'steps_offset': 0\n",
        "model_id = model_base\n",
        "\n",
        "params = {\n",
        "    \"prompt\":prompt,\n",
        "    \"prompt_2\":prompt2,\n",
        "    \"height\":height,\n",
        "    \"width\":width,\n",
        "    \"negative_prompt\":negative_prompt,\n",
        "    \"negative_prompt_2\":negative_prompt2,\n",
        "    \"guidance_scale\":scale,\n",
        "    \"num_inference_steps\":steps\n",
        "}\n",
        "\n",
        "save_dir = './samples_sdxl'\n",
        "if not os.path.exists(save_dir):\n",
        "    os.mkdir(save_dir)\n",
        "\n",
        "seed = random.randint(0, 1000000000000000000)\n",
        "generator = torch.Generator(device='cuda').manual_seed(seed)\n",
        "\n",
        "scheduler = DPMSolverMultistepScheduler.from_pretrained(\n",
        "    model_id,\n",
        "    subfolder=\"scheduler\",\n",
        "    **schedulers[scheduler_name][1],\n",
        ")\n",
        "pipe.scheduler = scheduler\n",
        "clear_output(wait=True)\n",
        "\n",
        "sdxl_img = pipe(**params, generator=generator).images[0]\n",
        "\n",
        "sdxl_img.save(os.path.join(save_dir, f\"seed_{seed}_steps_{params['num_inference_steps']}_{scheduler_name}.png\"))\n",
        "\n",
        "# Unload lora\n",
        "pipe.unfuse_lora()\n",
        "pipe.unload_lora_weights()\n",
        "\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()\n",
        "torch.cuda.synchronize()\n",
        "\n",
        "#sdxl_img\n",
        "\n",
        "image_path = os.path.join(save_dir, f\"seed_{seed}_steps_{params['num_inference_steps']}_{scheduler_name}.png\")\n",
        "\n",
        "if os.path.isfile(image_path):\n",
        "    clear_output(wait=True)\n",
        "    display(Image(filename=image_path, width=700))\n",
        "else:\n",
        "    print(f\"File '{image_path}' not found\")\n"
      ],
      "metadata": {
        "id": "Xys787lpHRw9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##**Refine and show (optional)** { display-mode: \"form\" }\n",
        "#@markdown running this code might lead to insufficient memory\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()\n",
        "torch.cuda.synchronize()\n",
        "\n",
        "denoising_start = \"0.85\" #@param [0.7, 0.71, 0.72, 0.73, 0.74, 0.75, 0.76, 0.77, 0.78, 0.79, 0.8, 0.81, 0.82, 0.83, 0.84, 0.85, 0.86, 0.87, 0.88, 0.89, 0.9, 0.91, 0.92, 0.93, 0.94, 0.95, 0.96, 0.97, 0.98, 0.99]\n",
        "denoising_start = float(denoising_start)\n",
        "\n",
        "steps = 50 #@param {type:\"integer\"}\n",
        "\n",
        "image_refined = pipe_refiner(prompt=prompt, num_inference_steps=steps, denoising_start=denoising_start, image=sdxl_img).images[0]\n",
        "\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()\n",
        "torch.cuda.synchronize()\n",
        "\n",
        "image_path = os.path.join(save_dir, f\"seed_{seed}_steps_{params['num_inference_steps']}_{scheduler_name}engance.png\")\n",
        "\n",
        "image_refined.save(os.path.join(save_dir, f\"seed_{seed}_steps_{params['num_inference_steps']}_{scheduler_name}engance.png\"))\n",
        "\n",
        "\n",
        "if os.path.isfile(image_path):\n",
        "    clear_output(wait=True)\n",
        "    display(Image(filename=image_path, width=700))\n",
        "else:\n",
        "    print(f\"File '{image_path}' not found\")\n"
      ],
      "metadata": {
        "id": "0w11kO6QHo6g"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}