app.py

prod = False
port = 8080
show_options = False
if prod:
    port = 8081
    # show_options = False

import gc
import os
import random
import time

import gradio as gr
import numpy as np

# import imageio
import torch
from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DPMSolverMultistepScheduler,
    StableDiffusionControlNetPipeline,
)
from diffusers.models.attention_processor import AttnProcessor2_0
from PIL import Image
from preprocess import Preprocessor

MAX_SEED = np.iinfo(np.int32).max
API_KEY = os.environ.get("API_KEY", None)

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",
    },
    {
        "name": "The Matrix",
        "prompt": "Neon (atompunk world) retro cyberpunk background",
    },
]

styles = {k["name"]: (k["prompt"]) for k in style_list}
STYLE_NAMES = list(styles.keys())


def apply_style(style_name):
    if style_name in styles:
        p = styles.get(style_name, "boho chic")
    return p


css = """
h1 {
    text-align: center;
    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)