Update app.py

app.py

@@ -1,161 +1,265 @@
-import os
-os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+"""
+CPU-Friendly Training Script for GLM-4.5V CAD Generation
+Simplified version for testing and development
+"""
 
-import gradio as gr
 import torch
-from transformers import pipeline
+from datasets import load_dataset
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
+from peft import LoraConfig, get_peft_model, TaskType
 from PIL import Image
-import spaces
+import json
+import os
+from dataclasses import dataclass
+from typing import Dict, List
 
-# Global model storage
-models = {}
+# Simple configuration for CPU testing
+CONFIG = {
+    "base_model": "microsoft/DialoGPT-small",  # Small model for CPU testing
+    "dataset_name": "CADCODER/GenCAD-Code",
+    "output_dir": "./test-cad-model",
+    "max_samples": 50,  # Very small for CPU
+    "batch_size": 1,
+    "gradient_accumulation": 4,
+    "epochs": 1,
+    "learning_rate": 5e-5,
+    "max_length": 512
+}
 
-# Wrap the entire generation in try-except to handle Zero GPU issues
-@spaces.GPU(duration=120)
-def generate_on_gpu(image_array, model_name, prompt_text):
-    """Core GPU function - simplified to avoid context issues."""
-    pipe = pipeline(
-        "image-text-to-text",
-        model=model_name,
-        device_map="auto", 
-        torch_dtype=torch.float16,
-        trust_remote_code=True
-    )
+@dataclass
+class SimpleDataCollator:
+    """Simple data collator for text-only training."""
+    tokenizer: any
+    max_length: int = 512
     
-    # Convert array back to PIL
-    image = Image.fromarray(image_array)
+    def __call__(self, features: List[Dict]) -> Dict[str, torch.Tensor]:
+        # Extract texts
+        texts = [f["text"] for f in features]
+        
+        # Tokenize
+        batch = self.tokenizer(
+            texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=self.max_length
+        )
+        
+        # Create labels for causal LM
+        batch["labels"] = batch["input_ids"].clone()
+        batch["labels"][batch["labels"] == self.tokenizer.pad_token_id] = -100
+        
+        return batch
+
+def prepare_simple_dataset(dataset_name: str, max_samples: int = 50):
+    """Prepare a simplified text-only dataset for CPU training."""
+    print(f"📊 Loading dataset: {dataset_name}")
     
