Add microforge/pipeline.py

microforge/pipeline.py

@@ -0,0 +1,335 @@
+"""
+MicroForge Pipeline: End-to-End Generation and Editing
+=======================================================
+
+Unified pipeline for:
+- Text-to-image generation
+- Image-to-image editing (spatial concat, DreamLite-style)
+- Inpainting (masked spatial concat)
+- Super-resolution (low-res spatial concat)
+
+The key insight (from DreamLite): spatial concatenation preserves generation
+priors when adding editing capabilities. The same backbone handles all tasks
+by varying what goes into the "context" panel:
+- Generation: context = blank (zeros)
+- Editing: context = source image latent
+- Inpainting: context = masked source image latent
+- Super-res: context = upsampled low-res latent
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Dict, List, Tuple, Union
+
+from .vae import MicroForgeVAE
+from .backbone import MicroForgeBackbone
+from .planner import RecurrentLatentPlanner
+from .training import FlowMatchingScheduler
+
+
+class SimpleTextEncoder(nn.Module):
+    """
+    Lightweight text encoder for prototyping.
+    In production, replace with CLIP-L or a small LLM (Gemma-2B).
+
+    This uses a small transformer on learned token embeddings.
+    For the prototype, we support:
+    1. Random projection (for testing)
+    2. Simple learned embedding (for small-scale training)
+    """
+
+    def __init__(
+        self,
+        vocab_size: int = 8192,
+        max_seq_len: int = 77,
+        embed_dim: int = 768,
+        num_heads: int = 8,
+        num_layers: int = 4,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+
+        self.token_embed = nn.Embedding(vocab_size, embed_dim)
+        self.pos_embed = nn.Parameter(torch.randn(1, max_seq_len, embed_dim) * 0.02)
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim,
+            nhead=num_heads,
+            dim_feedforward=embed_dim * 4,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+
+        self.final_norm = nn.LayerNorm(embed_dim)
+        self.pool_proj = nn.Linear(embed_dim, embed_dim)
+
+    def forward(self, token_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            token_ids: [B, L] integer token IDs
+
+        Returns:
+            text_emb: [B, L, D] token-level embeddings
+            text_pooled: [B, D] pooled embedding
+        """
+        x = self.token_embed(token_ids) + self.pos_embed[:, :token_ids.shape[1], :]
+        x = self.encoder(x)
+        x = self.final_norm(x)
+
+        # Pool: mean of all tokens
+        pooled = x.mean(dim=1)
+        pooled = self.pool_proj(pooled)
+
+        return x, pooled
+
+    def encode_text_simple(self, batch_size: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Generate random text embeddings for testing.
+        """
+        text_emb = torch.randn(batch_size, 77, self.embed_dim, device=device)
+        text_pooled = torch.randn(batch_size, self.embed_dim, device=device)
+        return text_emb, text_pooled
+
+
+class MicroForgePipeline:
+    """
+    End-to-end MicroForge pipeline.
+
+    Supports:
+    - text2img: Generate image from text
+    - img2img: Edit image with text guidance
+    - inpaint: Fill masked region with text guidance
+    - super_res: Upscale image with text guidance
+
+    All tasks use the same backbone via spatial concatenation.
+    """
+
+    def __init__(
+        self,
+        vae: MicroForgeVAE,
+        backbone: MicroForgeBackbone,
+        text_encoder: SimpleTextEncoder,
+        planner: Optional[RecurrentLatentPlanner] = None,
+        device: str = 'cpu',
+    ):
+        self.vae = vae.eval()
+        self.backbone = backbone.eval()
+        self.text_encoder = text_encoder.eval()
+        self.planner = planner.eval() if planner is not None else None
+        self.device = torch.device(device)
+        self.scheduler = FlowMatchingScheduler()
+
+        # Move to device
+        self.vae.to(self.device)
+        self.backbone.to(self.device)
+        self.text_encoder.to(self.device)
+        if self.planner is not None:
+            self.planner.to(self.device)
+
+    @torch.no_grad()
+    def text2img(
+        self,
+        text_tokens: torch.Tensor,
+        height: int = 256,
+        width: int = 256,
+        num_steps: int = 20,
+        cfg_scale: float = 7.5,
+        seed: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        Generate image from text.
+
+        Args:
+            text_tokens: [B, L] token IDs
+            height, width: output image size
+            num_steps: denoising steps
+            cfg_scale: classifier-free guidance scale
+            seed: random seed
+
+        Returns:
+            images: [B, 3, H, W] generated images in [-1, 1]
+        """
+        if seed is not None:
+            torch.manual_seed(seed)
+
+        B = text_tokens.shape[0]
+
+        # Encode text
+        text_emb, text_pooled = self.text_encoder(text_tokens.to(self.device))
+
+        # Latent dimensions (32x spatial compression)
+        latent_h = height // 32
+        latent_w = width // 32
+        latent_c = self.vae.latent_channels
+
+        # Sample noise
+        noise = torch.randn(B, latent_c, latent_h, latent_w, device=self.device)
+
+        # Denoise
+        z_0 = self.scheduler.sample(
+            self.backbone, noise, text_emb, text_pooled,
+            num_steps=num_steps, cfg_scale=cfg_scale,
+            planner=self.planner,
+        )
+
+        # Decode
+        images = self.vae.decode(z_0)
+        return images.clamp(-1, 1)
+
+    @torch.no_grad()
+    def img2img(
+        self,
+        source_image: torch.Tensor,
+        text_tokens: torch.Tensor,
+        strength: float = 0.7,
+        num_steps: int = 20,
+        cfg_scale: float = 7.5,
+    ) -> torch.Tensor:
+        """
+        Edit image with text guidance using spatial concatenation.
+        The source image latent is concatenated width-wise with the target latent.
+
+        Args:
+            source_image: [B, 3, H, W] source image
+            text_tokens: [B, L] edit instruction tokens
+            strength: how much to change (0=no change, 1=full regen)
+            num_steps: denoising steps
+            cfg_scale: guidance scale
+
+        Returns:
+            edited_images: [B, 3, H, W]
+        """
+        B = source_image.shape[0]
+
+        # Encode text and source
+        text_emb, text_pooled = self.text_encoder(text_tokens.to(self.device))
+        source_latent = self.vae.get_latent(source_image.to(self.device))
+
+        # Create noised target (start from partial noise of source)
+        noise = torch.randn_like(source_latent)
+        t_start = torch.tensor([strength], device=self.device)
+        z_t = (1 - t_start) * source_latent + t_start * noise
+
+        # Spatial concatenation: [target | source] along width
+        # This doubles the width of the latent
+        # The backbone processes both together
+        timesteps = torch.linspace(strength, 0, num_steps + 1, device=self.device)
+
+        for i in range(num_steps):
+            t = timesteps[i]
+            t_next = timesteps[i + 1]
+            t_batch = torch.full((B,), t, device=self.device)
+
+            # Concat: [target_noised | source_clean]
+            z_concat = torch.cat([z_t, source_latent], dim=-1)  # Width concat
+
+            v_pred = self.backbone(z_concat, t_batch, text_emb, text_pooled)
+
+            # Only take the target half of the prediction
+            v_target = v_pred[..., :z_t.shape[-1]]
+
+            z_t = self.scheduler.euler_step(z_t, v_target, t.item(), t_next.item())
+
+        images = self.vae.decode(z_t)
+        return images.clamp(-1, 1)
+
+    @torch.no_grad()
+    def inpaint(
+        self,
+        image: torch.Tensor,
+        mask: torch.Tensor,
+        text_tokens: torch.Tensor,
+        num_steps: int = 20,
+        cfg_scale: float = 7.5,
+    ) -> torch.Tensor:
+        """
+        Inpaint masked region.
+
+        Args:
+            image: [B, 3, H, W] source image
+            mask: [B, 1, H, W] binary mask (1=inpaint region)
+            text_tokens: [B, L] description of what to fill
+            num_steps: denoising steps
+            cfg_scale: guidance scale
+
+        Returns:
+            inpainted: [B, 3, H, W]
+        """
+        B = image.shape[0]
+
+        text_emb, text_pooled = self.text_encoder(text_tokens.to(self.device))
+        source_latent = self.vae.get_latent(image.to(self.device))
+
+        # Downsample mask to latent size
+        latent_mask = F.interpolate(mask.float(), size=source_latent.shape[2:], mode='nearest')
+
+        # Masked source: zero out inpaint region
+        masked_source = source_latent * (1 - latent_mask)
+
+        # Generate in masked region
+        noise = torch.randn_like(source_latent)
+        z_t = noise
+
+        timesteps = torch.linspace(1, 0, num_steps + 1, device=self.device)
+
+        for i in range(num_steps):
+            t = timesteps[i]
+            t_next = timesteps[i + 1]
+            t_batch = torch.full((B,), t, device=self.device)
+
+            # Concat masked source as context
+            z_concat = torch.cat([z_t, masked_source], dim=-1)
+            v_pred = self.backbone(z_concat, t_batch, text_emb, text_pooled)
+            v_target = v_pred[..., :z_t.shape[-1]]
+
+            z_t = self.scheduler.euler_step(z_t, v_target, t.item(), t_next.item())
+
+            # Replace unmasked region with source
+            z_t = z_t * latent_mask + source_latent * (1 - latent_mask)
+
+        images = self.vae.decode(z_t)
+        return images.clamp(-1, 1)
+
+    def get_memory_estimate(self, height: int = 512, width: int = 512) -> Dict[str, float]:
+        """
+        Estimate memory usage in MB for given resolution.
+        """
+        # Model parameters
+        vae_params = sum(p.numel() for p in self.vae.parameters()) * 4 / 1e6
+        backbone_params = sum(p.numel() for p in self.backbone.parameters()) * 4 / 1e6
+        text_params = sum(p.numel() for p in self.text_encoder.parameters()) * 4 / 1e6
+        planner_params = 0
+        if self.planner is not None:
+            planner_params = sum(p.numel() for p in self.planner.parameters()) * 4 / 1e6
+
+        # Activation memory (rough estimate)
+        latent_h = height // 32
+        latent_w = width // 32
+        latent_size = latent_h * latent_w * self.vae.latent_channels * 4 / 1e6  # MB
+
+        return {
+            'vae_params_mb': vae_params,
+            'backbone_params_mb': backbone_params,
+            'text_encoder_params_mb': text_params,
+            'planner_params_mb': planner_params,
+            'total_params_mb': vae_params + backbone_params + text_params + planner_params,
+            'latent_size_mb': latent_size,
+            'estimated_inference_mb': (vae_params + backbone_params + text_params + planner_params) * 1.3,  # +30% overhead
+        }
+
+    def count_parameters(self) -> Dict[str, int]:
+        """Count parameters per module."""
+        return {
+            'vae': sum(p.numel() for p in self.vae.parameters()),
+            'backbone': sum(p.numel() for p in self.backbone.parameters()),
+            'text_encoder': sum(p.numel() for p in self.text_encoder.parameters()),
+            'planner': sum(p.numel() for p in self.planner.parameters()) if self.planner else 0,
+            'total': (
+                sum(p.numel() for p in self.vae.parameters()) +
+                sum(p.numel() for p in self.backbone.parameters()) +
+                sum(p.numel() for p in self.text_encoder.parameters()) +
+                (sum(p.numel() for p in self.planner.parameters()) if self.planner else 0)
+            ),
+        }