utils.py

from transformers import PretrainedConfig
from PIL import Image
import torch
import numpy as np
import PIL
import os
from tqdm.auto import tqdm
from diffusers.models.attention_processor import (
    AttnProcessor2_0,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)

device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

def myroll2d(a, delta_x, delta_y):
    h, w = a.shape[0],  a.shape[1]
    delta_x = -delta_x
    delta_y = -delta_y
    if isinstance(a, np.ndarray):
        b = np.zeros   ([h,w]).astype(np.uint8)
    elif isinstance(a, torch.Tensor):
        b = torch.zeros([h,w]).to(torch.uint8)
    if delta_x > 0:
        left_a = delta_x
        right_a = w
        left_b = 0
        right_b = w - delta_x
    else:
        left_a = 0
        right_a = w + delta_x
        left_b = -delta_x
        right_b =  w
    if delta_y > 0:
        top_a = delta_y
        bot_a = h
        top_b = 0
        bot_b = h-delta_y
    else:
        top_a = 0
        bot_a = h + delta_y
        top_b = -delta_y
        bot_b = h
    b[left_b: right_b, top_b: bot_b] = a[left_a: right_a, top_a: bot_a]
    return b

def import_model_class_from_model_name_or_path(
    pretrained_model_name_or_path: str, revision = None, subfolder: str = "text_encoder"
):
    text_encoder_config = PretrainedConfig.from_pretrained(
        pretrained_model_name_or_path, subfolder=subfolder, revision=revision
    )
    model_class = text_encoder_config.architectures[0]

    if model_class == "CLIPTextModel":
        from transformers import CLIPTextModel
        return CLIPTextModel
    elif model_class == "CLIPTextModelWithProjection":
        from transformers import CLIPTextModelWithProjection
        return CLIPTextModelWithProjection
    else:
        raise ValueError(f"{model_class} is not supported.")

@torch.no_grad()
def image2latent(image, vae = None, dtype=None):
    with torch.no_grad():
        if type(image) is Image or type(image) is PIL.PngImagePlugin.PngImageFile or type(image) is PIL.JpegImagePlugin.JpegImageFile:
            image = np.array(image)
        if type(image) is torch.Tensor and image.dim() == 4:
            latents = image
        else:
            image = torch.from_numpy(image).float() / 127.5 - 1
            image = image.permute(2, 0, 1).unsqueeze(0).to(device, dtype= dtype)
            latents = vae.encode(image).latent_dist.sample()
            latents = latents * vae.config.scaling_factor
    return latents

@torch.no_grad()
def latent2image(latents, return_type = 'np', vae = None):
    # needs_upcasting = vae.dtype == torch.float16 and vae.config.force_upcast
    needs_upcasting = True
    if needs_upcasting:
        upcast_vae(vae)
        latents = latents.to(next(iter(vae.post_quant_conv.parameters())).dtype)
    image = vae.decode(latents /vae.config.scaling_factor, return_dict=False)[0]
    
    if return_type == 'np':
        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).numpy()#[0]
        image = (image * 255).astype(np.uint8)
    if needs_upcasting:
        vae.to(dtype=torch.float16)
    return image

def upcast_vae(vae):
    dtype = vae.dtype
    vae.to(dtype=torch.float32)
    use_torch_2_0_or_xformers = isinstance(
        vae.decoder.mid_block.attentions[0].processor,
        (
            AttnProcessor2_0,
            XFormersAttnProcessor,
            LoRAXFormersAttnProcessor,
            LoRAAttnProcessor2_0,
        ),
    )
    # if xformers or torch_2_0 is used attention block does not need
    # to be in float32 which can save lots of memory
    if use_torch_2_0_or_xformers:
        vae.post_quant_conv.to(dtype)
        vae.decoder.conv_in.to(dtype)
        vae.decoder.mid_block.to(dtype)

