Fix: streaming dataset (no full download), step-based training loop

train.py

@@ -7,12 +7,15 @@ Flow Matching training objective (velocity prediction):
 - Loss: MSE(model(x_t, t), v)
 
 At inference: solve ODE from t=1 (noise) to t=0 (clean) using Euler steps.
+
+Dataset loading: Uses STREAMING mode by default — no full download needed!
+For small datasets (<500MB), set use_streaming=False for faster epoch iteration.
 """
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.utils.data import DataLoader, Dataset
+from torch.utils.data import DataLoader, Dataset, IterableDataset
 from torch.amp import autocast, GradScaler
 import math
 import os
@@ -37,6 +40,9 @@ class TrainConfig:
     dataset_config: str = ""
     image_column: str = "image"
     label_column: str = ""
+    use_streaming: bool = True     # KEY: streaming mode, no full download
+    max_samples: int = 0           # 0 = use all (only for non-streaming)
+    streaming_buffer: int = 1000   # Shuffle buffer for streaming
     
     # VAE
     vae_id: str = "black-forest-labs/FLUX.1-schnell"
@@ -51,7 +57,7 @@ class TrainConfig:
     learning_rate: float = 1e-4
     weight_decay: float = 0.01
     max_grad_norm: float = 2.0
-    num_epochs: int = 100
+    max_steps: int = 50000         # Train by steps, not epochs (better for streaming)
     warmup_steps: int = 1000
     ema_decay: float = 0.9999
     mixed_precision: bool = True
@@ -73,7 +79,7 @@ class TrainConfig:
     
     # System
     seed: int = 42
-    num_workers: int = 2
+    num_workers: int = 0           # 0 for streaming (required)
     pin_memory: bool = True
     compile_model: bool = False
     
@@ -99,9 +105,114 @@ def get_model_config(size: str, num_classes: int = 0, class_drop_prob: float = 0
     return cfg
 
 
+# =============================================================================
+# Dataset Loaders
+# =============================================================================
+
+class StreamingImageDataset(IterableDataset):
+    """
+    Streaming dataset — loads images on-the-fly from HuggingFace Hub.
+    NO full download needed. Starts training immediately.
+    
+    Perfect for large datasets (WikiArt, LAION, etc.) on Colab free tier.
+    """
+    def __init__(self, dataset_name, image_column="image", label_column="",
+                 image_size=256, split="train", dataset_config="",
+                 buffer_size=1000, seed=42):
+        super().__init__()
+        self.dataset_name = dataset_name
+        self.image_column = image_column
+        self.label_column = label_column
+        self.split = split
+        self.dataset_config = dataset_config
+        self.buffer_size = buffer_size
+        self.seed = seed
+        
+        from torchvision import transforms
+        self.transform = transforms.Compose([
+            transforms.Resize(image_size, interpolation=transforms.InterpolationMode.LANCZOS),
+            transforms.CenterCrop(image_size),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+        ])
+    
+    def _get_stream(self):
+        from datasets import load_dataset
+        kwargs = {}
+        if self.dataset_config:
+            kwargs["name"] = self.dataset_config
+        ds = load_dataset(self.dataset_name, split=self.split, streaming=True, **kwargs)
+        ds = ds.shuffle(seed=self.seed, buffer_size=self.buffer_size)
+        return iter(ds)
+    
+    def __iter__(self):
+        stream = self._get_stream()
+        for item in stream:
+            try:
+                img = item[self.image_column]
+                if img.mode != "RGB":
+                    img = img.convert("RGB")
+                img_tensor = self.transform(img)
+                label = -1
+                if self.label_column and self.label_column in item:
+                    label = item[self.label_column]
+                yield img_tensor, label
+            except Exception:
+                continue
+
+
+class MapImageDataset(Dataset):
+    """
+    Standard map-style dataset for small datasets that fit in memory.
+    Downloads once, then fast random access.
+    
+    Good for: Pokemon (95MB), Flowers (330MB), few-shot-art (510MB)
+    """
+    def __init__(self, dataset_name, image_column="image", label_column="",
+                 image_size=256, split="train", dataset_config="", max_samples=0):
+        super().__init__()
+        self.image_column = image_column
+        self.label_column = label_column
+        
+        from datasets import load_dataset
+        from torchvision import transforms
+        
+        kwargs = {}
+        if dataset_config:
+            kwargs["name"] = dataset_config
+        
+        print(f"Downloading {dataset_name}...")
+        self.dataset = load_dataset(dataset_name, split=split, **kwargs)
+        if max_samples > 0:
+            self.dataset = self.dataset.select(range(min(max_samples, len(self.dataset))))
+        print(f"  {len(self.dataset)} images loaded")
+        
+        self.transform = transforms.Compose([
+            transforms.Resize(image_size, interpolation=transforms.InterpolationMode.LANCZOS),
+            transforms.CenterCrop(image_size),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+        ])
+    
+    def __len__(self):
+        return len(self.dataset)
+    
+    def __getitem__(self, idx):
+        item = self.dataset[idx]
+        img = item[self.image_column]
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+        img = self.transform(img)
+        label = item[self.label_column] if self.label_column and self.label_column in item else -1
+        return img, label
+
+
+# =============================================================================
+# Training Utilities
+# =============================================================================
+
 class EMAModel:
     """Exponential Moving Average of model parameters."""
-    
     def __init__(self, model: nn.Module, decay: float = 0.9999):
         self.decay = decay
         self.shadow = {name: p.clone().detach() for name, p in model.named_parameters() if p.requires_grad}
@@ -132,41 +243,28 @@ class EMAModel:
 
 
 class FlowMatchingScheduler:
-    """
-    Flow Matching scheduler for training and sampling.
-    
-    Training: x_t = (1-t)*x_0 + t*ε, v_target = ε - x_0
-    Sampling: Euler ODE from t=1 (noise) to t=0 (clean)
-    """
+    """Flow Matching: x_t = (1-t)*x_0 + t*ε, v_target = ε - x_0"""
+    def __init__(self, min_t=0.001, max_t=0.999):
+        self.min_t, self.max_t = min_t, max_t
     
-    def __init__(self, min_t: float = 0.001, max_t: float = 0.999):
-        self.min_t = min_t
-        self.max_t = max_t
-    
-    def sample_timesteps(self, batch_size: int, device: torch.device) -> torch.Tensor:
+    def sample_timesteps(self, batch_size, device):
         return torch.rand(batch_size, device=device) * (self.max_t - self.min_t) + self.min_t
     
-    def add_noise(self, x0: torch.Tensor, noise: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-        t_expand = t.view(-1, 1, 1, 1)
-        return (1 - t_expand) * x0 + t_expand * noise
+    def add_noise(self, x0, noise, t):
+        t = t.view(-1, 1, 1, 1)
+        return (1 - t) * x0 + t * noise
     
-    def get_velocity_target(self, x0: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+    def get_velocity_target(self, x0, noise):
         return noise - x0
     
     @torch.no_grad()
-    def sample(
-        self, model: nn.Module, shape: tuple, device: torch.device,
-        num_steps: int = 50, class_labels: Optional[torch.Tensor] = None,
-        cfg_scale: float = 1.0, dtype: torch.dtype = torch.float32,
-    ) -> torch.Tensor:
+    def sample(self, model, shape, device, num_steps=50, class_labels=None,
+               cfg_scale=1.0, dtype=torch.float32):
         model.eval()
         x = torch.randn(shape, device=device, dtype=dtype)
         dt = 1.0 / num_steps
-        times = torch.linspace(1.0, dt, num_steps, device=device)
-        
-        for t_val in times:
+        for t_val in torch.linspace(1.0, dt, num_steps, device=device):
             t = torch.full((shape[0],), t_val.item(), device=device, dtype=dtype)
-            
             if cfg_scale > 1.0 and class_labels is not None:
                 with torch.amp.autocast('cuda', enabled=(dtype != torch.float32)):
                     v_cond = model(x, t, class_labels)
@@ -175,42 +273,39 @@ class FlowMatchingScheduler:
             else:
                 with torch.amp.autocast('cuda', enabled=(dtype != torch.float32)):
                     v = model(x, t, class_labels)
-            
             x = x - dt * v
-        
         return x
 
 
 def get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps):
-    """Cosine LR schedule with linear warmup."""
-    def lr_lambda(current_step):
-        if current_step < warmup_steps:
-            return float(current_step) / float(max(1, warmup_steps))
-        progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+    def lr_lambda(step):
+        if step < warmup_steps:
+            return step / max(1, warmup_steps)
+        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
         return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
     return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
 
 
 @torch.no_grad()
 def encode_images_with_vae(images, vae, scaling_factor, shift_factor):
-    """Encode pixel images to VAE latents."""
     images = images * 2.0 - 1.0
     latents = vae.encode(images).latent_dist.sample()
-    latents = (latents - shift_factor) * scaling_factor
-    return latents
+    return (latents - shift_factor) * scaling_factor
 
 
 @torch.no_grad()
 def decode_latents_with_vae(latents, vae, scaling_factor, shift_factor):
-    """Decode VAE latents to pixel images."""
     latents = latents / scaling_factor + shift_factor
     images = vae.decode(latents).sample
-    images = (images + 1.0) / 2.0
-    return images.clamp(0, 1)
+    return ((images + 1.0) / 2.0).clamp(0, 1)
+
 
+# =============================================================================
+# Main Training Loop
+# =============================================================================
 
 def train(config: TrainConfig):
-    """Main training loop."""
+    """Main training loop with streaming dataset support."""
     from model import LiquidGen
     
     torch.manual_seed(config.seed)
@@ -224,7 +319,7 @@ def train(config: TrainConfig):
     with open(os.path.join(config.output_dir, "config.json"), "w") as f:
         json.dump(asdict(config), f, indent=2)
     
-    # Load VAE
+    # Load VAE (frozen)
     print("Loading VAE...")
     from diffusers import AutoencoderKL
     vae_dtype = torch.float16 if config.vae_dtype == "float16" else torch.bfloat16
@@ -233,48 +328,39 @@ def train(config: TrainConfig):
     ).to(device).eval()
     for p in vae.parameters():
         p.requires_grad_(False)
+    print(f"VAE: {sum(p.numel() for p in vae.parameters())/1e6:.1f}M params (frozen)")
     
     # Load Dataset
-    print(f"Loading dataset: {config.dataset_name}")
-    from datasets import load_dataset
-    from torchvision import transforms
-    
-    ds_kwargs = {}
-    if config.dataset_config:
-        ds_kwargs["name"] = config.dataset_config
-    dataset = load_dataset(config.dataset_name, split="train", **ds_kwargs)
-    
-    transform = transforms.Compose([
-        transforms.Resize(config.image_size, interpolation=transforms.InterpolationMode.LANCZOS),
-        transforms.CenterCrop(config.image_size),
-        transforms.RandomHorizontalFlip(),
-        transforms.ToTensor(),
-    ])
-    
-    class ImageDataset(Dataset):
-        def __init__(self, hf_dataset, transform, image_col, label_col=""):
-            self.dataset = hf_dataset
-            self.transform = transform
-            self.image_col = image_col
-            self.label_col = label_col
-        
-        def __len__(self):
-            return len(self.dataset)
-        
-        def __getitem__(self, idx):
-            item = self.dataset[idx]
-            img = item[self.image_col]
-            if img.mode != "RGB":
-                img = img.convert("RGB")
-            img = self.transform(img)
-            label = item[self.label_col] if self.label_col and self.label_col in item else -1
-            return img, label
-    
-    train_dataset = ImageDataset(dataset, transform, config.image_column, config.label_column)
-    train_loader = DataLoader(
-        train_dataset, batch_size=config.batch_size, shuffle=True,
-        num_workers=config.num_workers, pin_memory=config.pin_memory, drop_last=True,
-    )
+    print(f"Loading dataset: {config.dataset_name} (streaming={config.use_streaming})")
+    if config.use_streaming:
+        train_dataset = StreamingImageDataset(
+            dataset_name=config.dataset_name,
+            image_column=config.image_column,
+            label_column=config.label_column,
+            image_size=config.image_size,
+            dataset_config=config.dataset_config,
+            buffer_size=config.streaming_buffer,
+            seed=config.seed,
+        )
+        train_loader = DataLoader(
+            train_dataset, batch_size=config.batch_size,
+            num_workers=0,  # Required for streaming
+            pin_memory=config.pin_memory,
+        )
+        print("  Streaming mode — no full download, starts immediately!")
+    else:
+        train_dataset = MapImageDataset(
+            dataset_name=config.dataset_name,
+            image_column=config.image_column,
+            label_column=config.label_column,
+            image_size=config.image_size,
+            dataset_config=config.dataset_config,
+            max_samples=config.max_samples,
+        )
+        train_loader = DataLoader(
+            train_dataset, batch_size=config.batch_size, shuffle=True,
+            num_workers=2, pin_memory=config.pin_memory, drop_last=True,
+        )
     
     # Create Model
     model_kwargs = get_model_config(config.model_size, config.num_classes, config.class_drop_prob)
@@ -286,30 +372,36 @@ def train(config: TrainConfig):
     
     optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate,
                                    weight_decay=config.weight_decay, betas=(0.9, 0.999))
-    total_steps = len(train_loader) * config.num_epochs // config.gradient_accumulation_steps
-    scheduler = get_cosine_schedule_with_warmup(optimizer, config.warmup_steps, total_steps)
+    scheduler = get_cosine_schedule_with_warmup(optimizer, config.warmup_steps, config.max_steps)
     ema = EMAModel(model, decay=config.ema_decay)
     scaler = GradScaler('cuda', enabled=config.mixed_precision)
     fm = FlowMatchingScheduler(min_t=config.min_timestep, max_t=config.max_timestep)
     
-    print(f"\nTraining: {total_steps} steps, effective batch {config.batch_size * config.gradient_accumulation_steps}")
+    print(f"\nTraining for {config.max_steps} steps")
+    print(f"Effective batch size: {config.batch_size * config.gradient_accumulation_steps}")
     
+    # Step-based training loop (works for both streaming and map datasets)
     global_step = 0
     loss_accum = 0.0
-    
-    for epoch in range(config.num_epochs):
-        model.train()
-        t_start = time.time()
-        
-        for batch_idx, (images, labels) in enumerate(train_loader):
+    accum_count = 0
+    model.train()
+    t_start = time.time()
+    
+    while global_step < config.max_steps:
+        for images, labels in train_loader:
+            if global_step >= config.max_steps:
+                break
+            
             images = images.to(device)
             labels = labels.to(device) if config.num_classes > 0 else None
             
+            # Encode to latents
             with torch.no_grad():
                 latents = encode_images_with_vae(
                     images.to(vae_dtype), vae, config.vae_scaling_factor, config.vae_shift_factor
                 ).float()
             
+            # Flow matching
             t = fm.sample_timesteps(latents.shape[0], device)
             noise = torch.randn_like(latents)
             x_t = fm.add_noise(latents, noise, t)
@@ -321,8 +413,9 @@ def train(config: TrainConfig):
             
             scaler.scale(loss).backward()
             loss_accum += loss.item()
+            accum_count += 1
             
-            if (batch_idx + 1) % config.gradient_accumulation_steps == 0:
+            if accum_count % config.gradient_accumulation_steps == 0:
                 scaler.unscale_(optimizer)
                 grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
                 scaler.step(optimizer)
@@ -332,17 +425,22 @@ def train(config: TrainConfig):
                 ema.update(model)
                 global_step += 1
                 
+                # Logging
                 if global_step % config.log_every_n_steps == 0:
                     avg_loss = loss_accum / config.log_every_n_steps
                     lr = optimizer.param_groups[0]["lr"]
-                    print(f"step={global_step} | epoch={epoch} | loss={avg_loss:.4f} | "
-                          f"grad_norm={grad_norm:.2f} | lr={lr:.2e}")
+                    elapsed = time.time() - t_start
+                    steps_per_sec = global_step / max(elapsed, 1)
+                    print(f"step={global_step} | loss={avg_loss:.4f} | "
+                          f"grad_norm={grad_norm:.2f} | lr={lr:.2e} | "
+                          f"steps/s={steps_per_sec:.2f} | elapsed={elapsed:.0f}s")
                     loss_accum = 0.0
                     
                     if math.isnan(avg_loss) or avg_loss > 100:
                         print("⚠️ Training diverged!")
                         return
                 
+                # Sample
                 if global_step % config.sample_every_n_steps == 0:
                     ema.apply(model)
                     model.eval()
@@ -356,23 +454,32 @@ def train(config: TrainConfig):
                                                           config.vae_scaling_factor, config.vae_shift_factor).float()
                     from torchvision.utils import save_image
                     save_image(sample_imgs, os.path.join(config.output_dir, "samples", f"step_{global_step:07d}.png"), nrow=2)
+                    print(f"  📸 Saved samples: step_{global_step:07d}.png")
                     ema.restore(model)
                     model.train()
                 
+                # Checkpoint
                 if global_step % config.save_every_n_steps == 0:
+                    ckpt_path = os.path.join(config.output_dir, "checkpoints", f"step_{global_step:07d}.pt")
                     torch.save({
                         "model": model.state_dict(), "ema": ema.state_dict(),
                         "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(),
-                        "global_step": global_step, "epoch": epoch, "config": asdict(config),
-                    }, os.path.join(config.output_dir, "checkpoints", f"step_{global_step:07d}.pt"))
-        
-        print(f"Epoch {epoch} complete | time={time.time()-t_start:.0f}s")
-    
-    torch.save({"model": model.state_dict(), "ema": ema.state_dict(), "config": asdict(config),
-                "global_step": global_step}, os.path.join(config.output_dir, "checkpoints", "final.pt"))
-    print(f"Training complete! Final model saved.")
+                        "global_step": global_step, "config": asdict(config),
+                    }, ckpt_path)
+                    print(f"  💾 Checkpoint: {ckpt_path}")
+    
+    # Final save
+    final_path = os.path.join(config.output_dir, "checkpoints", "final.pt")
+    torch.save({"model": model.state_dict(), "ema": ema.state_dict(),
+                "config": asdict(config), "global_step": global_step}, final_path)
+    elapsed = time.time() - t_start
+    print(f"\n🎉 Training complete! {global_step} steps in {elapsed/60:.1f} min")
+    print(f"   Final model: {final_path}")
 
 
 if __name__ == "__main__":
-    config = TrainConfig(model_size="small", image_size=256, batch_size=4, num_epochs=2)
+    config = TrainConfig(
+        model_size="small", image_size=256, batch_size=4,
+        max_steps=100, use_streaming=True,
+    )
     train(config)