AI text to video model from scratch
1.0.0
Sora de OpenAI, Stable Video Diffusion de Stability AI y muchos otros modelos de texto a video que han surgido o aparecerán en el futuro se encuentran entre las tendencias de IA más populares en 2024, siguiendo los modelos de lenguaje grande (LLM). En este blog, crearemos un modelo de texto a video a pequeña escala desde cero . Introduciremos un mensaje de texto y nuestro modelo entrenado generará un vídeo basado en ese mensaje. Este blog cubrirá todo, desde la comprensión de los conceptos teóricos hasta la codificación de toda la arquitectura y la generación del resultado final.
Como no tengo una GPU sofisticada, codifiqué la arquitectura a pequeña escala. A continuación se muestra una comparación del tiempo necesario para entrenar el modelo en diferentes procesadores:
| Vídeos de formación | Épocas | UPC | GPU A10 | GPU T4 |
|---|---|---|---|---|
| 10K | 30 | más de 3 horas | 1 hora | 1 hora 42 minutos |
| 30K | 30 | más de 6 horas | 1 hora 30 | 2 h 30 |
| 100K | 30 | - | 3-4 horas | 5-6 horas |
Obviamente, ejecutarlo en una CPU llevará mucho más tiempo para entrenar el modelo. Si necesita probar rápidamente cambios en el código y ver resultados, la CPU no es la mejor opción. Recomiendo usar una GPU T4 de Colab o Kaggle para un entrenamiento más eficiente y rápido.
Aquí está el enlace del blog que le guía sobre cómo crear Difusión estable desde cero: Codificación de Difusión estable desde cero
Seguiremos un enfoque similar al aprendizaje automático tradicional o a los modelos de aprendizaje profundo que se entrenan en un conjunto de datos y luego se prueban en datos invisibles. En el contexto de la conversión de texto a vídeo, digamos que tenemos un conjunto de datos de entrenamiento de 100.000 vídeos de perros buscando pelotas y gatos persiguiendo ratones. Entrenaremos nuestro modelo para generar videos de un gato buscando una pelota o un perro persiguiendo un ratón.
Aunque estos conjuntos de datos de entrenamiento están fácilmente disponibles en Internet, la potencia computacional requerida es extremadamente alta. Por lo tanto, trabajaremos con un conjunto de datos de video de objetos en movimiento generados a partir de código Python.
Usaremos la arquitectura GAN (Generative Adversarial Networks) para crear nuestro modelo en lugar del modelo de difusión que usa OpenAI Sora. Intenté utilizar el modelo de difusión, pero falló debido a requisitos de memoria, que están más allá de mi capacidad. Las GAN, por otro lado, son más fáciles y rápidas de entrenar y probar.
Usaremos POO (Programación orientada a objetos), por lo que debe tener un conocimiento básico de ella junto con las redes neuronales. El conocimiento de GAN (Generative Adversarial Networks) no es obligatorio, ya que aquí cubriremos su arquitectura.
| Tema | Enlace |
|---|---|
| POO | Enlace de vídeo |
| Teoría de las redes neuronales | Enlace de vídeo |
| Arquitectura GAN | Enlace de vídeo |
| Conceptos básicos de Python | Enlace de vídeo |
Comprender la arquitectura GAN es importante porque gran parte de nuestra arquitectura depende de ella. Exploremos qué es, sus componentes y más.
Generative Adversarial Network (GAN) es un modelo de aprendizaje profundo en el que dos redes neuronales compiten: una crea nuevos datos (como imágenes o música) a partir de un conjunto de datos determinado y la otra intenta saber si los datos son reales o falsos. Este proceso continúa hasta que los datos generados no se pueden distinguir del original.
Generar imágenes : las GAN crean imágenes realistas a partir de indicaciones de texto o modifican imágenes existentes, como mejorar la resolución o agregar color a fotografías en blanco y negro.
Aumento de datos : generan datos sintéticos para entrenar otros modelos de aprendizaje automático, como la creación de datos de transacciones fraudulentas para sistemas de detección de fraude.
Completar la información faltante : las GAN pueden completar los datos faltantes, como generar imágenes del subsuelo a partir de mapas del terreno para aplicaciones de energía.
Generar modelos 3D : convierten imágenes 2D en modelos 3D, útiles en campos como el cuidado de la salud para crear imágenes realistas de órganos para la planificación quirúrgica.
Consta de dos redes neuronales profundas: el generador y el discriminador . Estas redes se entrenan juntas en una configuración de confrontación, donde una genera nuevos datos y la otra evalúa si los datos son reales o falsos.
Aquí hay una descripción general simplificada de cómo funciona GAN:
Análisis del conjunto de entrenamiento : el generador analiza el conjunto de entrenamiento para identificar los atributos de los datos, mientras que el discriminador analiza de forma independiente los mismos datos para conocer sus atributos.
Modificación de datos : el generador agrega ruido (cambios aleatorios) a algunos atributos de los datos.
Paso de datos : los datos modificados se pasan al discriminador.
Cálculo de probabilidad : el discriminador calcula la probabilidad de que los datos generados provengan del conjunto de datos original.
Bucle de retroalimentación : el discriminador proporciona retroalimentación al generador, guiándolo para reducir el ruido aleatorio en el siguiente ciclo.
Entrenamiento adversario : el generador intenta maximizar los errores del discriminador, mientras que el discriminador intenta minimizar sus propios errores. A través de muchas iteraciones de capacitación, ambas redes mejoran y evolucionan.
Estado de equilibrio : el entrenamiento continúa hasta que el discriminador ya no puede distinguir entre datos reales y sintetizados, lo que indica que el generador ha aprendido con éxito a producir datos realistas. En este punto, el proceso de formación está completo.
imagen de la guía de AWS
Expliquemos el modelo GAN con un ejemplo de traducción de imagen a imagen, centrándonos en modificar un rostro humano.
Imagen de entrada : la entrada es una imagen real de un rostro humano.
Modificación de atributos : el generador modifica los atributos del rostro, como agregar gafas de sol a los ojos.
Imágenes generadas : el generador crea un conjunto de imágenes con gafas de sol agregadas.
Tarea del discriminador : El discriminador recibe una mezcla de imágenes reales (personas con gafas de sol) e imágenes generadas (rostros donde se agregaron gafas de sol).
Evaluación : El discriminador intenta diferenciar entre imágenes reales y generadas.
Bucle de retroalimentación : si el discriminador identifica correctamente imágenes falsas, el generador ajusta sus parámetros para producir imágenes más convincentes. Si el generador engaña con éxito al discriminador, el discriminador actualiza sus parámetros para mejorar su detección.
A través de este proceso de confrontación, ambas redes mejoran continuamente. El generador mejora en la creación de imágenes realistas y el discriminador mejora en la identificación de falsificaciones hasta que se alcanza el equilibrio, donde el discriminador ya no puede distinguir entre las imágenes reales y las generadas. En este punto, GAN ha aprendido con éxito a producir modificaciones realistas.
Instalar las bibliotecas necesarias es el primer paso para construir nuestro modelo de texto a video.
pip install -r requirements.txtTrabajaremos con una variedad de bibliotecas de Python. Importémoslas.
# 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 base64Ahora que hemos importado todas nuestras bibliotecas, el siguiente paso es definir los datos de entrenamiento que usaremos para entrenar nuestra arquitectura GAN.
Necesitamos tener al menos 10.000 vídeos como datos de entrenamiento. ¿Por qué? Bueno, porque probé con números más pequeños y los resultados fueron muy pobres, prácticamente nada que ver. La siguiente gran pregunta es: ¿de qué tratan estos vídeos? Nuestro conjunto de datos de vídeos de entrenamiento consta de un círculo que se mueve en diferentes direcciones con diferentes movimientos. Entonces, codifiquémoslo y generemos 10,000 videos para ver cómo se ve.
# 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 = 10Después de configurar algunos parámetros básicos, a continuación debemos definir las indicaciones de texto de nuestro conjunto de datos de entrenamiento en función de los videos de entrenamiento que se generarán.
# 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
]Hemos definido varios movimientos de nuestro círculo usando estas indicaciones. Ahora, necesitamos codificar algunas ecuaciones matemáticas para mover ese círculo según las indicaciones.
# 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 )La función anterior se utiliza para mover nuestro círculo para cada cuadro según la dirección seleccionada. Solo necesitamos ejecutar un bucle encima hasta la cantidad de videos veces para generar todos los videos.
# 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 )Una vez que ejecute el código anterior, generará todo nuestro conjunto de datos de entrenamiento. Así es como se ve la estructura de nuestros archivos de conjuntos de datos de entrenamiento.
Cada carpeta de vídeo de formación contiene sus fotogramas junto con su mensaje de texto. Echemos un vistazo a la muestra de nuestro conjunto de datos de entrenamiento.
En nuestro conjunto de datos de entrenamiento, no hemos incluido el movimiento del círculo que se mueve hacia arriba y luego hacia la derecha . Usaremos esto como nuestro mensaje de prueba para evaluar nuestro modelo entrenado con datos invisibles.
Un punto más importante a tener en cuenta es que nuestros datos de entrenamiento contienen muchas muestras donde los objetos se alejan de la escena o aparecen parcialmente frente a la cámara, similar a lo que hemos observado en los videos de demostración de OpenAI Sora.
La razón para incluir tales muestras en nuestros datos de entrenamiento es probar si nuestro modelo puede mantener la coherencia cuando el círculo entra en escena desde la misma esquina sin romper su forma.
Ahora que se han generado nuestros datos de entrenamiento, necesitamos convertir los videos de entrenamiento a tensores, que son el tipo de datos principal utilizado en marcos de aprendizaje profundo como PyTorch. Además, realizar transformaciones como la normalización ayuda a mejorar la convergencia y la estabilidad de la arquitectura de entrenamiento al escalar los datos a un rango más pequeño.
Tenemos que codificar una clase de conjunto de datos para tareas de texto a video, que pueda leer fotogramas de video y sus mensajes de texto correspondientes desde el directorio del conjunto de datos de entrenamiento, haciéndolos disponibles para su uso en 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 proceder a codificar la arquitectura, necesitamos normalizar nuestros datos de entrenamiento. Usaremos un tamaño de lote de 16 y mezclaremos los datos para introducir más aleatoriedad.
# 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 )Es posible que haya visto en la arquitectura de transformadores donde el punto de partida es convertir nuestra entrada de texto en incrustación para su posterior procesamiento en atención de múltiples cabezales. De manera similar, aquí tenemos que codificar una capa de incrustación de texto basada en la cual se llevará a cabo el entrenamiento de la arquitectura GAN en nuestros datos de incrustación. y tensor de imágenes.
# 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 )El tamaño del vocabulario se basará en nuestros datos de entrenamiento, que calcularemos más adelante. El tamaño de incrustación será 10. Si trabaja con un conjunto de datos más grande, también puede usar su propia elección de modelo de incrustación disponible en Hugging Face.
Ahora que ya sabemos qué hace el generador en las GAN, codifiquemos esta capa y luego comprendamos su contenido.
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