def prompt_to_emb_length_sdxl(prompt, tokenizer, text_encoder, length = None):
    text_input = tokenizer(
        [prompt],
        padding="max_length",
        max_length=length,
        truncation=True,
        return_tensors="pt",
    )
    prompt_embeds = text_encoder(text_input.input_ids.to(device),output_hidden_states=True)
    pooled_prompt_embeds = prompt_embeds[0]

    prompt_embeds = prompt_embeds.hidden_states[-2]
    bs_embed, seq_len, _ = prompt_embeds.shape
    prompt_embeds = prompt_embeds.view(bs_embed, seq_len, -1)
    pooled_prompt_embeds = pooled_prompt_embeds.view(bs_embed, -1)
    
    return  {"prompt_embeds": prompt_embeds, "pooled_prompt_embeds": pooled_prompt_embeds}




def prompt_to_emb_length_sd(prompt, tokenizer, text_encoder,  length = None):
    text_input = tokenizer(
        [prompt],
        padding="max_length",
        max_length=length,
        truncation=True,
        return_tensors="pt",
    )
    emb = text_encoder(text_input.input_ids.to(device))[0]
    return  emb 

def sdxl_prepare_input_decom(
    set_string_list,
    tokenizer,
    tokenizer_2,
    text_encoder_1,
    text_encoder_2,
    length = 20,
    bsz = 1,
    weight_dtype = torch.float32,
    resolution = 1024,
    normal_token_id_list = []
):
    encoder_hidden_states_list = []
    pooled_prompt_embeds = 0

    for m_idx in range(len(set_string_list)):
        prompt_embeds_list = []
        if ("#" in set_string_list[m_idx] or "$" in set_string_list[m_idx]) and m_idx not in normal_token_id_list :  ###
            out = prompt_to_emb_length_sdxl(
                set_string_list[m_idx], tokenizer, text_encoder_1, length = length
            )
        else:
            out = prompt_to_emb_length_sdxl(
                set_string_list[m_idx], tokenizer, text_encoder_1, length = 77
            )
            print(m_idx, set_string_list[m_idx])
        prompt_embeds, _ = out["prompt_embeds"].to(dtype=weight_dtype), out["pooled_prompt_embeds"].to(dtype=weight_dtype)
        prompt_embeds = prompt_embeds.repeat(bsz, 1, 1)
        prompt_embeds_list.append(prompt_embeds)
        if ("#" in set_string_list[m_idx] or "$" in set_string_list[m_idx]) and m_idx not in  normal_token_id_list:
            out = prompt_to_emb_length_sdxl(
                set_string_list[m_idx], tokenizer_2, text_encoder_2, length = length
            )
        else:
            out = prompt_to_emb_length_sdxl(
                set_string_list[m_idx], tokenizer_2, text_encoder_2, length = 77
            )
            print(m_idx, set_string_list[m_idx])

        prompt_embeds = out["prompt_embeds"].to(dtype=weight_dtype)
        pooled_prompt_embeds += out["pooled_prompt_embeds"].to(dtype=weight_dtype)
        prompt_embeds = prompt_embeds.repeat(bsz, 1, 1)
        prompt_embeds_list.append(prompt_embeds)
            
        encoder_hidden_states_list.append(torch.concat(prompt_embeds_list, dim=-1))
        
    add_text_embeds = pooled_prompt_embeds /len(set_string_list)
    target_size, original_size,crops_coords_top_left = (resolution,resolution),(resolution,resolution),(0,0)
    add_time_ids = list(original_size + crops_coords_top_left + target_size)

    add_time_ids = torch.tensor([add_time_ids], dtype=prompt_embeds.dtype,device = pooled_prompt_embeds.device) #[B,6]
    return encoder_hidden_states_list, add_text_embeds, add_time_ids

def sd_prepare_input_decom(
    set_string_list,
    tokenizer,
    text_encoder_1,
    length = 20,
    bsz = 1,
    weight_dtype = torch.float32,
    normal_token_id_list = []
):
    encoder_hidden_states_list = []
    for m_idx in range(len(set_string_list)):
        if ("#" in set_string_list[m_idx] or "$" in set_string_list[m_idx]) and m_idx not in normal_token_id_list :  ###
            encoder_hidden_states = prompt_to_emb_length_sd(
                set_string_list[m_idx], tokenizer, text_encoder_1, length = length
            )
        else:
            encoder_hidden_states = prompt_to_emb_length_sd(
                set_string_list[m_idx], tokenizer, text_encoder_1, length = 77
            )
            print(m_idx, set_string_list[m_idx])
        encoder_hidden_states = encoder_hidden_states.repeat(bsz, 1, 1)
        encoder_hidden_states_list.append(encoder_hidden_states.to(dtype=weight_dtype))
    return encoder_hidden_states_list


