AI text to video model from scratch
1.0.0
Sora da OpenAI, Stable Video Diffusion da Stability AI e muitos outros modelos de texto para vídeo que foram lançados ou aparecerão no futuro estão entre as tendências de IA mais populares em 2024, seguindo modelos de linguagem grande (LLMs). Neste blog, construiremos um modelo de texto para vídeo em pequena escala do zero . Inseriremos um prompt de texto e nosso modelo treinado gerará um vídeo com base nesse prompt. Este blog abordará tudo, desde a compreensão dos conceitos teóricos até a codificação de toda a arquitetura e geração do resultado final.
Como não tenho uma GPU sofisticada, codifiquei a arquitetura em pequena escala. Aqui está uma comparação do tempo necessário para treinar o modelo em diferentes processadores:
| Vídeos de treinamento | Épocas | CPU | GPU A10 | GPU T4 |
|---|---|---|---|---|
| 10K | 30 | mais de 3 horas | 1 hora | 1h42m |
| 30 mil | 30 | mais de 6 horas | 1h30 | 2h30 |
| 100 mil | 30 | - | 3-4 horas | 5-6 horas |
A execução em uma CPU obviamente levará muito mais tempo para treinar o modelo. Se você precisar testar rapidamente as alterações no código e ver os resultados, a CPU não é a melhor escolha. Recomendo usar uma GPU T4 da Colab ou Kaggle para um treinamento mais eficiente e rápido.
Aqui está o link do blog que orienta você sobre como criar uma difusão estável do zero: Codificando a difusão estável do zero
Seguiremos uma abordagem semelhante ao aprendizado de máquina tradicional ou modelos de aprendizado profundo que são treinados em um conjunto de dados e depois testados em dados invisíveis. No contexto de conversão de texto em vídeo, digamos que temos um conjunto de dados de treinamento de 100 mil vídeos de cães buscando bolas e gatos perseguindo ratos. Treinaremos nosso modelo para gerar vídeos de um gato buscando uma bola ou de um cachorro perseguindo um rato.
Embora tais conjuntos de dados de treinamento estejam facilmente disponíveis na Internet, o poder computacional necessário é extremamente alto. Portanto, trabalharemos com um conjunto de dados de vídeo de objetos em movimento gerados a partir de código Python.
Usaremos a arquitetura GAN (Generative Adversarial Networks) para criar nosso modelo em vez do modelo de difusão que OpenAI Sora usa. Tentei usar o modelo de difusão, mas ele travou devido a requisitos de memória, que estão além da minha capacidade. Os GANs, por outro lado, são mais fáceis e rápidos de treinar e testar.
Estaremos usando OOP (Programação Orientada a Objetos), então você deve ter um conhecimento básico dela junto com redes neurais. O conhecimento de GANs (Redes Adversariais Generativas) não é obrigatório, pois abordaremos sua arquitetura aqui.
| Tópico | Link |
|---|---|
| POO | Link do vídeo |
| Teoria das Redes Neurais | Link do vídeo |
| Arquitetura GAN | Link do vídeo |
| Noções básicas de Python | Link do vídeo |
Compreender a arquitetura GAN é importante porque grande parte da nossa arquitetura depende dela. Vamos explorar o que é, seus componentes e muito mais.
Generative Adversarial Network (GAN) é um modelo de aprendizagem profunda onde duas redes neurais competem: uma cria novos dados (como imagens ou música) a partir de um determinado conjunto de dados e a outra tenta dizer se os dados são reais ou falsos. Este processo continua até que os dados gerados sejam indistinguíveis do original.
Gerar imagens : GANs criam imagens realistas a partir de prompts de texto ou modificam imagens existentes, como melhorar a resolução ou adicionar cor a fotos em preto e branco.
Aumento de dados : Geram dados sintéticos para treinar outros modelos de aprendizado de máquina, como a criação de dados de transações fraudulentas para sistemas de detecção de fraudes.
Completar informações ausentes : GANs podem preencher dados ausentes, como gerar imagens subterrâneas a partir de mapas de terreno para aplicações de energia.
Gerar modelos 3D : eles convertem imagens 2D em modelos 3D, úteis em áreas como saúde para criar imagens realistas de órgãos para planejamento cirúrgico.
Consiste em duas redes neurais profundas: o gerador e o discriminador . Essas redes treinam juntas em uma configuração adversária, onde uma gera novos dados e a outra avalia se os dados são reais ou falsos.
Aqui está uma visão geral simplificada de como o GAN funciona:
Análise do conjunto de treinamento : o gerador analisa o conjunto de treinamento para identificar atributos de dados, enquanto o discriminador analisa independentemente os mesmos dados para aprender seus atributos.
Modificação de dados : O gerador adiciona ruído (alterações aleatórias) a alguns atributos dos dados.
Passagem de dados : Os dados modificados são então passados para o discriminador.
Cálculo de probabilidade : O discriminador calcula a probabilidade de os dados gerados serem do conjunto de dados original.
Loop de Feedback : O discriminador fornece feedback ao gerador, orientando-o para reduzir o ruído aleatório no próximo ciclo.
Treinamento Adversarial : O gerador tenta maximizar os erros do discriminador, enquanto o discriminador tenta minimizar seus próprios erros. Através de muitas iterações de treinamento, ambas as redes melhoram e evoluem.
Estado de equilíbrio : O treinamento continua até que o discriminador não consiga mais distinguir entre dados reais e sintetizados, indicando que o gerador aprendeu com sucesso a produzir dados realistas. Neste ponto, o processo de treinamento está concluído.
imagem do guia aws
Vamos explicar o modelo GAN com um exemplo de tradução imagem para imagem, com foco na modificação de um rosto humano.
Imagem de entrada : A entrada é uma imagem real de um rosto humano.
Modificação de atributos : O gerador modifica atributos do rosto, como adicionar óculos de sol aos olhos.
Imagens Geradas : O gerador cria um conjunto de imagens com óculos de sol adicionados.
Tarefa do Discriminador : O discriminador recebe uma mistura de imagens reais (pessoas com óculos de sol) e imagens geradas (rostos onde foram adicionados óculos de sol).
Avaliação : O discriminador tenta diferenciar entre imagens reais e geradas.
Loop de feedback : Se o discriminador identificar corretamente as imagens falsas, o gerador ajusta seus parâmetros para produzir imagens mais convincentes. Se o gerador enganar o discriminador com sucesso, o discriminador atualizará seus parâmetros para melhorar sua detecção.
Através deste processo adversário, ambas as redes melhoram continuamente. O gerador fica melhor na criação de imagens realistas, e o discriminador fica melhor na identificação de falsificações até que o equilíbrio seja alcançado, onde o discriminador não consegue mais diferenciar entre imagens reais e geradas. Neste ponto, o GAN aprendeu com sucesso a produzir modificações realistas.
Instalar as bibliotecas necessárias é o primeiro passo na construção do nosso modelo de texto para vídeo.
pip install -r requirements.txtEstaremos trabalhando com uma variedade de bibliotecas Python. Vamos importá-las.
# Operating System module for interacting with the operating system
import os
# Module for generating random numbers
import random
# Module for numerical operations
import numpy as np
# OpenCV library for image processing
import cv2
# Python Imaging Library for image processing
from PIL import Image , ImageDraw , ImageFont
# PyTorch library for deep learning
import torch
# Dataset class for creating custom datasets in PyTorch
from torch . utils . data import Dataset
# Module for image transformations
import torchvision . transforms as transforms
# Neural network module in PyTorch
import torch . nn as nn
# Optimization algorithms in PyTorch
import torch . optim as optim
# Function for padding sequences in PyTorch
from torch . nn . utils . rnn import pad_sequence
# Function for saving images in PyTorch
from torchvision . utils import save_image
# Module for plotting graphs and images
import matplotlib . pyplot as plt
# Module for displaying rich content in IPython environments
from IPython . display import clear_output , display , HTML
# Module for encoding and decoding binary data to text
import base64Agora que importamos todas as nossas bibliotecas, a próxima etapa é definir nossos dados de treinamento que usaremos para treinar nossa arquitetura GAN.
Precisamos ter pelo menos 10.000 vídeos como dados de treinamento. Por que? Bom, porque testei com números menores e os resultados foram muito ruins, praticamente nada para ver. A próxima grande questão é: sobre o que são esses vídeos? Nosso conjunto de dados de vídeo de treinamento consiste em um círculo movendo-se em direções diferentes com movimentos diferentes. Então, vamos codificá-lo e gerar 10.000 vídeos para ver como fica.
# Create a directory named 'training_dataset'
os . makedirs ( 'training_dataset' , exist_ok = True )
# Define the number of videos to generate for the dataset
num_videos = 10000
# Define the number of frames per video (1 Second Video)
frames_per_video = 10
# Define the size of each image in the dataset
img_size = ( 64 , 64 )
# Define the size of the shapes (Circle)
shape_size = 10após definir alguns parâmetros básicos, precisamos definir os prompts de texto do nosso conjunto de dados de treinamento com base em quais vídeos de treinamento serão gerados.
# Define text prompts and corresponding movements for circles
prompts_and_movements = [
( "circle moving down" , "circle" , "down" ), # Move circle downward
( "circle moving left" , "circle" , "left" ), # Move circle leftward
( "circle moving right" , "circle" , "right" ), # Move circle rightward
( "circle moving diagonally up-right" , "circle" , "diagonal_up_right" ), # Move circle diagonally up-right
( "circle moving diagonally down-left" , "circle" , "diagonal_down_left" ), # Move circle diagonally down-left
( "circle moving diagonally up-left" , "circle" , "diagonal_up_left" ), # Move circle diagonally up-left
( "circle moving diagonally down-right" , "circle" , "diagonal_down_right" ), # Move circle diagonally down-right
( "circle rotating clockwise" , "circle" , "rotate_clockwise" ), # Rotate circle clockwise
( "circle rotating counter-clockwise" , "circle" , "rotate_counter_clockwise" ), # Rotate circle counter-clockwise
( "circle shrinking" , "circle" , "shrink" ), # Shrink circle
( "circle expanding" , "circle" , "expand" ), # Expand circle
( "circle bouncing vertically" , "circle" , "bounce_vertical" ), # Bounce circle vertically
( "circle bouncing horizontally" , "circle" , "bounce_horizontal" ), # Bounce circle horizontally
( "circle zigzagging vertically" , "circle" , "zigzag_vertical" ), # Zigzag circle vertically
( "circle zigzagging horizontally" , "circle" , "zigzag_horizontal" ), # Zigzag circle horizontally
( "circle moving up-left" , "circle" , "up_left" ), # Move circle up-left
( "circle moving down-right" , "circle" , "down_right" ), # Move circle down-right
( "circle moving down-left" , "circle" , "down_left" ), # Move circle down-left
]Definimos vários movimentos do nosso círculo usando essas instruções. Agora, precisamos codificar algumas equações matemáticas para mover esse círculo com base nas instruções.
# defining function to create image with moving shape
def create_image_with_moving_shape ( size , frame_num , shape , direction ):
# Create a new RGB image with specified size and white background
img = Image . new ( 'RGB' , size , color = ( 255 , 255 , 255 ))
# Create a drawing context for the image
draw = ImageDraw . Draw ( img )
# Calculate the center coordinates of the image
center_x , center_y = size [ 0 ] // 2 , size [ 1 ] // 2
# Initialize position with center for all movements
position = ( center_x , center_y )
# Define a dictionary mapping directions to their respective position adjustments or image transformations
direction_map = {
# Adjust position downwards based on frame number
"down" : ( 0 , frame_num * 5 % size [ 1 ]),
# Adjust position to the left based on frame number
"left" : ( - frame_num * 5 % size [ 0 ], 0 ),
# Adjust position to the right based on frame number
"right" : ( frame_num * 5 % size [ 0 ], 0 ),
# Adjust position diagonally up and to the right
"diagonal_up_right" : ( frame_num * 5 % size [ 0 ], - frame_num * 5 % size [ 1 ]),
# Adjust position diagonally down and to the left
"diagonal_down_left" : ( - frame_num * 5 % size [ 0 ], frame_num * 5 % size [ 1 ]),
# Adjust position diagonally up and to the left
"diagonal_up_left" : ( - frame_num * 5 % size [ 0 ], - frame_num * 5 % size [ 1 ]),
# Adjust position diagonally down and to the right
"diagonal_down_right" : ( frame_num * 5 % size [ 0 ], frame_num * 5 % size [ 1 ]),
# Rotate the image clockwise based on frame number
"rotate_clockwise" : img . rotate ( frame_num * 10 % 360 , center = ( center_x , center_y ), fillcolor = ( 255 , 255 , 255 )),
# Rotate the image counter-clockwise based on frame number
"rotate_counter_clockwise" : img . rotate ( - frame_num * 10 % 360 , center = ( center_x , center_y ), fillcolor = ( 255 , 255 , 255 )),
# Adjust position for a bouncing effect vertically
"bounce_vertical" : ( 0 , center_y - abs ( frame_num * 5 % size [ 1 ] - center_y )),
# Adjust position for a bouncing effect horizontally
"bounce_horizontal" : ( center_x - abs ( frame_num * 5 % size [ 0 ] - center_x ), 0 ),
# Adjust position for a zigzag effect vertically
"zigzag_vertical" : ( 0 , center_y - frame_num * 5 % size [ 1 ]) if frame_num % 2 == 0 else ( 0 , center_y + frame_num * 5 % size [ 1 ]),
# Adjust position for a zigzag effect horizontally
"zigzag_horizontal" : ( center_x - frame_num * 5 % size [ 0 ], center_y ) if frame_num % 2 == 0 else ( center_x + frame_num * 5 % size [ 0 ], center_y ),
# Adjust position upwards and to the right based on frame number
"up_right" : ( frame_num * 5 % size [ 0 ], - frame_num * 5 % size [ 1 ]),
# Adjust position upwards and to the left based on frame number
"up_left" : ( - frame_num * 5 % size [ 0 ], - frame_num * 5 % size [ 1 ]),
# Adjust position downwards and to the right based on frame number
"down_right" : ( frame_num * 5 % size [ 0 ], frame_num * 5 % size [ 1 ]),
# Adjust position downwards and to the left based on frame number
"down_left" : ( - frame_num * 5 % size [ 0 ], frame_num * 5 % size [ 1 ])
}
# Check if direction is in the direction map
if direction in direction_map :
# Check if the direction maps to a position adjustment
if isinstance ( direction_map [ direction ], tuple ):
# Update position based on the adjustment
position = tuple ( np . add ( position , direction_map [ direction ]))
else : # If the direction maps to an image transformation
# Update the image based on the transformation
img = direction_map [ direction ]
# Return the image as a numpy array
return np . array ( img )A função acima é usada para mover nosso círculo para cada quadro com base na direção selecionada. Só precisamos executar um loop em cima dele até o número de vídeos para gerar todos os vídeos.
# Iterate over the number of videos to generate
for i in range ( num_videos ):
# Randomly choose a prompt and movement from the predefined list
prompt , shape , direction = random . choice ( prompts_and_movements )
# Create a directory for the current video
video_dir = f'training_dataset/video_ { i } '
os . makedirs ( video_dir , exist_ok = True )
# Write the chosen prompt to a text file in the video directory
with open ( f' { video_dir } /prompt.txt' , 'w' ) as f :
f . write ( prompt )
# Generate frames for the current video
for frame_num in range ( frames_per_video ):
# Create an image with a moving shape based on the current frame number, shape, and direction
img = create_image_with_moving_shape ( img_size , frame_num , shape , direction )
# Save the generated image as a PNG file in the video directory
cv2 . imwrite ( f' { video_dir } /frame_ { frame_num } .png' , img )Depois de executar o código acima, ele gerará todo o nosso conjunto de dados de treinamento. Esta é a aparência da estrutura de nossos arquivos de conjunto de dados de treinamento.
Cada pasta de vídeo de treinamento contém seus quadros junto com seu prompt de texto. Vamos dar uma olhada no exemplo do nosso conjunto de dados de treinamento.
Em nosso conjunto de dados de treinamento, não incluímos o movimento do círculo subindo e depois para a direita . Usaremos isso como nosso prompt de teste para avaliar nosso modelo treinado em dados não vistos.
Mais um ponto importante a ser observado é que nossos dados de treinamento contêm muitas amostras onde objetos se afastam da cena ou aparecem parcialmente na frente da câmera, semelhante ao que observamos nos vídeos de demonstração do OpenAI Sora.
A razão para incluir tais amostras em nossos dados de treinamento é testar se nosso modelo pode manter a consistência quando o círculo entra em cena desde o canto, sem quebrar sua forma.
Agora que nossos dados de treinamento foram gerados, precisamos converter os vídeos de treinamento em tensores, que são o principal tipo de dados usado em estruturas de aprendizado profundo como o PyTorch. Além disso, realizar transformações como a normalização ajuda a melhorar a convergência e a estabilidade da arquitetura de treinamento, dimensionando os dados para um intervalo menor.
Temos que codificar uma classe de conjunto de dados para tarefas de texto para vídeo, que pode ler quadros de vídeo e seus prompts de texto correspondentes do diretório do conjunto de dados de treinamento, disponibilizando-os para uso no PyTorch.
# Define a dataset class inheriting from torch.utils.data.Dataset
class TextToVideoDataset ( Dataset ):
def __init__ ( self , root_dir , transform = None ):
# Initialize the dataset with root directory and optional transform
self . root_dir = root_dir
self . transform = transform
# List all subdirectories in the root directory
self . video_dirs = [ os . path . join ( root_dir , d ) for d in os . listdir ( root_dir ) if os . path . isdir ( os . path . join ( root_dir , d ))]
# Initialize lists to store frame paths and corresponding prompts
self . frame_paths = []
self . prompts = []
# Loop through each video directory
for video_dir in self . video_dirs :
# List all PNG files in the video directory and store their paths
frames = [ os . path . join ( video_dir , f ) for f in os . listdir ( video_dir ) if f . endswith ( '.png' )]
self . frame_paths . extend ( frames )
# Read the prompt text file in the video directory and store its content
with open ( os . path . join ( video_dir , 'prompt.txt' ), 'r' ) as f :
prompt = f . read (). strip ()
# Repeat the prompt for each frame in the video and store in prompts list
self . prompts . extend ([ prompt ] * len ( frames ))
# Return the total number of samples in the dataset
def __len__ ( self ):
return len ( self . frame_paths )
# Retrieve a sample from the dataset given an index
def __getitem__ ( self , idx ):
# Get the path of the frame corresponding to the given index
frame_path = self . frame_paths [ idx ]
# Open the image using PIL (Python Imaging Library)
image = Image . open ( frame_path )
# Get the prompt corresponding to the given index
prompt = self . prompts [ idx ]
# Apply transformation if specified
if self . transform :
image = self . transform ( image )
# Return the transformed image and the prompt
return image , promptAntes de prosseguir com a codificação da arquitetura, precisamos normalizar nossos dados de treinamento. Usaremos um tamanho de lote de 16 e embaralharemos os dados para introduzir mais aleatoriedade.
# Define a set of transformations to be applied to the data
transform = transforms . Compose ([
transforms . ToTensor (), # Convert PIL Image or numpy.ndarray to tensor
transforms . Normalize (( 0.5 ,), ( 0.5 ,)) # Normalize image with mean and standard deviation
])
# Load the dataset using the defined transform
dataset = TextToVideoDataset ( root_dir = 'training_dataset' , transform = transform )
# Create a dataloader to iterate over the dataset
dataloader = torch . utils . data . DataLoader ( dataset , batch_size = 16 , shuffle = True )Você deve ter visto na arquitetura do transformador onde o ponto de partida é converter nossa entrada de texto em incorporação para processamento adicional em atenção multi-head semelhante aqui, temos que codificar uma camada de incorporação de texto com base na qual o treinamento da arquitetura GAN ocorrerá em nossos dados de incorporação e tensor de imagens.
# Define a class for text embedding
class TextEmbedding ( nn . Module ):
# Constructor method with vocab_size and embed_size parameters
def __init__ ( self , vocab_size , embed_size ):
# Call the superclass constructor
super ( TextEmbedding , self ). __init__ ()
# Initialize embedding layer
self . embedding = nn . Embedding ( vocab_size , embed_size )
# Define the forward pass method
def forward ( self , x ):
# Return embedded representation of input
return self . embedding ( x )O tamanho do vocabulário será baseado em nossos dados de treinamento, que calcularemos posteriormente. O tamanho de incorporação será 10. Se estiver trabalhando com um conjunto de dados maior, você também pode usar sua própria escolha de modelo de incorporação disponível no Hugging Face.
Agora que já sabemos o que o gerador faz nas GANs, vamos codificar essa camada e então entender seu conteúdo.
class Generator ( nn . Module ):
def __init__ ( self , text_embed_size ):
super ( Generator , self ). __init__ ()
# Fully connected layer that takes noise and text embedding as input
self . fc1 = nn . Linear ( 100 + text_embed_size , 256 * 8 * 8 )
# Transposed convolutional layers to upsample the input
self . deconv1 = nn . ConvTranspose2d ( 256 , 128 , 4 , 2 , 1 )
self . deconv2 = nn . ConvTranspose2d ( 128 , 64 , 4 , 2 , 1 )
self . deconv3 = nn . ConvTranspose2d ( 64 , 3 , 4 , 2 , 1 ) # Output has 3 channels for RGB images
# Activation functions
self . relu = nn . ReLU ( True ) # ReLU activation function
self . tanh = nn . Tanh () # Tanh activation function for final output
def forward ( self , noise , text_embed ):
# Concatenate noise and text embedding along the channel dimension
x = torch . cat (( noise , text_embed ), dim = 1 )
# Fully connected layer followed by reshaping to 4D tensor
x = self . fc1 ( x ). view ( - 1 , 256 , 8 , 8 )
# Upsampling through transposed convolution layers with ReLU activation
x = self . relu ( self . deconv1 ( x ))
x = self . relu ( self . deconv2 ( x ))
# Final layer with Tanh activation to ensure output values are between -1 and 1 (for images)
x
