import gradio as gr import numpy as np import random #import spaces #[uncomment to use ZeroGPU] import os from PIL import Image, ImageDraw, ImageFont import torch from PIL import Image from diffusers.utils import load_image from diffusers import DPMSolverSDEScheduler from diffusers import StableDiffusionXLImg2ImgPipeline, DPMSolverMultistepScheduler, AutoencoderTiny, StableDiffusionXLControlNetPipeline, ControlNetModel from diffusers.utils import load_image from diffusers.image_processor import IPAdapterMaskProcessor from torch import nn ### auxiliary functions def ip_guide(guide, pool): distances = [] cos = nn.CosineSimilarity(dim=1, eps=1e-6) for embed in pool: dist = cos(guide, embed.to('cuda')) distances.append(dist) ### find the indexes of the top 5 embeddings indexed_distances = list(enumerate(distances)) # Sort the list of pairs based on the scores sorted_distances = sorted(indexed_distances, key=lambda x: x[1]) # Extract the indexes of the lowest scores lowest_indexes = [index for index, score in sorted_distances[:5]] ### return the embeddings with lowest_indexes return [pool[i] for i in lowest_indexes], lowest_indexes def make_inpaint_condition(image, image_mask): image = np.array(image.convert("RGB")).astype(np.float32) / 255.0 image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0 assert image.shape[0:1] == image_mask.shape[0:1] image[image_mask > 0.5] = -1.0 # set as masked pixel image = np.expand_dims(image, 0).transpose(0, 3, 1, 2) image = torch.from_numpy(image) return image device = "cuda" if torch.cuda.is_available() else "cpu" model_repo_id = "stabilityai/sdxl-turbo" #Replace to the model you would like to use if torch.cuda.is_available(): torch_dtype = torch.float16 else: torch_dtype = torch.float32 processor_mask = IPAdapterMaskProcessor() controlnets = [ ControlNetModel.from_pretrained( "diffusers/controlnet-depth-sdxl-1.0",variant="fp16",use_safetensors=True,torch_dtype=torch.float16 ), ControlNetModel.from_pretrained( "diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16, use_safetensors=True,variant="fp16" ), ] ###load pipelines pipe_CN = StableDiffusionXLControlNetPipeline.from_pretrained("SG161222/RealVisXL_V5.0", torch_dtype=torch.float16,controlnet=[controlnets[0],controlnets[0],controlnets[1]], use_safetensors=True, variant='fp16') pipe_CN.vae = AutoencoderTiny.from_pretrained("madebyollin/taesdxl", torch_dtype=torch.float16) pipe_CN.scheduler=DPMSolverSDEScheduler.from_pretrained("SG161222/RealVisXL_V5.0",subfolder="scheduler",use_karras_sigmas=True) pipe_CN.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin") pipe_CN.to("cuda") ##############################load loras state_dict, network_alphas = StableDiffusionXLControlNetPipeline.lora_state_dict('Tonioesparza/ourhood_training_dreambooth_lora_2_0', weight_name='pytorch_lora_weights.safetensors') pipe_CN.load_lora_into_unet(state_dict, network_alphas, pipe_CN.unet, adapter_name='unet_ourhood') pipe_CN.load_lora_into_text_encoder(state_dict, network_alphas, pipe_CN.text_encoder, adapter_name='text_ourhood') pipe_CN.load_lora_into_text_encoder(state_dict, network_alphas, pipe_CN.text_encoder ,prefix='2', adapter_name='text_2_ourhood') pipe_CN.set_adapters(["unet_ourhood","text_ourhood","text_2_ourhood"], adapter_weights=[1.0, 1.0,1.0]) pipe_CN.fuse_lora() refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0",text_encoder_2=pipe_CN.text_encoder_2,vae=pipe_CN.vae,torch_dtype=torch.float16,use_safetensors=True,variant="fp16") refiner.to("cuda") ip_pool = torch.load("./embeds_cases_for_ip.pt") pool = list(ip_pool.values()) MAX_SEED = np.iinfo(np.int32).max MAX_IMAGE_SIZE = 1024 def ourhood_inference(prompt=str,num_inference_steps=int,scaffold=int,seed=int): ###pro_encode = pipe_cn.encode_text(prompt) ###ip_images init guide = pipe_CN.encode_prompt(prompt) closest,indexes = ip_guide(guide[2],pool) print( [list(ip_pool.keys())[i] for i in indexes]) ip_embeds = torch.cat([torch.unsqueeze(torch.zeros_like(closest[0]),0),torch.unsqueeze(closest[0],0)],0).to(dtype=torch.float16,device='cuda') pipe_CN.set_ip_adapter_scale([[0.5]]) prompt1 = 'A photograph, of an OurHood privacy booth, with a silken oak frame, hickory stained melange polyester fabric, in ' + prompt ### function has no formats defined scaff_dic={1:{'mask1':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/mask_in_square_2.png", 'mask2':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/mask_out_square_2.png", 'depth_image':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/mask_depth_noroof_square.png", 'canny_image':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/mask_depth_solo_square.png"}, 2:{'mask1':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/mask_in_C.png", 'mask2':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/mask_out_C.png", 'depth_image':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/depth_C.png", 'canny_image':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/canny_C_solo.png"}, 3:{'mask1':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/mask_in_B.png", 'mask2':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/mask_out_B.png", 'depth_image':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/depth_B.png", 'canny_image':"https://huggingface.co/Tonioesparza/ourhood_training_dreambooth_lora_2_0/resolve/main/canny_B_solo.png"}} ### mask init output_height = 1024 output_width = 1024 mask1 = load_image(scaff_dic[scaffold]['mask1']) mask2 = load_image(scaff_dic[scaffold]['mask2']) masks = processor_mask.preprocess([mask1], height=output_height, width=output_width) masks = [masks.reshape(1, masks.shape[0], masks.shape[2], masks.shape[3])] ###precomputed depth image depth_image = load_image(scaff_dic[scaffold]['depth_image']) canny_image = load_image(scaff_dic[scaffold]['canny_image']) masked_depth=make_inpaint_condition(depth_image,mask2) images_CN = [depth_image, canny_image] ### inference n_steps = num_inference_steps generator = torch.Generator(device="cuda").manual_seed(seed) results = pipe_CN( prompt=prompt1, ip_adapter_image_embeds = [ip_embeds], negative_prompt="deformed, ugly, wrong proportion, low res, worst quality, low quality,text,watermark", generator=generator, num_inference_steps=n_steps, num_images_per_prompt=1, denoising_end=0.95, image=[depth_image,masked_depth,canny_image], output_type="latent", control_guidance_start=[0.0, 0.35, 0.35], control_guidance_end=[0.35, 0.95, 0.95], controlnet_conditioning_scale=[0.35,0.95, 0.95], cross_attention_kwargs={"ip_adapter_masks": masks} ).images[0] image = refiner( prompt=prompt1, generator=generator, num_inference_steps=n_steps, denoising_start=0.95, image=results, ).images[0] return image #@spaces.GPU #[uncomment to use ZeroGPU] examples = [ "in a British museum, pavillion, masonry, high-tables and chairs", "in a high ceilinged atrium, glass front, plantwalls, concrete floor, furniture, golden hour", "in a colorful open office environment", " in a Nordic atrium environment"] css=""" #col-container { margin: 0 auto; max-width: 640px; } """ with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown(f""" # HB8-Ourhood inference test """) with gr.Row(): prompt = gr.Text( label="Setting prompt", show_label=False, max_lines=1, placeholder="Where do you want to show the Ourhood pod?", container=False, ) run_button = gr.Button("Run", scale=0) result = gr.Image(label="Result", show_label=False) with gr.Accordion("Advanced Settings", open=False): perspective = gr.Slider( label="perspective", minimum=1, maximum=3, step=1, value=1, ) seed = gr.Slider( label="Tracking number (seed)", minimum=0, maximum=MAX_SEED, step=1, value=0, ) with gr.Row(): num_inference_steps = gr.Slider( label="Detail steps", minimum=35, maximum=75, step=1, value=50, #Replace with defaults that work for your model ) gr.Examples( examples = examples, inputs = [prompt] ) gr.on( triggers=[run_button.click, prompt.submit], fn = ourhood_inference, inputs = [prompt, num_inference_steps, perspective, seed], outputs = [result] ) demo.queue().launch()