optimzie

anime_app.py

@@ -14,13 +14,15 @@ import numpy as np
 # import imageio
 from huggingface_hub import HfApi
 import torch
+import spaces
 from PIL import Image
 from diffusers import (
     ControlNetModel,
     DPMSolverMultistepScheduler,
     StableDiffusionControlNetPipeline,
-    AutoencoderKL,
+    # AutoencoderKL,
 )
+from controlnet_aux_local import NormalBaeDetector
 from diffusers.models.attention_processor import AttnProcessor2_0
 MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
@@ -30,7 +32,28 @@ print("loading everything")
 compiled = False
 api = HfApi()
 
-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
@@ -59,10 +82,10 @@ scheduler = DPMSolverMultistepScheduler.from_pretrained(
 # Stable Diffusion Pipeline URL
 base_model_url = "https://huggingface.co/broyang/hentaidigitalart_v20/blob/main/realcartoon3d_v15.safetensors"
 # base_model = "broyang/realcartoon3d"
-print('loading vae')
-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_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 pipe')
 pipe = StableDiffusionControlNetPipeline.from_single_file(
@@ -71,30 +94,21 @@ pipe = StableDiffusionControlNetPipeline.from_single_file(
     # load_safety_checker=True,
     controlnet=controlnet,
     scheduler=scheduler,
-    vae=vae,
+    # vae=vae,
     torch_dtype=torch.float16,
-)
-
-# pipe = StableDiffusionControlNetPipeline.from_pretrained(
-#     base_model,
-#     safety_checker=None,
-#     controlnet=controlnet,
-#     scheduler=scheduler,
-#     vae=vae,
-#     torch_dtype=torch.float16,
-# )
+).to("cuda")
 
-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="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="fcNeg-neg.pt", token="fcNeg-neg")
-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")
+# 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="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="fcNeg-neg.pt", token="fcNeg-neg")
+# 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
@@ -103,14 +117,10 @@ preprocessor = Preprocessor()
 
 
 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)
+torch.cuda.empty_cache()
+gc.collect()
+print(f"CUDA memory allocated: {torch.cuda.max_memory_allocated(device='cuda') / 1e9:.2f} GB")
+print("Model Compiled!")
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -160,19 +170,30 @@ def get_prompt(prompt, additional_prompt):
     return prompt
 
 css = """
-h1 {
+h1, h2, h3 {
     text-align: center;
-    display:block;
+    display: block;
 }
-h2 {
-    text-align: center;
-    display:block;
+footer {
+    visibility: hidden;
 }
-h3 {
-    text-align: center;
-    display:block;
+.gradio-container {
+    max-width: 900px !important;
+}
+.gr-image {
+    display: flex;
+    justify-content: center; 
+    align-items: center;
+    width: 100%;
+    height: 512px;
+    overflow: hidden;
+}
+.gr-image img {
+    width: 100%;
+    height: 100%; 
+    object-fit: cover;
+    object-position: center;
 }
-footer {visibility: hidden}
 """
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     #############################################################################
@@ -216,7 +237,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     with gr.Column():
         prompt = gr.Textbox(
             label="Description",
-            placeholder="Leave empty for a random output 👀",
+            placeholder="Enter a description (optional)",
         )
     # input image
     with gr.Row():
@@ -230,7 +251,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             )
             # run button
             with gr.Column():
-                run_button = gr.Button(value="Use this one", size=["lg"], visible=False)
+                run_button = gr.Button(value="Use this one", size="lg", visible=False)
         # output image
         with gr.Column():
             result = gr.Image(  
@@ -241,7 +262,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             )
             # Use this image button
             with gr.Column():
-                use_ai_button = gr.Button(value="Use this one", size=["lg"], visible=False)
+                use_ai_button = gr.Button(value="Use this one", size="lg", visible=False)
     config = [
         image,
         prompt,
@@ -308,11 +329,8 @@ def process_image(
     torch.cuda.synchronize()
     preprocess_start = time.time()
     print("processing image")
+    global preprocessor
     preprocessor.load("NormalBae")
-    global compiled
-    if not compiled:
-        print("Not Compiled")
-        compiled = True
         
     seed = random.randint(0, MAX_SEED)
     generator = torch.cuda.manual_seed(seed)

local_anime_app.py

@@ -1,8 +1,7 @@
 prod = True
-port = 8080
 show_options = False
 if prod:
-    port = 8081
+    port = 8082
     # show_options = False
 
 import os
@@ -81,6 +80,11 @@ if gr.NO_RELOAD:
     pipe.to("cuda")
 
     print("---------------Loaded controlnet pipeline---------------") 
+    pipe.unet.set_attn_processor(AttnProcessor2_0())
+    torch.cuda.empty_cache()
+    gc.collect()
+    print(f"CUDA memory allocated: {torch.cuda.max_memory_allocated(device='cuda') / 1e9:.2f} GB")
+    print("Model Compiled!")
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -129,19 +133,30 @@ def get_prompt(prompt, additional_prompt):
     return prompt
 
 css = """
-h1 {
+h1, h2, h3 {
     text-align: center;
-    display:block;
+    display: block;
 }
-h2 {
-    text-align: center;
-    display:block;
+footer {
+    visibility: hidden;
 }
-h3 {
-    text-align: center;
-    display:block;
+.gradio-container {
+    max-width: 900px !important;
+}
+.gr-image {
+    display: flex;
+    justify-content: center; 
+    align-items: center;
+    width: 100%;
+    height: 512px;
+    overflow: hidden;
+}
+.gr-image img {
+    width: 100%;
+    height: 100%; 
+    object-fit: cover;
+    object-position: center;
 }
-footer {visibility: hidden}
 """
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     #############################################################################
@@ -277,10 +292,6 @@ def process_image(
     preprocess_start = time.time()
     print("processing image")
     preprocessor.load("NormalBae")
-    global compiled
-    if not compiled:
-        print("Not Compiled")
-        compiled = True
         
     seed = random.randint(0, MAX_SEED)
     generator = torch.cuda.manual_seed(seed)