app.py

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,num=3):
    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[:num]]

    ### 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

def find_token_sequence_in_pre_tokenized(input_string, other_string,pipe):
    # Load the tokenizer
    tokenizer = pipe.tokenizer

    # Tokenize the input string
    input_tokens = tokenizer.tokenize(input_string)

    # Tokenize the other string

    pre_tokenized_tokens = tokenizer.tokenize(other_string)
    # Find matching token sequences and their indexes
    matching_sequences = []
    input_length = len(input_tokens)
    for i in range(len(pre_tokenized_tokens) - input_length + 1):
        if pre_tokenized_tokens[i:i + input_length] == input_tokens:
            matching_sequences.append((pre_tokenized_tokens[i:i + input_length], i))

    return matching_sequences


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
pipe_CN.load_lora_weights('Tonioesparza/ourhood_training_dreambooth_lora_2_0',
                                   weight_name='pytorch_lora_weights.safetensors')
#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

slingshot = torch.load("./slingshot.pt")

def ourhood_inference(prompt=str, num_inference_steps=int, scaffold=int, seed=int, cases_strength=float , cases_scope=int ):
    ###pro_encode = pipe_cn.encode_text(prompt) ###ip_images init

    condition = 'both'

    guide = pipe_CN.encode_prompt(prompt)

    closest, indexes = ip_guide(guide[2], pool,cases_scope)

    ### torch.mean de los indexes

    ip_means = torch.mean(torch.stack([pool[i] for i in indexes]), dim=0)

    print([list(ip_pool.keys())[i] for i in indexes])

    ip_embeds = torch.cat([torch.unsqueeze(torch.zeros_like(closest[0]), 0), torch.unsqueeze(ip_means, 0)], 0).to(
        dtype=torch.float16, device='cuda')

    pipe_CN.set_ip_adapter_scale([[cases_strength]])

    prompt1 = 'A photograph, of an OurHood privacy booth, with a silken oak frame, hickory stained melange polyester fabric, in ' + prompt

    ### prompt encoding

    text_inputs = pipe_CN.tokenizer(
        prompt1,
        padding="max_length",
        max_length=pipe_CN.tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt",
    )

    text_input_ids = text_inputs.input_ids

    prompt_embeds_1 = pipe_CN.text_encoder(text_input_ids.to('cuda'), output_hidden_states=True)

    prompt_embeds_1 = prompt_embeds_1.hidden_states[-2]

    ###embed prompt encoding 2

    prompt_embeds_2 = pipe_CN.text_encoder_2(text_input_ids.to('cuda'), output_hidden_states=True)

    pooled_prompt_embeds_2 = prompt_embeds_2[0]

    prompt_embeds_2 = prompt_embeds_2.hidden_states[-2]

    #### substraction

    if condition == 'both':

        matches = find_token_sequence_in_pre_tokenized('ourhood privacy booth', prompt1, pipe_CN)

        items = []

        for match in matches:
            for w in range(len(match[0])):
                items.append(match[1] + w)

        for it in items:
            prompt_embeds_2[0][it] = prompt_embeds_2[0][it] + slingshot['b'].to('cuda')

        pooled_prompt_embeds = pooled_prompt_embeds_2 + slingshot['b'].to('cuda')

    elif condition == 'pooled':

        pooled_prompt_embeds = pooled_prompt_embeds_2 + slingshot['b'].to('cuda')

    elif condition == 'embeds':

        matches = find_token_sequence_in_pre_tokenized('ourhood privacy booth', prompt1, pipe_CN)

        items = []

        for match in matches:
            for w in range(len(match[0])):
                items.append(match[1] + w)

        for it in items:
            prompt_embeds_2[0][it] = prompt_embeds_2[0][it] + slingshot['b'].to('cuda')

    ### concatenation

    prompt_embeds = torch.cat([prompt_embeds_1, prompt_embeds_2], dim=-1)

    ### create negative embeds text encoder 1

    negative_prompt = "deformed, ugly, wrong proportion, low res, worst quality, low quality,text,watermark"

    max_length = prompt_embeds.shape[1]

    uncond_input = pipe_CN.tokenizer(
        negative_prompt,
        padding="max_length",
        max_length=max_length,
        truncation=True,
        return_tensors="pt",
    )

    uncond_input_ids = uncond_input.input_ids

    negative_prompt_embeds_1 = pipe_CN.text_encoder(
        uncond_input_ids.to('cuda'),
        output_hidden_states=True,
    )

    negative_prompt_embeds_1 = negative_prompt_embeds_1.hidden_states[-2]

    ### create negative embeds text encoder 2

    negative_prompt_embeds_2 = pipe_CN.text_encoder_2(
        uncond_input_ids.to('cuda'),
        output_hidden_states=True,
    )

    negative_pooled_prompt_embeds = negative_prompt_embeds_2[0]

    negative_prompt_embeds_2 = negative_prompt_embeds_2.hidden_states[-2]

    ### negative concatenation

    negative_prompt_embeds = torch.cat([negative_prompt_embeds_1, negative_prompt_embeds_2], dim=-1)

    ### 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_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds_2,
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        ip_adapter_image_embeds=[ip_embeds],
        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,
            )
            

            cases_strength = gr.Slider(
                label="Brand strenght",
                minimum=0.0,
                maximum=1.0,
                step=0.05,
                value=0.5,
            )

            cases_scope = gr.Slider(
                label="Brand scope",
                minimum=1,
                maximum=10,
                step=1,
                value=1,
            )
            

            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,cases_strength,cases_scope],
        outputs = [result]
    )

demo.queue().launch()