local and empty cache on hf

app.py

@@ -364,7 +364,6 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     # image processing
     @gr.on(triggers=[image.upload, prompt.submit, run_button.click], inputs=config, outputs=result, show_progress="minimal")
     def auto_process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
-        preprocessor.load("NormalBae")
         return process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
     
     # AI image processing
@@ -434,7 +433,7 @@ def process_image(
     ).images[0]
     print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------")
     # torch.cuda.synchronize()
-    # torch.cuda.empty_cache()
+    torch.cuda.empty_cache()
     return results
 
 if prod:

local_app.py

@@ -10,6 +10,8 @@ import random
 import time
 import gradio as gr
 import numpy as np
+# import spaces
+# import imageio
 import gc
 import torch
 from PIL import Image
@@ -17,85 +19,105 @@ from diffusers import (
     ControlNetModel,
     DPMSolverMultistepScheduler,
     StableDiffusionControlNetPipeline,
-    AutoencoderKL,
+    # 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)
+# os.environ['HF_HOME'] = '/data/.huggingface'
 
 print("CUDA version:", torch.version.cuda)
 print("loading everything")
 compiled = False
 
-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")
+class Preprocessor:
+    MODEL_ID = "lllyasviel/Annotators"
 
-    # 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,
-    )
+    def __init__(self):
+        self.model = None
+        self.name = ""
 
-    # 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"
+    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
 
-    # print('loading vae')
-    # vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda")
-    # vae.to(memory_format=torch.channels_last) 
+    def __call__(self, image: Image.Image, **kwargs) -> Image.Image:
+        return self.model(image, **kwargs)
 
-    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,
-    )
+# torch.cuda.max_memory_allocated(device="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="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")
+# 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")
 
-    print("loading preprocessor")
-    from preprocess import Preprocessor
-    preprocessor = Preprocessor()
-    preprocessor.load("NormalBae")
+# 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,
+)
 
-    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!")
+# 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"
+
+# print('loading vae')
+# 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(
+    base_model_url,
+    safety_checker=None,
+    controlnet=controlnet,
+    scheduler=scheduler,
+    # vae=vae,
+    torch_dtype=torch.float16,
+).to("cuda")
+
+print("loading 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---------------") 
+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:
@@ -227,7 +249,7 @@ footer {
     visibility: hidden;
 }
 .gradio-container {
-    max-width: 900px !important;
+    max-width: 1100px !important;
 }
 .gr-image {
     display: flex;
@@ -255,14 +277,14 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
                 label="Image resolution",
                 minimum=256,
                 maximum=1024,
-                value=768,
+                value=512,
                 step=256,
             )
             preprocess_resolution = gr.Slider(
                 label="Preprocess resolution",
                 minimum=128,
                 maximum=1024,
-                value=768,
+                value=512,
                 step=1,
             )
             num_steps = gr.Slider(
@@ -282,11 +304,13 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
                 value="EasyNegativeV2, fcNeg, (badhandv4:1.4), (worst quality, low quality, bad quality, normal quality:2.0), (bad hands, missing fingers, extra fingers:2.0)",
             )
     #############################################################################
+    # input text
     with gr.Column():
         prompt = gr.Textbox(
-            label="Custom Prompt",
-            placeholder="boho chic",
+            label="Custom Design",
+            placeholder="Enter a description (optional)",
         )
+    # design options
     with gr.Row(visible=True):
         style_selection = gr.Radio(
             show_label=True,
@@ -297,24 +321,24 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             label="Design Styles",
         )
     # input image
-    with gr.Row():
-        with gr.Column():
+    with gr.Row(equal_height=True):
+        with gr.Column(scale=1, min_width=300):
             image = gr.Image(
                 label="Input",
                 sources=["upload"],
                 show_label=True,
                 mirror_webcam=True,
-                format="webp",
+                type="pil",
             )
             # run button
             with gr.Column():
                 run_button = gr.Button(value="Use this one", size="lg", visible=False)
         # output image
-        with gr.Column():
+        with gr.Column(scale=1, min_width=300):
             result = gr.Image(  
                 label="Output",
                 interactive=False,
-                format="webp",
+                type="pil",
                 show_share_button= False,
             )
             # Use this image button
@@ -333,28 +357,22 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
         guidance_scale,
         seed,
     ]
-        
+    
     with gr.Row():
         helper_text = gr.Markdown("## Tap and hold (on mobile) to save the image.", visible=True)
-        
+    
     # image processing
     @gr.on(triggers=[image.upload, prompt.submit, run_button.click], inputs=config, outputs=result, show_progress="minimal")
     def auto_process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
         return process_image(image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
-
-    # # AI Image Processing
-    # @gr.on(triggers=[use_ai_button.click], inputs=config, outputs=result, show_progress="minimal")
-    # def submit(result, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
-    #     return process_image(result, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
-        
+    
+    # AI image processing
     @gr.on(triggers=[use_ai_button.click], inputs=[result] + config, outputs=[image, result], show_progress="minimal")
     def submit(previous_result, image, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
         # First, yield the previous result to update the input image immediately
         yield previous_result, gr.update()
-            
         # Then, process the new input image
         new_result = process_image(previous_result, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
-            
         # Finally, yield the new result
         yield previous_result, new_result
 
@@ -362,12 +380,13 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     @gr.on(triggers=[image.upload, use_ai_button.click, run_button.click], inputs=None, outputs=[run_button, use_ai_button], show_progress="hidden")
     def turn_buttons_off():
         return gr.update(visible=False), gr.update(visible=False)
-        
+    
     # Turn on buttons when processing is complete
     @gr.on(triggers=[result.change], inputs=None, outputs=[use_ai_button, run_button], show_progress="hidden")
     def turn_buttons_on():
         return gr.update(visible=True), gr.update(visible=True)
 
+# @spaces.GPU(duration=12)
 @torch.inference_mode()
 def process_image(
     image,
@@ -386,11 +405,9 @@ def process_image(
     preprocess_start = time.time()
     print("processing image")
 
-    # global preprocessor
-    # preprocessor.load("NormalBae")
-    
     seed = random.randint(0, MAX_SEED)
     generator = torch.cuda.manual_seed(seed)
+    preprocessor.load("NormalBae")
     control_image = preprocessor(
         image=image,
         image_resolution=image_resolution,
@@ -415,9 +432,8 @@ def process_image(
         image=control_image,
     ).images[0]
     print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------")
-    # results.save(os.path.join("/data", "temp_image.jpg"))
     # torch.cuda.synchronize()
-    # torch.cuda.empty_cache()
+    torch.cuda.empty_cache()
     return results
 
 if prod: