|
|
from __future__ import annotations
|
|
|
import gc
|
|
|
import numpy as np
|
|
|
import PIL.Image
|
|
|
import torch
|
|
|
from diffusers import (
|
|
|
ControlNetModel,
|
|
|
DiffusionPipeline,
|
|
|
StableDiffusionControlNetPipeline,
|
|
|
UniPCMultistepScheduler,
|
|
|
)
|
|
|
|
|
|
from preprocessor import Preprocessor
|
|
|
from settings import *
|
|
|
|
|
|
|
|
|
class Model:
|
|
|
def __init__(self, base_model_id: str = "runwayml/stable-diffusion-v1-5", task_name: str = "lineart"):
|
|
|
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
|
self.base_model_id = ""
|
|
|
self.task_name = ""
|
|
|
self.pipe = self.load_pipe(base_model_id, task_name)
|
|
|
self.preprocessor = Preprocessor()
|
|
|
|
|
|
def load_pipe(self, base_model_id: str, task_name) -> DiffusionPipeline:
|
|
|
if (
|
|
|
base_model_id == self.base_model_id
|
|
|
and task_name == self.task_name
|
|
|
and hasattr(self, "pipe")
|
|
|
and self.pipe is not None
|
|
|
):
|
|
|
return self.pipe
|
|
|
controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float16)
|
|
|
pipe = StableDiffusionControlNetPipeline.from_pretrained(
|
|
|
base_model_id, safety_checker=None, controlnet=controlnet, torch_dtype=torch.float16
|
|
|
)
|
|
|
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
|
|
|
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
|
|
|
if self.pipe is not None and hasattr(self.pipe, "controlnet"):
|
|
|
del self.pipe.controlnet
|
|
|
torch.cuda.empty_cache()
|
|
|
gc.collect()
|
|
|
controlnet = ControlNetModel.from_pretrained(model_id, torch_dtype=torch.float16)
|
|
|
controlnet.to(self.device)
|
|
|
torch.cuda.empty_cache()
|
|
|
gc.collect()
|
|
|
self.pipe.controlnet = controlnet
|
|
|
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,
|
|
|
control_image: PIL.Image.Image,
|
|
|
) -> list[PIL.Image.Image]:
|
|
|
generator = torch.Generator().manual_seed(randomize_seed)
|
|
|
return self.pipe(
|
|
|
prompt=prompt + ' ' + a_prompt,
|
|
|
negative_prompt=n_prompt,
|
|
|
guidance_scale=guidance_scale,
|
|
|
num_images_per_prompt=DEFAULT_NUM_IMAGES,
|
|
|
num_inference_steps=num_steps,
|
|
|
generator=generator,
|
|
|
image=control_image,
|
|
|
).images
|
|
|
|
|
|
def process_lineart(
|
|
|
self,
|
|
|
image: np.ndarray,
|
|
|
) -> list[PIL.Image.Image]:
|
|
|
if image is None:
|
|
|
raise ValueError
|
|
|
|
|
|
else:
|
|
|
|
|
|
self.preprocessor.load("Lineart")
|
|
|
control_image = self.preprocessor(
|
|
|
image=image,
|
|
|
image_resolution=DEFAULT_IMAGE_RESOLUTION,
|
|
|
detect_resolution=preprocess_resolution,
|
|
|
)
|
|
|
self.load_controlnet_weight("lineart")
|
|
|
results = self.run_pipe(
|
|
|
control_image=control_image
|
|
|
)
|
|
|
return [control_image] + results |
