Spaces:
Running
Running
from __future__ import annotations | |
import gc | |
import numpy as np | |
import PIL.Image | |
import torch | |
from controlnet_aux.util import HWC3 | |
from diffusers import ( | |
UniPCMultistepScheduler, | |
) | |
from unet import UNet2DConditionModelEx | |
from pipeline import StableDiffusionControlLoraV3Pipeline | |
from cv_utils import resize_image | |
from preprocessor import Preprocessor | |
from settings import MAX_IMAGE_RESOLUTION, MAX_NUM_IMAGES | |
from collections import OrderedDict | |
CONTROL_LORA_V3_MODEL_IDS = OrderedDict([ | |
("Openpose", "sd-control-lora-v3-pose-half-rank128-conv_in-rank128"), | |
("Canny", "sd-control-lora-v3-canny-half_skip_attn-rank16-conv_in-rank64"), | |
("segmentation", "sd-control-lora-v3-segmentation-half_skip_attn-rank128-conv_in-rank128"), | |
("depth", "sd-control-lora-v3-depth-half-rank8-conv_in-rank128"), | |
("Normal", "sd-control-lora-v3-normal-half-rank32-conv_in-rank128"), | |
("Tile", "sd-control-lora-v3-tile-half_skip_attn-rank16-conv_in-rank64"), | |
]) | |
class Model: | |
def __init__(self, base_model_id: str = "SG161222/Realistic_Vision_V4.0_noVAE", task_name: str = "Canny"): | |
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
self.base_model_id = "" | |
self.task_name = "" | |
self.pipe: StableDiffusionControlLoraV3Pipeline = self.load_pipe(base_model_id, task_name) | |
self.preprocessor = Preprocessor() | |
# preload | |
preprocessor = self.preprocessor | |
preprocessor.load("Openpose") | |
preprocessor.load("Canny") | |
preprocessor.load("OneFormer"); preprocessor.load("UPerNet") # segmentation | |
preprocessor.load("DPT") # depth | |
preprocessor.load("Midas") # normal (old) | |
def load_pipe(self, base_model_id: str, task_name) -> StableDiffusionControlLoraV3Pipeline: | |
if ( | |
base_model_id == self.base_model_id | |
and hasattr(self, "pipe") | |
and self.pipe is not None | |
): | |
unet: UNet2DConditionModelEx = self.pipe.unet | |
unet.set_adapter(task_name) | |
return self.pipe | |
unet: UNet2DConditionModelEx = UNet2DConditionModelEx.from_pretrained( | |
base_model_id, subfolder="unet", torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32 | |
) | |
unet.add_extra_conditions(["Placeholder"]) | |
pipe: StableDiffusionControlLoraV3Pipeline = StableDiffusionControlLoraV3Pipeline.from_pretrained( | |
base_model_id, safety_checker=None, unet=unet, torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32 | |
) | |
for _task_name, subfolder in CONTROL_LORA_V3_MODEL_IDS.items(): | |
pipe.load_lora_weights("HighCWu/control-lora-v3", adapter_name=_task_name, subfolder=subfolder) | |
pipe.unet.set_adapter(task_name) | |
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) | |
if self.device.type == "cuda": | |
pipe.enable_xformers_memory_efficient_attention() | |
pipe.to(self.device) | |
torch.cuda.empty_cache() | |
gc.collect() | |
self.base_model_id = base_model_id | |
self.task_name = task_name | |
return pipe | |
def set_base_model(self, base_model_id: str) -> str: | |
if not base_model_id or base_model_id == self.base_model_id: | |
return self.base_model_id | |
del self.pipe | |
if self.device.type == "cuda": | |
torch.cuda.empty_cache() | |
gc.collect() | |
try: | |
self.pipe = self.load_pipe(base_model_id, self.task_name) | |
except Exception: | |
self.pipe = self.load_pipe(self.base_model_id, self.task_name) | |
return self.base_model_id | |
def load_controlnet_weight(self, task_name: str) -> None: | |
if task_name == self.task_name: | |
return | |
unet: UNet2DConditionModelEx = self.pipe.unet | |
unet.set_adapter(task_name) | |
self.task_name = task_name | |
def get_prompt(self, prompt: str, additional_prompt: str) -> str: | |
if not prompt: | |
prompt = additional_prompt | |
else: | |
prompt = f"{prompt}, {additional_prompt}" | |
return prompt | |
def run_pipe( | |
self, | |
prompt: str, | |
negative_prompt: str, | |
control_image: PIL.Image.Image, | |
num_images: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
) -> list[PIL.Image.Image]: | |
def run(): | |
generator = torch.Generator().manual_seed(seed) | |
return self.pipe( | |
prompt=prompt, | |
negative_prompt=negative_prompt, | |
guidance_scale=guidance_scale, | |
num_images_per_prompt=num_images, | |
num_inference_steps=num_steps, | |
generator=generator, | |
image=control_image, | |
).images | |
if self.device.type == "cuda": | |
run = torch.autocast("cuda")(run) | |
return run() | |
def process_canny( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
low_threshold: int, | |
high_threshold: int, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
self.preprocessor.load("Canny") | |
control_image = self.preprocessor( | |
image=image, low_threshold=low_threshold, high_threshold=high_threshold, detect_resolution=image_resolution | |
) | |
self.load_controlnet_weight("Canny") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_mlsd( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
preprocess_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
value_threshold: float, | |
distance_threshold: float, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
self.preprocessor.load("MLSD") | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
thr_v=value_threshold, | |
thr_d=distance_threshold, | |
) | |
self.load_controlnet_weight("MLSD") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_scribble( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
preprocess_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
preprocessor_name: str, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
if preprocessor_name == "None": | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
elif preprocessor_name == "HED": | |
self.preprocessor.load(preprocessor_name) | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
scribble=False, | |
) | |
elif preprocessor_name == "PidiNet": | |
self.preprocessor.load(preprocessor_name) | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
safe=False, | |
) | |
self.load_controlnet_weight("scribble") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_scribble_interactive( | |
self, | |
image_and_mask: dict[str, np.ndarray], | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
) -> list[PIL.Image.Image]: | |
if image_and_mask is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
image = image_and_mask["mask"] | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
self.load_controlnet_weight("scribble") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_softedge( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
preprocess_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
preprocessor_name: str, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
if preprocessor_name == "None": | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
elif preprocessor_name in ["HED", "HED safe"]: | |
safe = "safe" in preprocessor_name | |
self.preprocessor.load("HED") | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
scribble=safe, | |
) | |
elif preprocessor_name in ["PidiNet", "PidiNet safe"]: | |
safe = "safe" in preprocessor_name | |
self.preprocessor.load("PidiNet") | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
safe=safe, | |
) | |
else: | |
raise ValueError | |
self.load_controlnet_weight("softedge") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_openpose( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
preprocess_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
preprocessor_name: str, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
if preprocessor_name == "None": | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
else: | |
self.preprocessor.load("Openpose") | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
hand_and_face=True, | |
) | |
self.load_controlnet_weight("Openpose") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_segmentation( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
preprocess_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
preprocessor_name: str, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
if preprocessor_name == "None": | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
else: | |
self.preprocessor.load(preprocessor_name) | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
) | |
self.load_controlnet_weight("segmentation") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_depth( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
preprocess_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
preprocessor_name: str, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
if preprocessor_name == "None": | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
else: | |
self.preprocessor.load(preprocessor_name) | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
) | |
self.load_controlnet_weight("depth") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_normal_old( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
preprocess_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
preprocessor_name: str, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
if preprocessor_name == "None": | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
else: | |
self.preprocessor.load("Midas") | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
depth_and_normal=True | |
) | |
self.load_controlnet_weight("Normal") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_normal( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
preprocess_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
preprocessor_name: str, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
if preprocessor_name == "None": | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
else: | |
self.preprocessor.load("NormalBae") | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
) | |
self.load_controlnet_weight("NormalBae") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_lineart( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
preprocess_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
preprocessor_name: str, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
if preprocessor_name in ["None", "None (anime)"]: | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
elif preprocessor_name in ["Lineart", "Lineart coarse"]: | |
coarse = "coarse" in preprocessor_name | |
self.preprocessor.load("Lineart") | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
coarse=coarse, | |
) | |
elif preprocessor_name == "Lineart (anime)": | |
self.preprocessor.load("LineartAnime") | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
detect_resolution=preprocess_resolution, | |
) | |
if "anime" in preprocessor_name: | |
self.load_controlnet_weight("lineart_anime") | |
else: | |
self.load_controlnet_weight("lineart") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_shuffle( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
preprocessor_name: str, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
if preprocessor_name == "None": | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
else: | |
self.preprocessor.load(preprocessor_name) | |
control_image = self.preprocessor( | |
image=image, | |
image_resolution=image_resolution, | |
) | |
self.load_controlnet_weight("shuffle") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_ip2p( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
self.load_controlnet_weight("ip2p") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |
def process_tile( | |
self, | |
image: np.ndarray, | |
prompt: str, | |
additional_prompt: str, | |
negative_prompt: str, | |
num_images: int, | |
image_resolution: int, | |
num_steps: int, | |
guidance_scale: float, | |
seed: int, | |
) -> list[PIL.Image.Image]: | |
if image is None: | |
raise ValueError | |
if image_resolution > MAX_IMAGE_RESOLUTION: | |
raise ValueError | |
if num_images > MAX_NUM_IMAGES: | |
raise ValueError | |
image = HWC3(image) | |
image = resize_image(image, resolution=image_resolution) | |
control_image = PIL.Image.fromarray(image) | |
self.load_controlnet_weight("Tile") | |
results = self.run_pipe( | |
prompt=self.get_prompt(prompt, additional_prompt), | |
negative_prompt=negative_prompt, | |
control_image=control_image, | |
num_images=num_images, | |
num_steps=num_steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
) | |
return [control_image] + results | |