| | 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 |
