Update to SD-Gcode end-to-end diffusion model

README.md

@@ -1,42 +1,41 @@
 ---
 title: dcode
-emoji: ✏️
-colorFrom: gray
-colorTo: green
+emoji: ✒️
+colorFrom: green
+colorTo: blue
 sdk: gradio
-sdk_version: "4.44.0"
+sdk_version: 5.9.1
 app_file: app.py
 pinned: false
 license: mit
-hardware: t4-small
-short_description: Text to Polargraph Gcode via Latent Diffusion
 ---
 
 # dcode
 
-Generate polargraph-compatible gcode from text prompts using latent diffusion.
+**Text -> Polargraph Gcode via Stable Diffusion**
 
-## How it works
+Single end-to-end diffusion model that converts text prompts directly to polargraph-compatible gcode.
 
-1. **Text → Latent**: Stable Diffusion generates a latent representation from your text prompt
-2. **Latent → Gcode**: Custom transformer decoder converts the latent to gcode commands
-3. **Validation**: Coordinates are clamped to machine bounds
+## Architecture
 
-## Usage
+```
+text prompt -> CLIP text encoder -> UNet diffusion -> latent [4,64,64] -> GcodeDecoder -> gcode tokens
+```
 
-1. Enter a prompt (e.g., "line drawing of a cat")
-2. Adjust diffusion steps and guidance scale
-3. Click Generate
-4. View preview and copy gcode
+All components post-trained end-to-end on 175,952 image-gcode pairs.
 
-## Model
+## Machine Specs
 
-- Base: Stable Diffusion 2.1
-- Decoder: 6-layer transformer trained on 175k image-gcode pairs
-- Final loss: 0.107
+- Work area: 841mm x 1189mm (A0)
+- Bounds: X [-420.5, 420.5], Y [-594.5, 594.5]
+- Pen servo: 40 deg (down), 90 deg (up)
 
 ## Links
 
