Update app.py

app.py

@@ -1,149 +1,458 @@
-import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import snapshot_download
 import os
+import sys
 import json
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Define model and checkpoint paths
-MODEL_REPO = "microsoft/CADFusion"
-CHECKPOINT_REVISION = "main"
-CHECKPOINT_SUBFOLDER = "exp/model_ckpt/v1_1"
-LOCAL_CHECKPOINT_DIR = "./model_ckpt/v1_1"
-FALLBACK_MODEL = "meta-llama/Llama-2-7b"
+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
 
-# Ensure local checkpoint directory exists
-os.makedirs(LOCAL_CHECKPOINT_DIR, exist_ok=True)
+# Add the src directory to Python path for imports
+sys.path.insert(0, './src')
 
-# Download checkpoint files
 try:
-    logger.info("Downloading checkpoint files...")
-    snapshot_download(
-        repo_id=MODEL_REPO,
-        revision=CHECKPOINT_REVISION,
-        allow_patterns=f"{CHECKPOINT_SUBFOLDER}/*",
-        local_dir=LOCAL_CHECKPOINT_DIR,
-        local_dir_use_symlinks=False
+    from transformers import (
+        AutoTokenizer, 
+        AutoModelForCausalLM,
+        LlamaTokenizer,
+        LlamaForCausalLM
     )
-    logger.info("Checkpoint files downloaded successfully.")
-except Exception as e:
-    logger.error(f"Error downloading checkpoint files: {str(e)}")
-    raise e
+    from huggingface_hub import snapshot_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
 
-# Load model and tokenizer
-try:
-    logger.info("Loading tokenizer from local checkpoint...")
-    tokenizer = AutoTokenizer.from_pretrained(
-        LOCAL_CHECKPOINT_DIR,
-        trust_remote_code=True
-    )
-    logger.info("Loading model from local checkpoint...")
-    model = AutoModelForCausalLM.from_pretrained(
-        LOCAL_CHECKPOINT_DIR,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        trust_remote_code=True
-    )
-    logger.info("Model and tokenizer loaded successfully from local checkpoint.")
-except Exception as e:
-    logger.error(f"Error loading from local checkpoint: {str(e)}")
-    logger.info(f"Falling back to {FALLBACK_MODEL}...")
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(
-            FALLBACK_MODEL,
-            trust_remote_code=True
-        )
-        model = AutoModelForCausalLM.from_pretrained(
-            FALLBACK_MODEL,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            trust_remote_code=True
-        )
-        logger.info(f"Fallback model {FALLBACK_MODEL} loaded successfully.")
-    except Exception as fallback_e:
-        logger.error(f"Error loading fallback model: {str(fallback_e)}")
-        raise fallback_e
+class CADFusionModel:
+    def __init__(self, model_path: str = "microsoft/CADFusion", version: str = "v1_1"):
+        """
+        Initialize the CADFusion model
+        
+        Args:
+            model_path: Path to the model on Hugging Face Hub
+            version: Model version (v1_0 or v1_1)
+        """
+        self.model_path = model_path
+        self.version = version
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        print(f"🚀 Initializing CADFusion {version} on {self.device}")
+        
+        # Download model if not already present
+        self.model_dir = self._download_model()
+        
+        # Initialize tokenizer and model
+        self.tokenizer = None
+        self.model = None
+        self._load_model()
+        
+        # CAD sequence processing utilities
+        self.max_sequence_length = 512
+        
+    def _download_model(self) -> str:
+        """Download the model from Hugging Face Hub"""
+        try:
+            cache_dir = "./model_cache"
+            model_dir = snapshot_download(
+                repo_id=self.model_path,
+                revision=self.version,
+                cache_dir=cache_dir,
+                token=os.getenv("HF_TOKEN")  # Use HF token if available
+            )
+            print(f"✅ Model downloaded to: {model_dir}")
+            return model_dir
+        except Exception as e:
+            print(f"❌ Error downloading model: {e}")
+            # Fallback to local directory structure
+            return f"./{self.version}"
+    
+    def _load_model(self):
+        """Load the tokenizer and model"""
+        try:
+            # Try loading as LLaMA model first (CADFusion is based on LLaMA)
+            model_files = list(Path(self.model_dir).glob("*.bin")) + list(Path(self.model_dir).glob("*.safetensors"))
+            
+            if model_files:
+                print(f"📦 Loading model from {self.model_dir}")
+                
+                # Load tokenizer
+                self.tokenizer = AutoTokenizer.from_pretrained(
+                    self.model_dir,
+                    trust_remote_code=True,
+                    padding_side="left"
+                )
+                
+                # Ensure pad token exists
+                if self.tokenizer.pad_token is None:
+                    self.tokenizer.pad_token = self.tokenizer.eos_token
+                
+                # Load model
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_dir,
+                    torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
+                    device_map="auto" if self.device.type == "cuda" else None,
+                    trust_remote_code=True
+                )
+                
+                if self.device.type != "cuda":
+                    self.model = self.model.to(self.device)
+                
+                self.model.eval()
+                print("✅ Model loaded successfully")
+                
+            else:
+                raise FileNotFoundError("No model files found")
+                
+        except Exception as e:
+            print(f"❌ Error loading model: {e}")
+            print("📝 Using 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)
+            
+            # Tokenize input
+            inputs = self.tokenizer(
+                processed_text,
+                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
+                )
+            
+            # Decode output
+            generated_sequence = self.tokenizer.decode(
+                outputs[0], 
+                skip_special_tokens=True
+            )
+            
+            # Extract the generated part (remove input prompt)
+            if processed_text in generated_sequence:
+                generated_part = generated_sequence.replace(processed_text, "").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": "Sketch('xy') -> Rectangle(0, 0, 10, 10) -> Extrude(10)",
+            "cylinder": "Sketch('xy') -> Circle(0, 0, 5) -> Extrude(15)",
+            "sphere": "Sketch('xy') -> Circle(0, 0, 5) -> Revolve(360)",
+            "bracket": "Sketch('xy') -> Rectangle(0, 0, 20, 10) -> Extrude(5) -> Sketch('top') -> Circle(15, 5, 2) -> Cut(5)"
+        }
+        
+        text_lower = text.lower()
+        for key, sequence in demo_sequences.items():
+            if key in text_lower:
+                return sequence
+        
+        # Default sequence
+        return "Sketch('xy') -> Rectangle(0, 0, 10, 10) -> Extrude(5)"
+
+# Global model instance
+model = None
 
