Update app.py

app.py

@@ -1,502 +1,109 @@
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
-import sys
 import json
-import torch
-import gradio as gr
-import numpy as np
-from PIL import Image
-from pathlib import Path
-import tempfile
-import subprocess
-import shutil
-from typing import Optional, List, Dict, Any
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Add the src directory to Python path for imports
-sys.path.insert(0, './src')
+# Define model and checkpoint paths
+MODEL_PATH = "microsoft/CADFusion"
+REVISION = "2687619"  # Use commit hash from the document
 
+# Load model and tokenizer
 try:
-    from transformers import (
-        AutoTokenizer, 
-        AutoModelForCausalLM,
-        LlamaTokenizer,
-        LlamaForCausalLM
+    logger.info("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_PATH,
+        revision=REVISION,
+        trust_remote_code=True
     )
-    from huggingface_hub import snapshot_download, hf_hub_download
-    print("✅ Successfully imported transformers and huggingface_hub")
-except ImportError as e:
-    print(f"❌ Import error: {e}")
-    print("Installing required packages...")
-    subprocess.run([sys.executable, "-m", "pip", "install", "transformers", "huggingface_hub", "torch", "accelerate"])
-    from transformers import AutoTokenizer, AutoModelForCausalLM
-    from huggingface_hub import snapshot_download, hf_hub_download
-
-class CADFusionModel:
-    def __init__(self, model_path: str = "microsoft/CADFusion", revision: str = "main"):
-        """
-        Initialize the CADFusion model
-        
-        Args:
-            model_path: Path to the model on Hugging Face Hub
-            revision: Model revision/branch (use 'main' instead of version numbers)
-        """
-        self.model_path = model_path
-        self.revision = revision
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
-        print(f"🚀 Initializing CADFusion from {model_path}@{revision} on {self.device}")
-        
-        # Initialize tokenizer and model
-        self.tokenizer = None
-        self.model = None
-        self._load_model()
-        
-        # CAD sequence processing utilities
-        self.max_sequence_length = 512
-        
-    def _load_model(self):
-        """Load the tokenizer and model directly from Hugging Face Hub"""
-        try:
-            print(f"📦 Loading model from {self.model_path}")
-            
-            # Load tokenizer
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_path,
-                revision=self.revision,
-                trust_remote_code=True,
-                padding_side="left",
-                token=os.getenv("HF_TOKEN")  # Use HF token if available
-            )
-            
-            # Ensure pad token exists
-            if self.tokenizer.pad_token is None:
-                self.tokenizer.pad_token = self.tokenizer.eos_token
-            
-            # Load model with appropriate dtype based on device
-            model_kwargs = {
-                "revision": self.revision,
-                "trust_remote_code": True,
-                "torch_dtype": torch.float16 if self.device.type == "cuda" else torch.float32,
-                "token": os.getenv("HF_TOKEN")
-            }
-            
-            # Add device mapping for CUDA
-            if self.device.type == "cuda":
-                model_kwargs["device_map"] = "auto"
-                model_kwargs["low_cpu_mem_usage"] = True
-            
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_path,
-                **model_kwargs
-            )
-            
-            # Move to device if not using device_map
-            if self.device.type != "cuda":
-                self.model = self.model.to(self.device)
-            
-            self.model.eval()
-            print("✅ Model loaded successfully")
-            
-        except Exception as e:
-            print(f"❌ Error loading model: {e}")
-            print("📝 Setting up placeholder model for demo purposes")
-            self._setup_placeholder_model()
-    
-    def _setup_placeholder_model(self):
-        """Setup a placeholder model for demo purposes"""
-        print("⚠️ Setting up placeholder model")
-        # This is a fallback when the actual model can't be loaded
-        self.model = None
-        self.tokenizer = None
-    
-    def preprocess_text(self, text: str) -> str:
-        """Preprocess input text for CAD generation"""
-        # Basic text cleaning and formatting
-        text = text.strip()
-        if not text:
-            return "Generate a simple 3D object"
-        
-        # Add any specific preprocessing for CAD descriptions
-        if not any(word in text.lower() for word in ['create', 'design', 'make', 'generate', 'build']):
-            text = f"Create a {text}"
-        
-        return text
-    
-    def generate_cad_sequence(self, text: str, max_length: int = 512, temperature: float = 0.7) -> Dict[str, Any]:
-        """
-        Generate CAD parametric sequence from text description
-        
-        Args:
-            text: Text description of the CAD object
-            max_length: Maximum sequence length
-            temperature: Generation temperature
-            
-        Returns:
-            Dictionary containing the generated sequence and metadata
-        """
-        try:
-            if self.model is None or self.tokenizer is None:
-                # Return placeholder response
-                return {
-                    "success": False,
-                    "message": "Model not loaded - showing demo output",
-                    "sequence": self._generate_demo_sequence(text),
-                    "text_input": text,
-                    "parameters": {
-                        "max_length": max_length,
-                        "temperature": temperature
-                    }
-                }
-            
-            # Preprocess input text
-            processed_text = self.preprocess_text(text)
-            
-            # Add special formatting for CADFusion if needed
-            # CADFusion may expect specific prompt formatting
-            prompt = f"Design a CAD model: {processed_text}\nCAD sequence:"
-            
-            # Tokenize input
-            inputs = self.tokenizer(
-                prompt,
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=256
-            ).to(self.device)
-            
-            # Generate sequence
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    inputs.input_ids,
-                    attention_mask=inputs.attention_mask,
-                    max_length=max_length,
-                    temperature=temperature,
-                    do_sample=True,
-                    top_p=0.9,
-                    top_k=50,
-                    pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id,
-                    repetition_penalty=1.1
-                )
-            
-            # Decode output
-            generated_sequence = self.tokenizer.decode(
-                outputs[0], 
-                skip_special_tokens=True
-            )
-            
-            # Extract the generated part (remove input prompt)
-            if "CAD sequence:" in generated_sequence:
-                generated_part = generated_sequence.split("CAD sequence:")[-1].strip()
-            elif prompt in generated_sequence:
-                generated_part = generated_sequence.replace(prompt, "").strip()
-            else:
-                generated_part = generated_sequence
-            
-            return {
-                "success": True,
-                "sequence": generated_part,
-                "full_output": generated_sequence,
-                "text_input": processed_text,
-                "parameters": {
-                    "max_length": max_length,
-                    "temperature": temperature
-                }
-            }
-            
-        except Exception as e:
-            print(f"❌ Generation error: {e}")
-            return {
-                "success": False,
-                "message": f"Generation failed: {str(e)}",
-                "sequence": self._generate_demo_sequence(text),
-                "text_input": text
-            }
-    
-    def _generate_demo_sequence(self, text: str) -> str:
-        """Generate a demo CAD sequence for demonstration purposes"""
-        # This is a simplified demo sequence based on the input text
-        demo_sequences = {
-            "cube": "NewSketch().Rectangle(0, 0, 10, 10).Extrude(10)",
-            "cylinder": "NewSketch().Circle(0, 0, 5).Extrude(15)",
-            "sphere": "NewSketch().Circle(0, 0, 5).Revolve(360, [0, 0, 1])",
-            "bracket": "NewSketch().Rectangle(0, 0, 20, 10).Extrude(5).NewSketch('top').Circle(15, 5, 2).Cut(5)",
-            "hole": "NewSketch().Rectangle(0, 0, 15, 8).Extrude(4).NewSketch('top').Circle(7.5, 4, 1.5).Cut(4)",
-            "gear": "NewSketch().Circle(0, 0, 10).Extrude(3).NewSketch('top').Circle(0, 0, 2).Cut(3)",
-            "pipe": "NewSketch().Circle(0, 0, 8).Extrude(20).NewSketch('top').Circle(0, 0, 6).Cut(20)",
-            "bolt": "NewSketch().Circle(0, 0, 4).Extrude(15).NewSketch('top').RegularPolygon(6, 0, 0, 6).Extrude(3)"
-        }
-        
-        text_lower = text.lower()
-        for key, sequence in demo_sequences.items():
-            if key in text_lower:
-                return sequence
-        
-        # Default sequence for rectangular objects
-        return "NewSketch().Rectangle(0, 0, 10, 10).Extrude(5)"
-
-# Global model instance
-model = None
-
-def initialize_model():
-    """Initialize the global model instance"""
-    global model
-    if model is None:
-        print("🔄 Initializing CADFusion model...")
-        try:
-            model = CADFusionModel()
-            if model.model is None:
-                print("⚠️ Model loaded in demo mode - using simulated responses")
-            else:
-                print("✅ Model loaded successfully!")
-        except Exception as e:
-            print(f"❌ Failed to initialize model: {e}")
-            print("🔄 Creating fallback demo model...")
-            model = CADFusionModel()
-    return model
+    logger.info("Loading model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_PATH,
+        revision=REVISION,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True
+    )
+    logger.info("Model and tokenizer loaded successfully.")
+except Exception as e:
+    logger.error(f"Error loading model or tokenizer: {e}")
+    raise Exception(f"Failed to load model from {MODEL_PATH} with revision {REVISION}. Please check the repository and revision ID.")
 