-    messages = [{
-        "role": "user",
-        "content": [
-            {"type": "image", "image": image},
-            {"type": "text", "text": prompt_text}
+    try:
+        # Load small subset
+        dataset = load_dataset(dataset_name, split=f"train[:{max_samples}]")
+        
+        def create_text_examples(examples):
+            """Convert to text-only format."""
+            texts = []
+            
+            for i in range(len(examples["code"])):
+                # Create simple prompt-response format
+                text = f"Generate CADQuery code:\n{examples['code'][i]}<|endoftext|>"
+                texts.append(text)
+            
+            return {"text": texts}
+        
+        # Process dataset
+        dataset = dataset.map(
+            create_text_examples,
+            batched=True,
+            remove_columns=dataset.column_names
+        )
+        
+        print(f"✅ Dataset prepared: {len(dataset)} samples")
+        return dataset
+        
+    except Exception as e:
+        print(f"❌ Dataset loading failed: {e}")
+        
+        # Create dummy dataset for testing
+        print("🔄 Creating dummy dataset for testing...")
+        dummy_codes = [
+            "import cadquery as cq\nresult = cq.Workplane('XY').box(10, 10, 5)",
+            "import cadquery as cq\nresult = cq.Workplane('XY').cylinder(5, 10)",
+            "import cadquery as cq\nresult = cq.Workplane('XY').box(20, 15, 8).fillet(2)",
         ]
-    }]
+        
+        texts = [f"Generate CADQuery code:\n{code}<|endoftext|>" for code in dummy_codes]
+        
+        from datasets import Dataset
+        dataset = Dataset.from_dict({"text": texts * (max_samples // 3 + 1)})
+        dataset = dataset.select(range(max_samples))
+        
+        print(f"✅ Dummy dataset created: {len(dataset)} samples")
+        return dataset
+
+def setup_simple_model(model_name: str):
+    """Set up a simple model for CPU training."""
+    print(f"🔧 Loading model: {model_name}")
     
-    result = pipe(messages, max_new_tokens=512, temperature=0.7)
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
     
-    if isinstance(result, list) and len(result) > 0:
-        return result[0].get("generated_text", str(result))
-    else:
-        return str(result)
-
-def generate_code(image, model_choice, prompt_style):
-    """Wrapper function that handles the UI logic."""
-    if image is None:
-        return "❌ Please upload an image first."
-    
-    # Prompts
-    prompts = {
-        "Simple": "Generate CADQuery Python code for this 3D model:",
-        "Detailed": "Analyze this 3D CAD model and generate Python CADQuery code.\n\nRequirements:\n- Import cadquery as cq\n- Store result in 'result' variable\n- Use proper CADQuery syntax\n\nCode:",
-        "Chain-of-Thought": "Analyze this 3D CAD model step by step:\n\nStep 1: Identify the basic geometry\nStep 2: Note any features\nStep 3: Generate clean CADQuery Python code\n\n```python\nimport cadquery as cq\n\n# Generated code:"
-    }
+    # Add pad token if missing
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
     
-    try:
-        # Model mapping
-        model_map = {
-            "GLM-4.5V-AWQ": "QuantTrio/GLM-4.5V-AWQ",
-            "GLM-4.5V-FP8": "zai-org/GLM-4.5V-FP8",
-            "GLM-4.5V": "zai-org/GLM-4.5V"
-        }
-        
-        model_name = model_map[model_choice]
-        prompt_text = prompts[prompt_style]
-        
-        # Convert PIL to array for passing
-        import numpy as np
-        image_array = np.array(image)
-        
-        # Call GPU function
-        generated_text = generate_on_gpu(image_array, model_name, prompt_text)
-        
-        # Extract code
-        code = generated_text.strip()
-        if "```python" in code:
-            start = code.find("```python") + 9
-            end = code.find("```", start)
-            if end > start:
-                code = code[start:end].strip()
-        
-        if "import cadquery" not in code:
-            code = "import cadquery as cq\n\n" + code
-        
-        return f"""## 🎯 Generated CADQuery Code
-
-```python
-{code}
-```
+    # Load model for CPU
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float32,  # Use float32 for CPU
+        device_map="cpu"
+    )
+    
+    # Simple LoRA config for CPU
+    lora_config = LoraConfig(
+        r=8,  # Small rank for CPU
+        lora_alpha=16,
+        lora_dropout=0.1,
+        bias="none",
+        task_type=TaskType.CAUSAL_LM,
+        target_modules=["c_attn", "c_proj"]  # DialoGPT modules
+    )
+    
+    # Apply LoRA
+    model = get_peft_model(model, lora_config)
+    
+    # Print parameters
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"💡 Trainable: {trainable_params:,} ({100 * trainable_params / total_params:.2f}%)")
+    
+    return model, tokenizer
 
-## 📊 Info
-- **Model**: {model_choice}
-- **Prompt**: {prompt_style}
-- **Device**: GPU
+def train_simple_model(model, tokenizer, dataset, config):
+    """Train the model with simple settings."""
+    print("🏋️ Starting CPU training...")
+    
+    # Training arguments for CPU
+    training_args = TrainingArguments(
+        output_dir=config["output_dir"],
+        per_device_train_batch_size=config["batch_size"],
+        gradient_accumulation_steps=config["gradient_accumulation"],
+        num_train_epochs=config["epochs"],
+        learning_rate=config["learning_rate"],
+        warmup_steps=10,
+        logging_steps=5,
+        save_steps=100,
+        evaluation_strategy="no",
+        save_total_limit=1,
+        remove_unused_columns=False,
+        report_to="none",
+        fp16=False,  # No FP16 on CPU
+        dataloader_pin_memory=False,
+        use_cpu=True
+    )
+    
+    # Data collator
+    data_collator = SimpleDataCollator(
+        tokenizer=tokenizer,
+        max_length=config["max_length"]
+    )
+    
+    # Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset,
+        data_collator=data_collator,
+        tokenizer=tokenizer
+    )
+    
+    # Train
+    print("⏳ Training will take a few minutes on CPU...")
+    trainer.train()
+    
+    # Save
+    trainer.save_model()
+    tokenizer.save_pretrained(config["output_dir"])
+    
+    print(f"✅ Training complete! Model saved to {config['output_dir']}")
+    return trainer
 
-## 🔧 Usage
-```bash
-pip install cadquery
-python your_script.py
-```
-"""
+def test_simple_model(model_path: str):
+    """Test the trained model."""
+    print(f"🧪 Testing model: {model_path}")
+    
+    try:
+        # Load model
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        model = AutoModelForCausalLM.from_pretrained(model_path)
+        
+        # Test generation
+        prompt = "Generate CADQuery code:"
+        inputs = tokenizer.encode(prompt, return_tensors="pt")
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs,
+                max_new_tokens=100,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
+        generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        print("🎯 Generated:")
+        print(generated)
+        return generated
         
     except Exception as e:
-        return f"❌ **Generation Failed**: {str(e)[:500]}"
-
-def system_info():
-    """Get system info."""
-    info = f"""## 🖥️ System Information
-
-- **CUDA Available**: {torch.cuda.is_available()}
-- **CUDA Devices**: {torch.cuda.device_count() if torch.cuda.is_available() else 0}
-- **PyTorch Version**: {torch.__version__}
-- **Device**: {"GPU" if torch.cuda.is_available() else "CPU"}
-"""
-    return info
+        print(f"❌ Testing failed: {e}")
+        return str(e)
 
-# Create interface WITHOUT ssr_mode
-with gr.Blocks(title="GLM-4.5V CAD Generator", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # 🔧 GLM-4.5V CAD Generator
-    
-    Generate CADQuery Python code from 3D CAD model images using GLM-4.5V models!
-    
-    **Models**: GLM-4.5V-AWQ (fastest) | GLM-4.5V-FP8 (balanced) | GLM-4.5V (best quality)
-    """)
-    
-    with gr.Tab("🚀 Generate"):
-        with gr.Row():
-            with gr.Column():
-                image_input = gr.Image(type="pil", label="Upload CAD Model Image")
-                model_choice = gr.Dropdown(
-                    choices=["GLM-4.5V-AWQ", "GLM-4.5V-FP8", "GLM-4.5V"],
-                    value="GLM-4.5V-AWQ",
-                    label="Select Model"
-                )
-                prompt_style = gr.Dropdown(
-                    choices=["Simple", "Detailed", "Chain-of-Thought"],
-                    value="Chain-of-Thought", 
-                    label="Prompt Style"
-                )
-                generate_btn = gr.Button("🚀 Generate CADQuery Code", variant="primary")
-            
-            with gr.Column():
-                output = gr.Markdown("Upload an image and click Generate!")
-        
-        generate_btn.click(
-            fn=generate_code,
-            inputs=[image_input, model_choice, prompt_style],
-            outputs=output
-        )
+def main():
+    """Main training pipeline for CPU."""
+    print("🚀 Starting CPU Training Pipeline")
+    print("=" * 50)
     
-    with gr.Tab("⚙️ System"):
-        info_display = gr.Markdown()
-        refresh_btn = gr.Button("🔄 Refresh")
+    try:
+        # 1. Prepare dataset
+        print("\n📊 Step 1: Preparing dataset...")
+        dataset = prepare_simple_dataset(CONFIG["dataset_name"], CONFIG["max_samples"])
         
-        demo.load(fn=system_info, outputs=info_display)
-        refresh_btn.click(fn=system_info, outputs=info_display)
+        # 2. Setup model
+        print("\n🔧 Step 2: Setting up model...")
+        model, tokenizer = setup_simple_model(CONFIG["base_model"])
+        
+        # 3. Train
+        print("\n🏋️ Step 3: Training...")
+        trainer = train_simple_model(model, tokenizer, dataset, CONFIG)
+        
+        # 4. Test
+        print("\n🧪 Step 4: Testing...")
+        test_simple_model(CONFIG["output_dir"])
+        
+        print("\n🎉 Pipeline complete!")
+        print(f"Model saved to: {CONFIG['output_dir']}")
+        
+        return True
+        
+    except Exception as e:
+        print(f"\n❌ Pipeline failed: {e}")
+        return False
 
 if __name__ == "__main__":
-    print("🚀 Starting GLM-4.5V CAD Generator...")
-    print(f"CUDA available: {torch.cuda.is_available()}")
-    demo.queue().launch(ssr_mode=False)  # Disable SSR mode
+    success = main()
+    
+    if success:
+        print("\n📝 Next steps:")
+        print("1. Check the generated model in ./test-cad-model/")
+        print("2. Run test_simple_model() to generate more examples")
+        print("3. Once working, move to GPU version")
+    else:
+        print("\n🔧 Troubleshooting:")
+        print("1. Check internet connection for dataset download")
+        print("2. Ensure you have enough disk space")
+        print("3. Try reducing max_samples to 10")