Spaces:
Sleeping
Sleeping
| # preprocess.py | |
| ```py | |
| import PIL.Image | |
| import torch, gc | |
| from controlnet_aux_local import NormalBaeDetector#, CannyDetector | |
| 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: PIL.Image.Image, **kwargs) -> PIL.Image.Image: | |
| return self.model(image, **kwargs) | |
| ``` | |
| # app.py | |
| ```py | |
| 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 | |
| from PIL import Image | |
| from diffusers import ( | |
| ControlNetModel, | |
| DPMSolverMultistepScheduler, | |
| StableDiffusionControlNetPipeline, | |
| # 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", | |
| 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 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 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, | |
| ): | |
| preprocess_start = time.time() | |
| print("processing image") | |
| 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, | |
| ) | |
| 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) | |
| print(f"\n-------------------------Preprocess done in: {preprocess_time:.2f} seconds-------------------------") | |
| 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] | |
| print(f"\n-------------------------Inference done in: {time.time() - start:.2f} seconds-------------------------") | |
| torch.cuda.empty_cache() | |
| # upload block | |
| timestamp = int(time.time()) | |
| img_path = f"{timestamp}.jpg" | |
| results_path = f"{timestamp}_out.jpg" | |
| imageio.imsave(img_path, image) | |
| imageio.imsave(results_path, results) | |
| 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 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) | |
| ``` | |
| # .aidigestignore | |
| ``` | |
| controlnet_aux_local/normalbae/* | |
| requirements.txt | |
| win.requirements.txt | |
| web.html | |
| client.py | |
| local_app.py | |
| README.md | |
| Dockerfile | |
| .gitignore | |
| .gitattributes | |
| ``` | |
| # controlnet_aux_local/util.py | |
| ```py | |
| import os | |
| import random | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| annotator_ckpts_path = os.path.join(os.path.dirname(__file__), 'ckpts') | |
| def HWC3(x): | |
| assert x.dtype == np.uint8 | |
| if x.ndim == 2: | |
| x = x[:, :, None] | |
| assert x.ndim == 3 | |
| H, W, C = x.shape | |
| assert C == 1 or C == 3 or C == 4 | |
| if C == 3: | |
| return x | |
| if C == 1: | |
| return np.concatenate([x, x, x], axis=2) | |
| if C == 4: | |
| color = x[:, :, 0:3].astype(np.float32) | |
| alpha = x[:, :, 3:4].astype(np.float32) / 255.0 | |
| y = color * alpha + 255.0 * (1.0 - alpha) | |
| y = y.clip(0, 255).astype(np.uint8) | |
| return y | |
| def make_noise_disk(H, W, C, F): | |
| noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C)) | |
| noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC) | |
| noise = noise[F: F + H, F: F + W] | |
| noise -= np.min(noise) | |
| noise /= np.max(noise) | |
| if C == 1: | |
| noise = noise[:, :, None] | |
| return noise | |
| def nms(x, t, s): | |
| x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s) | |
| f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8) | |
| f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8) | |
| f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8) | |
| f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8) | |
| y = np.zeros_like(x) | |
| for f in [f1, f2, f3, f4]: | |
| np.putmask(y, cv2.dilate(x, kernel=f) == x, x) | |
| z = np.zeros_like(y, dtype=np.uint8) | |
| z[y > t] = 255 | |
| return z | |
| def min_max_norm(x): | |
| x -= np.min(x) | |
| x /= np.maximum(np.max(x), 1e-5) | |
| return x | |
| def safe_step(x, step=2): | |
| y = x.astype(np.float32) * float(step + 1) | |
| y = y.astype(np.int32).astype(np.float32) / float(step) | |
| return y | |
| def img2mask(img, H, W, low=10, high=90): | |
| assert img.ndim == 3 or img.ndim == 2 | |
| assert img.dtype == np.uint8 | |
| if img.ndim == 3: | |
| y = img[:, :, random.randrange(0, img.shape[2])] | |
| else: | |
| y = img | |
| y = cv2.resize(y, (W, H), interpolation=cv2.INTER_CUBIC) | |
| if random.uniform(0, 1) < 0.5: | |
| y = 255 - y | |
| return y < np.percentile(y, random.randrange(low, high)) | |
| def resize_image(input_image, resolution): | |
| H, W, C = input_image.shape | |
| H = float(H) | |
| W = float(W) | |
| k = float(resolution) / min(H, W) | |
| H *= k | |
| W *= k | |
| H = int(np.round(H / 64.0)) * 64 | |
| W = int(np.round(W / 64.0)) * 64 | |
| img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA) | |
| return img | |
| def torch_gc(): | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| torch.cuda.ipc_collect() | |
| def ade_palette(): | |
| """ADE20K palette that maps each class to RGB values.""" | |
| return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50], | |
| [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255], | |
| [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7], | |
| [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82], | |
| [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3], | |
| [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255], | |
| [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220], | |
| [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224], | |
| [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255], | |
| [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7], | |
| [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153], | |
| [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255], | |
| [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0], | |
| [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255], | |
| [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255], | |
| [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255], | |
| [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0], | |
| [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0], | |
| [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255], | |
| [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255], | |
| [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20], | |
| [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255], | |
| [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255], | |
| [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255], | |
| [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0], | |
| [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0], | |
| [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255], | |
| [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112], | |
| [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160], | |
| [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163], | |
| [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0], | |
| [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0], | |
| [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255], | |
| [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204], | |
| [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255], | |
| [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255], | |
| [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194], | |
| [102, 255, 0], [92, 0, 255]] | |
| ``` | |
| # controlnet_aux_local/processor.py | |
| ```py | |
| """ | |
| This file contains a Processor that can be used to process images with controlnet aux processors | |
| """ | |
| import io | |
| import logging | |
| from typing import Dict, Optional, Union | |
| from PIL import Image | |
| from controlnet_aux_local import (CannyDetector, ContentShuffleDetector, HEDdetector, | |
| LeresDetector, LineartAnimeDetector, | |
| LineartDetector, MediapipeFaceDetector, | |
| MidasDetector, MLSDdetector, NormalBaeDetector, | |
| OpenposeDetector, PidiNetDetector, ZoeDetector, | |
| DWposeDetector) | |
| LOGGER = logging.getLogger(__name__) | |
| MODELS = { | |
| # checkpoint models | |
| 'scribble_hed': {'class': HEDdetector, 'checkpoint': True}, | |
| 'softedge_hed': {'class': HEDdetector, 'checkpoint': True}, | |
| 'scribble_hedsafe': {'class': HEDdetector, 'checkpoint': True}, | |
| 'softedge_hedsafe': {'class': HEDdetector, 'checkpoint': True}, | |
| 'depth_midas': {'class': MidasDetector, 'checkpoint': True}, | |
| 'mlsd': {'class': MLSDdetector, 'checkpoint': True}, | |
| 'openpose': {'class': OpenposeDetector, 'checkpoint': True}, | |
| 'openpose_face': {'class': OpenposeDetector, 'checkpoint': True}, | |
| 'openpose_faceonly': {'class': OpenposeDetector, 'checkpoint': True}, | |
| 'openpose_full': {'class': OpenposeDetector, 'checkpoint': True}, | |
| 'openpose_hand': {'class': OpenposeDetector, 'checkpoint': True}, | |
| 'dwpose': {'class': DWposeDetector, 'checkpoint': True}, | |
| 'scribble_pidinet': {'class': PidiNetDetector, 'checkpoint': True}, | |
| 'softedge_pidinet': {'class': PidiNetDetector, 'checkpoint': True}, | |
| 'scribble_pidsafe': {'class': PidiNetDetector, 'checkpoint': True}, | |
| 'softedge_pidsafe': {'class': PidiNetDetector, 'checkpoint': True}, | |
| 'normal_bae': {'class': NormalBaeDetector, 'checkpoint': True}, | |
| 'lineart_coarse': {'class': LineartDetector, 'checkpoint': True}, | |
| 'lineart_realistic': {'class': LineartDetector, 'checkpoint': True}, | |
| 'lineart_anime': {'class': LineartAnimeDetector, 'checkpoint': True}, | |
| 'depth_zoe': {'class': ZoeDetector, 'checkpoint': True}, | |
| 'depth_leres': {'class': LeresDetector, 'checkpoint': True}, | |
| 'depth_leres++': {'class': LeresDetector, 'checkpoint': True}, | |
| # instantiate | |
| 'shuffle': {'class': ContentShuffleDetector, 'checkpoint': False}, | |
| 'mediapipe_face': {'class': MediapipeFaceDetector, 'checkpoint': False}, | |
| 'canny': {'class': CannyDetector, 'checkpoint': False}, | |
| } | |
| MODEL_PARAMS = { | |
| 'scribble_hed': {'scribble': True}, | |
| 'softedge_hed': {'scribble': False}, | |
| 'scribble_hedsafe': {'scribble': True, 'safe': True}, | |
| 'softedge_hedsafe': {'scribble': False, 'safe': True}, | |
| 'depth_midas': {}, | |
| 'mlsd': {}, | |
| 'openpose': {'include_body': True, 'include_hand': False, 'include_face': False}, | |
| 'openpose_face': {'include_body': True, 'include_hand': False, 'include_face': True}, | |
| 'openpose_faceonly': {'include_body': False, 'include_hand': False, 'include_face': True}, | |
| 'openpose_full': {'include_body': True, 'include_hand': True, 'include_face': True}, | |
| 'openpose_hand': {'include_body': False, 'include_hand': True, 'include_face': False}, | |
| 'dwpose': {}, | |
| 'scribble_pidinet': {'safe': False, 'scribble': True}, | |
| 'softedge_pidinet': {'safe': False, 'scribble': False}, | |
| 'scribble_pidsafe': {'safe': True, 'scribble': True}, | |
| 'softedge_pidsafe': {'safe': True, 'scribble': False}, | |
| 'normal_bae': {}, | |
| 'lineart_realistic': {'coarse': False}, | |
| 'lineart_coarse': {'coarse': True}, | |
| 'lineart_anime': {}, | |
| 'canny': {}, | |
| 'shuffle': {}, | |
| 'depth_zoe': {}, | |
| 'depth_leres': {'boost': False}, | |
| 'depth_leres++': {'boost': True}, | |
| 'mediapipe_face': {}, | |
| } | |
| CHOICES = f"Choices for the processor are {list(MODELS.keys())}" | |
| class Processor: | |
| def __init__(self, processor_id: str, params: Optional[Dict] = None) -> None: | |
| """Processor that can be used to process images with controlnet aux processors | |
| Args: | |
| processor_id (str): processor name, options are 'hed, midas, mlsd, openpose, | |
| pidinet, normalbae, lineart, lineart_coarse, lineart_anime, | |
| canny, content_shuffle, zoe, mediapipe_face | |
| params (Optional[Dict]): parameters for the processor | |
| """ | |
| LOGGER.info(f"Loading {processor_id}") | |
| if processor_id not in MODELS: | |
| raise ValueError(f"{processor_id} is not a valid processor id. Please make sure to choose one of {', '.join(MODELS.keys())}") | |
| self.processor_id = processor_id | |
| self.processor = self.load_processor(self.processor_id) | |
| # load default params | |
| self.params = MODEL_PARAMS[self.processor_id] | |
| # update with user params | |
| if params: | |
| self.params.update(params) | |
| def load_processor(self, processor_id: str) -> 'Processor': | |
| """Load controlnet aux processors | |
| Args: | |
| processor_id (str): processor name | |
| Returns: | |
| Processor: controlnet aux processor | |
| """ | |
| processor = MODELS[processor_id]['class'] | |
| # check if the proecssor is a checkpoint model | |
| if MODELS[processor_id]['checkpoint']: | |
| processor = processor.from_pretrained("lllyasviel/Annotators") | |
| else: | |
| processor = processor() | |
| return processor | |
| def __call__(self, image: Union[Image.Image, bytes], | |
| to_pil: bool = True) -> Union[Image.Image, bytes]: | |
| """processes an image with a controlnet aux processor | |
| Args: | |
| image (Union[Image.Image, bytes]): input image in bytes or PIL Image | |
| to_pil (bool): whether to return bytes or PIL Image | |
| Returns: | |
| Union[Image.Image, bytes]: processed image in bytes or PIL Image | |
| """ | |
| # check if bytes or PIL Image | |
| if isinstance(image, bytes): | |
| image = Image.open(io.BytesIO(image)).convert("RGB") | |
| processed_image = self.processor(image, **self.params) | |
| if to_pil: | |
| return processed_image | |
| else: | |
| output_bytes = io.BytesIO() | |
| processed_image.save(output_bytes, format='JPEG') | |
| return output_bytes.getvalue() | |
| ``` | |
| # controlnet_aux_local/__init__.py | |
| ```py | |
| __version__ = "0.0.8" | |
| # from .hed import HEDdetector | |
| # from .leres import LeresDetector | |
| # from .lineart import LineartDetector | |
| # from .lineart_anime import LineartAnimeDetector | |
| # from .midas import MidasDetector | |
| # from .mlsd import MLSDdetector | |
| from .normalbae import NormalBaeDetector | |
| # from .open_pose import OpenposeDetector | |
| # from .pidi import PidiNetDetector | |
| # from .zoe import ZoeDetector | |
| # from .canny import CannyDetector | |
| # from .mediapipe_face import MediapipeFaceDetector | |
| # from .segment_anything import SamDetector | |
| # from .shuffle import ContentShuffleDetector | |
| # from .dwpose import DWposeDetector | |
| ``` | |