ysharma/style-aligned-controlnet
21
1{2 "cells": [3 {4 "cell_type": "markdown",5 "id": "f86ede39-8d9f-4da9-bc12-955f2fddd484",6 "metadata": {7 "pycharm": {8 "name": "#%% md\n"9 }10 },11 "source": [12 "## Copyright 2023 Google LLC"13 ]14 },15 {16 "cell_type": "code",17 "execution_count": null,18 "id": "3f3cbf47-a52b-48b1-9bd3-3435f92f2174",19 "metadata": {20 "pycharm": {21 "name": "#%%\n"22 }23 },24 "outputs": [],25 "source": [26 "# Copyright 2023 Google LLC\n",27 "#\n",28 "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",29 "# you may not use this file except in compliance with the License.\n",30 "# You may obtain a copy of the License at\n",31 "#\n",32 "# http://www.apache.org/licenses/LICENSE-2.0\n",33 "#\n",34 "# Unless required by applicable law or agreed to in writing, software\n",35 "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",36 "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",37 "# See the License for the specific language governing permissions and\n",38 "# limitations under the License."39 ]40 },41 {42 "cell_type": "markdown",43 "id": "22de629b-581f-4335-9e7b-f73221d8dbcb",44 "metadata": {45 "pycharm": {46 "name": "#%% md\n"47 }48 },49 "source": [50 "# ControlNet depth with StyleAligned over SDXL"51 ]52 },53 {54 "cell_type": "code",55 "execution_count": null,56 "id": "486b7ebb-c483-4bf0-ace8-f8092c2d1f23",57 "metadata": {58 "pycharm": {59 "name": "#%%\n"60 }61 },62 "outputs": [],63 "source": [64 "from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL\n",65 "from diffusers.utils import load_image\n",66 "from transformers import DPTImageProcessor, DPTForDepthEstimation\n",67 "import torch\n",68 "import mediapy\n",69 "import sa_handler\n",70 "import pipeline_calls"71 ]72 },73 {74 "cell_type": "code",75 "execution_count": null,76 "id": "2a7e85e7-b5cf-45b2-946a-5ba1e4923586",77 "metadata": {78 "pycharm": {79 "name": "#%%\n"80 }81 },82 "outputs": [],83 "source": [84 "# init models\n",85 "\n",86 "depth_estimator = DPTForDepthEstimation.from_pretrained(\"Intel/dpt-hybrid-midas\").to(\"cuda\")\n",87 "feature_processor = DPTImageProcessor.from_pretrained(\"Intel/dpt-hybrid-midas\")\n",88 "\n",89 "controlnet = ControlNetModel.from_pretrained(\n",90 " \"diffusers/controlnet-depth-sdxl-1.0\",\n",91 " variant=\"fp16\",\n",92 " use_safetensors=True,\n",93 " torch_dtype=torch.float16,\n",94 ").to(\"cuda\")\n",95 "vae = AutoencoderKL.from_pretrained(\"madebyollin/sdxl-vae-fp16-fix\", torch_dtype=torch.float16).to(\"cuda\")\n",96 "pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(\n",97 " \"stabilityai/stable-diffusion-xl-base-1.0\",\n",98 " controlnet=controlnet,\n",99 " vae=vae,\n",100 " variant=\"fp16\",\n",101 " use_safetensors=True,\n",102 " torch_dtype=torch.float16,\n",103 ").to(\"cuda\")\n",104 "pipeline.enable_model_cpu_offload()\n",105 "\n",106 "sa_args = sa_handler.StyleAlignedArgs(share_group_norm=False,\n",107 " share_layer_norm=False,\n",108 " share_attention=True,\n",109 " adain_queries=True,\n",110 " adain_keys=True,\n",111 " adain_values=False,\n",112 " )\n",113 "handler = sa_handler.Handler(pipeline)\n",114 "handler.register(sa_args, )"115 ]116 },117 {118 "cell_type": "code",119 "execution_count": null,120 "id": "94ca26b4-9061-4012-9400-8d97ef212d87",121 "metadata": {122 "pycharm": {123 "name": "#%%\n"124 }125 },126 "outputs": [],127 "source": [128 "# get depth maps\n",129 "\n",130 "image = load_image(\"./example_image/train.png\")\n",131 "depth_image1 = pipeline_calls.get_depth_map(image, feature_processor, depth_estimator)\n",132 "depth_image2 = load_image(\"./example_image/sun.png\").resize((1024, 1024))\n",133 "mediapy.show_images([depth_image1, depth_image2])"134 ]135 },136 {137 "cell_type": "code",138 "execution_count": null,139 "id": "c8f56fe4-559f-49ff-a2d8-460dcfeb56a0",140 "metadata": {141 "pycharm": {142 "name": "#%%\n"143 }144 },145 "outputs": [],146 "source": [147 "# run ControlNet depth with StyleAligned\n",148 "\n",149 "reference_prompt = \"a poster in flat design style\"\n",150 "target_prompts = [\"a train in flat design style\", \"the sun in flat design style\"]\n",151 "controlnet_conditioning_scale = 0.8\n",152 "num_images_per_prompt = 3 # adjust according to VRAM size\n",153 "latents = torch.randn(1 + num_images_per_prompt, 4, 128, 128).to(pipeline.unet.dtype)\n",154 "for deph_map, target_prompt in zip((depth_image1, depth_image2), target_prompts):\n",155 " latents[1:] = torch.randn(num_images_per_prompt, 4, 128, 128).to(pipeline.unet.dtype)\n",156 " images = pipeline_calls.controlnet_call(pipeline, [reference_prompt, target_prompt],\n",157 " image=deph_map,\n",158 " num_inference_steps=50,\n",159 " controlnet_conditioning_scale=controlnet_conditioning_scale,\n",160 " num_images_per_prompt=num_images_per_prompt,\n",161 " latents=latents)\n",162 " \n",163 " mediapy.show_images([images[0], deph_map] + images[1:], titles=[\"reference\", \"depth\"] + [f'result {i}' for i in range(1, len(images))])\n"164 ]165 },166 {167 "cell_type": "code",168 "execution_count": null,169 "id": "437ba4bd-6243-486b-8ba5-3b7cd661d53a",170 "metadata": {171 "pycharm": {172 "name": "#%%\n"173 }174 },175 "outputs": [],176 "source": []177 }178 ],179 "metadata": {180 "kernelspec": {181 "display_name": "Python 3 (ipykernel)",182 "language": "python",183 "name": "python3"184 },185 "language_info": {186 "codemirror_mode": {187 "name": "ipython",188 "version": 3189 },190 "file_extension": ".py",191 "mimetype": "text/x-python",192 "name": "python",193 "nbconvert_exporter": "python",194 "pygments_lexer": "ipython3",195 "version": "3.11.5"196 }197 },198 "nbformat": 4,199 "nbformat_minor": 5200}