heh

app.py

@@ -5,26 +5,23 @@ if prod:
     port = 8081
     # show_options = False
 
-import gc
 import os
+import gc
 import random
 import time
-
 import gradio as gr
 import numpy as np
-
 # import imageio
 import torch
+from PIL import Image
 from diffusers import (
-    AutoencoderKL,
     ControlNetModel,
     DPMSolverMultistepScheduler,
     StableDiffusionControlNetPipeline,
+    AutoencoderKL,
 )
 from diffusers.models.attention_processor import AttnProcessor2_0
-from PIL import Image
 from preprocess import Preprocessor
-
 MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
 
@@ -34,7 +31,7 @@ compiled = False
 # api = HfApi()
 
 import spaces
-
+    
 preprocessor = Preprocessor()
 preprocessor.load("NormalBae")
 
@@ -49,7 +46,7 @@ if gr.NO_RELOAD:
         torch_dtype=torch.float16,
         attn_implementation="flash_attention_2",
     ).to("cuda")
-
+    
     # Scheduler
     scheduler = DPMSolverMultistepScheduler.from_pretrained(
         "runwayml/stable-diffusion-v1-5",
@@ -69,10 +66,10 @@ if gr.NO_RELOAD:
     # 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)
-
+    vae.to(memory_format=torch.channels_last) 
+    
     pipe = StableDiffusionControlNetPipeline.from_single_file(
         base_model_url,
         # safety_checker=None,
@@ -83,57 +80,23 @@ if gr.NO_RELOAD:
         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.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")
-
+    
     # experimental speedup?
     # pipe.compile()
     # torch.cuda.empty_cache()
     # gc.collect()
-    print("---------------Loaded controlnet pipeline---------------")
+    print("---------------Loaded controlnet pipeline---------------") 
 
     @spaces.GPU(duration=12)
     def init(pipe):
@@ -141,42 +104,21 @@ if gr.NO_RELOAD:
         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:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-
 def get_additional_prompt():
     prompt = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
     top = ["tank top", "blouse", "button up shirt", "sweater", "corset top"]
-    bottom = [
-        "short skirt",
-        "athletic shorts",
-        "jean shorts",
-        "pleated skirt",
-        "short skirt",
-        "leggings",
-        "high-waisted shorts",
-    ]
-    accessory = [
-        "knee-high boots",
-        "gloves",
-        "Thigh-high stockings",
-        "Garter belt",
-        "choker",
-        "necklace",
-        "headband",
-        "headphones",
-    ]
+    bottom = ["short skirt", "athletic shorts", "jean shorts", "pleated skirt", "short skirt", "leggings", "high-waisted shorts"]
+    accessory = ["knee-high boots", "gloves", "Thigh-high stockings", "Garter belt", "choker", "necklace", "headband", "headphones"]
     return f"{prompt}, {random.choice(top)}, {random.choice(bottom)}, {random.choice(accessory)}, score_9"
     # outfit = ["schoolgirl outfit", "playboy outfit", "red dress", "gala dress", "cheerleader outfit", "nurse outfit", "Kimono"]
 
-
 def get_prompt(prompt, additional_prompt):
     interior = "design-style interior designed (interior space), captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length, tungsten white balance, (sharp focus), professional photography, high-resolution, 8k, Pulitzer Prize-winning"
     default = "hyperrealistic photography,extremely detailed,(intricate details),unity 8k wallpaper,ultra detailed"
@@ -197,22 +139,9 @@ def get_prompt(prompt, additional_prompt):
     abg = "(1girl, asian body covered in words, words on body, tattoos of (words) on body),(masterpiece, best quality),medium breasts,(intricate details),unity 8k wallpaper,ultra detailed,(pastel colors),beautiful and aesthetic,see-through (clothes),detailed,solo"
     # shibari = "extremely detailed, hyperrealistic photography, earrings, blushing, lace choker, tattoo, medium hair, score_9, HDA_Shibari"
     shibari2 = "octane render, highly detailed, volumetric, HDA_Shibari"
-
+    
     if prompt == "":
-        girls = [
-            randomize,
-            pet_play,
-            bondage,
-            lab_girl,
-            athleisure,
-            atompunk,
-            maid,
-            nundress,
-            naked_hoodie,
-            abg,
-            shibari2,
-            ahegao2,
-        ]
+        girls = [randomize, pet_play, bondage, lab_girl, athleisure, atompunk, maid, nundress, naked_hoodie, abg, shibari2, ahegao2]
         prompts_nsfw = [abg, shibari2, ahegao2]
         prompt = f"{random.choice(girls)}"
         prompt = f"boho chic"
@@ -222,11 +151,19 @@ def get_prompt(prompt, additional_prompt):
         # prompt = default2
     return f"{prompt} f{additional_prompt}"
 
-
 style_list = [
-    {"name": "None", "prompt": ""},
-    {"name": "Minimalistic", "prompt": "Minimalistic"},
-    {"name": "Boho Chic", "prompt": "boho chic"},
+    {
+        "name": "None",
+        "prompt": ""
+    },
+    {
+        "name": "Minimalistic",
+        "prompt": "Minimalistic"
+    },
+    {
+        "name": "Boho Chic",
+        "prompt": "boho chic"
+    },
     {
         "name": "Saudi Prince Gold",
         "prompt": "saudi prince gold",
@@ -271,22 +208,17 @@ style_list = [
         "name": "Beach",
         "prompt": "Beach",
     },
-    {
-        "name": "The Matrix",
-        "prompt": "Neon (atompunk world) retro cyberpunk background",
-    },
 ]
 
 styles = {k["name"]: (k["prompt"]) for k in style_list}
 STYLE_NAMES = list(styles.keys())
 
-
 def apply_style(style_name):
     if style_name in styles:
         p = styles.get(style_name, "boho chic")
     return p
 
-
+    
 css = """
 h1 {
     text-align: center;
@@ -334,7 +266,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
             a_prompt = gr.Textbox(
                 label="Additional prompt",
-                value="design-style interior designed (interior space), captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length, tungsten white balance, (sharp focus), professional photography, high-resolution, 8k, Pulitzer Prize-winning",
+                value = "design-style interior designed (interior space), captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length, tungsten white balance, (sharp focus), professional photography, high-resolution, 8k, Pulitzer Prize-winning"
             )
             n_prompt = gr.Textbox(
                 label="Negative prompt",
@@ -343,10 +275,7 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
     #############################################################################
     # input text
     with gr.Row():
-        gr.Text(
-            label="Interior Design Style Examples",
-            value="Eclectic, Maximalist, Bohemian, Scandinavian, Minimalist, Rustic, Modern Farmhouse, Contemporary, Luxury, Airbnb, Boho Chic, Midcentury Modern, Art Deco, Zen, Beach, Neoclassical, Industrial, Biophilic, Eco-friendly, Hollywood Glam, Parisian White, Saudi Prince Gold, French Country, Monster Energy Drink, Cyberpunk, Vaporwave, Baroque, etc.\n\nPro tip: add a color to customize it! You can also describe the furniture type.",
-        )
+        gr.Text(label="Interior Design Style Examples", value="Eclectic, Maximalist, Bohemian, Scandinavian, Minimalist, Rustic, Modern Farmhouse, Contemporary, Luxury, Airbnb, Boho Chic, Midcentury Modern, Art Deco, Zen, Beach, Neoclassical, Industrial, Biophilic, Eco-friendly, Hollywood Glam, Parisian White, Saudi Prince Gold, French Country, Monster Energy Drink, Cyberpunk, Vaporwave, Baroque, etc.\n\nPro tip: add a color to customize it! You can also describe the furniture type.")
     with gr.Column():
         prompt = gr.Textbox(
             label="Custom Prompt",
@@ -376,17 +305,15 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
                 run_button = gr.Button(value="Use this one", size=["lg"], visible=False)
         # output image
         with gr.Column():
-            result = gr.Image(
+            result = gr.Image(  
                 label="Output",
                 interactive=False,
                 format="webp",
-                show_share_button=False,
+                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,
         style_selection,
@@ -400,89 +327,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
-        )
-
+        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,
-        )
-
+    @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(
-        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,
-        )
+    @gr.on(triggers=[use_ai_button.click], inputs=config, outputs=result, show_progress="minimal")
+    def submit(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)
 
     # Change input to result
-    @gr.on(
-        triggers=[use_ai_button.click],
-        inputs=None,
-        outputs=image,
-        show_progress="hidden",
-    )
+    @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")
@@ -491,28 +351,17 @@ with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
         except FileNotFoundError:
             print("No AI Image Available")
             return None
-
+        
     # 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",
-    )
+    @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",
-    )
+    @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=10)
 @torch.inference_mode()
 def process_image(
@@ -527,7 +376,7 @@ def process_image(
     num_steps,
     guidance_scale,
     seed,
-    progress=gr.Progress(track_tqdm=True),
+    progress=gr.Progress(track_tqdm=True)
 ):
     torch.cuda.synchronize()
     preprocess_start = time.time()
@@ -539,7 +388,7 @@ def process_image(
     if not compiled:
         print("Not Compiled")
         compiled = True
-
+    
     seed = random.randint(0, MAX_SEED)
     generator = torch.cuda.manual_seed(seed)
     control_image = preprocessor(
@@ -549,17 +398,10 @@ def process_image(
     )
     preprocess_time = time.time() - preprocess_start
     if style_selection is not None or style_selection != "None":
-        prompt = (
-            "Photo from Pinterest of "
-            + apply_style(style_selection)
-            + " "
-            + prompt
-            + " "
-            + a_prompt
-        )
+        prompt = "Photo from Pinterest of " + apply_style(style_selection) + " " + prompt + " " + a_prompt
     else:
-        prompt = str(get_prompt(prompt, a_prompt))
-    negative_prompt = str(n_prompt)
+        prompt=str(get_prompt(prompt, a_prompt))
+    negative_prompt=str(n_prompt)
     print(prompt)
     start = time.time()
     results = pipe(
@@ -573,22 +415,18 @@ def process_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-------------------------"
-    )
-
+    print(f"\n-------------------------Preprocess done in: {preprocess_time:.2f} seconds-------------------------")    
+    print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------")
+    
     # timestamp = int(time.time())
-    # if not os.path.exists("./outputs"):
+    #if not os.path.exists("./outputs"):
     #    os.makedirs("./outputs")
     # img_path = f"./{timestamp}.jpg"
     # results_path = f"./{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,
@@ -607,9 +445,7 @@ def process_image(
     # )
 
     return results
-
-
 if prod:
     demo.queue(max_size=20).launch(server_name="localhost", server_port=port)
 else:
-    demo.queue(api_open=False).launch(show_api=False)
+    demo.queue(api_open=False).launch(show_api=False)