local

app.py

@@ -1,435 +1,451 @@
-prod = False
-port = 8080
-show_options = False
-if prod:
-    port = 8081
-    # show_options = False
-
-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 (
-    ControlNetModel,
-    DPMSolverMultistepScheduler,
-    StableDiffusionControlNetPipeline,
-)
-from diffusers.models.attention_processor import AttnProcessor2_0
-from preprocess import Preprocessor
-MAX_SEED = np.iinfo(np.int32).max
-API_KEY = os.environ.get("API_KEY", None)
-
-print("CUDA version:", torch.version.cuda)
-print("loading pipe")
-compiled = False
-# api = HfApi()
-
-import spaces
-    
-preprocessor = Preprocessor()
-preprocessor.load("NormalBae")
-
-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")
-    
-    # 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"
-
-    pipe = StableDiffusionControlNetPipeline.from_single_file(
-        base_model_url,
-        # safety_checker=None,
-        # load_safety_checker=True,
-        controlnet=controlnet,
-        scheduler=scheduler,
-        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")
-    
-    # experimental speedup?
-    # pipe.compile()
-    # torch.cuda.empty_cache()
-    # gc.collect()
-    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)
-
-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"]
-    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"
-    default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,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"
-    lab_girl = "hyperrealistic photography, extremely detailed, shy assistant wearing minidress boots and gloves, laboratory background, score_9, 1girl"
-    pet_play = "hyperrealistic photography, extremely detailed, playful, blush, glasses, collar, score_9, HDA_pet_play"
-    bondage = "hyperrealistic photography, extremely detailed, submissive, glasses, score_9, HDA_Bondage"
-    # ahegao = "((invisible clothing)), hyperrealistic photography,exposed vagina,sexy,nsfw,HDA_Ahegao"
-    ahegao2 = "(invisiblebodypaint),rating_newd,HDA_Ahegao"
-    athleisure = "hyperrealistic photography, extremely detailed, 1girl athlete, exhausted embarrassed sweaty,outdoors, ((athleisure clothing)), score_9"
-    atompunk = "((atompunk world)), hyperrealistic photography, extremely detailed, short hair, bodysuit, glasses, neon cyberpunk background, score_9"
-    maid = "hyperrealistic photography, extremely detailed, shy, blushing, score_9, pastel background, HDA_unconventional_maid"
-    nundress = "hyperrealistic photography, extremely detailed, shy, blushing, fantasy background, score_9, HDA_NunDress"
-    naked_hoodie = "hyperrealistic photography, extremely detailed, medium hair, cityscape, (neon lights), score_9, HDA_NakedHoodie"
-    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]
-        prompts_nsfw = [abg, shibari2, ahegao2]
-        prompt = f"{random.choice(girls)}"
-        prompt = f"boho chic"
-        # print(f"-------------{preset}-------------")
-    else:
-        prompt = f"Photo from Pinterest of {prompt} {interior}"
-        # prompt = default2
-    return f"{prompt} f{additional_prompt}"
-
-style_list = [
-    {
-        "name": "None",
-        "prompt": ""
-    },
-    {
-        "name": "Minimalistic",
-        "prompt": "Minimalistic"
-    },
-    {
-        "name": "Boho Chic",
-        "prompt": "boho chic"
-    },
-    {
-        "name": "Saudi Prince Gold",
-        "prompt": "saudi prince gold",
-    },
-    {
-        "name": "Modern Farmhouse",
-        "prompt": "modern farmhouse",
-    },
-    {
-        "name": "Neoclassical",
-        "prompt": "Neoclassical",
-    },
-    {
-        "name": "Eclectic",
-        "prompt": "Eclectic",
-    },
-    {
-        "name": "Parisian White",
-        "prompt": "Parisian White",
-    },
-    {
-        "name": "Hollywood Glam",
-        "prompt": "Hollywood Glam",
-    },
-    {
-        "name": "Scandinavian",
-        "prompt": "Scandinavian",
-    },
-    {
-        "name": "Japanese",
-        "prompt": "Japanese",
-    },
-    {
-        "name": "Texas Cowboy",
-        "prompt": "Texas Cowboy",
-    },
-]
-
-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;
-    display:block;
-}
-h2 {
-    text-align: center;
-    display:block;
-}
-h3 {
-    text-align: center;
-    display:block;
-}
-.gradio-container{max-width: 1200px !important}
-footer {visibility: hidden}
-"""
-with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
-    #############################################################################
-    with gr.Row():
-        with gr.Accordion("Advanced options", open=show_options, visible=show_options):
-            num_images = gr.Slider(
-                label="Images", minimum=1, maximum=4, value=1, step=1
-            )
-            image_resolution = gr.Slider(
-                label="Image resolution",
-                minimum=256,
-                maximum=1024,
-                value=768,
-                step=256,
-            )
-            preprocess_resolution = gr.Slider(
-                label="Preprocess resolution",
-                minimum=128,
-                maximum=1024,
-                value=768,
-                step=1,
-            )
-            num_steps = gr.Slider(
-                label="Number of steps", minimum=1, maximum=100, value=12, step=1
-            )  # 20/4.5 or 12 without lora, 4 with lora
-            guidance_scale = gr.Slider(
-                label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
-            )  # 5 without lora, 2 with lora
-            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
-            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"
-            )
-            n_prompt = gr.Textbox(
-                label="Negative prompt",
-                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.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.")
-    with gr.Column():
-        prompt = gr.Textbox(
-            label="Description",
-            placeholder="boho chic",
-        )
-    with gr.Row(visible=True):
-        style_selection = gr.Radio(
-            show_label=True,
-            container=True,
-            interactive=True,
-            choices=STYLE_NAMES,
-            value="None",
-            label="Design Styles",
-        )
-    # input image
-    with gr.Row():
-        with gr.Column():
-            image = gr.Image(
-                label="Input",
-                sources=["upload"],
-                show_label=True,
-                mirror_webcam=True,
-                format="webp",
-            )
-            # run button
-            with gr.Column():
-                run_button = gr.Button(value="Use this one", size=["lg"], visible=False)
-        # output image
-        with gr.Column():
-            result = gr.Image(  
-                label="Output",
-                interactive=False,
-                format="webp",
-                show_share_button= False,
-            )
-            # Use this image button
-            with gr.Column():
-                use_ai_button = gr.Button(value="Use this one", size=["lg"], visible=False)
-    config = [
-        image,
-        style_selection,
-        prompt,
-        a_prompt,
-        n_prompt,
-        num_images,
-        image_resolution,
-        preprocess_resolution,
-        num_steps,
-        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(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")
-    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
-        
-    # 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():
-        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=10)
-@torch.inference_mode()
-def 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)
-):
-    preprocess_start = time.time()
-    print("processing image")
-    preprocessor.load("NormalBae")
-    # preprocessor.load("Canny") #20 steps, 9 guidance, 512, 512
-
-    global compiled
-    if not compiled:
-        print("Not Compiled")
-        compiled = True
-    
-    seed = random.randint(0, MAX_SEED)
-    generator = torch.cuda.manual_seed(seed)
-    control_image = preprocessor(
-        image=image,
-        image_resolution=image_resolution,
-        detect_resolution=preprocess_resolution,
-    )
-    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
-    else:
-        prompt=str(get_prompt(prompt, a_prompt))
-    negative_prompt=str(n_prompt)
-    print(prompt)
-    start = time.time()
-    results = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_images_per_prompt=num_images,
-        num_inference_steps=num_steps,
-        generator=generator,
-        image=control_image,
-    ).images[0]
-    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-------------------------")
-    
-    # timestamp = int(time.time())
-    #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,
-    #     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,
-    # )
-
-    return results
-if prod:
-    demo.queue(max_size=20).launch(server_name="localhost", server_port=port)
-else:
+prod = False
+port = 8080
+show_options = False
+if prod:
+    port = 8081
+    # show_options = False
+
+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 (
+    ControlNetModel,
+    DPMSolverMultistepScheduler,
+    StableDiffusionControlNetPipeline,
+    AutoencoderKL,
+)
+from diffusers.models.attention_processor import AttnProcessor2_0
+from preprocess import Preprocessor
+MAX_SEED = np.iinfo(np.int32).max
+API_KEY = os.environ.get("API_KEY", None)
+
+print("CUDA version:", torch.version.cuda)
+print("loading pipe")
+compiled = False
+# api = HfApi()
+
+import spaces
+    
+preprocessor = Preprocessor()
+preprocessor.load("NormalBae")
+
+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")
+    
+    # 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"
+    
+    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,
+    )
+
+    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---------------") 
+
+    @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:
+        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"]
+    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"
+    default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,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"
+    lab_girl = "hyperrealistic photography, extremely detailed, shy assistant wearing minidress boots and gloves, laboratory background, score_9, 1girl"
+    pet_play = "hyperrealistic photography, extremely detailed, playful, blush, glasses, collar, score_9, HDA_pet_play"
+    bondage = "hyperrealistic photography, extremely detailed, submissive, glasses, score_9, HDA_Bondage"
+    # ahegao = "((invisible clothing)), hyperrealistic photography,exposed vagina,sexy,nsfw,HDA_Ahegao"
+    ahegao2 = "(invisiblebodypaint),rating_newd,HDA_Ahegao"
+    athleisure = "hyperrealistic photography, extremely detailed, 1girl athlete, exhausted embarrassed sweaty,outdoors, ((athleisure clothing)), score_9"
+    atompunk = "((atompunk world)), hyperrealistic photography, extremely detailed, short hair, bodysuit, glasses, neon cyberpunk background, score_9"
+    maid = "hyperrealistic photography, extremely detailed, shy, blushing, score_9, pastel background, HDA_unconventional_maid"
+    nundress = "hyperrealistic photography, extremely detailed, shy, blushing, fantasy background, score_9, HDA_NunDress"
+    naked_hoodie = "hyperrealistic photography, extremely detailed, medium hair, cityscape, (neon lights), score_9, HDA_NakedHoodie"
+    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]
+        prompts_nsfw = [abg, shibari2, ahegao2]
+        prompt = f"{random.choice(girls)}"
+        prompt = f"boho chic"
+        # print(f"-------------{preset}-------------")
+    else:
+        prompt = f"Photo from Pinterest of {prompt} {interior}"
+        # prompt = default2
+    return f"{prompt} f{additional_prompt}"
+
+style_list = [
+    {
+        "name": "None",
+        "prompt": ""
+    },
+    {
+        "name": "Minimalistic",
+        "prompt": "Minimalistic"
+    },
+    {
+        "name": "Boho Chic",
+        "prompt": "boho chic"
+    },
+    {
+        "name": "Saudi Prince Gold",
+        "prompt": "saudi prince gold",
+    },
+    {
+        "name": "Modern Farmhouse",
+        "prompt": "modern farmhouse",
+    },
+    {
+        "name": "Neoclassical",
+        "prompt": "Neoclassical",
+    },
+    {
+        "name": "Eclectic",
+        "prompt": "Eclectic",
+    },
+    {
+        "name": "Parisian White",
+        "prompt": "Parisian White",
+    },
+    {
+        "name": "Hollywood Glam",
+        "prompt": "Hollywood Glam",
+    },
+    {
+        "name": "Scandinavian",
+        "prompt": "Scandinavian",
+    },
+    {
+        "name": "Japanese",
+        "prompt": "Japanese",
+    },
+    {
+        "name": "Texas Cowboy",
+        "prompt": "Texas Cowboy",
+    },
+    {
+        "name": "Midcentury Modern",
+        "prompt": "Midcentury Modern",
+    },
+    {
+        "name": "Beach",
+        "prompt": "Beach",
+    },
+]
+
+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;
+    display:block;
+}
+h2 {
+    text-align: center;
+    display:block;
+}
+h3 {
+    text-align: center;
+    display:block;
+}
+.gradio-container{max-width: 1200px !important}
+footer {visibility: hidden}
+"""
+with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
+    #############################################################################
+    with gr.Row():
+        with gr.Accordion("Advanced options", open=show_options, visible=show_options):
+            num_images = gr.Slider(
+                label="Images", minimum=1, maximum=4, value=1, step=1
+            )
+            image_resolution = gr.Slider(
+                label="Image resolution",
+                minimum=256,
+                maximum=1024,
+                value=512,
+                step=256,
+            )
+            preprocess_resolution = gr.Slider(
+                label="Preprocess resolution",
+                minimum=128,
+                maximum=1024,
+                value=512,
+                step=1,
+            )
+            num_steps = gr.Slider(
+                label="Number of steps", minimum=1, maximum=100, value=15, step=1
+            )  # 20/4.5 or 12 without lora, 4 with lora
+            guidance_scale = gr.Slider(
+                label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
+            )  # 5 without lora, 2 with lora
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            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"
+            )
+            n_prompt = gr.Textbox(
+                label="Negative prompt",
+                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.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.")
+    with gr.Column():
+        prompt = gr.Textbox(
+            label="Custom Prompt",
+            placeholder="boho chic",
+        )
+    with gr.Row(visible=True):
+        style_selection = gr.Radio(
+            show_label=True,
+            container=True,
+            interactive=True,
+            choices=STYLE_NAMES,
+            value="None",
+            label="Design Styles",
+        )
+    # input image
+    with gr.Row():
+        with gr.Column():
+            image = gr.Image(
+                label="Input",
+                sources=["upload"],
+                show_label=True,
+                mirror_webcam=True,
+                format="webp",
+            )
+            # run button
+            with gr.Column():
+                run_button = gr.Button(value="Use this one", size=["lg"], visible=False)
+        # output image
+        with gr.Column():
+            result = gr.Image(  
+                label="Output",
+                interactive=False,
+                format="webp",
+                show_share_button= False,
+            )
+            # Use this image button
+            with gr.Column():
+                use_ai_button = gr.Button(value="Use this one", size=["lg"], visible=False)
+    config = [
+        image,
+        style_selection,
+        prompt,
+        a_prompt,
+        n_prompt,
+        num_images,
+        image_resolution,
+        preprocess_resolution,
+        num_steps,
+        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(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")
+    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
+        
+    # 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():
+        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=10)
+@torch.inference_mode()
+def 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)
+):
+    torch.cuda.synchronize()
+    preprocess_start = time.time()
+    print("processing image")
+    preprocessor.load("NormalBae")
+    # preprocessor.load("Canny") #20 steps, 9 guidance, 512, 512
+
+    global compiled
+    if not compiled:
+        print("Not Compiled")
+        compiled = True
+    
+    seed = random.randint(0, MAX_SEED)
+    generator = torch.cuda.manual_seed(seed)
+    control_image = preprocessor(
+        image=image,
+        image_resolution=image_resolution,
+        detect_resolution=preprocess_resolution,
+    )
+    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
+    else:
+        prompt=str(get_prompt(prompt, a_prompt))
+    negative_prompt=str(n_prompt)
+    print(prompt)
+    start = time.time()
+    results = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=num_images,
+        num_inference_steps=num_steps,
+        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-------------------------")
+    
+    # timestamp = int(time.time())
+    #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,
+    #     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,
+    # )
+
+    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)

