optimized!

app.py

@@ -24,77 +24,82 @@ MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
 
 print("CUDA version:", torch.version.cuda)
-print("loading pipe")
+print("loading everything")
 compiled = False
-from preprocess import Preprocessor
-preprocessor = Preprocessor()
 
 import spaces
 
-if gr.NO_RELOAD:
-    torch.cuda.max_memory_allocated(device="cuda")
-    # Controlnet Normal
-    model_id = "lllyasviel/control_v11p_sd15_normalbae"
-    print("initializing controlnet")
-    controlnet = ControlNetModel.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16,
-        attn_implementation="flash_attention_2",
-    ).to("cuda")
+torch.cuda.max_memory_allocated(device="cuda")
+# Controlnet Normal
+model_id = "lllyasviel/control_v11p_sd15_normalbae"
+print("initializing controlnet")
+controlnet = ControlNetModel.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16,
+    attn_implementation="flash_attention_2",
+).to("cuda")
 
-    # Scheduler
-    scheduler = DPMSolverMultistepScheduler.from_pretrained(
-        "runwayml/stable-diffusion-v1-5",
-        solver_order=2,
-        subfolder="scheduler",
-        use_karras_sigmas=True,
-        final_sigmas_type="sigma_min",
-        algorithm_type="sde-dpmsolver++",
-        prediction_type="epsilon",
-        thresholding=False,
-        denoise_final=True,
-        device_map="cuda",
-        torch_dtype=torch.float16,
-    )
+# Scheduler
+scheduler = DPMSolverMultistepScheduler.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    solver_order=2,
+    subfolder="scheduler",
+    use_karras_sigmas=True,
+    final_sigmas_type="sigma_min",
+    algorithm_type="sde-dpmsolver++",
+    prediction_type="epsilon",
+    thresholding=False,
+    denoise_final=True,
+    device_map="cuda",
+    torch_dtype=torch.float16,
+)
 
-    # Stable Diffusion Pipeline URL
-    # base_model_url = "https://huggingface.co/broyang/hentaidigitalart_v20/blob/main/realcartoon3d_v15.safetensors"
-    base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors"
-    vae_url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors"
+# Stable Diffusion Pipeline URL
+# base_model_url = "https://huggingface.co/broyang/hentaidigitalart_v20/blob/main/realcartoon3d_v15.safetensors"
+base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors"
+vae_url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors"
 
-    vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda")
-    vae.to(memory_format=torch.channels_last) 
+print('loading vae')
+vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda")
+vae.to(memory_format=torch.channels_last) 
 
-    pipe = StableDiffusionControlNetPipeline.from_single_file(
-        base_model_url,
-        # safety_checker=None,
-        # load_safety_checker=True,
-        controlnet=controlnet,
-        scheduler=scheduler,
-        vae=vae,
-        torch_dtype=torch.float16,
-    )
+print('loading pipe')
+pipe = StableDiffusionControlNetPipeline.from_single_file(
+    base_model_url,
+    safety_checker=None,
+    # load_safety_checker=True,
+    controlnet=controlnet,
+    scheduler=scheduler,
+    vae=vae,
+    torch_dtype=torch.float16,
+)
+
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="EasyNegativeV2.safetensors", token="EasyNegativeV2",)
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="badhandv4.pt", token="badhandv4")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="fcNeg-neg.pt", token="fcNeg-neg")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Ahegao.pt", token="HDA_Ahegao")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Bondage.pt", token="HDA_Bondage")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_pet_play.pt", token="HDA_pet_play")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_unconventional maid.pt", token="HDA_unconventional_maid")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NakedHoodie.pt", token="HDA_NakedHoodie")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NunDress.pt", token="HDA_NunDress")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Shibari.pt", token="HDA_Shibari")
+pipe.to("cuda")
+
+print("loading preprocessor")
+from preprocess import Preprocessor
+preprocessor = Preprocessor()
+preprocessor.load("NormalBae")
 
-    pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="EasyNegativeV2.safetensors", token="EasyNegativeV2",)
-    pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="badhandv4.pt", token="badhandv4")
-    pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="fcNeg-neg.pt", token="fcNeg-neg")
-    pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Ahegao.pt", token="HDA_Ahegao")
-    pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Bondage.pt", token="HDA_Bondage")
-    pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_pet_play.pt", token="HDA_pet_play")
-    pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_unconventional maid.pt", token="HDA_unconventional_maid")
-    pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NakedHoodie.pt", token="HDA_NakedHoodie")
-    pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NunDress.pt", token="HDA_NunDress")
-    pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Shibari.pt", token="HDA_Shibari")
-    pipe.to("cuda")
-    print("---------------Loaded controlnet pipeline---------------") 
+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())
+    print("Model Compiled!")
+init(pipe)
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:

preprocess.py

@@ -1,12 +1,7 @@
-# import numpy as np
 import PIL.Image
 import torch, gc
 from controlnet_aux_local import NormalBaeDetector#, CannyDetector
 
-# from controlnet_aux.util import HWC3
-# import cv2
-# from cv_utils import resize_image
-
 class Preprocessor:
     MODEL_ID = "lllyasviel/Annotators"
 
@@ -21,14 +16,10 @@ class Preprocessor:
             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()
+            self.name = name
         else:
             raise ValueError
-
-        
-        self.name = name
+        return
 
     def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
         return self.model(image, **kwargs)