Upload 5 files

handler.py

@@ -0,0 +1,36 @@
+from typing import  Dict, Any
+import torch
+import base64
+from io import BytesIO
+from model import Model
+from PIL import Image
+# set device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+if device.type != 'cuda':
+    raise ValueError("need to run on GPU")
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # load the optimized model
+        self.model = Model()
+
+
+    def __call__(self, data: Any) -> Any:
+        """
+        Args:
+            data (:obj:):
+                includes the input data and the parameters for the inference.
+        Return:
+            A :obj:`dict`:. base64 encoded image
+        """
+        inputs = data.pop("image", data)
+        
+        image = Image.open(BytesIO(base64.b64decode(inputs)))
+
+        # run inference pipeline
+        _, res = self.model.process_lineart(image = image) 
+        
+            
+        # encoding image as base 64 is done by the default toolkit
+        return res

model.py

@@ -0,0 +1,115 @@
+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

preprocessor.py

@@ -0,0 +1,27 @@
+import gc
+import PIL.Image
+import torch
+from controlnet_aux import LineartDetector
+
+class Preprocessor:
+    MODEL_ID = "lllyasviel/Annotators"
+
+    def __init__(self):
+        self.model = None
+        self.name = ""
+
+    def load(self, name: str) -> None:
+        if name == self.name:
+            return
+        if name == "Lineart":
+            self.model = LineartDetector.from_pretrained(self.MODEL_ID)
+
+        else:
+            raise ValueError
+        torch.cuda.empty_cache()
+        gc.collect()
+        self.name = name
+
+    def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
+            return self.model(image, **kwargs)
+

requirements.txt

@@ -0,0 +1,13 @@
+accelerate==0.21.0
+controlnet_aux==0.0.6
+diffusers==0.18.2
+einops==0.6.1
+gradio==3.45.2
+huggingface-hub==0.16.4
+mediapipe==0.10.1
+opencv-python-headless==4.8.0.74
+safetensors==0.3.1
+torch==2.0.1
+torchvision==0.15.2
+transformers==4.30.2
+xformers==0.0.20

settings.py

@@ -0,0 +1,17 @@
+DEFAULT_MODEL_ID = "runwayml/stable-diffusion-v1-5"
+DEFAULT_NUM_IMAGES = 1
+MAX_IMAGE_RESOLUTION = 768
+DEFAULT_IMAGE_RESOLUTION = 768
+preprocess_resolution = 512
+num_steps = 20
+guidance_scale = 9
+randomize_seed = 0
+
+task_name = "lineart"
+model_id = "lllyasviel/control_v11p_sd15_lineart"
+prompt = "Architecture, Building, Realistic, 3D Rendering, 2D Elevation, Professional."
+
+a_prompt = "best quality, extremely detailed"
+n_prompt = "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
+preprocessor_name = 'lineart'
+