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)