-def generate_cad(
-    text_input: str,
-    max_length: int = 512,
-    temperature: float = 0.7
-) -> tuple:
-    """
-    Gradio interface function for CAD generation
-    
-    Returns:
-        Tuple of (generated_sequence, status_message, parameters_info)
-    """
+# Function to generate CAD model from text description
+def generate_cad_model(text_description):
     try:
-        # Initialize model if needed
-        global model
-        if model is None:
-            model = initialize_model()
-        
-        # Validate inputs
-        if not text_input or not text_input.strip():
-            return "Please provide a text description.", "❌ Error: Empty input", "No parameters"
-        
-        # Generate CAD sequence
-        result = model.generate_cad_sequence(
-            text_input,
-            max_length=max_length,
-            temperature=temperature
+        if not text_description.strip():
+            return "Error: Please provide a valid text description."
+        
+        # Tokenize input
+        inputs = tokenizer(text_description, return_tensors="pt").to(model.device)
+        
+        # Generate output
+        outputs = model.generate(
+            **inputs,
+            max_length=512,
+            num_return_sequences=1,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
         )
         
-        # Format output
-        if result["success"]:
-            status = "✅ Generation successful"
-            sequence = result["sequence"]
-        else:
-            status = f"⚠️ {result.get('message', 'Generation failed')}"
-            sequence = result["sequence"]
-        
-        # Format parameters info
-        params = result.get("parameters", {})
-        param_info = f"Max Length: {params.get('max_length', max_length)}, Temperature: {params.get('temperature', temperature)}"
-        
-        return sequence, status, param_info
+        # Decode output
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
+        # Parse generated text to extract CAD model data (assuming JSON-like output)
+        try:
+            cad_data = json.loads(generated_text)
+            return json.dumps(cad_data, indent=2)
+        except json.JSONDecodeError:
+            return generated_text  # Return raw text if JSON parsing fails
     except Exception as e:
-        error_msg = f"❌ Error: {str(e)}"
-        return "Generation failed", error_msg, "No parameters"
+        logger.error(f"Error during generation: {e}")
+        return f"Error: {str(e)}"
 
+# Gradio interface
 def create_gradio_interface():
-    """Create the Gradio interface"""
-    
-    # Custom CSS for better styling
-    css = """
-    .gradio-container {
-        font-family: 'Arial', sans-serif;
-    }
-    .gr-button-primary {
-        background: linear-gradient(45deg, #1e3a8a, #3b82f6);
-        border: none;
-    }
-    .gr-panel {
-        border-radius: 8px;
-        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
-    }
-    .title-container {
-        text-align: center;
-        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-        padding: 2rem;
-        border-radius: 10px;
-        margin-bottom: 2rem;
-        color: white;
-    }
-    """
-    
-    with gr.Blocks(css=css, title="CADFusion - Text to CAD Generation") as interface:
-        
-        # Header
-        with gr.HTML():
-            gr.HTML("""
-            <div class="title-container">
-                <h1>🔧 CADFusion - Text to CAD Generation</h1>
-                <p>Convert natural language descriptions into CAD parametric sequences using Microsoft's CADFusion model.</p>
-            </div>
-            """)
-        
-        gr.Markdown("""
-        **Model**: microsoft/CADFusion (based on LLaMA-3-8B)  
-        **Paper**: [Text-to-CAD Generation Through Infusing Visual Feedback in Large Language Models](https://arxiv.org/abs/2501.19054)  
-        **Repository**: [GitHub](https://github.com/microsoft/CADFusion)
-        """)
+    with gr.Blocks() as demo:
+        gr.Markdown("# CADFusion: Text-to-CAD Generation")
+        gr.Markdown("Enter a textual description of the CAD model you want to generate. For example: 'A 3D model of a chair with four legs and a curved backrest.'")
         
         with gr.Row():
-            with gr.Column(scale=2):
-                # Input section
-                gr.Markdown("### 📝 Input")
+            with gr.Column():
                 text_input = gr.Textbox(
-                    label="CAD Description",
-                    placeholder="Describe the CAD object you want to create (e.g., 'Create a cylindrical bracket with mounting holes')",
-                    lines=4,
-                    value="Create a rectangular bracket with two circular mounting holes"
-                )
-                
-                # Parameters section
-                gr.Markdown("### ⚙️ Generation Parameters")
-                with gr.Row():
-                    max_length = gr.Slider(
-                        label="Max Length",
-                        minimum=128,
-                        maximum=1024,
-                        value=512,
-                        step=64,
-                        info="Maximum length of generated sequence"
-                    )
-                    temperature = gr.Slider(
-                        label="Temperature",
-                        minimum=0.1,
-                        maximum=1.5,
-                        value=0.7,
-                        step=0.1,
-                        info="Generation randomness (lower = more deterministic)"
-                    )
-                
-                # Generate button
-                generate_btn = gr.Button(
-                    "🚀 Generate CAD Sequence",
-                    variant="primary",
-                    size="lg"
+                    label="Text Description",
+                    placeholder="Enter your CAD model description here...",
+                    lines=5
                 )
+                submit_button = gr.Button("Generate CAD Model")
             
-            with gr.Column(scale=3):
-                # Output section
-                gr.Markdown("### 🎯 Generated CAD Sequence")
-                sequence_output = gr.Textbox(
-                    label="Parametric Sequence",
-                    lines=10,
-                    interactive=False,
-                    placeholder="Generated CAD sequence will appear here..."
+            with gr.Column():
+                output_text = gr.Textbox(
+                    label="Generated CAD Model (JSON or Text)",
+                    placeholder="Generated output will appear here...",
+                    lines=10
                 )
-                
-                status_output = gr.Textbox(
-                    label="Status",
-                    lines=1,
-                    interactive=False
-                )
-                
-                params_output = gr.Textbox(
-                    label="Parameters Used",
-                    lines=1,
-                    interactive=False
-                )
-        
-        # Examples section
-        gr.Markdown("### 💡 Example Prompts")
-        examples = gr.Examples(
-            examples=[
-                ["Create a cylindrical rod with a square base"],
-                ["Design a mounting bracket with four holes"],
-                ["Make a simple cube with rounded corners"],
-                ["Create a T-shaped connector piece"],
-                ["Design a gear wheel with 12 teeth"],
-                ["Make a pipe elbow joint at 90 degrees"],
-                ["Create a hexagonal bolt head"],
-                ["Design a simple housing enclosure"],
-                ["Create a rectangular plate with center hole"],
-                ["Design a cylindrical bearing housing"]
-            ],
-            inputs=[text_input],
-            label="Click on any example to try it out"
-        )
-        
-        # Information section
-        with gr.Accordion("ℹ️ About CADFusion", open=False):
-            gr.Markdown("""
-            ### Model Overview
-            
-            CADFusion is a state-of-the-art text-to-CAD generation model that:
-            - Uses visual feedback to enhance LLM performance
-            - Generates parametric sequences for CAD modeling
-            - Supports complex 3D object descriptions
-            - Based on alternating sequential and visual learning stages
-            
-            ### Training Approach
-            - **Sequential Learning**: Fine-tuning LLM with paired text-CAD data
-            - **Visual Feedback**: Using vision-language models to improve generation quality
-            - **Alternating Training**: 9 rounds of SL and VF stages for optimal performance
-            
-            ### Usage Tips
-            - Be specific about shapes, dimensions, and features
-            - Use technical CAD terminology when possible
-            - Mention materials or constraints if relevant
-            - Start with simple descriptions and add complexity gradually
-            
-            ### Model Specifications
-            - **Base Model**: LLaMA-3-8B
-            - **Training Data**: SkexGen dataset with human annotations
-            - **License**: MIT License
-            - **Intended Use**: Research and educational purposes
-            
-            ### Performance
-            CADFusion significantly outperforms baselines like GPT-4o and Text2CAD:
-            - **VLM Score**: 8.96 (vs 5.13 for GPT-4o, 2.01 for Text2CAD)
-            - **Better**: Generation diversity, visual quality, and technical accuracy
-            """)
         
-        # Connect the generate button to the function
-        generate_btn.click(
-            fn=generate_cad,
-            inputs=[text_input, max_length, temperature],
-            outputs=[sequence_output, status_output, params_output],
-            show_progress=True
+        submit_button.click(
+            fn=generate_cad_model,
+            inputs=text_input,
+            outputs=output_text
         )
         
-        # Auto-generate on example selection
-        examples.click(
-            fn=generate_cad,
-            inputs=[text_input, max_length, temperature],
-            outputs=[sequence_output, status_output, params_output],
-            show_progress=True
-        )
-    
-    return interface
-
-def main():
-    """Main function to run the Gradio app"""
-    print("===== Application Startup at {} =====".format(
-        __import__('datetime').datetime.now().strftime('%Y-%m-%d %H:%M:%S')
-    ))
-    print("🌟 Starting CADFusion Gradio App")
-    
-    # Initialize model
-    print("🔄 Initializing model...")
-    initialize_model()
-    
-    # Create and launch interface
-    interface = create_gradio_interface()
+        gr.Markdown("""
+        **Note**: 
+        - CADFusion is for research purposes only. Generated models may not be technically accurate and require validation.
+        - Ensure descriptions are clear and specific for best results.
+        - For more details, visit the [CADFusion GitHub repo](https://github.com/microsoft/CADFusion).
+        """)
     
-    # Launch configuration
-    interface.launch(
-        server_name="0.0.0.0",  # Allow external access
-        server_port=7860,       # Standard Gradio port
-        share=False,            # Set to True for public sharing
-        debug=False,            # Disable debug mode in production
-        show_error=True,        # Show errors in interface
-        quiet=False             # Show startup logs
-    )
+    return demo
 
+# Launch Gradio app
 if __name__ == "__main__":
-    main()
+    demo = create_gradio_interface()
+    demo.launch(server_name="0.0.0.0", server_port=7860)