Improve inference: prompt enhancement, better sampling, repetition penalty

app.py

@@ -456,6 +456,21 @@ def gcode_to_svg(gcode: str) -> str:
 # GENERATION
 # ============================================================================
 
+def enhance_prompt(prompt: str) -> str:
+    """Enhance prompt for better SD line drawing generation."""
+    prompt = prompt.strip().lower()
+    
+    # Skip if already detailed
+    if any(x in prompt for x in ["drawing", "sketch", "line", "illustration"]):
+        enhanced = prompt
+    else:
+        enhanced = f"a simple line drawing of {prompt}"
+    
+    # Add style suffixes for better SD output
+    enhanced += ", black ink on white paper, single continuous line, minimalist sketch, vector art style"
+    return enhanced
+
+
 @spaces.GPU
 def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, guidance: float):
     """Generate gcode from text prompt."""
@@ -471,10 +486,16 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
         dtype = m["dtype"]
         is_v3 = m.get("is_v3", False)
         
+        # Enhance prompt for better line drawing generation
+        enhanced = enhance_prompt(prompt)
+        print(f"Enhanced prompt: {enhanced}")
+        
         # Text -> Latent via SD diffusion
         with torch.no_grad():
+            # Use negative prompt to avoid unwanted styles
             result = pipe(
-                prompt,
+                enhanced,
+                negative_prompt="color, shading, gradient, photorealistic, 3d, complex, detailed texture",
                 num_inference_steps=num_steps,
                 guidance_scale=guidance,
                 output_type="latent",
@@ -499,27 +520,52 @@ def generate(prompt: str, temperature: float, max_tokens: int, num_steps: int, g
             
             max_gen = min(max_tokens, gcode_decoder.config.max_seq_len - 1)
             
+            # Track generated content for repetition detection
+            recent_tokens = []
+            repetition_window = 50
+            
             for step in range(max_gen):
                 logits = gcode_decoder(latent, input_ids)
                 next_logits = logits[:, -1, :] / temperature
                 
-                # Top-p sampling
-                sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
+                # Repetition penalty - reduce probability of recent tokens
+                if recent_tokens:
+                    for token_id in set(recent_tokens[-repetition_window:]):
+                        next_logits[:, token_id] *= 0.7
+                
+                # Top-k + Top-p sampling for better coherence
+                top_k = 50
+                top_p = 0.85
+                
+                # Top-k filtering
+                top_k_logits, top_k_indices = torch.topk(next_logits, top_k, dim=-1)
+                
+                # Top-p filtering within top-k
+                sorted_logits, sorted_idx = torch.sort(top_k_logits, descending=True, dim=-1)
                 cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-                sorted_indices_to_remove = cumulative_probs > 0.9
+                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
+                sorted_logits[sorted_indices_to_remove] = float('-inf')
                 
-                for b in range(batch_size):
-                    next_logits[b, sorted_indices[b, sorted_indices_to_remove[b]]] = float('-inf')
+                probs = torch.softmax(sorted_logits, dim=-1)
+                sampled_idx = torch.multinomial(probs, num_samples=1)
                 
-                probs = torch.softmax(next_logits, dim=-1)
-                next_token = torch.multinomial(probs, num_samples=1)
+                # Map back to vocabulary indices
+                next_token = top_k_indices.gather(-1, sorted_idx.gather(-1, sampled_idx))
                 input_ids = torch.cat([input_ids, next_token], dim=1)
+                recent_tokens.append(next_token.item())
                 
                 # Check EOS
                 if next_token.item() == gcode_tokenizer.eos_token_id:
                     break
+                
+                # Early stop on excessive repetition
+                if len(recent_tokens) > 20:
+                    last_20 = recent_tokens[-20:]
+                    if len(set(last_20)) < 5:  # Less than 5 unique tokens in last 20
+                        print("Stopping due to repetition")
+                        break
             
             print(f"Generated {input_ids.shape[1]} tokens")
             gcode = gcode_tokenizer.decode(input_ids[0], skip_special_tokens=True)
@@ -659,22 +705,25 @@ 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, 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")
+                temperature = gr.Slider(0.3, 1.2, value=0.6, label="temperature", step=0.1)
+                max_tokens = gr.Slider(256, 2048, value=1536, step=256, label="max tokens")
+                num_steps = gr.Slider(20, 50, value=30, step=5, label="diffusion steps")
+                guidance = gr.Slider(5.0, 20.0, value=12.0, step=0.5, label="guidance")
             
             generate_btn = gr.Button("generate", variant="secondary")
             
             gr.Examples(
                 examples=[
-                    ["a line drawing of a horse"],
-                    ["portrait sketch"],
-                    ["geometric shapes"],
+                    ["horse"],
+                    ["cat face"],
+                    ["spiral"],
+                    ["star"],
+                    ["tree"],
+                    ["flower"],
                 ],
                 inputs=prompt,
                 label=None,
-                examples_per_page=3,
+                examples_per_page=6,
             )
         
         with gr.Column(scale=2):