testing new preprocess

app.py

@@ -21,7 +21,6 @@ from diffusers import (
     AutoencoderKL,
 )
 from diffusers.models.attention_processor import AttnProcessor2_0
-from preprocess import Preprocessor
 MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
 
@@ -31,9 +30,6 @@ compiled = False
 # api = HfApi()
 
 import spaces
-    
-preprocessor = Preprocessor()
-preprocessor.load("NormalBae")
 
 if gr.NO_RELOAD:
     torch.cuda.max_memory_allocated(device="cuda")
@@ -98,13 +94,16 @@ if gr.NO_RELOAD:
     # gc.collect()
     print("---------------Loaded controlnet pipeline---------------") 
 
-    # @spaces.GPU(duration=12)
-    # def init(pipe):
-    #     pipe.enable_xformers_memory_efficient_attention()
-    #     pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-    #     pipe.unet.set_attn_processor(AttnProcessor2_0())
-    #     print("Model Compiled!")
-    # init(pipe)
+    @spaces.GPU(duration=12)
+    def init(pipe):
+        pipe.enable_xformers_memory_efficient_attention()
+        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+        pipe.unet.set_attn_processor(AttnProcessor2_0())
+        from preprocess import Preprocessor
+        preprocessor = Preprocessor()
+        preprocessor.load("NormalBae")
+        print("Model Compiled!")
+    init(pipe)
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:

local_app.py

@@ -21,7 +21,6 @@ from diffusers import (
     AutoencoderKL,
 )
 from diffusers.models.attention_processor import AttnProcessor2_0
-from local_preprocess import Preprocessor
 MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
 
@@ -30,6 +29,7 @@ print("loading pipe")
 compiled = False
 
 if gr.NO_RELOAD:
+    from local_preprocess import Preprocessor
     preprocessor = Preprocessor()
     preprocessor.load("NormalBae")
     torch.cuda.max_memory_allocated(device="cuda")

local_preprocess.py

@@ -10,21 +10,6 @@ from controlnet_aux import NormalBaeDetector
 
 class Preprocessor:
     MODEL_ID = "lllyasviel/Annotators"
-    
-    # def resize_image(input_image, resolution, interpolation=None):
-    #     H, W, C = input_image.shape
-    #     H = float(H)
-    #     W = float(W)
-    #     k = float(resolution) / max(H, W)
-    #     H *= k
-    #     W *= k
-    #     H = int(np.round(H / 64.0)) * 64
-    #     W = int(np.round(W / 64.0)) * 64
-    #     if interpolation is None:
-    #         interpolation = cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA
-    #     img = cv2.resize(input_image, (W, H), interpolation=interpolation)
-    #     return img
-
 
     def __init__(self):
         self.model = None
@@ -36,33 +21,13 @@ class Preprocessor:
         elif name == "NormalBae":
             print("Loading NormalBae")
             self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
+            torch.cuda.empty_cache()
+            gc.collect()
         # elif name == "Canny":
         #     self.model = CannyDetector()
         else:
             raise ValueError
-        torch.cuda.empty_cache()
-        gc.collect()
-        
         self.name = name
 
     def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
-        # if self.name == "Canny":
-        #     if "detect_resolution" in kwargs:
-        #         detect_resolution = kwargs.pop("detect_resolution")
-        #         image = np.array(image)
-        #         image = HWC3(image)
-        #         image = resize_image(image, resolution=detect_resolution)
-        #     image = self.model(image, **kwargs)
-        #     return PIL.Image.fromarray(image)
-        # elif self.name == "Midas":
-        #     detect_resolution = kwargs.pop("detect_resolution", 512)
-        #     image_resolution = kwargs.pop("image_resolution", 512)
-        #     image = np.array(image)
-        #     image = HWC3(image)
-        #     image = resize_image(image, resolution=detect_resolution)
-        #     image = self.model(image, **kwargs)
-        #     image = HWC3(image)
-        #     image = resize_image(image, resolution=image_resolution)
-        #     return PIL.Image.fromarray(image)
-        # else:
         return self.model(image, **kwargs)

preprocess.py

@@ -1,7 +1,8 @@
 # import numpy as np
 import PIL.Image
-# import torch
+import torch, gc
 from controlnet_aux_local import NormalBaeDetector#, CannyDetector
+import spaces
 
 # from controlnet_aux.util import HWC3
 # import cv2
@@ -9,21 +10,6 @@ from controlnet_aux_local import NormalBaeDetector#, CannyDetector
 
 class Preprocessor:
     MODEL_ID = "lllyasviel/Annotators"
-    
-    # def resize_image(input_image, resolution, interpolation=None):
-    #     H, W, C = input_image.shape
-    #     H = float(H)
-    #     W = float(W)
-    #     k = float(resolution) / max(H, W)
-    #     H *= k
-    #     W *= k
-    #     H = int(np.round(H / 64.0)) * 64
-    #     W = int(np.round(W / 64.0)) * 64
-    #     if interpolation is None:
-    #         interpolation = cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA
-    #     img = cv2.resize(input_image, (W, H), interpolation=interpolation)
-    #     return img
-
 
     def __init__(self):
         self.model = None
@@ -35,33 +21,15 @@ class Preprocessor:
         elif name == "NormalBae":
             print("Loading NormalBae")
             self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
+            torch.cuda.empty_cache()
+            gc.collect()
         # elif name == "Canny":
         #     self.model = CannyDetector()
         else:
             raise ValueError
-        # torch.cuda.empty_cache()
-        # gc.collect()
+
         
         self.name = name
 
     def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
-        # if self.name == "Canny":
-        #     if "detect_resolution" in kwargs:
-        #         detect_resolution = kwargs.pop("detect_resolution")
-        #         image = np.array(image)
-        #         image = HWC3(image)
-        #         image = resize_image(image, resolution=detect_resolution)
-        #     image = self.model(image, **kwargs)
-        #     return PIL.Image.fromarray(image)
-        # elif self.name == "Midas":
-        #     detect_resolution = kwargs.pop("detect_resolution", 512)
-        #     image_resolution = kwargs.pop("image_resolution", 512)
-        #     image = np.array(image)
-        #     image = HWC3(image)
-        #     image = resize_image(image, resolution=detect_resolution)
-        #     image = self.model(image, **kwargs)
-        #     image = HWC3(image)
-        #     image = resize_image(image, resolution=image_resolution)
-        #     return PIL.Image.fromarray(image)
-        # else:
         return self.model(image, **kwargs)