Update app.py

app.py

@@ -21,32 +21,29 @@ try:
         LlamaTokenizer,
         LlamaForCausalLM
     )
-    from huggingface_hub import snapshot_download
+    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
+    from huggingface_hub import snapshot_download, hf_hub_download
 
 class CADFusionModel:
-    def __init__(self, model_path: str = "microsoft/CADFusion", version: str = "v1_1"):
+    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
-            version: Model version (v1_0 or v1_1)
+            revision: Model revision/branch (use 'main' instead of version numbers)
         """
         self.model_path = model_path
-        self.version = version
+        self.revision = revision
         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()
+        print(f"🚀 Initializing CADFusion from {model_path}@{revision} on {self.device}")
         
         # Initialize tokenizer and model
         self.tokenizer = None
@@ -56,63 +53,52 @@ class CADFusionModel:
         # CAD sequence processing utilities
         self.max_sequence_length = 512
         
-    def _download_model(self) -> str:
-        """Download the model from Hugging Face Hub"""
+    def _load_model(self):
+        """Load the tokenizer and model directly 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,
+            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
             )
-            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")
-                
+            # 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("📝 Using placeholder model for demo purposes")
+            print("📝 Setting up placeholder model for demo purposes")
             self._setup_placeholder_model()
     
     def _setup_placeholder_model(self):
@@ -164,9 +150,13 @@ class CADFusionModel:
             # 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(
-                processed_text,
+                prompt,
                 return_tensors="pt",
                 padding=True,
                 truncation=True,
@@ -184,7 +174,8 @@ class CADFusionModel:
                     top_p=0.9,
                     top_k=50,
                     pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    repetition_penalty=1.1
                 )
             
             # Decode output
@@ -194,8 +185,10 @@ class CADFusionModel:
             )
             
             # Extract the generated part (remove input prompt)
-            if processed_text in generated_sequence:
-                generated_part = generated_sequence.replace(processed_text, "").strip()
+            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
             
@@ -223,10 +216,14 @@ class CADFusionModel:
         """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)"
+            "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()
@@ -234,8 +231,8 @@ class CADFusionModel:
             if key in text_lower:
                 return sequence
         
-        # Default sequence
-        return "Sketch('xy') -> Rectangle(0, 0, 10, 10) -> Extrude(5)"
+        # Default sequence for rectangular objects
+        return "NewSketch().Rectangle(0, 0, 10, 10).Extrude(5)"
 
 # Global model instance
 model = None
@@ -310,18 +307,31 @@ def create_gradio_interface():
         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
-        gr.Markdown("""
-        # 🔧 CADFusion - Text to CAD Generation
-        
-        Convert natural language descriptions into CAD parametric sequences using Microsoft's CADFusion model.
+        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>
+            """)
         
-        **Model**: microsoft/CADFusion v1.1  
-        **Paper**: [Text-to-CAD Generation Through Infusing Visual Feedback in Large Language Models](https://arxiv.org/abs/2501.19054)
+        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.Row():
@@ -331,8 +341,8 @@ def create_gradio_interface():
                 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=3,
-                    value="Create a simple rectangular bracket with two circular holes"
+                    lines=4,
+                    value="Create a rectangular bracket with two circular mounting holes"
                 )
                 
                 # Parameters section
@@ -367,7 +377,7 @@ def create_gradio_interface():
                 gr.Markdown("### 🎯 Generated CAD Sequence")
                 sequence_output = gr.Textbox(
                     label="Parametric Sequence",
-                    lines=8,
+                    lines=10,
                     interactive=False,
                     placeholder="Generated CAD sequence will appear here..."
                 )
@@ -395,33 +405,47 @@ def create_gradio_interface():
                 ["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"]
+                ["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
-        gr.Markdown("""
-        ### ℹ️ 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
-        """)
+        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(
@@ -430,11 +454,22 @@ def create_gradio_interface():
             outputs=[sequence_output, status_output, params_output],
             show_progress=True
         )
+        
+        # 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
@@ -449,11 +484,10 @@ def main():
         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
+        debug=False,            # Disable debug mode in production
         show_error=True,        # Show errors in interface
         quiet=False             # Show startup logs
     )
 
 if __name__ == "__main__":
-    main()
-
+    main()