preprocessor class

app.py

@@ -10,6 +10,7 @@ import random
 import time
 import gradio as gr
 import numpy as np
+import spaces
 # import imageio
 import torch
 from PIL import Image
@@ -19,7 +20,8 @@ from diffusers import (
     StableDiffusionControlNetPipeline,
     AutoencoderKL,
 )
-# from diffusers.models.attention_processor import AttnProcessor2_0
+from controlnet_aux_local import NormalBaeDetector
+
 MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
 
@@ -27,7 +29,27 @@ print("CUDA version:", torch.version.cuda)
 print("loading everything")
 compiled = False
 
-import spaces
+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
+        elif name == "NormalBae":
+            print("Loading NormalBae")
+            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
+            torch.cuda.empty_cache()
+            self.name = name
+        else:
+            raise ValueError
+        return
+
+    def __call__(self, image: Image.Image, **kwargs) -> Image.Image:
+        return self.model(image, **kwargs)
 
 torch.cuda.max_memory_allocated(device="cuda")
 # Controlnet Normal
@@ -87,7 +109,6 @@ pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Shi
 pipe.to("cuda")
 
 print("loading preprocessor")
-from preprocess import Preprocessor
 preprocessor = Preprocessor()
 # preprocessor.load("NormalBae")