Update app.py

app.py

@@ -1,21 +1,23 @@
 import gradio as gr
 import numpy as np
 import random
-#import spaces #[uncomment to use ZeroGPU]
+# 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 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):
+def ip_guide(guide, pool,num=3):
     distances = []
     cos = nn.CosineSimilarity(dim=1, eps=1e-6)
     for embed in pool:
@@ -26,7 +28,7 @@ def ip_guide(guide, pool):
     # 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]]
+    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
@@ -42,9 +44,28 @@ def make_inpaint_condition(image, image_mask):
     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
+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
@@ -54,32 +75,40 @@ else:
 processor_mask = IPAdapterMaskProcessor()
 controlnets = [
     ControlNetModel.from_pretrained(
-        "diffusers/controlnet-depth-sdxl-1.0",variant="fp16",use_safetensors=True,torch_dtype=torch.float16
+        "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"
+        "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 = 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.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.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 = 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")
@@ -89,37 +118,148 @@ 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):
+slingshot = torch.load("./slingshot.pt")
 
-###pro_encode = pipe_cn.encode_text(prompt) ###ip_images init
+def ourhood_inference(prompt=str, num_inference_steps=int, scaffold=int, seed=int, cases_strength = 0.5, cases_scope = 5):
+    ###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)
+    closest, indexes = ip_guide(guide[2], pool,cases_scope)
 
-    print( [list(ip_pool.keys())[i] for i in indexes])
+    ### torch.mean de los 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')
+    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]])
 
-    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
+    ### 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
@@ -130,38 +270,38 @@ def ourhood_inference(prompt=str,num_inference_steps=int,scaffold=int,seed=int):
     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
+    ###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)
+    masked_depth = make_inpaint_condition(depth_image, mask2)
 
     images_CN = [depth_image, canny_image]
 
-
-### inference
+    ### 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",
+        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],
+        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],
+        controlnet_conditioning_scale=[0.35, 0.95, 0.95],
         cross_attention_kwargs={"ip_adapter_masks": masks}
     ).images[0]
 
-
     image = refiner(
         prompt=prompt1,
         generator=generator,
@@ -227,6 +367,22 @@ with gr.Blocks(css=css) as demo:
                 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():