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", "lllyasviel/control_v11f1p_sd15_depth"), ("NormalBae", "sd-control-lora-v3-normal-half-rank32-conv_in-rank128"), ("depth", "sd-control-lora-v3-depth-half-rank8-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() 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 ) unet.add_extra_conditions(["Placeholder"]) pipe: StableDiffusionControlLoraV3Pipeline = StableDiffusionControlLoraV3Pipeline.from_pretrained( base_model_id, safety_checker=None, unet=unet, torch_dtype=torch.float16 ) 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.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 # @torch.autocast("cuda") 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]: 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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 @torch.inference_mode() 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