-# Function to generate CAD model from text description
-def generate_cad_model(text_description):
+def initialize_model():
+    """Initialize the global model instance"""
+    global model
+    if model is None:
+        print("🔄 Initializing CADFusion model...")
+        model = CADFusionModel()
+    return model
+
+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)
+    """
     try:
-        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
+        # 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
         )
         
-        # Decode output
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # 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
         
-        # 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:
-        logger.error(f"Error during generation: {str(e)}")
-        return f"Error: {str(e)}"
+        error_msg = f"❌ Error: {str(e)}"
+        return "Generation failed", error_msg, "No parameters"
 
-# Gradio interface
 def create_gradio_interface():
-    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.'")
+    """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);
+    }
+    """
+    
+    with gr.Blocks(css=css, title="CADFusion - Text to CAD Generation") as interface:
+        
+        # Header
+        gr.Markdown("""
+        # 🔧 CADFusion - Text to CAD Generation
+        
+        Convert natural language descriptions into CAD parametric sequences using Microsoft's CADFusion model.
+        
+        **Model**: microsoft/CADFusion v1.1  
+        **Paper**: [Text-to-CAD Generation Through Infusing Visual Feedback in Large Language Models](https://arxiv.org/abs/2501.19054)
+        """)
         
         with gr.Row():
-            with gr.Column():
+            with gr.Column(scale=2):
+                # Input section
+                gr.Markdown("### 📝 Input")
                 text_input = gr.Textbox(
-                    label="Text Description",
-                    placeholder="Enter your CAD model description here...",
-                    lines=5
+                    label="CAD Description",
+                    placeholder="Describe the CAD object you want to create (e.g., 'Create a cylindrical bracket with mounting holes')",
+                    lines=3,
+                    value="Create a simple rectangular bracket with two circular 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"
                 )
-                submit_button = gr.Button("Generate CAD Model")
             
-            with gr.Column():
-                output_text = gr.Textbox(
-                    label="Generated CAD Model (JSON or Text)",
-                    placeholder="Generated output will appear here...",
-                    lines=10
+            with gr.Column(scale=3):
+                # Output section
+                gr.Markdown("### 🎯 Generated CAD Sequence")
+                sequence_output = gr.Textbox(
+                    label="Parametric Sequence",
+                    lines=8,
+                    interactive=False,
+                    placeholder="Generated CAD sequence will appear here..."
+                )
+                
+                status_output = gr.Textbox(
+                    label="Status",
+                    lines=1,
+                    interactive=False
+                )
+                
+                params_output = gr.Textbox(
+                    label="Parameters Used",
+                    lines=1,
+                    interactive=False
                 )
         
-        submit_button.click(
-            fn=generate_cad_model,
-            inputs=text_input,
-            outputs=output_text
+        # 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"]
+            ],
+            inputs=[text_input],
+            label="Click on any example to try it out"
         )
         
+        # Information section
         gr.Markdown("""
-        **Note**: 
-        - This deployment may use a fallback model (Llama-2-7b) due to issues with the CADFusion v1_1 checkpoint.
-        - CADFusion is for research purposes only. Generated models may not be technically accurate and require validation.
-        - For full CADFusion functionality, follow the setup instructions in the [CADFusion GitHub repo](https://github.com/microsoft/CADFusion).
-        - Contact Shizhao Sun (shizsu@microsoft.com) for checkpoint access or issues.
+        ### ℹ️ About CADFusion
+        
+        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
+        
+        **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 Info**:
+        - Version: v1.1 (9 rounds of alternate training)
+        - Base Model: LLaMA architecture
+        - Training Data: SkexGen dataset with human annotations
         """)
+        
+        # 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
+        )
     
-    return demo
+    return interface
+
+def main():
+    """Main function to run the Gradio app"""
+    print("🌟 Starting CADFusion Gradio App")
+    
+    # Initialize model
+    print("🔄 Initializing model...")
+    initialize_model()
+    
+    # Create and launch interface
+    interface = create_gradio_interface()
+    
+    # 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=True,             # Enable debug mode
+        show_error=True,        # Show errors in interface
+        quiet=False             # Show startup logs
+    )
 
-# Launch Gradio app
 if __name__ == "__main__":
-    try:
-        demo = create_gradio_interface()
-        demo.launch()
-    except Exception as e:
-        logger.error(f"Error launching Gradio app: {str(e)}")
-        raise e
+    main()