Create text2video_model.py

text2video_model.py

@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class TextEncoder(nn.Module):
+    def __init__(self, vocab_size, embed_dim, hidden_dim):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.transformer = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(embed_dim, nhead=8),
+            num_layers=6
+        )
+        
+    def forward(self, text):
+        x = self.embedding(text)
+        return self.transformer(x)
+
+class VideoGenerator(nn.Module):
+    def __init__(self, latent_dim, num_frames, frame_size):
+        super().__init__()
+        self.latent_dim = latent_dim
+        self.num_frames = num_frames
+        
+        self.generator = nn.Sequential(
+            nn.ConvTranspose3d(latent_dim, 512, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm3d(512),
+            nn.ReLU(),
+            nn.ConvTranspose3d(512, 256, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm3d(256),
+            nn.ReLU(),
+            nn.ConvTranspose3d(256, 128, kernel_size=4, stride=2, padding=1),
+            nn.BatchNorm3d(128),
+            nn.ReLU(),
+            nn.ConvTranspose3d(128, 3, kernel_size=4, stride=2, padding=1),
+            nn.Tanh()
+        )
+    
+    def forward(self, z):
+        return self.generator(z)
+
+class Text2VideoModel(nn.Module):
+    def __init__(self, vocab_size, embed_dim, latent_dim, num_frames, frame_size):
+        super().__init__()
+        self.text_encoder = TextEncoder(vocab_size, embed_dim, hidden_dim=512)
+        self.video_generator = VideoGenerator(latent_dim, num_frames, frame_size)
+        self.latent_mapper = nn.Linear(embed_dim, latent_dim * num_frames)
+        
+    def forward(self, text):
+        text_features = self.text_encoder(text)
+        latent_vector = self.latent_mapper(text_features.mean(dim=1))
+        latent_video = latent_vector.view(-1, self.video_generator.latent_dim, 1, 1, 1)
+        generated_video = self.video_generator(latent_video)
+        return generated_video