File size: 7,104 Bytes
19b3da3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 |
from typing import List
import cv2
import numpy as np
import torch
from controlnet_aux import OpenposeDetector
from diffusers import (
ControlNetModel,
DiffusionPipeline,
StableDiffusionControlNetPipeline,
UniPCMultistepScheduler,
)
from PIL import Image
from tqdm import gui
from internals.data.result import Result
from internals.pipelines.commons import AbstractPipeline
from internals.util.cache import clear_cuda_and_gc
from internals.util.commons import download_image
class ControlNet(AbstractPipeline):
__current_task_name = ""
def load(self, model_dir: str):
# we will load canny by default
self.load_canny()
# controlnet pipeline for canny and pose
pipe = DiffusionPipeline.from_pretrained(
model_dir,
controlnet=self.controlnet,
torch_dtype=torch.float16,
custom_pipeline="stable_diffusion_controlnet_img2img",
).to("cuda")
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()
pipe.enable_xformers_memory_efficient_attention()
self.pipe = pipe
# controlnet pipeline for tile upscaler
pipe2 = StableDiffusionControlNetPipeline(**pipe.components).to("cuda")
pipe2.scheduler = UniPCMultistepScheduler.from_config(pipe2.scheduler.config)
pipe2.enable_xformers_memory_efficient_attention()
self.pipe2 = pipe2
def load_canny(self):
if self.__current_task_name == "canny":
return
canny = ControlNetModel.from_pretrained(
"lllyasviel/control_v11p_sd15_canny", torch_dtype=torch.float16
).to("cuda")
self.__current_task_name = "canny"
self.controlnet = canny
if hasattr(self, "pipe"):
self.pipe.controlnet = canny
if hasattr(self, "pipe2"):
self.pipe2.controlnet = canny
clear_cuda_and_gc()
def load_pose(self):
if self.__current_task_name == "pose":
return
pose = ControlNetModel.from_pretrained(
"lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16
).to("cuda")
self.__current_task_name = "pose"
self.controlnet = pose
if hasattr(self, "pipe"):
self.pipe.controlnet = pose
if hasattr(self, "pipe2"):
self.pipe2.controlnet = pose
clear_cuda_and_gc()
def load_tile_upscaler(self):
if self.__current_task_name == "tile_upscaler":
return
tile_upscaler = ControlNetModel.from_pretrained(
"lllyasviel/control_v11f1e_sd15_tile", torch_dtype=torch.float16
).to("cuda")
self.__current_task_name = "tile_upscaler"
self.controlnet = tile_upscaler
if hasattr(self, "pipe"):
self.pipe.controlnet = tile_upscaler
if hasattr(self, "pipe2"):
self.pipe2.controlnet = tile_upscaler
clear_cuda_and_gc()
def cleanup(self):
self.pipe.controlnet = None
self.pipe2.controlnet = None
self.controlnet = None
self.__current_task_name = ""
clear_cuda_and_gc()
@torch.inference_mode()
def process_canny(
self,
prompt: List[str],
imageUrl: str,
seed: int,
steps: int,
negative_prompt: List[str],
guidance_scale: float,
height: int,
width: int,
):
if self.__current_task_name != "canny":
raise Exception("ControlNet is not loaded with canny model")
torch.manual_seed(seed)
init_image = download_image(imageUrl).resize((width, height))
init_image = self.__canny_detect_edge(init_image)
result = self.pipe2.__call__(
prompt=prompt,
image=init_image,
guidance_scale=guidance_scale,
num_images_per_prompt=1,
negative_prompt=negative_prompt,
num_inference_steps=steps,
height=height,
width=width,
)
return Result.from_result(result)
@torch.inference_mode()
def process_pose(
self,
prompt: List[str],
image: List[Image.Image],
seed: int,
steps: int,
guidance_scale: float,
negative_prompt: List[str],
height: int,
width: int,
):
if self.__current_task_name != "pose":
raise Exception("ControlNet is not loaded with pose model")
torch.manual_seed(seed)
result = self.pipe2.__call__(
prompt=prompt,
image=image,
num_images_per_prompt=1,
num_inference_steps=steps,
negative_prompt=negative_prompt,
guidance_scale=guidance_scale,
height=height,
width=width,
)
return Result.from_result(result)
@torch.inference_mode()
def process_tile_upscaler(
self,
imageUrl: str,
prompt: str,
negative_prompt: str,
steps: int,
seed: int,
height: int,
width: int,
resize_dimension: int,
guidance_scale: float,
):
if self.__current_task_name != "tile_upscaler":
raise Exception("ControlNet is not loaded with tile_upscaler model")
torch.manual_seed(seed)
init_image = download_image(imageUrl).resize((width, height))
condition_image = self.__resize_for_condition_image(
init_image, resize_dimension
)
result = self.pipe.__call__(
image=condition_image,
prompt=prompt,
controlnet_conditioning_image=condition_image,
num_inference_steps=steps,
negative_prompt=negative_prompt,
height=condition_image.size[1],
width=condition_image.size[0],
strength=1.0,
guidance_scale=guidance_scale,
)
return Result.from_result(result)
def detect_pose(self, imageUrl: str) -> Image.Image:
detector = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
image = download_image(imageUrl)
image = detector.__call__(image, hand_and_face=True)
return image
def __canny_detect_edge(self, image: Image.Image) -> Image.Image:
image_array = np.array(image)
low_threshold = 100
high_threshold = 200
image_array = cv2.Canny(image_array, low_threshold, high_threshold)
image_array = image_array[:, :, None]
image_array = np.concatenate([image_array, image_array, image_array], axis=2)
canny_image = Image.fromarray(image_array)
return canny_image
def __resize_for_condition_image(self, image: Image.Image, resolution: int):
input_image = image.convert("RGB")
W, H = input_image.size
k = float(resolution) / min(W, H)
H *= k
W *= k
H = int(round(H / 64.0)) * 64
W = int(round(W / 64.0)) * 64
img = input_image.resize((W, H), resample=Image.LANCZOS)
return img
|