-- [Model](https://huggingface.co/twarner/dcode-latent-gcode)
-- [Dataset](https://huggingface.co/datasets/twarner/dcode-polargraph-gcode)
 - [GitHub](https://github.com/Twarner491/dcode)
+- [Model](https://huggingface.co/twarner/dcode-sd-gcode)
+- [Dataset](https://huggingface.co/datasets/twarner/dcode-polargraph-gcode)
+
+## License
+
+MIT

app.py

@@ -1,115 +1,188 @@
-"""dcode Gradio Space - Text to Gcode via Latent Diffusion."""
+"""dcode Gradio Space - Text to Gcode via SD-Gcode Diffusion."""
 
 import re
+import os
+import json
 import gradio as gr
 import torch
+import torch.nn as nn
 from pathlib import Path
 
 # Machine limits
 BOUNDS = {"left": -420.5, "right": 420.5, "top": 594.5, "bottom": -594.5}
 
-# Model caches
-_generator = None
+# Model cache
+_model = None
 
 
-def get_generator():
-    """Load and cache the latent-gcode generator."""
-    global _generator
-    if _generator is None:
-        from diffusers import StableDiffusionPipeline, AutoencoderKL
+class GcodeDecoderConfig:
+    """Configuration for gcode decoder."""
+    def __init__(
+        self,
+        latent_channels: int = 4,
+        latent_size: int = 64,
+        hidden_size: int = 768,
+        num_layers: int = 6,
+        num_heads: int = 12,
+        vocab_size: int = 32128,
+        max_seq_len: int = 1024,
+        dropout: float = 0.1,
+    ):
+        self.latent_channels = latent_channels
+        self.latent_size = latent_size
+        self.latent_dim = latent_channels * 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
+
+
+class GcodeDecoder(nn.Module):
+    """Transformer decoder: SD latent -> gcode tokens."""
+    
+    def __init__(self, config: GcodeDecoderConfig):
+        super().__init__()
+        self.config = config
+        
+        self.latent_proj = nn.Sequential(
+            nn.Linear(config.latent_dim, config.hidden_size * 4),
+            nn.GELU(),
+            nn.Linear(config.hidden_size * 4, config.hidden_size * 16),
+            nn.LayerNorm(config.hidden_size * 16),
+        )
+        
+        self.token_embed = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.pos_embed = nn.Embedding(config.max_seq_len, config.hidden_size)
+        
+        decoder_layer = nn.TransformerDecoderLayer(
+            d_model=config.hidden_size,
+            nhead=config.num_heads,
+            dim_feedforward=config.hidden_size * 4,
+            dropout=config.dropout,
+            activation='gelu',
+            batch_first=True,
+            norm_first=True,
+        )
+        self.decoder = nn.TransformerDecoder(decoder_layer, config.num_layers)
+        
+        self.ln_f = nn.LayerNorm(config.hidden_size)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.token_embed.weight
+        
+    def forward(self, latent: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len = input_ids.shape
+        device = input_ids.device
+        
+        latent_flat = latent.view(batch_size, -1)
+        memory = self.latent_proj(latent_flat)
+        memory = memory.view(batch_size, 16, self.config.hidden_size)
+        
+        positions = torch.arange(seq_len, device=device)
+        x = self.token_embed(input_ids) + self.pos_embed(positions)
+        
+        causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=device)
+        
+        x = self.decoder(x, memory, tgt_mask=causal_mask)
+        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)
+
+
+def get_model():
+    """Load and cache the SD-Gcode model."""
+    global _model
+    if _model is None:
+        from diffusers import StableDiffusionPipeline
         from transformers import AutoTokenizer
-        import torch.nn as nn
+        from huggingface_hub import hf_hub_download
         
         device = "cuda" if torch.cuda.is_available() else "cpu"
         dtype = torch.float16 if device == "cuda" else torch.float32
         
-        print("Loading Stable Diffusion pipeline...")
-        # Use SD 1.5 which is more reliably available
-        pipe = StableDiffusionPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5",
-            torch_dtype=dtype,
-            safety_checker=None,
-            use_safetensors=True,
-        ).to(device)
+        print("Loading SD-Gcode model...")
         
-        print("Loading gcode decoder...")
-        from huggingface_hub import hf_hub_download
-        
-        # Download model files
-        model_path = hf_hub_download("twarner/dcode-latent-gcode", "pytorch_model.bin")
-        config_path = hf_hub_download("twarner/dcode-latent-gcode", "config.json")
+        # Download config and weights
+        config_path = hf_hub_download("twarner/dcode-sd-gcode", "config.json")
+        weights_path = hf_hub_download("twarner/dcode-sd-gcode", "pytorch_model.bin")
         
-        import json
         with open(config_path) as f:
             config = json.load(f)
         
-        # Load tokenizer
-        tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
-        
-        # Build decoder model
-        class LatentProjector(nn.Module):
-            def __init__(self, latent_dim, hidden_size):
-                super().__init__()
-                self.proj = nn.Sequential(
-                    nn.Linear(latent_dim, hidden_size * 2),
-                    nn.GELU(),
-                    nn.Linear(hidden_size * 2, hidden_size),
-                    nn.LayerNorm(hidden_size),
-                )
-            def forward(self, x):
-                return self.proj(x)
-        
-        class GcodeDecoder(nn.Module):
-            def __init__(self, hidden_size, vocab_size, num_layers, num_heads, max_seq_len):
-                super().__init__()
-                self.embed = nn.Embedding(vocab_size, hidden_size)
-                self.pos_embed = nn.Embedding(max_seq_len, hidden_size)
-                layer = nn.TransformerDecoderLayer(hidden_size, num_heads, hidden_size * 4, batch_first=True)
-                self.decoder = nn.TransformerDecoder(layer, num_layers)
-                self.head = nn.Linear(hidden_size, vocab_size)
-                self.max_seq_len = max_seq_len
-                
-            def forward(self, tgt, memory, tgt_mask=None):
-                pos = torch.arange(tgt.size(1), device=tgt.device)
-                x = self.embed(tgt) + self.pos_embed(pos)
-                x = self.decoder(x, memory, tgt_mask=tgt_mask)
-                return self.head(x)
-        
-        # Initialize models
-        latent_dim = 4 * 64 * 64
-        hidden_size = config.get("hidden_size", 512)
-        vocab_size = tokenizer.vocab_size
-        num_layers = config.get("num_layers", 6)
-        num_heads = config.get("num_heads", 8)
-        max_seq_len = config.get("max_seq_len", 1024)
+        # 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(
+            sd_model_id,
+            torch_dtype=dtype,
+            safety_checker=None,
+        ).to(device)
         
-        projector = LatentProjector(latent_dim, hidden_size).to(device, dtype)
-        decoder = GcodeDecoder(hidden_size, vocab_size, num_layers, num_heads, max_seq_len).to(device, dtype)
+        # 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)
         
         # Load weights
-        state_dict = torch.load(model_path, map_location=device)
-        
-        proj_state = {k.replace("projector.", ""): v for k, v in state_dict.items() if k.startswith("projector.")}
-        dec_state = {k.replace("decoder.", ""): v for k, v in state_dict.items() if k.startswith("decoder.")}
+        state_dict = torch.load(weights_path, map_location=device)
         
-        projector.load_state_dict(proj_state)
-        decoder.load_state_dict(dec_state)
+        # Extract decoder weights
+        decoder_state = {k.replace("gcode_decoder.", ""): v for k, v in state_dict.items() 
+                        if k.startswith("gcode_decoder.")}
+        gcode_decoder.load_state_dict(decoder_state)
+        gcode_decoder.eval()
         
-        projector.eval()
-        decoder.eval()
+        # Gcode tokenizer
+        gcode_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
         
-        _generator = {
+        _model = {
             "pipe": pipe,
-            "projector": projector,
-            "decoder": decoder,
-            "tokenizer": tokenizer,
+            "gcode_decoder": gcode_decoder,
+            "gcode_tokenizer": gcode_tokenizer,
             "device": device,
             "dtype": dtype,
-            "max_seq_len": max_seq_len,
+            "num_inference_steps": config.get("num_inference_steps", 20),
         }
-        print("Models loaded!")
+        print("Model loaded!")
     
-    return _generator
+    return _model
 
 
 def validate_gcode(gcode: str) -> str:
@@ -220,21 +293,19 @@ def gcode_to_svg(gcode: str) -> str:
 
 
 def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, guidance: float):
-    """Generate gcode from text prompt via latent diffusion."""
+    """Generate gcode from text prompt via SD-Gcode diffusion."""
     if not prompt or not prompt.strip():
         return "Enter a prompt to generate gcode", gcode_to_svg("")
 
     try:
-        gen = get_generator()
-        pipe = gen["pipe"]
-        projector = gen["projector"]
-        decoder = gen["decoder"]
-        tokenizer = gen["tokenizer"]
-        device = gen["device"]
-        dtype = gen["dtype"]
-        max_seq_len = gen["max_seq_len"]
+        m = get_model()
+        pipe = m["pipe"]
+        gcode_decoder = m["gcode_decoder"]
+        gcode_tokenizer = m["gcode_tokenizer"]
+        device = m["device"]
+        dtype = m["dtype"]
         
-        # 1. Text -> Latent via Stable Diffusion
+        # 1. Text -> Latent via full SD diffusion
         with torch.no_grad():
             result = pipe(
                 prompt,
@@ -242,37 +313,22 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
                 guidance_scale=guidance,
                 output_type="latent",
             )
-            latent = result.images  # [1, 4, 64, 64]
+            latent = result.images.to(dtype)  # [1, 4, 64, 64]
         
-        # 2. Latent -> Gcode via decoder
+        # 2. Latent -> Gcode via trained decoder
         with torch.no_grad():
-            # Flatten and project latent
-            latent_flat = latent.view(1, -1).to(dtype)  # [1, 4*64*64]
-            memory = projector(latent_flat).unsqueeze(1)  # [1, 1, hidden]
-            
-            # Autoregressive decoding
-            bos_id = tokenizer.bos_token_id or tokenizer.pad_token_id
-            eos_id = tokenizer.eos_token_id
-            
-            tokens = torch.tensor([[bos_id]], device=device)
-            
-            for _ in range(min(max_tokens, max_seq_len - 1)):
-                logits = decoder(tokens, memory)
-                next_logits = logits[:, -1, :] / temperature
-                probs = torch.softmax(next_logits, dim=-1)
-                next_token = torch.multinomial(probs, 1)
-                tokens = torch.cat([tokens, next_token], dim=1)
-                
-                if next_token.item() == eos_id:
-                    break
-            
-            gcode = tokenizer.decode(tokens[0], skip_special_tokens=True)
+            gcode = gcode_decoder.generate(
+                latent,
+                gcode_tokenizer,
+                max_length=max_tokens,
+                temperature=temperature,
+            )
         
         gcode = validate_gcode(gcode)
         line_count = len(gcode.split("\n"))
         svg = gcode_to_svg(gcode)
         
-        gcode_with_header = f"; dcode output - {line_count} lines\n; Prompt: {prompt}\n; Machine validated\n\n{gcode}"
+        gcode_with_header = f"; dcode SD-Gcode output - {line_count} lines\n; Prompt: {prompt}\n; Machine validated\n\n{gcode}"
         return gcode_with_header, svg
         
     except Exception as e:
@@ -291,11 +347,11 @@ custom_css = """
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="emerald")) as demo:
     gr.Markdown("""
     # dcode
-    **Text → Polargraph Gcode via Latent Diffusion**
+    **Text -> Polargraph Gcode via Stable Diffusion**
     
-    Uses Stable Diffusion to generate latents from text, then decodes to machine gcode.
+    Single end-to-end diffusion model: text -> CLIP -> UNet -> latent -> gcode decoder -> gcode
     
-    [GitHub](https://github.com/Twarner491/dcode) | [Model](https://huggingface.co/twarner/dcode-latent-gcode) | [Dataset](https://huggingface.co/datasets/twarner/dcode-polargraph-gcode)
+    [GitHub](https://github.com/Twarner491/dcode) | [Model](https://huggingface.co/twarner/dcode-sd-gcode) | [Dataset](https://huggingface.co/datasets/twarner/dcode-polargraph-gcode)
     """)
     
     with gr.Row():
@@ -307,11 +363,11 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="emerald")) as d
             )
             
             with gr.Row():
-                temperature = gr.Slider(0.5, 1.5, value=0.9, label="Temperature")
+                temperature = gr.Slider(0.5, 1.5, value=0.8, label="Temperature")
                 max_tokens = gr.Slider(256, 1024, value=512, step=128, label="Max Tokens")
             
             with gr.Row():
-                num_steps = gr.Slider(10, 50, value=25, step=5, label="Diffusion Steps")
+                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 Scale")
             
             generate_btn = gr.Button("Generate", variant="primary", size="lg")
@@ -338,7 +394,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="emerald")) as d
     
     gr.Markdown("""
     ---
-    **Machine Bounds**: X: ±420.5mm, Y: ±594.5mm | Pen servo: 40° (down) / 90° (up) | **License**: MIT
+    **Machine Bounds**: X: +/-420.5mm, Y: +/-594.5mm | Pen servo: 40 deg (down) / 90 deg (up) | **License**: MIT
     """)
     
     generate_btn.click(

requirements.txt

@@ -1,7 +1,7 @@
-gradio>=4.44.0
-torch>=2.0
-transformers>=4.36
-diffusers>=0.25
-accelerate>=0.25
-huggingface_hub>=0.20
-safetensors>=0.4
+gradio>=4.44.1
+gradio_client==1.3.0
+torch
+diffusers
+transformers
+accelerate
+huggingface_hub