Support v3 decoder architecture with CNN projection

app.py

@@ -16,7 +16,120 @@ BOUNDS = {"left": -420.5, "right": 420.5, "top": 594.5, "bottom": -594.5}
 _model = None
 
 
-class GcodeDecoderConfig:
+# ============================================================================
+# V3 DECODER ARCHITECTURE
+# ============================================================================
+
+class GcodeDecoderConfigV3:
+    """Config for v3 decoder architecture."""
+    
+    def __init__(
+        self,
+        latent_channels: int = 4,
+        latent_size: int = 64,
+        hidden_size: int = 1024,
+        num_layers: int = 12,
+        num_heads: int = 16,
+        vocab_size: int = 8192,
+        max_seq_len: int = 2048,
+        dropout: float = 0.1,
+        ffn_mult: int = 4,
+    ):
+        self.latent_channels = latent_channels
+        self.latent_size = latent_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.vocab_size = vocab_size
+        self.max_seq_len = max_seq_len
+        self.dropout = dropout
+        self.ffn_mult = ffn_mult
+
+
+class CNNLatentProjector(nn.Module):
+    """CNN-based latent projector preserving spatial structure."""
+    
+    def __init__(self, config: GcodeDecoderConfigV3):
+        super().__init__()
+        
+        self.cnn = nn.Sequential(
+            nn.Conv2d(config.latent_channels, 64, 3, stride=2, padding=1),
+            nn.LayerNorm([64, 32, 32]),
+            nn.GELU(),
+            nn.Conv2d(64, 128, 3, stride=2, padding=1),
+            nn.LayerNorm([128, 16, 16]),
+            nn.GELU(),
+            nn.Conv2d(128, 256, 3, stride=2, padding=1),
+            nn.LayerNorm([256, 8, 8]),
+            nn.GELU(),
+            nn.Conv2d(256, config.hidden_size, 3, stride=2, padding=1),
+            nn.LayerNorm([config.hidden_size, 4, 4]),
+            nn.GELU(),
+        )
+        
+        self.num_memory_tokens = 16
+        self.memory_pos = nn.Parameter(torch.randn(1, self.num_memory_tokens, config.hidden_size) * 0.02)
+    
+    def forward(self, latent: torch.Tensor) -> torch.Tensor:
+        B = latent.shape[0]
+        x = self.cnn(latent)
+        x = x.view(B, x.shape[1], -1).transpose(1, 2)
+        x = x + self.memory_pos.expand(B, -1, -1)
+        return x
+
+
+class GcodeDecoderV3(nn.Module):
+    """Large transformer decoder for gcode generation (v3)."""
+    
+    def __init__(self, config: GcodeDecoderConfigV3):
+        super().__init__()
+        self.config = config
+        
+        self.latent_proj = CNNLatentProjector(config)
+        self.token_embed = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.pos_embed = nn.Embedding(config.max_seq_len, config.hidden_size)
+        self.embed_drop = nn.Dropout(config.dropout)
+        
+        self.layers = nn.ModuleList([
+            nn.TransformerDecoderLayer(
+                d_model=config.hidden_size,
+                nhead=config.num_heads,
+                dim_feedforward=config.hidden_size * config.ffn_mult,
+                dropout=config.dropout,
+                activation='gelu',
+                batch_first=True,
+                norm_first=True,
+            )
+            for _ in range(config.num_layers)
+        ])
+        
+        self.ln_f = nn.LayerNorm(config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        
+    def forward(self, latent: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
+        B, seq_len = input_ids.shape
+        device = input_ids.device
+        dtype = latent.dtype
+        
+        memory = self.latent_proj(latent)
+        positions = torch.arange(seq_len, device=device)
+        x = self.token_embed(input_ids) + self.pos_embed(positions)
+        x = self.embed_drop(x)
+        
+        causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=device, dtype=dtype)
+        
+        for layer in self.layers:
+            x = layer(x, memory, tgt_mask=causal_mask)
+        
+        x = self.ln_f(x)
+        return self.lm_head(x)
+
+
+# ============================================================================
+# V2 DECODER ARCHITECTURE (for backwards compatibility)
+# ============================================================================
+
+class GcodeDecoderConfigV2:
     def __init__(
         self,
         latent_channels: int = 4,
@@ -39,8 +152,8 @@ class GcodeDecoderConfig:
         self.dropout = dropout
 
 
-class GcodeDecoder(nn.Module):
-    def __init__(self, config: GcodeDecoderConfig):
+class GcodeDecoderV2(nn.Module):
+    def __init__(self, config: GcodeDecoderConfigV2):
         super().__init__()
         self.config = config
         
@@ -54,7 +167,6 @@ class GcodeDecoder(nn.Module):
         self.token_embed = nn.Embedding(config.vocab_size, config.hidden_size)
         self.pos_embed = nn.Embedding(config.max_seq_len, config.hidden_size)
         
-        # Individual layers (matches v2 training architecture)
         self.layers = nn.ModuleList([
             nn.TransformerDecoderLayer(
                 d_model=config.hidden_size,
@@ -84,7 +196,6 @@ class GcodeDecoder(nn.Module):
         positions = torch.arange(seq_len, device=device)
         x = self.token_embed(input_ids) + self.pos_embed(positions)
         
-        # Causal mask must match dtype for attention
         causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=device, dtype=dtype)
         
         for layer in self.layers:
@@ -92,43 +203,18 @@ class GcodeDecoder(nn.Module):
         
         x = self.ln_f(x)
         return self.lm_head(x)
-    
-    @torch.no_grad()
-    def generate(self, latent, tokenizer, max_length=512, temperature=0.8, top_p=0.9):
-        device = latent.device
-        batch_size = latent.shape[0]
-        
-        input_ids = torch.full((batch_size, 1), tokenizer.pad_token_id, dtype=torch.long, device=device)
-        
-        for _ in range(max_length - 1):
-            logits = self(latent, input_ids)
-            next_logits = logits[:, -1, :] / temperature
-            
-            sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
-            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-            sorted_indices_to_remove = cumulative_probs > top_p
-            sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
-            sorted_indices_to_remove[:, 0] = False
-            
-            for b in range(batch_size):
-                next_logits[b, sorted_indices[b, sorted_indices_to_remove[b]]] = float('-inf')
-            
-            probs = torch.softmax(next_logits, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1)
-            input_ids = torch.cat([input_ids, next_token], dim=1)
-            
-            if next_token.item() == tokenizer.eos_token_id:
-                break
-        
-        return tokenizer.decode(input_ids[0], skip_special_tokens=True)
 
 
+# ============================================================================
+# MODEL LOADING
+# ============================================================================
+
 def get_model():
-    """Load and cache the SD-Gcode model with full finetuned weights."""
+    """Load and cache the SD-Gcode model."""
     global _model
     if _model is None:
         from diffusers import StableDiffusionPipeline
-        from transformers import AutoTokenizer
+        from transformers import AutoTokenizer, PreTrainedTokenizerFast
         from huggingface_hub import hf_hub_download
         
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -143,7 +229,13 @@ def get_model():
         with open(config_path) as f:
             config = json.load(f)
         
-        # Load SD pipeline (we'll replace weights with finetuned ones)
+        # Determine model version
+        gcode_cfg = config.get("gcode_decoder", {})
+        is_v3 = gcode_cfg.get("ffn_mult") is not None or gcode_cfg.get("hidden_size", 768) >= 1024
+        
+        print(f"Model version: {'v3' if is_v3 else 'v2'}")
+        
+        # Load SD pipeline
         sd_model_id = config.get("sd_model_id", "runwayml/stable-diffusion-v1-5")
         print(f"Loading SD from {sd_model_id}...")
         pipe = StableDiffusionPipeline.from_pretrained(
@@ -152,58 +244,52 @@ def get_model():
             safety_checker=None,
         ).to(device)
         
-        # Build gcode decoder
-        gcode_cfg = config.get("gcode_decoder", {})
-        decoder_config = GcodeDecoderConfig(
-            latent_channels=gcode_cfg.get("latent_channels", 4),
-            latent_size=gcode_cfg.get("latent_size", 64),
-            hidden_size=gcode_cfg.get("hidden_size", 768),
-            num_layers=gcode_cfg.get("num_layers", 6),
-            num_heads=gcode_cfg.get("num_heads", 12),
-            vocab_size=gcode_cfg.get("vocab_size", 32128),
-            max_seq_len=gcode_cfg.get("max_seq_len", 1024),
-        )
-        gcode_decoder = GcodeDecoder(decoder_config).to(device, dtype)
+        # Build decoder based on version
+        if is_v3:
+            decoder_config = GcodeDecoderConfigV3(
+                latent_channels=gcode_cfg.get("latent_channels", 4),
+                latent_size=gcode_cfg.get("latent_size", 64),
+                hidden_size=gcode_cfg.get("hidden_size", 1024),
+                num_layers=gcode_cfg.get("num_layers", 12),
+                num_heads=gcode_cfg.get("num_heads", 16),
+                vocab_size=gcode_cfg.get("vocab_size", 8192),
+                max_seq_len=gcode_cfg.get("max_seq_len", 2048),
+                ffn_mult=gcode_cfg.get("ffn_mult", 4),
+            )
+            gcode_decoder = GcodeDecoderV3(decoder_config).to(device, dtype)
+        else:
+            decoder_config = GcodeDecoderConfigV2(
+                latent_channels=gcode_cfg.get("latent_channels", 4),
+                latent_size=gcode_cfg.get("latent_size", 64),
+                hidden_size=gcode_cfg.get("hidden_size", 768),
+                num_layers=gcode_cfg.get("num_layers", 6),
+                num_heads=gcode_cfg.get("num_heads", 12),
+                vocab_size=gcode_cfg.get("vocab_size", 32128),
+                max_seq_len=gcode_cfg.get("max_seq_len", 1024),
+            )
+            gcode_decoder = GcodeDecoderV2(decoder_config).to(device, dtype)
         
-        # Load ALL finetuned weights
+        # Load weights
         print("Loading finetuned weights...")
         state_dict = torch.load(weights_path, map_location=device, weights_only=False)
         
-        # Debug: print all key prefixes
-        prefixes = set(k.split(".")[0] for k in state_dict.keys())
-        print(f"State dict prefixes: {prefixes}")
-        print(f"Sample keys: {list(state_dict.keys())[:5]}")
-        
-        # Load text encoder weights
+        # Load SD components if present
         text_encoder_state = {k.replace("text_encoder.", ""): v for k, v in state_dict.items() 
                              if k.startswith("text_encoder.")}
         if text_encoder_state:
             pipe.text_encoder.load_state_dict(text_encoder_state, strict=False)
             print(f"Loaded {len(text_encoder_state)} text encoder weights")
         
-        # Load UNet weights
         unet_state = {k.replace("unet.", ""): v for k, v in state_dict.items() 
                      if k.startswith("unet.")}
         if unet_state:
             pipe.unet.load_state_dict(unet_state, strict=False)
             print(f"Loaded {len(unet_state)} UNet weights")
         
-        # Load gcode decoder weights
+        # Load decoder weights
         decoder_state = {k.replace("gcode_decoder.", ""): v for k, v in state_dict.items() 
                         if k.startswith("gcode_decoder.")}
         if decoder_state:
-            # Check what keys the model expects vs what we have
-            model_keys = set(gcode_decoder.state_dict().keys())
-            ckpt_keys = set(decoder_state.keys())
-            missing = model_keys - ckpt_keys
-            extra = ckpt_keys - model_keys
-            print(f"Decoder: model expects {len(model_keys)} keys, checkpoint has {len(ckpt_keys)}")
-            if missing:
-                print(f"Missing keys: {list(missing)[:5]}")
-            if extra:
-                print(f"Extra keys: {list(extra)[:5]}")
-            
-            # Try loading with strict=True to see errors
             try:
                 gcode_decoder.load_state_dict(decoder_state, strict=True)
                 print(f"Loaded {len(decoder_state)} decoder weights (strict)")
@@ -211,13 +297,19 @@ def get_model():
                 print(f"Strict load failed: {e}")
                 gcode_decoder.load_state_dict(decoder_state, strict=False)
                 print(f"Loaded {len(decoder_state)} decoder weights (non-strict)")
-        else:
-            print("WARNING: No gcode_decoder weights found!")
         
         gcode_decoder.eval()
         
-        # Gcode tokenizer
-        gcode_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+        # Load gcode tokenizer
+        try:
+            # Try loading custom tokenizer
+            tokenizer_path = hf_hub_download("twarner/dcode-sd-gcode", "gcode_tokenizer/tokenizer.json")
+            gcode_tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_path)
+            print("Loaded custom gcode tokenizer")
+        except Exception:
+            # Fallback to T5 tokenizer
+            gcode_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
+            print("Using fallback T5 tokenizer")
         
         _model = {
             "pipe": pipe,
@@ -226,12 +318,17 @@ def get_model():
             "device": device,
             "dtype": dtype,
             "num_inference_steps": config.get("num_inference_steps", 20),
+            "is_v3": is_v3,
         }
         print("Model loaded!")
     
     return _model
 
 
+# ============================================================================
+# GCODE PROCESSING
+# ============================================================================
+
 def validate_gcode(gcode: str) -> str:
     """Clamp coordinates to machine bounds."""
     lines = []
@@ -268,13 +365,15 @@ def gcode_to_svg(gcode: str) -> str:
     x, y = 0.0, 0.0
     pen_down = False
 
-    # Split on newlines OR command boundaries (for single-line gcode)
+    # Split on newlines, newline tokens, or command boundaries
     lines = []
+    # Replace newline tokens with actual newlines
+    gcode = gcode.replace("<newline>", "\n")
+    
     for line in gcode.replace(";", "\n;").split("\n"):
         line = line.strip()
         if not line:
             continue
-        # Split on G/M commands
         parts = re.split(r'(?=[GM]\d)', line)
         for part in parts:
             part = part.strip()
@@ -316,7 +415,6 @@ def gcode_to_svg(gcode: str) -> str:
     h = BOUNDS["top"] - BOUNDS["bottom"]
     padding = 20
     
-    # Dark mode compatible SVG
     svg = f'''<svg xmlns="http://www.w3.org/2000/svg" 
                   viewBox="{BOUNDS["left"] - padding} {-BOUNDS["top"] - padding} {w + 2*padding} {h + 2*padding}" 
                   style="width: 100%; height: 480px; border: 1px solid var(--border, #e0e0e0); border-radius: 4px;">
@@ -354,6 +452,10 @@ def gcode_to_svg(gcode: str) -> str:
     return svg
 
 
+# ============================================================================
+# GENERATION
+# ============================================================================
+
 @spaces.GPU
 def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, guidance: float):
     """Generate gcode from text prompt."""
@@ -367,6 +469,7 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
         gcode_tokenizer = m["gcode_tokenizer"]
         device = m["device"]
         dtype = m["dtype"]
+        is_v3 = m.get("is_v3", False)
         
         # Text -> Latent via SD diffusion
         with torch.no_grad():
@@ -377,25 +480,26 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
                 output_type="latent",
             )
             latent = result.images.to(dtype)
-            print(f"Latent shape: {latent.shape}, dtype: {latent.dtype}, device: {latent.device}")
-            print(f"Latent stats: min={latent.min():.3f}, max={latent.max():.3f}, mean={latent.mean():.3f}")
-            print(f"Decoder dtype: {next(gcode_decoder.parameters()).dtype}, device: {next(gcode_decoder.parameters()).device}")
+            print(f"Latent shape: {latent.shape}, dtype: {latent.dtype}")
         
-        # Latent -> Gcode via trained decoder (with debug)
+        # Latent -> Gcode via trained decoder
         with torch.no_grad():
             batch_size = latent.shape[0]
-            # Start with semicolon (gcode comment start) instead of pad
-            # Gcode files start with "; Source: ..." 
-            start_tokens = gcode_tokenizer.encode(";", add_special_tokens=False)
-            print(f"Start tokens for ';': {start_tokens}")
-            if start_tokens:
-                start_id = start_tokens[0]
+            
+            # Start token
+            if is_v3:
+                # V3 uses custom tokenizer with BOS
+                start_id = gcode_tokenizer.bos_token_id or 0
             else:
-                start_id = gcode_tokenizer.pad_token_id
+                # V2 uses semicolon as start
+                start_tokens = gcode_tokenizer.encode(";", add_special_tokens=False)
+                start_id = start_tokens[0] if start_tokens else gcode_tokenizer.pad_token_id
+            
             input_ids = torch.tensor([[start_id]], dtype=torch.long, device=device)
             
-            generated_tokens = []
-            for step in range(min(max_tokens, 1024) - 1):
+            max_gen = min(max_tokens, gcode_decoder.config.max_seq_len - 1)
+            
+            for step in range(max_gen):
                 logits = gcode_decoder(latent, input_ids)
                 next_logits = logits[:, -1, :] / temperature
                 
@@ -413,23 +517,17 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
                 next_token = torch.multinomial(probs, num_samples=1)
                 input_ids = torch.cat([input_ids, next_token], dim=1)
                 
-                token_id = next_token.item()
-                generated_tokens.append(token_id)
-                
-                # Debug first few tokens
-                if step < 5:
-                    token_str = gcode_tokenizer.decode([token_id])
-                    # Check logits distribution
-                    top5_vals, top5_ids = torch.topk(logits[0, -1, :], 5)
-                    top5_tokens = [gcode_tokenizer.decode([i.item()]) for i in top5_ids]
-                    print(f"Step {step}: token_id={token_id}, token='{token_str}', top5={list(zip(top5_tokens, top5_vals.tolist()))}")
-                
-                if token_id == gcode_tokenizer.eos_token_id:
-                    print(f"Hit EOS at step {step}")
+                # Check EOS
+                if next_token.item() == gcode_tokenizer.eos_token_id:
                     break
             
-            print(f"Generated {len(generated_tokens)} tokens")
+            print(f"Generated {input_ids.shape[1]} tokens")
             gcode = gcode_tokenizer.decode(input_ids[0], skip_special_tokens=True)
+            
+            # Post-process for v3: restore newlines
+            if is_v3:
+                gcode = gcode.replace("<newline>", "\n")
+            
             print(f"Decoded gcode length: {len(gcode)} chars")
         
         gcode = validate_gcode(gcode)
@@ -445,7 +543,10 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
         return f"; Error: {e}", gcode_to_svg("")
 
 
-# Minimal monochrome CSS with dark mode
+# ============================================================================
+# UI
+# ============================================================================
+
 css = """
 @import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;500&display=swap');
 
@@ -550,7 +651,7 @@ with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
             
             with gr.Accordion("settings", open=False):
                 temperature = gr.Slider(0.5, 1.5, value=0.8, label="temperature", step=0.1)
-                max_tokens = gr.Slider(256, 1024, value=512, step=128, label="max tokens")
+                max_tokens = gr.Slider(256, 2048, value=1024, step=256, label="max tokens")
                 num_steps = gr.Slider(10, 50, value=20, step=5, label="diffusion steps")
                 guidance = gr.Slider(1.0, 15.0, value=7.5, step=0.5, label="guidance")