app.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e24bf5be8b15309a5f50ad7c94d94dfb5fab0bff4f0baea1f1a67af5cc3f925
+size 13317

app/app.py

@@ -0,0 +1,451 @@
+prod = False
+port = 8080
+show_options = False
+if prod:
+    port = 8081
+    # show_options = False
+
+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 (
+    ControlNetModel,
+    DPMSolverMultistepScheduler,
+    StableDiffusionControlNetPipeline,
+    AutoencoderKL,
+)
+from diffusers.models.attention_processor import AttnProcessor2_0
+from preprocess import Preprocessor
+MAX_SEED = np.iinfo(np.int32).max
+API_KEY = os.environ.get("API_KEY", None)
+
+print("CUDA version:", torch.version.cuda)
+print("loading pipe")
+compiled = False
+# api = HfApi()
+
+import spaces
+    
+preprocessor = Preprocessor()
+preprocessor.load("NormalBae")
+
+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")
+    
+    # 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"
+    
+    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,
+    )
+
+    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---------------") 
+
+    @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:
+        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"]
+    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"
+    default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,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"
+    lab_girl = "hyperrealistic photography, extremely detailed, shy assistant wearing minidress boots and gloves, laboratory background, score_9, 1girl"
+    pet_play = "hyperrealistic photography, extremely detailed, playful, blush, glasses, collar, score_9, HDA_pet_play"
+    bondage = "hyperrealistic photography, extremely detailed, submissive, glasses, score_9, HDA_Bondage"
+    # ahegao = "((invisible clothing)), hyperrealistic photography,exposed vagina,sexy,nsfw,HDA_Ahegao"
+    ahegao2 = "(invisiblebodypaint),rating_newd,HDA_Ahegao"
+    athleisure = "hyperrealistic photography, extremely detailed, 1girl athlete, exhausted embarrassed sweaty,outdoors, ((athleisure clothing)), score_9"
+    atompunk = "((atompunk world)), hyperrealistic photography, extremely detailed, short hair, bodysuit, glasses, neon cyberpunk background, score_9"
+    maid = "hyperrealistic photography, extremely detailed, shy, blushing, score_9, pastel background, HDA_unconventional_maid"
+    nundress = "hyperrealistic photography, extremely detailed, shy, blushing, fantasy background, score_9, HDA_NunDress"
+    naked_hoodie = "hyperrealistic photography, extremely detailed, medium hair, cityscape, (neon lights), score_9, HDA_NakedHoodie"
+    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]
+        prompts_nsfw = [abg, shibari2, ahegao2]
+        prompt = f"{random.choice(girls)}"
+        prompt = f"boho chic"
+        # print(f"-------------{preset}-------------")
+    else:
+        prompt = f"Photo from Pinterest of {prompt} {interior}"
+        # prompt = default2
+    return f"{prompt} f{additional_prompt}"
+
+style_list = [
+    {
+        "name": "None",
+        "prompt": ""
+    },
+    {
+        "name": "Minimalistic",
+        "prompt": "Minimalistic"
+    },
+    {
+        "name": "Boho Chic",
+        "prompt": "boho chic"
+    },
+    {
+        "name": "Saudi Prince Gold",
+        "prompt": "saudi prince gold",
+    },
+    {
+        "name": "Modern Farmhouse",
+        "prompt": "modern farmhouse",
+    },
+    {
+        "name": "Neoclassical",
+        "prompt": "Neoclassical",
+    },
+    {
+        "name": "Eclectic",
+        "prompt": "Eclectic",
+    },
+    {
+        "name": "Parisian White",
+        "prompt": "Parisian White",
+    },
+    {
+        "name": "Hollywood Glam",
+        "prompt": "Hollywood Glam",
+    },
+    {
+        "name": "Scandinavian",
+        "prompt": "Scandinavian",
+    },
+    {
+        "name": "Japanese",
+        "prompt": "Japanese",
+    },
+    {
+        "name": "Texas Cowboy",
+        "prompt": "Texas Cowboy",
+    },
+    {
+        "name": "Midcentury Modern",
+        "prompt": "Midcentury Modern",
+    },
+    {
+        "name": "Beach",
+        "prompt": "Beach",
+    },
+]
+
+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;
+    display:block;
+}
+h2 {
+    text-align: center;
+    display:block;
+}
+h3 {
+    text-align: center;
+    display:block;
+}
+.gradio-container{max-width: 1200px !important}
+footer {visibility: hidden}
+"""
+with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
+    #############################################################################
+    with gr.Row():
+        with gr.Accordion("Advanced options", open=show_options, visible=show_options):
+            num_images = gr.Slider(
+                label="Images", minimum=1, maximum=4, value=1, step=1
+            )
+            image_resolution = gr.Slider(
+                label="Image resolution",
+                minimum=256,
+                maximum=1024,
+                value=512,
+                step=256,
+            )
+            preprocess_resolution = gr.Slider(
+                label="Preprocess resolution",
+                minimum=128,
+                maximum=1024,
+                value=512,
+                step=1,
+            )
+            num_steps = gr.Slider(
+                label="Number of steps", minimum=1, maximum=100, value=15, step=1
+            )  # 20/4.5 or 12 without lora, 4 with lora
+            guidance_scale = gr.Slider(
+                label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
+            )  # 5 without lora, 2 with lora
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            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"
+            )
+            n_prompt = gr.Textbox(
+                label="Negative prompt",
+                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.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.")
+    with gr.Column():
+        prompt = gr.Textbox(
+            label="Custom Prompt",
+            placeholder="boho chic",
+        )
+    with gr.Row(visible=True):
+        style_selection = gr.Radio(
+            show_label=True,
+            container=True,
+            interactive=True,
+            choices=STYLE_NAMES,
+            value="None",
+            label="Design Styles",
+        )
+    # input image
+    with gr.Row():
+        with gr.Column():
+            image = gr.Image(
+                label="Input",
+                sources=["upload"],
+                show_label=True,
+                mirror_webcam=True,
+                format="webp",
+            )
+            # run button
+            with gr.Column():
+                run_button = gr.Button(value="Use this one", size=["lg"], visible=False)
+        # output image
+        with gr.Column():
+            result = gr.Image(  
+                label="Output",
+                interactive=False,
+                format="webp",
+                show_share_button= False,
+            )
+            # Use this image button
+            with gr.Column():
+                use_ai_button = gr.Button(value="Use this one", size=["lg"], visible=False)
+    config = [
+        image,
+        style_selection,
+        prompt,
+        a_prompt,
+        n_prompt,
+        num_images,
+        image_resolution,
+        preprocess_resolution,
+        num_steps,
+        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(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")
+    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
+        
+    # 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():
+        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=10)
+@torch.inference_mode()
+def 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)
+):
+    torch.cuda.synchronize()
+    preprocess_start = time.time()
+    print("processing image")
+    preprocessor.load("NormalBae")
+    # preprocessor.load("Canny") #20 steps, 9 guidance, 512, 512
+
+    global compiled
+    if not compiled:
+        print("Not Compiled")
+        compiled = True
+    
+    seed = random.randint(0, MAX_SEED)
+    generator = torch.cuda.manual_seed(seed)
+    control_image = preprocessor(
+        image=image,
+        image_resolution=image_resolution,
+        detect_resolution=preprocess_resolution,
+    )
+    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
+    else:
+        prompt=str(get_prompt(prompt, a_prompt))
+    negative_prompt=str(n_prompt)
+    print(prompt)
+    start = time.time()
+    results = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=num_images,
+        num_inference_steps=num_steps,
+        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-------------------------")
+    
+    # timestamp = int(time.time())
+    #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,
+    #     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,
+    # )
+
+    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)

app/local_app.py

@@ -0,0 +1,455 @@
+prod = True
+port = 8080
+show_options = False
+if prod:
+    port = 8081
+    # show_options = False
+
+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 (
+    ControlNetModel,
+    DPMSolverMultistepScheduler,
+    StableDiffusionControlNetPipeline,
+    AutoencoderKL,
+)
+from diffusers.models.attention_processor import AttnProcessor2_0
+from local_preprocess import Preprocessor
+MAX_SEED = np.iinfo(np.int32).max
+API_KEY = os.environ.get("API_KEY", None)
+
+print("CUDA version:", torch.version.cuda)
+print("loading pipe")
+compiled = False
+    
+preprocessor = Preprocessor()
+preprocessor.load("NormalBae")
+
+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")
+    
+    # Scheduler
+    scheduler = DPMSolverMultistepScheduler.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0", 
+        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",
+        attn_implementation="flash_attention_2",   
+    )
+
+    # 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"
+    base_model_id = "Lykon/absolute-reality-1.81"
+    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) 
+    
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        base_model_id,
+        safety_checker=None,
+        controlnet=controlnet,
+        scheduler=scheduler,
+        vae=vae,
+        torch_dtype=torch.float16,
+    ).to("cuda")
+    
+    # pipe = StableDiffusionControlNetPipeline.from_single_file(
+    #     base_model_url,
+    #     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")
+    
+    # experimental speedup?
+    # pipe.compile()
+    # torch.cuda.empty_cache()
+    # gc.collect()
+    print("---------------Loaded controlnet pipeline---------------") 
+
+    # @spaces.GPU(duration=12)
+    # pipe.enable_xformers_memory_efficient_attention()
+    # pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+    # pipe.unet.set_attn_processor(AttnProcessor2_0())
+    torch.cuda.empty_cache()
+    gc.collect()
+    print("Model Compiled!")
+    
+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"]
+    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"
+    default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,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"
+    lab_girl = "hyperrealistic photography, extremely detailed, shy assistant wearing minidress boots and gloves, laboratory background, score_9, 1girl"
+    pet_play = "hyperrealistic photography, extremely detailed, playful, blush, glasses, collar, score_9, HDA_pet_play"
+    bondage = "hyperrealistic photography, extremely detailed, submissive, glasses, score_9, HDA_Bondage"
+    # ahegao = "((invisible clothing)), hyperrealistic photography,exposed vagina,sexy,nsfw,HDA_Ahegao"
+    ahegao2 = "(invisiblebodypaint),rating_newd,HDA_Ahegao"
+    athleisure = "hyperrealistic photography, extremely detailed, 1girl athlete, exhausted embarrassed sweaty,outdoors, ((athleisure clothing)), score_9"
+    atompunk = "((atompunk world)), hyperrealistic photography, extremely detailed, short hair, bodysuit, glasses, neon cyberpunk background, score_9"
+    maid = "hyperrealistic photography, extremely detailed, shy, blushing, score_9, pastel background, HDA_unconventional_maid"
+    nundress = "hyperrealistic photography, extremely detailed, shy, blushing, fantasy background, score_9, HDA_NunDress"
+    naked_hoodie = "hyperrealistic photography, extremely detailed, medium hair, cityscape, (neon lights), score_9, HDA_NakedHoodie"
+    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]
+        prompts_nsfw = [abg, shibari2, ahegao2]
+        prompt = f"{random.choice(girls)}"
+        prompt = f"boho chic"
+        # print(f"-------------{preset}-------------")
+    else:
+        prompt = f"Photo from Pinterest of {prompt} {interior}"
+        # prompt = default2
+    return f"{prompt} f{additional_prompt}"
+
+style_list = [
+    {
+        "name": "None",
+        "prompt": ""
+    },
+    {
+        "name": "Minimalistic",
+        "prompt": "Minimalistic"
+    },
+    {
+        "name": "Boho Chic",
+        "prompt": "boho chic"
+    },
+    {
+        "name": "Saudi Prince Gold",
+        "prompt": "saudi prince gold",
+    },
+    {
+        "name": "Modern Farmhouse",
+        "prompt": "modern farmhouse",
+    },
+    {
+        "name": "Neoclassical",
+        "prompt": "Neoclassical",
+    },
+    {
+        "name": "Eclectic",
+        "prompt": "Eclectic",
+    },
+    {
+        "name": "Parisian White",
+        "prompt": "Parisian White",
+    },
+    {
+        "name": "Hollywood Glam",
+        "prompt": "Hollywood Glam",
+    },
+    {
+        "name": "Scandinavian",
+        "prompt": "Scandinavian",
+    },
+    {
+        "name": "Japanese",
+        "prompt": "Japanese",
+    },
+    {
+        "name": "Texas Cowboy",
+        "prompt": "Texas Cowboy",
+    },
+    {
+        "name": "Midcentury Modern",
+        "prompt": "Midcentury Modern",
+    },
+    {
+        "name": "Beach",
+        "prompt": "Beach",
+    },
+]
+
+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;
+    display:block;
+}
+h2 {
+    text-align: center;
+    display:block;
+}
+h3 {
+    text-align: center;
+    display:block;
+}
+.gradio-container{max-width: 1200px !important}
+footer {visibility: hidden}
+"""
+with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
+    #############################################################################
+    with gr.Row():
+        with gr.Accordion("Advanced options", open=show_options, visible=show_options):
+            num_images = gr.Slider(
+                label="Images", minimum=1, maximum=4, value=1, step=1
+            )
+            image_resolution = gr.Slider(
+                label="Image resolution",
+                minimum=256,
+                maximum=1024,
+                value=512,
+                step=256,
+            )
+            preprocess_resolution = gr.Slider(
+                label="Preprocess resolution",
+                minimum=128,
+                maximum=1024,
+                value=512,
+                step=1,
+            )
+            num_steps = gr.Slider(
+                label="Number of steps", minimum=1, maximum=100, value=15, step=1
+            )  # 20/4.5 or 12 without lora, 4 with lora
+            guidance_scale = gr.Slider(
+                label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
+            )  # 5 without lora, 2 with lora
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            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"
+            )
+            n_prompt = gr.Textbox(
+                label="Negative prompt",
+                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.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.")
+    with gr.Column():
+        prompt = gr.Textbox(
+            label="Custom Prompt",
+            placeholder="boho chic",
+        )
+    with gr.Row(visible=True):
+        style_selection = gr.Radio(
+            show_label=True,
+            container=True,
+            interactive=True,
+            choices=STYLE_NAMES,
+            value="None",
+            label="Design Styles",
+        )
+    # input image
+    with gr.Row():
+        with gr.Column():
+            image = gr.Image(
+                label="Input",
+                sources=["upload"],
+                show_label=True,
+                mirror_webcam=True,
+                format="webp",
+            )
+            # run button
+            with gr.Column():
+                run_button = gr.Button(value="Use this one", size=["lg"], visible=False)
+        # output image
+        with gr.Column():
+            result = gr.Image(  
+                label="Output",
+                interactive=False,
+                format="webp",
+                show_share_button= False,
+            )
+            # Use this image button
+            with gr.Column():
+                use_ai_button = gr.Button(value="Use this one", size=["lg"], visible=False)
+    config = [
+        image,
+        style_selection,
+        prompt,
+        a_prompt,
+        n_prompt,
+        num_images,
+        image_resolution,
+        preprocess_resolution,
+        num_steps,
+        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(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")
+    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
+        
+    # 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():
+        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,
+    style_selection,
+    prompt,
+    a_prompt,
+    n_prompt,
+    num_images,
+    image_resolution,
+    preprocess_resolution,
+    num_steps,
+    guidance_scale,
+    seed,
+    progress=gr.Progress(track_tqdm=True)
+):
+    torch.cuda.synchronize()
+    preprocess_start = time.time()
+    print("processing image")
+    preprocessor.load("NormalBae")
+    # preprocessor.load("Canny") #20 steps, 9 guidance, 512, 512
+    
+    global compiled
+    if not compiled:
+        print("Not Compiled")
+        compiled = True
+        
+    seed = random.randint(0, MAX_SEED)
+    generator = torch.cuda.manual_seed(seed)
+    control_image = preprocessor(
+        image=image,
+        image_resolution=image_resolution,
+        detect_resolution=preprocess_resolution,
+    )
+    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
+    else:
+        prompt=str(get_prompt(prompt, a_prompt))
+    negative_prompt=str(n_prompt)
+    print(prompt)
+    start = time.time()
+    results = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=num_images,
+        num_inference_steps=num_steps,
+        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-------------------------")
+    
+    # timestamp = int(time.time())
+    #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,
+    #     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,
+    # )
+
+    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)

app/local_preprocess.py

@@ -0,0 +1,69 @@
+# import numpy as np
+import PIL.Image
+import torch
+import gc
+# from controlnet_aux_local import NormalBaeDetector#, CannyDetector
+from controlnet_aux import NormalBaeDetector
+
+# from controlnet_aux.util import HWC3
+# import cv2
+# from cv_utils import resize_image
+
+class Preprocessor:
+    MODEL_ID = "lllyasviel/Annotators"
+    
+    # def resize_image(input_image, resolution, interpolation=None):
+    #     H, W, C = input_image.shape
+    #     H = float(H)
+    #     W = float(W)
+    #     k = float(resolution) / max(H, W)
+    #     H *= k
+    #     W *= k
+    #     H = int(np.round(H / 64.0)) * 64
+    #     W = int(np.round(W / 64.0)) * 64
+    #     if interpolation is None:
+    #         interpolation = cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA
+    #     img = cv2.resize(input_image, (W, H), interpolation=interpolation)
+    #     return img
+
+
+    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")
+        # elif name == "Canny":
+        #     self.model = CannyDetector()
+        else:
+            raise ValueError
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        self.name = name
+
+    def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
+        # if self.name == "Canny":
+        #     if "detect_resolution" in kwargs:
+        #         detect_resolution = kwargs.pop("detect_resolution")
+        #         image = np.array(image)
+        #         image = HWC3(image)
+        #         image = resize_image(image, resolution=detect_resolution)
+        #     image = self.model(image, **kwargs)
+        #     return PIL.Image.fromarray(image)
+        # elif self.name == "Midas":
+        #     detect_resolution = kwargs.pop("detect_resolution", 512)
+        #     image_resolution = kwargs.pop("image_resolution", 512)
+        #     image = np.array(image)
+        #     image = HWC3(image)
+        #     image = resize_image(image, resolution=detect_resolution)
+        #     image = self.model(image, **kwargs)
+        #     image = HWC3(image)
+        #     image = resize_image(image, resolution=image_resolution)
+        #     return PIL.Image.fromarray(image)
+        # else:
+        return self.model(image, **kwargs)

app/preprocess.py

@@ -0,0 +1,67 @@
+# import numpy as np
+import PIL.Image
+# import torch
+from controlnet_aux import NormalBaeDetector#, CannyDetector
+
+# from controlnet_aux.util import HWC3
+# import cv2
+# from cv_utils import resize_image
+
+class Preprocessor:
+    MODEL_ID = "lllyasviel/Annotators"
+    
+    # def resize_image(input_image, resolution, interpolation=None):
+    #     H, W, C = input_image.shape
+    #     H = float(H)
+    #     W = float(W)
+    #     k = float(resolution) / max(H, W)
+    #     H *= k
+    #     W *= k
+    #     H = int(np.round(H / 64.0)) * 64
+    #     W = int(np.round(W / 64.0)) * 64
+    #     if interpolation is None:
+    #         interpolation = cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA
+    #     img = cv2.resize(input_image, (W, H), interpolation=interpolation)
+    #     return img
+
+
+    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")
+        # elif name == "Canny":
+        #     self.model = CannyDetector()
+        else:
+            raise ValueError
+        # torch.cuda.empty_cache()
+        # gc.collect()
+        
+        self.name = name
+
+    def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
+        # if self.name == "Canny":
+        #     if "detect_resolution" in kwargs:
+        #         detect_resolution = kwargs.pop("detect_resolution")
+        #         image = np.array(image)
+        #         image = HWC3(image)
+        #         image = resize_image(image, resolution=detect_resolution)
+        #     image = self.model(image, **kwargs)
+        #     return PIL.Image.fromarray(image)
+        # elif self.name == "Midas":
+        #     detect_resolution = kwargs.pop("detect_resolution", 512)
+        #     image_resolution = kwargs.pop("image_resolution", 512)
+        #     image = np.array(image)
+        #     image = HWC3(image)
+        #     image = resize_image(image, resolution=detect_resolution)
+        #     image = self.model(image, **kwargs)
+        #     image = HWC3(image)
+        #     image = resize_image(image, resolution=image_resolution)
+        #     return PIL.Image.fromarray(image)
+        # else:
+        return self.model(image, **kwargs)

app/requirements.txt

@@ -0,0 +1,12 @@
+torch
+torchvision
+diffusers
+einops
+huggingface-hub
+mediapipe
+opencv-python-headless
+safetensors
+transformers
+xformers
+accelerate
+imageio

app/win.requirements.txt

@@ -0,0 +1,17 @@
+torch
+torchvision
+torchaudio
+--index-url https://download.pytorch.org/whl/cu121
+
+diffusers
+einops
+gradio
+gradio-client
+mediapipe
+opencv-python-headless
+safetensors
+transformers
+xformers
+accelerate
+imageio
+controlnet_aux

local_app.py

@@ -0,0 +1,455 @@
+prod = True
+port = 8080
+show_options = False
+if prod:
+    port = 8081
+    # show_options = False
+
+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 (
+    ControlNetModel,
+    DPMSolverMultistepScheduler,
+    StableDiffusionControlNetPipeline,
+    AutoencoderKL,
+)
+from diffusers.models.attention_processor import AttnProcessor2_0
+from local_preprocess import Preprocessor
+MAX_SEED = np.iinfo(np.int32).max
+API_KEY = os.environ.get("API_KEY", None)
+
+print("CUDA version:", torch.version.cuda)
+print("loading pipe")
+compiled = False
+    
+preprocessor = Preprocessor()
+preprocessor.load("NormalBae")
+
+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")
+    
+    # Scheduler
+    scheduler = DPMSolverMultistepScheduler.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0", 
+        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",
+        attn_implementation="flash_attention_2",   
+    )
+
+    # 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"
+    base_model_id = "Lykon/absolute-reality-1.81"
+    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) 
+    
+    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+        base_model_id,
+        safety_checker=None,
+        controlnet=controlnet,
+        scheduler=scheduler,
+        vae=vae,
+        torch_dtype=torch.float16,
+    ).to("cuda")
+    
+    # pipe = StableDiffusionControlNetPipeline.from_single_file(
+    #     base_model_url,
+    #     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")
+    
+    # experimental speedup?
+    # pipe.compile()
+    # torch.cuda.empty_cache()
+    # gc.collect()
+    print("---------------Loaded controlnet pipeline---------------") 
+
+    # @spaces.GPU(duration=12)
+    # pipe.enable_xformers_memory_efficient_attention()
+    # pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+    # pipe.unet.set_attn_processor(AttnProcessor2_0())
+    torch.cuda.empty_cache()
+    gc.collect()
+    print("Model Compiled!")
+    
+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"]
+    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"
+    default2 = f"professional 3d model {prompt},octane render,highly detailed,volumetric,dramatic lighting,hyperrealistic photography,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"
+    lab_girl = "hyperrealistic photography, extremely detailed, shy assistant wearing minidress boots and gloves, laboratory background, score_9, 1girl"
+    pet_play = "hyperrealistic photography, extremely detailed, playful, blush, glasses, collar, score_9, HDA_pet_play"
+    bondage = "hyperrealistic photography, extremely detailed, submissive, glasses, score_9, HDA_Bondage"
+    # ahegao = "((invisible clothing)), hyperrealistic photography,exposed vagina,sexy,nsfw,HDA_Ahegao"
+    ahegao2 = "(invisiblebodypaint),rating_newd,HDA_Ahegao"
+    athleisure = "hyperrealistic photography, extremely detailed, 1girl athlete, exhausted embarrassed sweaty,outdoors, ((athleisure clothing)), score_9"
+    atompunk = "((atompunk world)), hyperrealistic photography, extremely detailed, short hair, bodysuit, glasses, neon cyberpunk background, score_9"
+    maid = "hyperrealistic photography, extremely detailed, shy, blushing, score_9, pastel background, HDA_unconventional_maid"
+    nundress = "hyperrealistic photography, extremely detailed, shy, blushing, fantasy background, score_9, HDA_NunDress"
+    naked_hoodie = "hyperrealistic photography, extremely detailed, medium hair, cityscape, (neon lights), score_9, HDA_NakedHoodie"
+    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]
+        prompts_nsfw = [abg, shibari2, ahegao2]
+        prompt = f"{random.choice(girls)}"
+        prompt = f"boho chic"
+        # print(f"-------------{preset}-------------")
+    else:
+        prompt = f"Photo from Pinterest of {prompt} {interior}"
+        # prompt = default2
+    return f"{prompt} f{additional_prompt}"
+
+style_list = [
+    {
+        "name": "None",
+        "prompt": ""
+    },
+    {
+        "name": "Minimalistic",
+        "prompt": "Minimalistic"
+    },
+    {
+        "name": "Boho Chic",
+        "prompt": "boho chic"
+    },
+    {
+        "name": "Saudi Prince Gold",
+        "prompt": "saudi prince gold",
+    },
+    {
+        "name": "Modern Farmhouse",
+        "prompt": "modern farmhouse",
+    },
+    {
+        "name": "Neoclassical",
+        "prompt": "Neoclassical",
+    },
+    {
+        "name": "Eclectic",
+        "prompt": "Eclectic",
+    },
+    {
+        "name": "Parisian White",
+        "prompt": "Parisian White",
+    },
+    {
+        "name": "Hollywood Glam",
+        "prompt": "Hollywood Glam",
+    },
+    {
+        "name": "Scandinavian",
+        "prompt": "Scandinavian",
+    },
+    {
+        "name": "Japanese",
+        "prompt": "Japanese",
+    },
+    {
+        "name": "Texas Cowboy",
+        "prompt": "Texas Cowboy",
+    },
+    {
+        "name": "Midcentury Modern",
+        "prompt": "Midcentury Modern",
+    },
+    {
+        "name": "Beach",
+        "prompt": "Beach",
+    },
+]
+
+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;
+    display:block;
+}
+h2 {
+    text-align: center;
+    display:block;
+}
+h3 {
+    text-align: center;
+    display:block;
+}
+.gradio-container{max-width: 1200px !important}
+footer {visibility: hidden}
+"""
+with gr.Blocks(theme="bethecloud/storj_theme", css=css) as demo:
+    #############################################################################
+    with gr.Row():
+        with gr.Accordion("Advanced options", open=show_options, visible=show_options):
+            num_images = gr.Slider(
+                label="Images", minimum=1, maximum=4, value=1, step=1
+            )
+            image_resolution = gr.Slider(
+                label="Image resolution",
+                minimum=256,
+                maximum=1024,
+                value=512,
+                step=256,
+            )
+            preprocess_resolution = gr.Slider(
+                label="Preprocess resolution",
+                minimum=128,
+                maximum=1024,
+                value=512,
+                step=1,
+            )
+            num_steps = gr.Slider(
+                label="Number of steps", minimum=1, maximum=100, value=15, step=1
+            )  # 20/4.5 or 12 without lora, 4 with lora
+            guidance_scale = gr.Slider(
+                label="Guidance scale", minimum=0.1, maximum=30.0, value=5.5, step=0.1
+            )  # 5 without lora, 2 with lora
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            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"
+            )
+            n_prompt = gr.Textbox(
+                label="Negative prompt",
+                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.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.")
+    with gr.Column():
+        prompt = gr.Textbox(
+            label="Custom Prompt",
+            placeholder="boho chic",
+        )
+    with gr.Row(visible=True):
+        style_selection = gr.Radio(
+            show_label=True,
+            container=True,
+            interactive=True,
+            choices=STYLE_NAMES,
+            value="None",
+            label="Design Styles",
+        )
+    # input image
+    with gr.Row():
+        with gr.Column():
+            image = gr.Image(
+                label="Input",
+                sources=["upload"],
+                show_label=True,
+                mirror_webcam=True,
+                format="webp",
+            )
+            # run button
+            with gr.Column():
+                run_button = gr.Button(value="Use this one", size=["lg"], visible=False)
+        # output image
+        with gr.Column():
+            result = gr.Image(  
+                label="Output",
+                interactive=False,
+                format="webp",
+                show_share_button= False,
+            )
+            # Use this image button
+            with gr.Column():
+                use_ai_button = gr.Button(value="Use this one", size=["lg"], visible=False)
+    config = [
+        image,
+        style_selection,
+        prompt,
+        a_prompt,
+        n_prompt,
+        num_images,
+        image_resolution,
+        preprocess_resolution,
+        num_steps,
+        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(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")
+    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
+        
+    # 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():
+        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,
+    style_selection,
+    prompt,
+    a_prompt,
+    n_prompt,
+    num_images,
+    image_resolution,
+    preprocess_resolution,
+    num_steps,
+    guidance_scale,
+    seed,
+    progress=gr.Progress(track_tqdm=True)
+):
+    torch.cuda.synchronize()
+    preprocess_start = time.time()
+    print("processing image")
+    preprocessor.load("NormalBae")
+    # preprocessor.load("Canny") #20 steps, 9 guidance, 512, 512
+    
+    global compiled
+    if not compiled:
+        print("Not Compiled")
+        compiled = True
+        
+    seed = random.randint(0, MAX_SEED)
+    generator = torch.cuda.manual_seed(seed)
+    control_image = preprocessor(
+        image=image,
+        image_resolution=image_resolution,
+        detect_resolution=preprocess_resolution,
+    )
+    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
+    else:
+        prompt=str(get_prompt(prompt, a_prompt))
+    negative_prompt=str(n_prompt)
+    print(prompt)
+    start = time.time()
+    results = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_images_per_prompt=num_images,
+        num_inference_steps=num_steps,
+        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-------------------------")
+    
+    # timestamp = int(time.time())
+    #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,
+    #     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,
+    # )
+
+    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)

local_preprocess.py

@@ -0,0 +1,69 @@
+# import numpy as np
+import PIL.Image
+import torch
+import gc
+# from controlnet_aux_local import NormalBaeDetector#, CannyDetector
+from controlnet_aux import NormalBaeDetector
+
+# from controlnet_aux.util import HWC3
+# import cv2
+# from cv_utils import resize_image
+
+class Preprocessor:
+    MODEL_ID = "lllyasviel/Annotators"
+    
+    # def resize_image(input_image, resolution, interpolation=None):
+    #     H, W, C = input_image.shape
+    #     H = float(H)
+    #     W = float(W)
+    #     k = float(resolution) / max(H, W)
+    #     H *= k
+    #     W *= k
+    #     H = int(np.round(H / 64.0)) * 64
+    #     W = int(np.round(W / 64.0)) * 64
+    #     if interpolation is None:
+    #         interpolation = cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA
+    #     img = cv2.resize(input_image, (W, H), interpolation=interpolation)
+    #     return img
+
+
+    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")
+        # elif name == "Canny":
+        #     self.model = CannyDetector()
+        else:
+            raise ValueError
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        self.name = name
+
+    def __call__(self, image: PIL.Image.Image, **kwargs) -> PIL.Image.Image:
+        # if self.name == "Canny":
+        #     if "detect_resolution" in kwargs:
+        #         detect_resolution = kwargs.pop("detect_resolution")
+        #         image = np.array(image)
+        #         image = HWC3(image)
+        #         image = resize_image(image, resolution=detect_resolution)
+        #     image = self.model(image, **kwargs)
+        #     return PIL.Image.fromarray(image)
+        # elif self.name == "Midas":
+        #     detect_resolution = kwargs.pop("detect_resolution", 512)
+        #     image_resolution = kwargs.pop("image_resolution", 512)
+        #     image = np.array(image)
+        #     image = HWC3(image)
+        #     image = resize_image(image, resolution=detect_resolution)
+        #     image = self.model(image, **kwargs)
+        #     image = HWC3(image)
+        #     image = resize_image(image, resolution=image_resolution)
+        #     return PIL.Image.fromarray(image)
+        # else:
+        return self.model(image, **kwargs)

preprocess.py

@@ -64,4 +64,4 @@ class Preprocessor:
         #     image = resize_image(image, resolution=image_resolution)
         #     return PIL.Image.fromarray(image)
         # else:
-        return self.model(image, **kwargs)
+        return self.model(image, **kwargs)

requirements.txt

@@ -1,12 +1,12 @@
-torch
-torchvision
-diffusers
-einops
-huggingface-hub
-mediapipe
-opencv-python-headless
-safetensors
-transformers
-xformers
-accelerate
+torch
+torchvision
+diffusers
+einops
+huggingface-hub
+mediapipe
+opencv-python-headless
+safetensors
+transformers
+xformers
+accelerate
 imageio

win.requirements.txt

@@ -0,0 +1,17 @@
+torch
+torchvision
+torchaudio
+--index-url https://download.pytorch.org/whl/cu121
+
+diffusers
+einops
+gradio
+gradio-client
+mediapipe
+opencv-python-headless
+safetensors
+transformers
+xformers
+accelerate
+imageio
+controlnet_aux