Update train_full_pipeline.py: full training pipeline with real datasets

train_full_pipeline.py

@@ -0,0 +1,690 @@
+#!/usr/bin/env python3
+"""
+Full Pipeline: Train VQ-VAE → Tokenize OpenVid → Train LLM → Push to EeshaAI/zeeb
+=================================================================================
+Runs on HuggingFace Spaces (free CPU tier, 16GB RAM).
+
+Phase 1: Train VQ-VAE on COCO 2017 images (118K real images, streaming)
+Phase 2: Stream 10K clips from OpenVid-1M → tokenize via trained VQ-VAE → save integers
+Phase 3: Fine-tune OLMo 2 1B with LoRA on 10K tokenized samples
+Phase 4: Push trained model to EeshaAI/zeeb
+"""
+
+import os
+import sys
+import json
+import time
+import gc
+import threading
+import traceback
+import numpy as np
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset, IterableDataset
+
+# ============================================================================
+# CONFIGURATION
+# ============================================================================
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+REPO_ID = "eeshaAI/zeeb"
+MODEL_NAME = "allenai/OLMo-2-0425-1B-Instruct"
+CODEBOOK_SIZE = 1024
+CODEBOOK_DIM = 256
+LATENT_DIM = 256
+VIDEO_START = "<video_start>"
+VIDEO_END = "<video_end>"
+VIDEO_PAD = "<video_pad>"
+
+# VQ-VAE training
+VQ_VAE_EPOCHS = 5
+VQ_VAE_LR = 1e-3
+VQ_VAE_BATCH = 32
+VQ_VAE_IMG_SIZE = 128  # resize images to 128x128
+
+# Dataset preparation
+NUM_OPENVID_CLIPS = 10000
+TOKENS_PER_CLIP = 128  # number of visual tokens per video clip
+
+# LLM training
+NUM_EPOCHS = 3
+LORA_R = 4
+LORA_ALPHA = 8
+LORA_DROPOUT = 0.05
+LEARNING_RATE = 1e-4
+BATCH_SIZE = 1
+MAX_SEQ_LEN = 384
+GRADIENT_ACCUMULATION = 4
+
+LOG_FILE = "/tmp/pipeline_log.txt"
+
+
+# ============================================================================
+# LOGGING
+# ============================================================================
+class Logger:
+    def __init__(self, path):
+        self.path = path
+        self.lock = threading.Lock()
+        with open(path, "w") as f:
+            f.write("🚀 Zeeb Full Pipeline Starting...\n\n")
+
+    def log(self, msg):
+        with self.lock:
+            with open(self.path, "a") as f:
+                f.write(msg)
+                f.flush()
+            print(msg, end="", flush=True)
+
+
+# ============================================================================
+# VQ-VAE MODEL
+# ============================================================================
+class Encoder(nn.Module):
+    def __init__(self, in_channels=3, latent_dim=LATENT_DIM):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(in_channels, 64, 4, stride=2, padding=1),   # → 64x64
+            nn.ReLU(),
+            nn.Conv2d(64, 128, 4, stride=2, padding=1),            # → 32x32
+            nn.ReLU(),
+            nn.Conv2d(128, 256, 4, stride=2, padding=1),           # → 16x16
+            nn.ReLU(),
+            nn.Conv2d(256, latent_dim, 4, stride=2, padding=1),    # → 8x8
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class VectorQuantizer(nn.Module):
+    def __init__(self, codebook_size=CODEBOOK_SIZE, codebook_dim=CODEBOOK_DIM, commitment_cost=0.25):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.commitment_cost = commitment_cost
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+        self.codebook.weight.data.uniform_(-1.0 / codebook_size, 1.0 / codebook_size)
+
+    def forward(self, z):
+        # z: [B, H, W, C] (channels last)
+        B, H, W, C = z.shape
+        z_flat = z.reshape(-1, C)
+
+        # Find nearest codebook entry
+        dist = (z_flat.unsqueeze(1) - self.codebook.weight.unsqueeze(0)).pow(2).sum(-1)
+        indices = dist.argmin(dim=1)
+
+        z_q = self.codebook(indices).reshape(B, H, W, C)
+
+        # Losses
+        commitment_loss = F.mse_loss(z_flat, z_q.reshape(-1, C).detach())
+        codebook_loss = F.mse_loss(z_q.reshape(-1, C), z_flat.detach())
+        loss = codebook_loss + self.commitment_cost * commitment_loss
+
+        # Straight-through estimator
+        z_q_st = z + (z_q - z).detach()
+
+        return z_q_st, loss, indices.reshape(B, H, W)
+
+
+class Decoder(nn.Module):
+    def __init__(self, out_channels=3, latent_dim=LATENT_DIM):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.ConvTranspose2d(latent_dim, 256, 4, stride=2, padding=1),  # → 16x16
+            nn.ReLU(),
+            nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1),         # → 32x32
+            nn.ReLU(),
+            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1),          # → 64x64
+            nn.ReLU(),
+            nn.ConvTranspose2d(64, out_channels, 4, stride=2, padding=1), # → 128x128
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class VQVAE(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = Encoder()
+        self.quantizer = VectorQuantizer()
+        self.proj_in = nn.Linear(LATENT_DIM, CODEBOOK_DIM)
+        self.proj_out = nn.Linear(CODEBOOK_DIM, LATENT_DIM)
+        self.decoder = Decoder()
+
+    def forward(self, x):
+        z = self.encoder(x)                              # [B, C, H, W]
+        z = z.permute(0, 2, 3, 1)                        # [B, H, W, C]
+        z = self.proj_in(z)                               # [B, H, W, codebook_dim]
+        z_q, vq_loss, indices = self.quantizer(z)
+        z_q = self.proj_out(z_q)                          # [B, H, W, latent_dim]
+        z_q = z_q.permute(0, 3, 1, 2)                    # [B, C, H, W]
+        recon = self.decoder(z_q)
+        return recon, vq_loss, indices
+
+    def encode(self, x):
+        """Encode image to token indices."""
+        z = self.encoder(x)
+        z = z.permute(0, 2, 3, 1)
+        z = self.proj_in(z)
+        _, _, indices = self.quantizer(z)
+        return indices  # [B, H, W]
+
+    def decode_tokens(self, token_ids, grid_h=8, grid_w=8):
+        """Decode token IDs back to image."""
+        if isinstance(token_ids, list):
+            token_ids = torch.tensor(token_ids, dtype=torch.long)
+        token_ids = token_ids[:grid_h * grid_w]
+        if len(token_ids) < grid_h * grid_w:
+            token_ids = torch.cat([token_ids, torch.zeros(grid_h * grid_w - len(token_ids), dtype=torch.long)])
+
+        z_q = self.quantizer.codebook(token_ids)   # [H*W, D]
+        z_q = self.proj_out(z_q)                    # [H*W, latent_dim]
+        z_q = z_q.reshape(1, grid_h, grid_w, -1).permute(0, 3, 1, 2)
+        return self.decoder(z_q)
+
+
+# ============================================================================
+# PHASE 1: TRAIN VQ-VAE ON COCO IMAGES
+# ============================================================================
+def train_vq_vae(logger: Logger) -> VQVAE:
+    """Train VQ-VAE on COCO 2017 images (streaming, so no massive download)."""
+    logger.log("=" * 60 + "\n")
+    logger.log("PHASE 1: Training VQ-VAE on COCO 2017 images\n")
+    logger.log("=" * 60 + "\n\n")
+
+    from datasets import load_dataset
+    from torchvision import transforms
+
+    # Load COCO in streaming mode
+    logger.log("📦 Loading COCO 2017 dataset (streaming)...\n")
+    coco = load_dataset("HuggingFaceM4/COCO", split="train", streaming=True, trust_remote_code=True)
+
+    # Image transforms
+    transform = transforms.Compose([
+        transforms.Resize((VQ_VAE_IMG_SIZE, VQ_VAE_IMG_SIZE)),
+        transforms.ToTensor(),  # [0, 1]
+    ])
+
+    class COCOStreamDataset(IterableDataset):
+        def __init__(self, hf_dataset, transform, max_samples=50000):
+            self.dataset = hf_dataset
+            self.transform = transform
+            self.max_samples = max_samples
+
+        def __iter__(self):
+            count = 0
+            for item in self.dataset:
+                if count >= self.max_samples:
+                    break
+                try:
+                    img = item["image"]
+                    if img.mode != "RGB":
+                        img = img.convert("RGB")
+                    tensor = self.transform(img)
+                    yield tensor
+                    count += 1
+                except Exception:
+                    continue
+
+    dataset = COCOStreamDataset(coco, transform, max_samples=50000)
+    dataloader = DataLoader(dataset, batch_size=VQ_VAE_BATCH, num_workers=0)
+
+    # Initialize model
+    model = VQVAE()
+    n_params = sum(p.numel() for p in model.parameters()) / 1e6
+    logger.log(f"✅ VQ-VAE initialized: {n_params:.1f}M parameters\n")
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=VQ_VAE_LR)
+    model.train()
+
+    for epoch in range(VQ_VAE_EPOCHS):
+        epoch_loss = 0.0
+        epoch_recon = 0.0
+        epoch_vq = 0.0
+        num_batches = 0
+        start_time = time.time()
+
+        for batch_idx, batch in enumerate(dataloader):
+            recon, vq_loss, _ = model(batch)
+            recon_loss = F.mse_loss(recon, batch)
+            loss = recon_loss + vq_loss
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+            epoch_recon += recon_loss.item()
+            epoch_vq += vq_loss.item()
+            num_batches += 1
+
+            if batch_idx % 50 == 0 and batch_idx > 0:
+                avg = epoch_loss / num_batches
+                avg_r = epoch_recon / num_batches
+                avg_v = epoch_vq / num_batches
+                logger.log(f"  Epoch {epoch+1}/{VQ_VAE_EPOCHS} | Batch {batch_idx} | "
+                          f"Loss: {avg:.4f} (recon: {avg_r:.4f}, vq: {avg_v:.4f})\n")
+
+            del recon, vq_loss, loss
+            if batch_idx % 200 == 0:
+                gc.collect()
+
+        elapsed = time.time() - start_time
+        avg_loss = epoch_loss / max(num_batches, 1)
+        logger.log(f"\n📈 Epoch {epoch+1} done. Avg Loss: {avg_loss:.4f} | "
+                  f"Batches: {num_batches} | Time: {elapsed:.0f}s\n\n")
+
+    # Save
+    torch.save(model.state_dict(), "vq_vae_real.pt")
+    logger.log("✅ VQ-VAE saved to vq_vae_real.pt\n\n")
+    return model
+
+
+# ============================================================================
+# PHASE 2: TOKENIZE OPENVID-1M DATASET
+# ============================================================================
+def tokenize_openvid(logger: Logger, vq_vae: Optional[VQVAE] = None):
+    """Stream OpenVid-1M, tokenize videos with VQ-VAE, save tokenized data."""
+    logger.log("=" * 60 + "\n")
+    logger.log("PHASE 2: Tokenizing OpenVid-1M dataset (10K clips)\n")
+    logger.log("=" * 60 + "\n\n")
+
+    # Load VQ-VAE if not provided
+    if vq_vae is None:
+        if os.path.exists("vq_vae_real.pt"):
+            vq_vae = VQVAE()
+            vq_vae.load_state_dict(torch.load("vq_vae_real.pt", map_location="cpu", weights_only=False))
+            logger.log("✅ Loaded trained VQ-VAE from vq_vae_real.pt\n")
+        else:
+            logger.log("❌ No trained VQ-VAE found! Run Phase 1 first.\n")
+            return None
+
+    vq_vae.eval()
+
+    from datasets import load_dataset
+
+    logger.log("📦 Loading OpenVid-1M dataset (streaming)...\n")
+    try:
+        dataset = load_dataset("NJU-PCALab/OpenVid-1M", split="train", streaming=True, trust_remote_code=True)
+    except Exception as e:
+        logger.log(f"⚠️ OpenVid-1M load error: {e}\n")
+        logger.log("🔄 Trying alternative: WebVid-2M...\n")
+        try:
+            dataset = load_dataset("tmpdump/webvid10m", split="train", streaming=True, trust_remote_code=True)
+        except Exception as e2:
+            logger.log(f"⚠️ WebVid load error: {e2}\n")
+            logger.log("🔄 Falling back to COCO captions (image-only, but much more data)...\n")
+            return _tokenize_coco_fallback(logger, vq_vae)
+
+    # Tokenize clips
+    tokenized_data = []
+    count = 0
+    errors = 0
+
+    for item in dataset:
+        if count >= NUM_OPENVID_CLIPS:
+            break
+
+        try:
+            # Get text caption
+            caption = ""
+            for key in ["caption", "text", "description", "title"]:
+                if key in item and item[key]:
+                    caption = item[key]
+                    break
+
+            if not caption:
+                caption = f"video clip {count}"
+
+            # Get video frames
+            video = item.get("video", None)
+            if video is None:
+                errors += 1
+                continue
+
+            # Process video frames
+            import io
+            from PIL import Image
+
+            frames = []
+            if hasattr(video, 'read'):
+                # It's bytes
+                pass
+
+            # Try to extract frames
+            if isinstance(video, dict) and "bytes" in video:
+                video_bytes = video["bytes"]
+            elif isinstance(video, bytes):
+                video_bytes = video
+            else:
+                errors += 1
+                continue
+
+            # Use imageio or decord to extract frames
+            try:
+                import imageio
+                reader = imageio.get_reader(io.BytesIO(video_bytes), format='mp4')
+                for i, frame in enumerate(reader):
+                    if i >= 4:  # Take first 4 frames
+                        break
+                    img = Image.fromarray(frame).convert("RGB").resize((128, 128))
+                    frames.append(np.array(img))
+                reader.close()
+            except Exception:
+                errors += 1
+                continue
+
+            if not frames:
+                errors += 1
+                continue
+
+            # Tokenize frames through VQ-VAE
+            from torchvision import transforms
+            transform = transforms.ToTensor()
+            all_tokens = []
+
+            for frame in frames:
+                img_tensor = transform(Image.fromarray(frame)).unsqueeze(0)
+                with torch.no_grad():
+                    tokens = vq_vae.encode(img_tensor)
+                    all_tokens.extend(tokens.flatten().tolist())
+
+            # Truncate/pad to fixed length
+            all_tokens = all_tokens[:TOKENS_PER_CLIP]
+            while len(all_tokens) < TOKENS_PER_CLIP:
+                all_tokens.append(0)
+
+            tokenized_data.append({
+                "text_prompt": caption,
+                "video_tokens": all_tokens,
+            })
+
+            count += 1
+            if count % 100 == 0:
+                logger.log(f"  Tokenized {count}/{NUM_OPENVID_CLIPS} clips (errors: {errors})\n")
+
+        except Exception as e:
+            errors += 1
+            if errors <= 3:
+                logger.log(f"  ⚠️ Error on item: {e}\n")
+            continue
+
+    if not tokenized_data:
+        logger.log("❌ No clips tokenized from OpenVid-1M! Falling back to COCO captions.\n")
+        return _tokenize_coco_fallback(logger, vq_vae)
+
+    # Save
+    with open("tokenized_dataset.json", "w") as f:
+        json.dump(tokenized_data, f)
+
+    logger.log(f"\n✅ Tokenized {len(tokenized_data)} clips saved to tokenized_dataset.json\n")
+    logger.log(f"   Errors: {errors}\n\n")
+    return tokenized_data
+
+
+def _tokenize_coco_fallback(logger: Logger, vq_vae: VQVAE):
+    """Fallback: tokenize COCO captions as image-text pairs."""
+    logger.log("📦 Using COCO captions as image-text pairs (50K samples)...\n")
+
+    from datasets import load_dataset
+    from torchvision import transforms
+    from PIL import Image
+
+    coco = load_dataset("HuggingFaceM4/COCO", split="train", streaming=True, trust_remote_code=True)
+    transform = transforms.Compose([
+        transforms.Resize((VQ_VAE_IMG_SIZE, VQ_VAE_IMG_SIZE)),
+        transforms.ToTensor(),
+    ])
+
+    vq_vae.eval()
+    tokenized_data = []
+    count = 0
+
+    for item in coco:
+        if count >= 50000:
+            break
+
+        try:
+            img = item["image"]
+            if img.mode != "RGB":
+                img = img.convert("RGB")
+
+            caption = ""
+            if "caption" in item:
+                caption = item["caption"] if isinstance(item["caption"], str) else item["caption"][0]
+            elif "text" in item:
+                caption = item["text"]
+            if not caption:
+                caption = f"image {count}"
+
+            img_tensor = transform(img).unsqueeze(0)
+            with torch.no_grad():
+                tokens = vq_vae.encode(img_tensor)
+                flat_tokens = tokens.flatten().tolist()
+
+            # Truncate/pad
+            flat_tokens = flat_tokens[:TOKENS_PER_CLIP]
+            while len(flat_tokens) < TOKENS_PER_CLIP:
+                flat_tokens.append(0)
+
+            tokenized_data.append({
+                "text_prompt": caption,
+                "video_tokens": flat_tokens,
+            })
+
+            count += 1
+            if count % 1000 == 0:
+                logger.log(f"  Tokenized {count}/50000 images\n")
+                # Save checkpoint periodically
+                if count % 10000 == 0:
+                    with open("tokenized_dataset.json", "w") as f:
+                        json.dump(tokenized_data, f)
+                    logger.log(f"  💾 Checkpoint saved ({len(tokenized_data)} samples)\n")
+
+        except Exception:
+            continue
+
+    # Final save
+    with open("tokenized_dataset.json", "w") as f:
+        json.dump(tokenized_data, f)
+
+    logger.log(f"\n✅ Tokenized {len(tokenized_data)} images saved to tokenized_dataset.json\n\n")
+    return tokenized_data
+
+
+# ============================================================================
+# PHASE 3: TRAIN LLM WITH LORA
+# ============================================================================
+def train_llm(logger: Logger):
+    """Fine-tune OLMo 2 1B with LoRA on tokenized data."""
+    logger.log("=" * 60 + "\n")
+    logger.log("PHASE 3: Fine-tuning OLMo 2 1B + LoRA\n")
+    logger.log("=" * 60 + "\n\n")
+
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from peft import LoraConfig, get_peft_model, TaskType
+
+    # Load data
+    data_path = "tokenized_dataset.json"
+    if not os.path.exists(data_path):
+        logger.log("❌ No tokenized dataset found! Run Phase 2 first.\n")
+        return
+
+    with open(data_path) as f:
+        data = json.load(f)
+    logger.log(f"📊 Loaded {len(data)} training samples\n")
+
+    # Tokenizer
+    logger.log("📦 Loading OLMo 2 1B tokenizer...\n")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Model
+    logger.log("📦 Loading model (fp32, CPU)...\n")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, trust_remote_code=True, torch_dtype=torch.float32
+    )
+    logger.log(f"✅ Model loaded. Original vocab: {len(tokenizer)}\n")
+
+    # Expand vocab
+    logger.log(f"🔤 Adding {CODEBOOK_SIZE} visual tokens...\n")
+    visual_tokens = [VIDEO_START, VIDEO_END, VIDEO_PAD]
+    for i in range(CODEBOOK_SIZE):
+        visual_tokens.append(f"<v_{i}>")
+    tokenizer.add_tokens(visual_tokens)
+    model.resize_token_embeddings(len(tokenizer))
+    logger.log(f"✅ New vocab: {len(tokenizer)}\n")
+
+    # LoRA
+    logger.log(f"🔧 Applying LoRA (r={LORA_R})...\n")
+    lora_config = LoraConfig(
+        r=LORA_R, lora_alpha=LORA_ALPHA,
+        target_modules=["q_proj", "v_proj"],
+        lora_dropout=LORA_DROPOUT, bias="none",
+        task_type=TaskType.CAUSAL_LM,
+    )
+    model = get_peft_model(model, lora_config)
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    logger.log(f"✅ LoRA: {trainable:,} / {total:,} trainable ({100*trainable/total:.2f}%)\n")
+
+    # Dataset
+    class VideoTokenDataset(Dataset):
+        def __init__(self, data, max_tokens=TOKENS_PER_CLIP):
+            self.data = data
+            self.max_tokens = max_tokens
+
+        def __len__(self):
+            return len(self.data)
+
+        def __getitem__(self, idx):
+            item = self.data[idx]
+            prompt = item["text_prompt"]
+            tokens = item["video_tokens"][:self.max_tokens]
+            while len(tokens) < self.max_tokens:
+                tokens.append(0)
+            return {"prompt": prompt, "video_tokens": torch.tensor(tokens, dtype=torch.long)}
+
+    dataset = VideoTokenDataset(data)
+    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+    total_steps = NUM_EPOCHS * len(dataloader)
+    logger.log(f"📊 {len(dataset)} samples × {NUM_EPOCHS} epochs = {total_steps} steps\n\n")
+
+    # Train
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
+    model.train()
+    global_step = 0
+    running_loss = 0.0
+    start_time = time.time()
+
+    for epoch in range(NUM_EPOCHS):
+        epoch_loss = 0.0
+        num_batches = 0
+
+        for batch_idx, batch in enumerate(dataloader):
+            prompt = batch["prompt"][0]
+            video_tokens = batch["video_tokens"][0]
+
+            # Format: use 64 visual tokens per sample for memory
+            token_str = " ".join(f"<v_{t.item()}>" for t in video_tokens[:64])
+            text = f"Create a video of: {prompt} {VIDEO_START} {token_str} {VIDEO_END}"
+
+            inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=MAX_SEQ_LEN, padding="max_length")
+            outputs = model(**inputs, labels=inputs["input_ids"])
+            loss = outputs.loss / GRADIENT_ACCUMULATION
+            loss.backward()
+
+            if (batch_idx + 1) % GRADIENT_ACCUMULATION == 0 or (batch_idx + 1) == len(dataloader):
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                optimizer.step()
+                optimizer.zero_grad()
+
+            global_step += 1
+            batch_loss = loss.item() * GRADIENT_ACCUMULATION
+            epoch_loss += batch_loss
+            running_loss += batch_loss
+            num_batches += 1
+
+            if batch_idx % 100 == 0:
+                elapsed = time.time() - start_time
+                speed = global_step / elapsed if elapsed > 0 else 0
+                logger.log(f"  Epoch {epoch+1}/{NUM_EPOCHS} | Step {batch_idx+1}/{len(dataloader)} | "
+                          f"Loss: {batch_loss:.4f} | Avg: {epoch_loss/num_batches:.4f} | "
+                          f"Speed: {speed:.2f} steps/s\n")
+
+            del outputs, loss
+            gc.collect()
+
+        logger.log(f"\n📈 Epoch {epoch+1} done. Avg Loss: {epoch_loss/num_batches:.4f}\n\n")
+
+    total_time = time.time() - start_time
+    logger.log(f"✅ Training complete in {total_time:.0f}s ({total_time/60:.1f} min)\n")
+    logger.log(f"   Final avg loss: {running_loss/global_step:.4f}\n\n")
+
+    # Merge & save
+    logger.log("🔀 Merging LoRA → base model...\n")
+    model = model.merge_and_unload()
+
+    save_dir = "./trained_model"
+    model.save_pretrained(save_dir, safe_serialization=True)
+    tokenizer.save_pretrained(save_dir)
+
+    # Also save VQ-VAE
+    if os.path.exists("vq_vae_real.pt"):
+        import shutil
+        shutil.copy("vq_vae_real.pt", f"{save_dir}/vq_vae_final.pt")
+
+    # Copy tokenized dataset
+    if os.path.exists("tokenized_dataset.json"):
+        import shutil
+        shutil.copy("tokenized_dataset.json", f"{save_dir}/tokenized_dataset.json")
+
+    logger.log("✅ Model saved locally.\n")
+
+    # Push
+    logger.log(f"🚀 Pushing to {REPO_ID}...\n")
+    from huggingface_hub import HfApi
+    api = HfApi(token=HF_TOKEN)
+    try:
+        api.create_repo(repo_id=REPO_ID, repo_type="model", exist_ok=True)
+    except:
+        pass
+    api.upload_folder(folder_path=save_dir, repo_id=REPO_ID, repo_type="model",
+                      commit_message=f"LoRA OLMo 2 1B (r={LORA_R}, {NUM_EPOCHS} epochs, {len(data)} samples)")
+    logger.log(f"✅ Pushed to https://huggingface.co/{REPO_ID}\n\n")
+
+
+# ============================================================================
+# MAIN PIPELINE
+# ============================================================================
+def run_pipeline(log_path: str = LOG_FILE):
+    logger = Logger(log_path)
+
+    try:
+        # Phase 1: Train VQ-VAE
+        vq_vae = train_vq_vae(logger)
+        gc.collect()
+
+        # Phase 2: Tokenize dataset
+        tokenize_openvid(logger, vq_vae)
+        gc.collect()
+
+        # Phase 3: Train LLM
+        train_llm(logger)
+
+        logger.log("\n🎉 FULL PIPELINE COMPLETE!\n")
+    except Exception as e:
+        logger.log(f"\n❌ PIPELINE ERROR: {e}\n")
+        logger.log(traceback.format_exc())
+
+
+# CLI
+if __name__ == "__main__":
+    run_pipeline()