def load_mask (input_folder):
    np_mask_dtype = 'uint8'
    mask_np_list = []
    mask_label_list = []
    files = [
        file_name for file_name in os.listdir(input_folder) \
        if "mask" in file_name and ".npy" in file_name \
        and "_" in file_name and "Edited"  not in file_name 
    ]
    files = sorted(files, key = lambda x: int(x.split("_")[0][4:]))

    for idx, file_name in enumerate(files):
        if "mask" in file_name and ".npy" in file_name and "_" in file_name \
            and "Edited"  not in file_name:
            mask_np =  np.load(os.path.join(input_folder, file_name)).astype(np_mask_dtype) 
            mask_np_list.append(mask_np)  
            mask_label = file_name.split("_")[1][:-4]
            mask_label_list.append(mask_label)
    mask_list = []
    for mask_np in mask_np_list:
        mask = torch.from_numpy(mask_np)
        mask_list.append(mask)
    try: 
        assert torch.all(sum(mask_list)==1)
    except:
        print("please check mask")
        # plt.imsave( "out_mask.png", mask_list_edit[0]) 
        import pdb; pdb.set_trace()
    return mask_list, mask_label_list

def load_image(image_path, left=0, right=0, top=0, bottom=0, size = 512):
    if type(image_path) is str:
        image = np.array(Image.open(image_path))[:, :, :3]
    else:
        image = image_path
    h, w, c = image.shape
    left = min(left, w-1)
    right = min(right, w - left - 1)
    top = min(top, h - left - 1)
    bottom = min(bottom, h - top - 1)
    image = image[top:h-bottom, left:w-right]
    h, w, c = image.shape
    if h < w:
        offset = (w - h) // 2
        image = image[:, offset:offset + h]
    elif w < h:
        offset = (h - w) // 2
        image = image[offset:offset + w]
    image = np.array(Image.fromarray(image).resize((size, size)))
    return image

def mask_union_torch(*masks):
    masks = [m.to(torch.float) for m in masks]
    res = sum(masks)>0
    return res

def load_mask_edit(input_folder):
    np_mask_dtype = 'uint8'
    mask_np_list = []
    mask_label_list = []

    files = [file_name for file_name in os.listdir(input_folder)  if "mask" in file_name and ".npy" in file_name and "_" in file_name and "Edited" in file_name and "-1" not in file_name]
    files = sorted(files, key = lambda x: int(x.split("_")[0][10:]))
    
    for idx, file_name in enumerate(files):
        if "mask" in file_name and ".npy" in file_name and "_" in file_name and "Edited" in file_name and "-1" not in file_name:
            mask_np =  np.load(os.path.join(input_folder, file_name)).astype(np_mask_dtype) 
            mask_np_list.append(mask_np)  
            mask_label = file_name.split("_")[1][:-4]
            # mask_label = mask_label.split("-")[0]
            mask_label_list.append(mask_label)
    mask_list = []
    for mask_np in mask_np_list:
        mask = torch.from_numpy(mask_np)
        mask_list.append(mask)
    try: 
        assert torch.all(sum(mask_list)==1)
    except:
        print("Make sure maskEdited is in the folder, if not, generate using the UI")
        import pdb; pdb.set_trace()
    return mask_list, mask_label_list

def save_images(images,filename, num_rows=1, offset_ratio=0.02):
    if type(images) is list:
        num_empty = len(images) % num_rows
    elif images.ndim == 4:
        num_empty = images.shape[0] % num_rows
    else:
        images = [images]
        num_empty = 0

    empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
    images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
    num_items = len(images)

    folder = os.path.dirname(filename)
    for i, image in enumerate(images):
        pil_img = Image.fromarray(image)
        name = filename.split("/")[-1]
        name = name.split(".")[-2]+"_{}".format(i) +"."+filename.split(".")[-1]
        pil_img.save(os.path.join(folder, name))
        print("saved to ", os.path.join(folder, name))