Spaces:
Runtime error
Runtime error
| import os | |
| import csv | |
| import torch | |
| import torch.nn as nn | |
| from torch.utils.data import Dataset, DataLoader | |
| import cv2 | |
| import numpy as np | |
| from wendas.models import ( | |
| TextEncoder, | |
| VideoGenerator, | |
| VideoDiscriminator | |
| ) | |
| # ---------------------------- | |
| # SETTINGS | |
| # ---------------------------- | |
| CSV_FILE = "video_captions.csv" | |
| OUTPUT_FILE = "video_generator.pt" | |
| EPOCHS = 20 | |
| BATCH_SIZE = 2 | |
| LEARNING_RATE = 0.0002 | |
| FRAMES = 8 | |
| IMAGE_SIZE = 64 | |
| NOISE_SIZE = 128 | |
| TEXT_SIZE = 128 | |
| DEVICE = torch.device( | |
| "cuda" if torch.cuda.is_available() else "cpu" | |
| ) | |
| # ---------------------------- | |
| # VIDEO DATASET | |
| # ---------------------------- | |
| class VideoDataset(Dataset): | |
| def __init__(self, csv_file): | |
| self.items = [] | |
| if not os.path.exists(csv_file): | |
| raise FileNotFoundError( | |
| f"Missing {csv_file}. Create it with columns: path,caption" | |
| ) | |
| with open(csv_file, "r", encoding="utf-8") as file: | |
| reader = csv.DictReader(file) | |
| for row in reader: | |
| path = row["path"] | |
| caption = row["caption"] | |
| if os.path.exists(path): | |
| self.items.append((path, caption)) | |
| else: | |
| print("Missing video:", path) | |
| if not self.items: | |
| raise RuntimeError( | |
| "No usable videos found. Check the paths in video_captions.csv." | |
| ) | |
| def __len__(self): | |
| return len(self.items) | |
| def __getitem__(self, index): | |
| path, caption = self.items[index] | |
| capture = cv2.VideoCapture(path) | |
| frames = [] | |
| while len(frames) < FRAMES: | |
| success, frame = capture.read() | |
| if not success: | |
| break | |
| frame = cv2.resize( | |
| frame, | |
| (IMAGE_SIZE, IMAGE_SIZE) | |
| ) | |
| frame = cv2.cvtColor( | |
| frame, | |
| cv2.COLOR_BGR2RGB | |
| ) | |
| frame = frame.astype(np.float32) / 127.5 - 1.0 | |
| frames.append(frame) | |
| capture.release() | |
| if not frames: | |
| raise RuntimeError(f"Could not read video: {path}") | |
| while len(frames) < FRAMES: | |
| frames.append(frames[-1]) | |
| frames = np.array(frames) | |
| # Frames: T, H, W, C | |
| # Output: C, T, H, W | |
| frames = torch.tensor(frames).permute(3, 0, 1, 2) | |
| return frames, caption | |
| # ---------------------------- | |
| # TEXT PROCESSING | |
| # ---------------------------- | |
| def make_vocabulary(dataset): | |
| words = {"<PAD>": 0, "<UNK>": 1} | |
| for _, caption in dataset.items: | |
| for word in caption.lower().split(): | |
| if word not in words: | |
| words[word] = len(words) | |
| return words | |
| def encode_caption(caption, vocabulary): | |
| tokens = [ | |
| vocabulary.get(word, vocabulary["<UNK>"]) | |
| for word in caption.lower().split() | |
| ] | |
| if not tokens: | |
| tokens = [vocabulary["<UNK>"]] | |
| return torch.tensor(tokens, dtype=torch.long) | |
| # ---------------------------- | |
| # TRAINING | |
| # ---------------------------- | |
| def main(): | |
| print("Wendas AI training starting...") | |
| print("Device:", DEVICE) | |
| dataset = VideoDataset(CSV_FILE) | |
| vocabulary = make_vocabulary(dataset) | |
| print("Videos found:", len(dataset)) | |
| print("Vocabulary size:", len(vocabulary)) | |
| # Models from wendas/models.py | |
| text_encoder = TextEncoder( | |
| vocab_size=len(vocabulary) | |
| ).to(DEVICE) | |
| generator = VideoGenerator().to(DEVICE) | |
| discriminator = VideoDiscriminator().to(DEVICE) | |
| optimizer_g = torch.optim.Adam( | |
| list(generator.parameters()) + | |
| list(text_encoder.parameters()), | |
| lr=LEARNING_RATE, | |
| betas=(0.5, 0.999) | |
| ) | |
| optimizer_d = torch.optim.Adam( | |
| discriminator.parameters(), | |
| lr=LEARNING_RATE, | |
| betas=(0.5, 0.999) | |
| ) | |
| loss_function = nn.BCELoss() | |
| for epoch in range(EPOCHS): | |
| generator.train() | |
| discriminator.train() | |
| text_encoder.train() | |
| total_g = 0 | |
| total_d = 0 | |
| for real_videos, captions in DataLoader( | |
| dataset, | |
| batch_size=BATCH_SIZE, | |
| shuffle=True, | |
| drop_last=True, | |
| collate_fn=lambda batch: ( | |
| torch.stack([item[0] for item in batch]), | |
| [item[1] for item in batch] | |
| ) | |
| ): | |
| real_videos = real_videos.to(DEVICE) | |
| # Encode captions | |
| encoded = [ | |
| encode_caption(caption, vocabulary) | |
| for caption in captions | |
| ] | |
| max_length = max(len(x) for x in encoded) | |
| padded = torch.zeros( | |
| len(encoded), | |
| max_length, | |
| dtype=torch.long | |
| ) | |
| for i, sequence in enumerate(encoded): | |
| padded[i, :len(sequence)] = sequence | |
| padded = padded.to(DEVICE) | |
| text_features = text_encoder(padded) | |
| batch_size = real_videos.size(0) | |
| real_labels = torch.ones( | |
| batch_size, 1, device=DEVICE | |
| ) | |
| fake_labels = torch.zeros( | |
| batch_size, 1, device=DEVICE | |
| ) | |
| # -------------------- | |
| # Train discriminator | |
| # -------------------- | |
| noise = torch.randn( | |
| batch_size, | |
| NOISE_SIZE, | |
| device=DEVICE | |
| ) | |
| with torch.no_grad(): | |
| fake_videos = generator( | |
| noise, | |
| text_features | |
| ) | |
| optimizer_d.zero_grad() | |
| real_output = discriminator( | |
| real_videos, | |
| text_features.detach() | |
| ) | |
| fake_output = discriminator( | |
| fake_videos.detach(), | |
| text_features.detach() | |
| ) | |
| loss_real = loss_function( | |
| real_output, | |
| real_labels | |
| ) | |
| loss_fake = loss_function( | |
| fake_output, | |
| fake_labels | |
| ) | |
| loss_d = (loss_real + loss_fake) / 2 | |
| loss_d.backward() | |
| optimizer_d.step() | |
| # -------------------- | |
| # Train generator | |
| # -------------------- | |
| optimizer_g.zero_grad() | |
| noise = torch.randn( | |
| batch_size, | |
| NOISE_SIZE, | |
| device=DEVICE | |
| ) | |
| text_features = text_encoder(padded) | |
| generated_videos = generator( | |
| noise, | |
| text_features | |
| ) | |
| generated_output = discriminator( | |
| generated_videos, | |
| text_features | |
| ) | |
| loss_g = loss_function( | |
| generated_output, | |
| real_labels | |
| ) | |
| loss_g.backward() | |
| optimizer_g.step() | |
| total_g += loss_g.item() | |
| total_d += loss_d.item() | |
| print( | |
| f"Epoch {epoch + 1}/{EPOCHS} | " | |
| f"Generator: {total_g:.4f} | " | |
| f"Discriminator: {total_d:.4f}" | |
| ) | |
| # Save progress after every epoch | |
| torch.save( | |
| generator.state_dict(), | |
| OUTPUT_FILE | |
| ) | |
| print("Training finished!") | |
| print("Saved:", OUTPUT_FILE) | |
| if __name__ == "__main__": | |
| main() | |