prod = False port = 8080 show_options = False if prod: port = 8081 # show_options = False import os import random import time import gradio as gr import numpy as np import spaces import imageio from huggingface_hub import HfApi import gc import torch import cv2 from PIL import Image from diffusers import ( ControlNetModel, DPMSolverMultistepScheduler, StableDiffusionControlNetPipeline, # StableDiffusionInpaintPipeline, # AutoencoderKL, ) from controlnet_aux_local import NormalBaeDetector MAX_SEED = np.iinfo(np.int32).max API_KEY = os.environ.get("API_KEY", None) # os.environ['HF_HOME'] = '/data/.huggingface' print("CUDA version:", torch.version.cuda) print("loading everything") compiled = False api = HfApi() class Preprocessor: MODEL_ID = "lllyasviel/Annotators" 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") torch.cuda.empty_cache() self.name = name else: raise ValueError return def __call__(self, image: Image.Image, **kwargs) -> Image.Image: return self.model(image, **kwargs) if gr.NO_RELOAD: # 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", # "stable-diffusion-v1-5/stable-diffusion-v1-5", "ashllay/stable-diffusion-v1-5-archive", solver_order=2, subfolder="scheduler", use_karras_sigmas=True, final_sigmas_type="sigma_min", algorithm_type="sde-dpmsolver++", prediction_type="epsilon", thresholding=False, denoise_final=True, device_map="cuda", torch_dtype=torch.float16, ) # Stable Diffusion Pipeline URL # base_model_url = "https://huggingface.co/broyang/hentaidigitalart_v20/blob/main/realcartoon3d_v15.safetensors" base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors" # vae_url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors" # print('loading vae') # vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda") # vae.to(memory_format=torch.channels_last) print('loading pipe') pipe = StableDiffusionControlNetPipeline.from_single_file( base_model_url, safety_checker=None, controlnet=controlnet, scheduler=scheduler, # vae=vae, torch_dtype=torch.float16, ).to("cuda") # print('loading inpainting pipe') # inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained( # "runwayml/stable-diffusion-inpainting", # torch_dtype=torch.float16, # ).to("cuda") print("loading preprocessor") preprocessor = Preprocessor() preprocessor.load("NormalBae") pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="EasyNegativeV2.safetensors", token="EasyNegativeV2",) pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="badhandv4.pt", token="badhandv4") pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="fcNeg-neg.pt", token="fcNeg-neg") pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Ahegao.pt", token="HDA_Ahegao") pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Bondage.pt", token="HDA_Bondage") pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_pet_play.pt", token="HDA_pet_play") pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_unconventional maid.pt", token="HDA_unconventional_maid") pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NakedHoodie.pt", token="HDA_NakedHoodie") pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NunDress.pt", token="HDA_NunDress") pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Shibari.pt", token="HDA_Shibari") pipe.to("cuda") print("---------------Loaded controlnet pipeline---------------") torch.cuda.empty_cache() gc.collect() print(f"CUDA memory allocated: {torch.cuda.max_memory_allocated(device='cuda') / 1e9:.2f} GB") print("Model Compiled!") def generate_furniture_mask(image, furniture_type): image_np = np.array(image) height, width = image_np.shape[:2] mask = np.zeros((height, width), dtype=np.uint8) if furniture_type == "sofa": cv2.rectangle(mask, (width//4, int(height*0.6)), (width*3//4, height), 255, -1) elif furniture_type == "table": cv2.rectangle(mask, (width//3, height//3), (width*2//3, height*2//3), 255, -1) elif furniture_type == "chair": cv2.circle(mask, (width*3//5, height*2//3), height//6, 255, -1) return Image.fromarray(mask) 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),tungsten white balance,captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length" 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": "Minimalist interior design,clean lines,neutral colors,uncluttered space,functional furniture,lots of natural light" }, { "name": "Boho", "prompt": "Bohemian chic interior,eclectic mix of patterns and textures,vintage furniture,plants,woven textiles,warm earthy colors" }, { "name": "Farmhouse", "prompt": "Modern farmhouse interior,rustic wood elements,shiplap walls,neutral color palette,industrial accents,cozy textiles" }, { "name": "Saudi Prince", "prompt": "Opulent gold interior,luxurious ornate furniture,crystal chandeliers,rich fabrics,marble floors,intricate Arabic patterns" }, { "name": "Neoclassical", "prompt": "Neoclassical interior design,elegant columns,ornate moldings,symmetrical layout,refined furniture,muted color palette" }, { "name": "Eclectic", "prompt": "Eclectic interior design,mix of styles and eras,bold color combinations,diverse furniture pieces,unique art objects" }, { "name": "Parisian", "prompt": "Parisian apartment interior,all-white color scheme,ornate moldings,herringbone wood floors,elegant furniture,large windows" }, { "name": "Hollywood", "prompt": "Hollywood Regency interior,glamorous and luxurious,bold colors,mirrored surfaces,velvet upholstery,gold accents" }, { "name": "Scandinavian", "prompt": "Scandinavian interior design,light wood tones,white walls,minimalist furniture,cozy textiles,hygge atmosphere" }, { "name": "Beach", "prompt": "Coastal beach house interior,light blue and white color scheme,weathered wood,nautical accents,sheer curtains,ocean view" }, { "name": "Japanese", "prompt": "Traditional Japanese interior,tatami mats,shoji screens,low furniture,zen garden view,minimalist decor,natural materials" }, { "name": "Midcentury Modern", "prompt": "Mid-century modern interior,1950s-60s style furniture,organic shapes,warm wood tones,bold accent colors,large windows" }, { "name": "Retro Futurism", "prompt": "Neon (atompunk world) retro cyberpunk background", }, { "name": "Texan", "prompt": "Western cowboy interior,rustic wood beams,leather furniture,cowhide rugs,antler chandeliers,southwestern patterns" }, { "name": "Matrix", "prompt": "Futuristic cyberpunk interior,neon accent lighting,holographic plants,sleek black surfaces,advanced gaming setup,transparent screens,Blade Runner inspired decor,high-tech minimalist furniture" } ] 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, "none") return p css = """ h1, h2, h3 { text-align: center; display: block; } footer { visibility: hidden; } .gradio-container { max-width: 1100px !important; } .gr-image { display: flex; justify-content: center; align-items: center; width: 100%; height: 512px; overflow: hidden; } .gr-image img { width: 100%; height: 100%; object-fit: cover; object-position: center; } """ 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), tungsten white balance, captured with a DSLR camera using f/10 aperture, 1/60 sec shutter speed, ISO 400, 20mm focal length" ) 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.Column(): prompt = gr.Textbox( label="Custom Design", placeholder="Enter a description (optional)", ) # design options 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(equal_height=True): with gr.Column(scale=1, min_width=300): image = gr.Image( label="Input", sources=["upload"], show_label=True, mirror_webcam=True, type="pil", ) # run button with gr.Column(): run_button = gr.Button(value="Use this one", size="lg", visible=False) # output image with gr.Column(scale=1, min_width=300): result = gr.Image( label="Output", interactive=False, type="pil", show_share_button= False, ) # Use this image button with gr.Column(): use_ai_button = gr.Button(value="Use this one", size="lg", visible=False) 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=[result] + config, outputs=[image, result], show_progress="minimal") def submit(previous_result, 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)): # First, yield the previous result to update the input image immediately yield previous_result, gr.update() # Then, process the new input image new_result = process_image(previous_result, style_selection, prompt, a_prompt, n_prompt, num_images, image_resolution, preprocess_resolution, num_steps, guidance_scale, seed) # Finally, yield the new result yield previous_result, new_result # Turn off buttons when processing @gr.on(triggers=[image.upload, use_ai_button.click, run_button.click], inputs=None, outputs=[run_button, use_ai_button], show_progress="hidden") def turn_buttons_off(): 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, ): seed = random.randint(0, MAX_SEED) generator = torch.cuda.manual_seed(seed) preprocessor.load("NormalBae") control_image = preprocessor( image=image, image_resolution=image_resolution, detect_resolution=preprocess_resolution, ) if style_selection is not None and style_selection != "None": prompt = f"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) # Generate the initial room image initial_result = pipe( prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_images_per_prompt=1, num_inference_steps=num_steps, generator=generator, image=control_image, ).images[0] # # Randomly choose whether to add furniture and which type # furniture_types = ["sofa", "table", "chair", "dresser", "bookshelf", "desk", "coffee table"] # furniture_type = random.choice(furniture_types) # furniture_mask = generate_furniture_mask(initial_result, furniture_type) # furniture_prompt = f"{prompt}, with a {furniture_type} in the style of {style_selection}" # print(furniture_prompt) # # Use the inpainting pipeline to add furniture # final_result = inpaint_pipe( # prompt=furniture_prompt, # image=initial_result, # mask_image=furniture_mask, # negative_prompt=negative_prompt, # num_inference_steps=num_steps, # guidance_scale=guidance_scale, # generator=generator, # ).images[0] # Save and upload results timestamp = int(time.time()) img_path = f"{timestamp}.jpg" results_path = f"{timestamp}_out.jpg" imageio.imsave(img_path, image) imageio.imsave(results_path, initial_result) api.upload_file( path_or_fileobj=img_path, path_in_repo=img_path, repo_id="broyang/interior-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/interior-ai-outputs", repo_type="dataset", token=API_KEY, run_as_future=True, ) return initial_result if prod: demo.queue(max_size=20).launch(server_name="localhost", server_port=port) else: demo.queue(api_open=False).launch(show_api=False)