592 lines
26 KiB
Python
592 lines
26 KiB
Python
from calendar import c
|
||
from operator import invert
|
||
from webbrowser import get
|
||
import torch
|
||
import random
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import re
|
||
from .load_models_utils import get_lora_dict
|
||
|
||
global lora_get
|
||
lora_get=get_lora_dict()
|
||
lora_lightning_list = lora_get["lightning_xl_lora"]
|
||
class SpatialAttnProcessor2_0(torch.nn.Module):
|
||
r"""
|
||
Attention processor for IP-Adapater for PyTorch 2.0.
|
||
Args:
|
||
hidden_size (`int`):
|
||
The hidden size of the attention layer.
|
||
cross_attention_dim (`int`):
|
||
The number of channels in the `encoder_hidden_states`.
|
||
text_context_len (`int`, defaults to 77):
|
||
The context length of the text features.
|
||
scale (`float`, defaults to 1.0):
|
||
the weight scale of image prompt.
|
||
"""
|
||
|
||
def __init__(self, hidden_size = None, cross_attention_dim=None,id_length = 4,device = "cuda",dtype = torch.float16):
|
||
super().__init__()
|
||
if not hasattr(F, "scaled_dot_product_attention"):
|
||
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
||
self.device = device
|
||
self.dtype = dtype
|
||
self.hidden_size = hidden_size
|
||
self.cross_attention_dim = cross_attention_dim
|
||
self.total_length = id_length + 1
|
||
self.id_length = id_length
|
||
self.id_bank = {}
|
||
|
||
def __call__(
|
||
self,
|
||
attn,
|
||
hidden_states,
|
||
encoder_hidden_states=None,
|
||
attention_mask=None,
|
||
temb=None):
|
||
# un_cond_hidden_states, cond_hidden_states = hidden_states.chunk(2)
|
||
# un_cond_hidden_states = self.__call2__(attn, un_cond_hidden_states,encoder_hidden_states,attention_mask,temb)
|
||
# 生成一个0到1之间的随机数
|
||
global total_count,attn_count,cur_step,mask256,mask1024,mask4096
|
||
global sa16, sa32, sa64
|
||
global write
|
||
if write:
|
||
self.id_bank[cur_step] = [hidden_states[:self.id_length], hidden_states[self.id_length:]]
|
||
else:
|
||
encoder_hidden_states = torch.cat(self.id_bank[cur_step][0],hidden_states[:1],self.id_bank[cur_step][1],hidden_states[1:])
|
||
# 判断随机数是否大于0.5
|
||
if cur_step <5:
|
||
hidden_states = self.__call2__(attn, hidden_states,encoder_hidden_states,attention_mask,temb)
|
||
else: # 256 1024 4096
|
||
random_number = random.random()
|
||
if cur_step <20:
|
||
rand_num = 0.3
|
||
else:
|
||
rand_num = 0.1
|
||
if random_number > rand_num:
|
||
if not write:
|
||
if hidden_states.shape[1] == 32* 32:
|
||
attention_mask = mask1024[mask1024.shape[0] // self.total_length * self.id_length:]
|
||
elif hidden_states.shape[1] ==16*16:
|
||
attention_mask = mask256[mask256.shape[0] // self.total_length * self.id_length:]
|
||
else:
|
||
attention_mask = mask4096[mask4096.shape[0] // self.total_length * self.id_length:]
|
||
else:
|
||
if hidden_states.shape[1] == 32* 32:
|
||
attention_mask = mask1024[:mask1024.shape[0] // self.total_length * self.id_length]
|
||
elif hidden_states.shape[1] ==16*16:
|
||
attention_mask = mask256[:mask256.shape[0] // self.total_length * self.id_length]
|
||
else:
|
||
attention_mask = mask4096[:mask4096.shape[0] // self.total_length * self.id_length]
|
||
hidden_states = self.__call1__(attn, hidden_states,encoder_hidden_states,attention_mask,temb)
|
||
else:
|
||
hidden_states = self.__call2__(attn, hidden_states,None,attention_mask,temb)
|
||
attn_count +=1
|
||
if attn_count == total_count:
|
||
attn_count = 0
|
||
cur_step += 1
|
||
mask256,mask1024,mask4096 = cal_attn_mask(self.total_length,self.id_length,sa16,sa32,sa64, device=self.device, dtype= self.dtype)
|
||
|
||
return hidden_states
|
||
def __call1__(
|
||
self,
|
||
attn,
|
||
hidden_states,
|
||
encoder_hidden_states=None,
|
||
attention_mask=None,
|
||
temb=None,
|
||
):
|
||
residual = hidden_states
|
||
if encoder_hidden_states is not None:
|
||
raise Exception("not implement")
|
||
if attn.spatial_norm is not None:
|
||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||
input_ndim = hidden_states.ndim
|
||
|
||
if input_ndim == 4:
|
||
total_batch_size, channel, height, width = hidden_states.shape
|
||
hidden_states = hidden_states.view(total_batch_size, channel, height * width).transpose(1, 2)
|
||
total_batch_size,nums_token,channel = hidden_states.shape
|
||
img_nums = total_batch_size//2
|
||
hidden_states = hidden_states.view(-1,img_nums,nums_token,channel).reshape(-1,img_nums * nums_token,channel)
|
||
|
||
batch_size, sequence_length, _ = hidden_states.shape
|
||
|
||
if attn.group_norm is not None:
|
||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||
|
||
query = attn.to_q(hidden_states)
|
||
|
||
if encoder_hidden_states is None:
|
||
encoder_hidden_states = hidden_states # B, N, C
|
||
else:
|
||
encoder_hidden_states = encoder_hidden_states.view(-1,self.id_length+1,nums_token,channel).reshape(-1,(self.id_length+1) * nums_token,channel)
|
||
|
||
key = attn.to_k(encoder_hidden_states)
|
||
value = attn.to_v(encoder_hidden_states)
|
||
|
||
|
||
inner_dim = key.shape[-1]
|
||
head_dim = inner_dim // attn.heads
|
||
|
||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||
|
||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||
|
||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||
# TODO: add support for attn.scale when we move to Torch 2.1
|
||
hidden_states = F.scaled_dot_product_attention(
|
||
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
||
)
|
||
|
||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||
hidden_states = hidden_states.to(query.dtype)
|
||
|
||
|
||
|
||
# linear proj
|
||
hidden_states = attn.to_out[0](hidden_states)
|
||
# dropout
|
||
hidden_states = attn.to_out[1](hidden_states)
|
||
|
||
# if input_ndim == 4:
|
||
# tile_hidden_states = tile_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||
|
||
# if attn.residual_connection:
|
||
# tile_hidden_states = tile_hidden_states + residual
|
||
|
||
if input_ndim == 4:
|
||
hidden_states = hidden_states.transpose(-1, -2).reshape(total_batch_size, channel, height, width)
|
||
if attn.residual_connection:
|
||
hidden_states = hidden_states + residual
|
||
hidden_states = hidden_states / attn.rescale_output_factor
|
||
|
||
return hidden_states
|
||
def __call2__(
|
||
self,
|
||
attn,
|
||
hidden_states,
|
||
encoder_hidden_states=None,
|
||
attention_mask=None,
|
||
temb=None):
|
||
residual = hidden_states
|
||
|
||
if attn.spatial_norm is not None:
|
||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||
|
||
input_ndim = hidden_states.ndim
|
||
|
||
if input_ndim == 4:
|
||
batch_size, channel, height, width = hidden_states.shape
|
||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||
|
||
batch_size, sequence_length, _ = (
|
||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||
)
|
||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||
|
||
if attn.group_norm is not None:
|
||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||
|
||
query = attn.to_q(hidden_states)
|
||
|
||
if encoder_hidden_states is None:
|
||
encoder_hidden_states = hidden_states
|
||
elif attn.norm_cross:
|
||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||
|
||
key = attn.to_k(encoder_hidden_states)
|
||
value = attn.to_v(encoder_hidden_states)
|
||
|
||
query = attn.head_to_batch_dim(query)
|
||
key = attn.head_to_batch_dim(key)
|
||
value = attn.head_to_batch_dim(value)
|
||
|
||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||
hidden_states = torch.bmm(attention_probs, value)
|
||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||
|
||
# linear proj
|
||
hidden_states = attn.to_out[0](hidden_states)
|
||
# dropout
|
||
hidden_states = attn.to_out[1](hidden_states)
|
||
|
||
if input_ndim == 4:
|
||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||
|
||
if attn.residual_connection:
|
||
hidden_states = hidden_states + residual
|
||
|
||
hidden_states = hidden_states / attn.rescale_output_factor
|
||
|
||
return hidden_states
|
||
|
||
|
||
def cal_attn_mask(total_length,id_length,sa16,sa32,sa64,device="cuda",dtype= torch.float16):
|
||
bool_matrix256 = torch.rand((1, total_length * 256),device = device,dtype = dtype) < sa16
|
||
bool_matrix1024 = torch.rand((1, total_length * 1024),device = device,dtype = dtype) < sa32
|
||
bool_matrix4096 = torch.rand((1, total_length * 4096),device = device,dtype = dtype) < sa64
|
||
bool_matrix256 = bool_matrix256.repeat(total_length,1)
|
||
bool_matrix1024 = bool_matrix1024.repeat(total_length,1)
|
||
bool_matrix4096 = bool_matrix4096.repeat(total_length,1)
|
||
for i in range(total_length):
|
||
bool_matrix256[i:i+1,id_length*256:] = False
|
||
bool_matrix1024[i:i+1,id_length*1024:] = False
|
||
bool_matrix4096[i:i+1,id_length*4096:] = False
|
||
bool_matrix256[i:i+1,i*256:(i+1)*256] = True
|
||
bool_matrix1024[i:i+1,i*1024:(i+1)*1024] = True
|
||
bool_matrix4096[i:i+1,i*4096:(i+1)*4096] = True
|
||
mask256 = bool_matrix256.unsqueeze(1).repeat(1,256,1).reshape(-1,total_length * 256)
|
||
mask1024 = bool_matrix1024.unsqueeze(1).repeat(1,1024,1).reshape(-1,total_length * 1024)
|
||
mask4096 = bool_matrix4096.unsqueeze(1).repeat(1,4096,1).reshape(-1,total_length * 4096)
|
||
return mask256,mask1024,mask4096
|
||
|
||
def cal_attn_mask_xl(total_length,id_length,sa32,sa64,height_s,width_s,device="cuda",dtype= torch.float16):
|
||
nums_1024 = (height_s // 32) * (width_s // 32)
|
||
nums_4096 = (height_s // 16) * (width_s // 16)
|
||
bool_matrix1024 = torch.rand((1, total_length * nums_1024),device = device,dtype = dtype) < sa32
|
||
bool_matrix4096 = torch.rand((1, total_length * nums_4096),device = device,dtype = dtype) < sa64
|
||
bool_matrix1024 = bool_matrix1024.repeat(total_length,1)
|
||
bool_matrix4096 = bool_matrix4096.repeat(total_length,1)
|
||
for i in range(total_length):
|
||
bool_matrix1024[i:i+1,id_length*nums_1024:] = False
|
||
bool_matrix4096[i:i+1,id_length*nums_4096:] = False
|
||
bool_matrix1024[i:i+1,i*nums_1024:(i+1)*nums_1024] = True
|
||
bool_matrix4096[i:i+1,i*nums_4096:(i+1)*nums_4096] = True
|
||
mask1024 = bool_matrix1024.unsqueeze(1).repeat(1,nums_1024,1).reshape(-1,total_length * nums_1024)
|
||
mask4096 = bool_matrix4096.unsqueeze(1).repeat(1,nums_4096,1).reshape(-1,total_length * nums_4096)
|
||
return mask1024,mask4096
|
||
|
||
|
||
def cal_attn_indice_xl_effcient_memory(total_length,id_length,sa32,sa64,height_s,width_s,device="cuda",dtype= torch.float16):
|
||
nums_1024 = (height_s // 32) * (width_s // 32)
|
||
nums_4096 = (height_s // 16) * (width_s // 16)
|
||
bool_matrix1024 = torch.rand((total_length,nums_1024),device = device,dtype = dtype) < sa32
|
||
bool_matrix4096 = torch.rand((total_length,nums_4096),device = device,dtype = dtype) < sa64
|
||
# 用nonzero()函数获取所有为True的值的索引
|
||
indices1024 = [torch.nonzero(bool_matrix1024[i], as_tuple=True)[0] for i in range(total_length)]
|
||
indices4096 = [torch.nonzero(bool_matrix4096[i], as_tuple=True)[0] for i in range(total_length)]
|
||
|
||
return indices1024,indices4096
|
||
|
||
|
||
class AttnProcessor(nn.Module):
|
||
r"""
|
||
Default processor for performing attention-related computations.
|
||
"""
|
||
def __init__(
|
||
self,
|
||
hidden_size=None,
|
||
cross_attention_dim=None,
|
||
):
|
||
super().__init__()
|
||
|
||
def __call__(
|
||
self,
|
||
attn,
|
||
hidden_states,
|
||
encoder_hidden_states=None,
|
||
attention_mask=None,
|
||
temb=None,
|
||
):
|
||
residual = hidden_states
|
||
|
||
if attn.spatial_norm is not None:
|
||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||
|
||
input_ndim = hidden_states.ndim
|
||
|
||
if input_ndim == 4:
|
||
batch_size, channel, height, width = hidden_states.shape
|
||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||
|
||
batch_size, sequence_length, _ = (
|
||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||
)
|
||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||
|
||
if attn.group_norm is not None:
|
||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||
|
||
query = attn.to_q(hidden_states)
|
||
|
||
if encoder_hidden_states is None:
|
||
encoder_hidden_states = hidden_states
|
||
elif attn.norm_cross:
|
||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||
|
||
key = attn.to_k(encoder_hidden_states)
|
||
value = attn.to_v(encoder_hidden_states)
|
||
|
||
query = attn.head_to_batch_dim(query)
|
||
key = attn.head_to_batch_dim(key)
|
||
value = attn.head_to_batch_dim(value)
|
||
|
||
attention_probs = attn.get_attention_scores(query, key, attention_mask)
|
||
hidden_states = torch.bmm(attention_probs, value)
|
||
hidden_states = attn.batch_to_head_dim(hidden_states)
|
||
|
||
# linear proj
|
||
hidden_states = attn.to_out[0](hidden_states)
|
||
# dropout
|
||
hidden_states = attn.to_out[1](hidden_states)
|
||
|
||
if input_ndim == 4:
|
||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||
|
||
if attn.residual_connection:
|
||
hidden_states = hidden_states + residual
|
||
|
||
hidden_states = hidden_states / attn.rescale_output_factor
|
||
|
||
return hidden_states
|
||
|
||
|
||
class AttnProcessor2_0(torch.nn.Module):
|
||
r"""
|
||
Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
|
||
"""
|
||
def __init__(
|
||
self,
|
||
hidden_size=None,
|
||
cross_attention_dim=None,
|
||
):
|
||
super().__init__()
|
||
if not hasattr(F, "scaled_dot_product_attention"):
|
||
raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
|
||
|
||
def __call__(
|
||
self,
|
||
attn,
|
||
hidden_states,
|
||
encoder_hidden_states=None,
|
||
attention_mask=None,
|
||
temb=None,
|
||
):
|
||
residual = hidden_states
|
||
|
||
if attn.spatial_norm is not None:
|
||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
||
|
||
input_ndim = hidden_states.ndim
|
||
|
||
if input_ndim == 4:
|
||
batch_size, channel, height, width = hidden_states.shape
|
||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||
|
||
batch_size, sequence_length, _ = (
|
||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||
)
|
||
|
||
if attention_mask is not None:
|
||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||
# scaled_dot_product_attention expects attention_mask shape to be
|
||
# (batch, heads, source_length, target_length)
|
||
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
|
||
|
||
if attn.group_norm is not None:
|
||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
||
|
||
query = attn.to_q(hidden_states)
|
||
|
||
if encoder_hidden_states is None:
|
||
encoder_hidden_states = hidden_states
|
||
elif attn.norm_cross:
|
||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
||
|
||
key = attn.to_k(encoder_hidden_states)
|
||
value = attn.to_v(encoder_hidden_states)
|
||
|
||
inner_dim = key.shape[-1]
|
||
head_dim = inner_dim // attn.heads
|
||
|
||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||
|
||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||
|
||
# the output of sdp = (batch, num_heads, seq_len, head_dim)
|
||
# TODO: add support for attn.scale when we move to Torch 2.1
|
||
hidden_states = F.scaled_dot_product_attention(
|
||
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
|
||
)
|
||
|
||
hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||
hidden_states = hidden_states.to(query.dtype)
|
||
|
||
# linear proj
|
||
hidden_states = attn.to_out[0](hidden_states)
|
||
# dropout
|
||
hidden_states = attn.to_out[1](hidden_states)
|
||
|
||
if input_ndim == 4:
|
||
hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
|
||
|
||
if attn.residual_connection:
|
||
hidden_states = hidden_states + residual
|
||
|
||
hidden_states = hidden_states / attn.rescale_output_factor
|
||
|
||
return hidden_states
|
||
|
||
|
||
def is_torch2_available():
|
||
return hasattr(F, "scaled_dot_product_attention")
|
||
|
||
|
||
# 将列表转换为字典的函数
|
||
def character_to_dict(general_prompt,lora,add_trigger_words):
|
||
character_dict = {}
|
||
generate_prompt_arr = general_prompt.splitlines()
|
||
if lora:
|
||
if lora not in lora_lightning_list:
|
||
generate_prompt_arr = [item + add_trigger_words for item in generate_prompt_arr]
|
||
#print(prompts)
|
||
character_index_dict = {}
|
||
invert_character_index_dict = {}
|
||
character_list = []
|
||
for ind,string in enumerate(generate_prompt_arr):
|
||
# 分割字符串寻找key和value
|
||
start = string.find('[')
|
||
end = string.find(']')
|
||
if start != -1 and end != -1:
|
||
key = string[start:end+1]
|
||
value = string[end+1:]
|
||
if "#" in value:
|
||
value = value.rpartition('#')[0]
|
||
if key in character_dict:
|
||
raise f"duplicate character descirption:{key}"
|
||
character_dict[key] = value
|
||
character_list.append(key)
|
||
|
||
|
||
return character_dict,character_list
|
||
|
||
def get_id_prompt_index(character_dict,id_prompts):
|
||
replace_id_prompts = []
|
||
character_index_dict = {}
|
||
invert_character_index_dict = {}
|
||
for ind,id_prompt in enumerate(id_prompts):
|
||
for key in character_dict.keys():
|
||
if key in id_prompt:
|
||
if key not in character_index_dict:
|
||
character_index_dict[key] = []
|
||
character_index_dict[key].append(ind)
|
||
invert_character_index_dict[ind] = key
|
||
replace_id_prompts.append(id_prompt.replace(key,character_dict[key]))
|
||
|
||
return character_index_dict,invert_character_index_dict,replace_id_prompts
|
||
|
||
def get_cur_id_list(real_prompt,character_dict,character_index_dict):
|
||
list_arr = []
|
||
for keys in character_index_dict.keys():
|
||
if keys in real_prompt:
|
||
list_arr = list_arr + character_index_dict[keys]
|
||
real_prompt = real_prompt.replace(keys,character_dict[keys])
|
||
return list_arr,real_prompt
|
||
|
||
def process_original_prompt(character_dict,prompts,id_length,img_mode):
|
||
replace_prompts = []
|
||
character_index_dict = {}
|
||
invert_character_index_dict = {}
|
||
for ind,prompt in enumerate(prompts):
|
||
for key in character_dict.keys():
|
||
if key in prompt:
|
||
if key not in character_index_dict:
|
||
character_index_dict[key] = []
|
||
character_index_dict[key].append(ind)
|
||
if ind not in invert_character_index_dict:
|
||
invert_character_index_dict[ind] = []
|
||
invert_character_index_dict[ind].append(key)
|
||
cur_prompt = prompt
|
||
if ind in invert_character_index_dict:
|
||
for key in invert_character_index_dict[ind]:
|
||
cur_prompt = cur_prompt.replace(key,character_dict[key] + " ")
|
||
replace_prompts.append(cur_prompt)
|
||
ref_index_dict = {}
|
||
ref_totals = []
|
||
for character_key in character_index_dict.keys():
|
||
if character_key not in character_index_dict:
|
||
raise f"{character_key} not have prompt description, please remove it"
|
||
index_list = character_index_dict[character_key]
|
||
index_list = [index for index in index_list if len(invert_character_index_dict[index]) == 1]
|
||
# if not img_mode:
|
||
# if len(index_list) < id_length:
|
||
# raise f"{character_key} not have enough prompt description, need no less than {id_length}, but you give {len(index_list)}"
|
||
ref_index_dict[character_key] = index_list[:id_length]
|
||
ref_totals = ref_totals + index_list[:id_length]
|
||
return character_index_dict,invert_character_index_dict,replace_prompts,ref_index_dict,ref_totals
|
||
|
||
#character_index_dict：{'[Taylor]': [0, 3], '[sam]': [1, 2]},if 1 role {'[Taylor]': [0, 1, 2]}
|
||
#invert_character_index_dict:{0: ['[Taylor]'], 1: ['[sam]'], 2: ['[sam]'], 3: ['[Taylor]']},if 1 role {0: ['[Taylor]'], 1: ['[Taylor]'], 2: ['[Taylor]']}
|
||
#ref_indexs_dict:{'[Taylor]': [0, 3], '[sam]': [1, 2]},if 1 role {'[Taylor]': [0]}
|
||
#ref_totals: [0, 3, 1, 2] if 1 role [0]
|
||
|
||
|
||
def process_original_text(role_dict,prompts):
|
||
'''
|
||
|
||
Args:
|
||
role_dict: # character_dict:{'[Taylor]': ' a woman img, wearing a white T-shirt, blue loose hair.'},character_list:['[Taylor]'] 1role
|
||
#character_dict:{'[Taylor]': ' a woman img, wearing a white T-shirt, blue loose hair.', '[Lecun]': ' a man img,wearing a suit,black hair.'} 2role
|
||
prompts: #[ list of scense txt]
|
||
|
||
Returns:
|
||
1 role
|
||
role_index_dict:{'[Taylor]': [0, 1]},
|
||
invert_role_index_dict:{0: ['[Taylor]'], 1: ['[Taylor]']},
|
||
replace_prompts:[' a woman img, wearing a white T-shirt, blue loose hair. wake up in the bed ;', ' a woman img, wearing a white T-shirt, blue loose hair. have breakfast by the window;'],
|
||
ref_role_index_dict:{'[Taylor]': [0]},
|
||
ref_role_totals:[0]
|
||
2 role
|
||
role_index_dict:{'[Taylor]': [0, 1], '[Lecun]': [2, 3]},
|
||
invert_role_index_dict:{0: ['[Taylor]'], 1: ['[Taylor]'], 2: ['[Lecun]'], 3: ['[Lecun]']},
|
||
replace_prompts:[' a woman img, wearing a white T-shirt, blue loose hair. wake up in the bed ;', ' a woman img, wearing a white T-shirt, blue loose hair. have breakfast by the window;', ' a man img,wearing a suit,black hair. driving a car;', ' a man img,wearing a suit,black hair. is working.'],
|
||
ref_role_index_dict:{'[Taylor]': [0, 1], '[Lecun]': [2, 3]},
|
||
ref_role_totals:[0, 1, 2, 3]
|
||
'''
|
||
replace_prompts = []
|
||
role_index_dict = {}
|
||
invert_role_index_dict = {}
|
||
for ind,prompt in enumerate(prompts):
|
||
for key in role_dict.keys(): # [Taylor],[Lecun]
|
||
if key in prompt: # [Taylor] a woman img
|
||
if key not in role_index_dict:
|
||
role_index_dict[key] = [] #character_index_dict:{'[Taylor]': [], '[Lecun]': []},
|
||
role_index_dict[key].append(ind) #character_index_dict:{'[Taylor]': [0, 1], '[Lecun]': [2, 3]},
|
||
if ind not in invert_role_index_dict:
|
||
invert_role_index_dict[ind] = [] #invert_character_index_dict:{0: []},
|
||
invert_role_index_dict[ind].append(key) # invert_character_index_dict:{0: ['[Taylor]'], 1: ['[Taylor]'], 2: ['[Lecun]'], 3: ['[Lecun]']},
|
||
cur_prompt = prompt
|
||
if ind in invert_role_index_dict:
|
||
for key in invert_role_index_dict[ind]:
|
||
cur_prompt = cur_prompt.replace(key,role_dict[key] + " ") #去掉[Taylor]
|
||
replace_prompts.append(cur_prompt)
|
||
ref_role_index_dict = {}
|
||
ref_role_totals = []
|
||
id_len=len(role_dict)
|
||
for character_key in role_index_dict.keys():
|
||
if character_key not in role_index_dict:
|
||
raise f"{character_key} not have prompt description, please remove it"
|
||
index_list = role_index_dict[character_key] #[0, 1],[3, 4] 如果2是NC
|
||
index_list = [index for index in index_list if len(invert_role_index_dict[index]) == 1]
|
||
# if not img_mode:
|
||
# if len(index_list) < id_length:
|
||
# raise f"{character_key} not have enough prompt description, need no less than {id_length}, but you give {len(index_list)}"
|
||
ref_role_index_dict[character_key] = index_list[:id_len]
|
||
ref_role_totals = ref_role_totals + index_list[:id_len]
|
||
return role_index_dict,invert_role_index_dict,replace_prompts,ref_role_index_dict,ref_role_totals
|
||
|
||
#character_index_dict：{'[Taylor]': [0, 3], '[sam]': [1, 2]},if 1 role {'[Taylor]': [0, 1, 2]}
|
||
#invert_character_index_dict:{0: ['[Taylor]'], 1: ['[sam]'], 2: ['[sam]'], 3: ['[Taylor]']},if 1 role {0: ['[Taylor]'], 1: ['[Taylor]'], 2: ['[Taylor]']}
|
||
#ref_indexs_dict:{'[Taylor]': [0, 3], '[sam]': [1, 2]},if 1 role {'[Taylor]': [0]}
|
||
#ref_totals: [0, 3, 1, 2] if 1 role [0]
|
||
|
||
def get_ref_character(real_prompt,character_dict):
|
||
list_arr = []
|
||
for keys in character_dict.keys():
|
||
if keys in real_prompt:
|
||
list_arr = list_arr + [keys]
|
||
return list_arr
|