117 lines
4.1 KiB
Python
117 lines
4.1 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
import torch.nn.functional as F
|
|
def get_generator(seed, device):
|
|
|
|
if seed is not None:
|
|
if isinstance(seed, list):
|
|
generator = [torch.Generator(device).manual_seed(seed_item) for seed_item in seed]
|
|
else:
|
|
generator = torch.Generator(device).manual_seed(seed)
|
|
else:
|
|
generator = None
|
|
|
|
return generator
|
|
|
|
def is_torch2_available():
|
|
return hasattr(F, "scaled_dot_product_attention")
|
|
|
|
# https://github.com/tencent-ailab/IP-Adapter/issues/54
|
|
# import cv2
|
|
# import numpy as np
|
|
# import insightface
|
|
# from insightface.app import FaceAnalysis
|
|
# from insightface.data import get_image as ins_get_image
|
|
# from insightface.utils import face_align
|
|
|
|
# app = FaceAnalysis(providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
|
|
# app.prepare(ctx_id=0, det_size=(640, 640))
|
|
# img = cv2.imread("person.png")
|
|
|
|
# faces = app.get(img)
|
|
# norm_face = face_align.norm_crop(img, landmark=faces[0].kps, image_size=224)
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
class MultiHeadAttention(nn.Module):
|
|
def __init__(self, d_model, num_heads):
|
|
super(MultiHeadAttention, self).__init__()
|
|
assert d_model % num_heads == 0, f'd_model={d_model}, numheads={num_heads}'
|
|
self.num_heads = num_heads
|
|
self.head_dim = d_model // num_heads
|
|
self.W_q = nn.Linear(d_model, d_model)
|
|
self.W_k = nn.Linear(d_model, d_model)
|
|
self.W_v = nn.Linear(d_model, d_model)
|
|
self.W_o = nn.Linear(d_model, d_model)
|
|
|
|
def split_heads(self, x, batch_size):
|
|
x = x.view(batch_size, -1, self.num_heads, self.head_dim)
|
|
return x.permute(0, 2, 1, 3)
|
|
|
|
def forward(self, query, key, value):
|
|
batch_size = query.shape[0]
|
|
Q = self.split_heads(self.W_q(query), batch_size)
|
|
K = self.split_heads(self.W_k(key), batch_size)
|
|
V = self.split_heads(self.W_v(value), batch_size)
|
|
|
|
scores = torch.matmul(Q, K.permute(0, 1, 3, 2)) / (self.head_dim**0.5)
|
|
attention_weights = torch.nn.functional.softmax(scores, dim=-1)
|
|
|
|
x = torch.matmul(attention_weights, V)
|
|
x = x.permute(0, 2, 1, 3).contiguous().view(batch_size, -1, self.num_heads * self.head_dim)
|
|
x = self.W_o(x)
|
|
return x
|
|
|
|
class TransformerLayer(nn.Module):
|
|
def __init__(self, d_model, num_heads):
|
|
super(TransformerLayer, self).__init__()
|
|
self.multi_head_attention = MultiHeadAttention(d_model, num_heads)
|
|
self.feed_forward = nn.Sequential(
|
|
nn.Linear(d_model, 4*d_model),
|
|
nn.ReLU(),
|
|
nn.Linear(4*d_model, d_model)
|
|
)
|
|
self.layer_norm1 = nn.LayerNorm(d_model)
|
|
self.layer_norm2 = nn.LayerNorm(d_model)
|
|
|
|
def forward(self, x):
|
|
attention_output = self.multi_head_attention(x, x, x)
|
|
x = self.layer_norm1(x + attention_output)
|
|
feed_forward_output = self.feed_forward(x)
|
|
x = self.layer_norm2(x + feed_forward_output)
|
|
return x
|
|
|
|
class Transformer(nn.Module):
|
|
def __init__(self, d_model, num_heads, num_layers):
|
|
super(Transformer, self).__init__()
|
|
self.num_layers = num_layers
|
|
self.embedding = nn.Linear(d_model, d_model)
|
|
self.layers = nn.ModuleList([TransformerLayer(d_model, num_heads) for _ in range(num_layers)])
|
|
|
|
def forward(self, x):
|
|
x = self.embedding(x)
|
|
for _ in range(self.num_layers):
|
|
x = self.layers[_](x)
|
|
return x
|
|
|
|
# Example usage:
|
|
# input_dim = 512 # Dimension of the input tensor
|
|
# num_heads = 8 # Number of attention heads
|
|
# num_layers = 3 # Number of transformer layers
|
|
|
|
# # Create an instance of the Transformer model
|
|
# model = Transformer(input_dim, num_heads, num_layers)
|
|
|
|
# # Test the model with a random input tensor (batch_size, sequence_length, d_model)
|
|
# batch_size, sequence_length = 16, 20
|
|
# input_tensor = torch.randn(batch_size, sequence_length, input_dim)
|
|
# output = model(input_tensor)
|
|
|
|
# print("Input shape:", input_tensor.shape)
|
|
# print("Output shape:", output.shape)
|
|
|
|
|