local and empty cache

anime_app.py

@@ -284,22 +284,6 @@ with gr.Blocks("bethecloud/storj_theme", css=css) as demo:
     def auto_process_image(image, 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, 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(image, 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, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
-
-    # # Change input to result
-    # @gr.on(triggers=[use_ai_button.click], inputs=None, outputs=image, show_progress="hidden")
-    # def update_input():
-    #     try:
-    #         print("Updating image to AI Temp Image")
-    #         ai_temp_image = Image.open("temp_image.jpg")
-    #         return ai_temp_image
-    #     except FileNotFoundError:
-    #         print("No AI Image Available")
-    #         return None
-
     @gr.on(triggers=[use_ai_button.click], inputs=[result] + config, outputs=[image, result], show_progress="minimal")
     def submit(previous_result, image, 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
@@ -362,10 +346,8 @@ def process_image(
         generator=generator,
         image=control_image,
     ).images[0]
-    # torch.cuda.synchronize()
-    # torch.cuda.empty_cache()
     print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------")
-    # results.save("temp_image.jpg")
+    torch.cuda.empty_cache()
     return results
 
 if prod:

local_anime_app.py

@@ -1,7 +1,8 @@
-prod = True
+prod = False
+port = 8080
 show_options = False
 if prod:
-    port = 8082
+    port = 8081
     # show_options = False
 
 import os
@@ -10,81 +11,116 @@ import random
 import time
 import gradio as gr
 import numpy as np
-import imageio
+# import imageio
 from huggingface_hub import HfApi
 import torch
+# import spaces
 from PIL import Image
 from diffusers import (
     ControlNetModel,
     DPMSolverMultistepScheduler,
     StableDiffusionControlNetPipeline,
+    # AutoencoderKL,
 )
+from controlnet_aux_local import NormalBaeDetector
+# from controlnet_aux 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)
 
 print("CUDA version:", torch.version.cuda)
-print("loading pipe")
+print("loading everything")
 compiled = False
-from preprocess import Preprocessor
-preprocessor = Preprocessor()
 api = HfApi()
 
-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",
-    )
+    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"
+    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
 
-    pipe = StableDiffusionControlNetPipeline.from_single_file(
-        base_model_url,
-        # safety_checker=None,
-        # load_safety_checker=True,
-        controlnet=controlnet,
-        scheduler=scheduler,
-        torch_dtype=torch.float16,
-    )
+    def __call__(self, image: Image.Image, **kwargs) -> Image.Image:
+        return self.model(image, **kwargs)
 
-    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")
+# 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,
+)
+
+# 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!")
 
-    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:
@@ -100,8 +136,9 @@ def get_additional_prompt():
     # outfit = ["schoolgirl outfit", "playboy outfit", "red dress", "gala dress", "cheerleader outfit", "nurse outfit", "Kimono"]
 
 def get_prompt(prompt, additional_prompt):
-    default = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
-    default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
+    default = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed,tungsten white balance"
+    # default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
+    default2 = f"hyperrealistic photography of {prompt},extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
     randomize = get_additional_prompt()
     # nude = "NSFW,((nude)),medium bare breasts,hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
     # bodypaint = "((fully naked with no clothes)),nude naked seethroughxray,invisiblebodypaint,rating_newd,NSFW"
@@ -141,7 +178,7 @@ footer {
     visibility: hidden;
 }
 .gradio-container {
-    max-width: 900px !important;
+    max-width: 1100px !important;
 }
 .gr-image {
     display: flex;
@@ -158,7 +195,7 @@ footer {
     object-position: center;
 }
 """
-with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
+with gr.Blocks("bethecloud/storj_theme", css=css) as demo:
     #############################################################################
     with gr.Row():
         with gr.Accordion("Advanced options", open=show_options, visible=show_options):
@@ -200,32 +237,32 @@ 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():
-        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)
+                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="Anime AI",
+                label="Output",
                 interactive=False,
-                format="webp",
+                type="pil",
                 show_share_button= False,
             )
             # 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,
@@ -247,22 +284,15 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     def auto_process_image(image, 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, 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(image, 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, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed)
-
-    # Change input to result
-    @gr.on(triggers=[use_ai_button.click], inputs=None, outputs=image, show_progress="hidden")
-    def update_input():
-        try:
-            print("Updating image to AI Temp Image")
-            ai_temp_image = Image.open("temp_image.jpg")
-            return ai_temp_image
-        except FileNotFoundError:
-            print("No AI Image Available")
-            return None
-        
+    @gr.on(triggers=[use_ai_button.click], inputs=[result] + config, outputs=[image, result], show_progress="minimal")
+    def submit(previous_result, image, 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, 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
+                
     # Turn off buttons when processing
     @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():
@@ -274,6 +304,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         return gr.update(visible=True), gr.update(visible=True)
 
 
+# @spaces.GPU(duration=12)
 @torch.inference_mode()
 def process_image(
     image,
@@ -288,13 +319,12 @@ def process_image(
     seed,
     progress=gr.Progress(track_tqdm=True)
 ):
-    torch.cuda.synchronize()
+    # torch.cuda.synchronize()
     preprocess_start = time.time()
     print("processing image")
-    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,
@@ -305,7 +335,7 @@ def process_image(
     custom_prompt=str(get_prompt(prompt, a_prompt))
     negative_prompt=str(n_prompt)
     print(f"{custom_prompt}")
-    
+    print(f"\n-------------------------Preprocess done in: {preprocess_time:.2f} seconds-------------------------")    
     start = time.time()
     results = pipe(
         prompt=custom_prompt,
@@ -316,36 +346,8 @@ def process_image(
         generator=generator,
         image=control_image,
     ).images[0]
-    torch.cuda.synchronize()
-    torch.cuda.empty_cache()
-    print(f"\n-------------------------Preprocess done in: {preprocess_time:.2f} seconds-------------------------")    
     print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------")
-    
-    # if not os.path.exists("./outputs"):
-    #     os.makedirs("./outputs")
-    # img_path = f"./outputs/{timestamp}.jpg"
-    # results_path = f"./outputs/{timestamp}_out_{prompt}.jpg"
-    # imageio.imsave(img_path, image)
-    # results.save(results_path)
-    results.save("temp_image.jpg")
-    
-    # api.upload_file(
-    #     path_or_fileobj=img_path,
-    #     path_in_repo=img_path,
-    #     repo_id="broyang/anime-ai-outputs",
-    #     repo_type="dataset",
-    #     token=API_KEY,
-    #     run_as_future=True,
-    # )
-    # api.upload_file(
-    #     path_or_fileobj=results_path,
-    #     path_in_repo=results_path,
-    #     repo_id="broyang/anime-ai-outputs",
-    #     repo_type="dataset",
-    #     token=API_KEY,
-    #     run_as_future=True,
-    # )
-
+    torch.cuda.empty_cache()
     return results
 
 if prod: