Update app.py

app.py

@@ -1,7 +1,488 @@
+#!/usr/bin/env python3
+"""
+Hugging Face Space App for GLM-4.5V CAD Generation
+Optimized for HF Space deployment with GPU runtime
+"""
+
 import gradio as gr
+import torch
+from transformers import pipeline, AutoProcessor, AutoModelForImageTextToText
+from PIL import Image
+import json
+import time
+import traceback
+import os
+
+# Global model storage
+models = {}
+model_status = {}
+
+def check_environment():
+    """Check the HF Space environment."""
+    info = {
+        "CUDA Available": torch.cuda.is_available(),
+        "CUDA Device Count": torch.cuda.device_count() if torch.cuda.is_available() else 0,
+        "Current Device": torch.cuda.current_device() if torch.cuda.is_available() else "CPU",
+        "GPU Name": torch.cuda.get_device_name() if torch.cuda.is_available() else "None",
+        "Python Version": os.sys.version,
+        "PyTorch Version": torch.__version__,
+        "Space ID": os.environ.get("SPACE_ID", "Unknown"),
+        "Space Author": os.environ.get("SPACE_AUTHOR_NAME", "Unknown")
+    }
+    return info
+
+def load_model(model_name: str):
+    """Load GLM model with error handling."""
+    if model_name in models:
+        return True, f"✅ {model_name} already loaded"
+    
+    try:
+        print(f"🔄 Loading {model_name}...")
+        model_status[model_name] = "Loading..."
+        
+        # Try pipeline approach first (simpler)
+        pipe = pipeline(
+            "image-text-to-text",
+            model=model_name,
+            device_map="auto",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            trust_remote_code=True
+        )
+        
+        models[model_name] = {"pipe": pipe, "type": "pipeline"}
+        model_status[model_name] = "✅ Loaded (Pipeline)"
+        print(f"✅ Successfully loaded {model_name} via pipeline")
+        return True, f"✅ {model_name} loaded successfully"
+        
+    except Exception as e:
+        print(f"❌ Pipeline failed for {model_name}: {e}")
+        
+        try:
+            # Fallback to direct loading
+            print(f"🔄 Trying direct loading for {model_name}...")
+            processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+            model = AutoModelForImageTextToText.from_pretrained(
+                model_name,
+                device_map="auto",
+                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                trust_remote_code=True
+            )
+            
+            models[model_name] = {
+                "processor": processor,
+                "model": model,
+                "type": "direct"
+            }
+            model_status[model_name] = "✅ Loaded (Direct)"
+            print(f"✅ Successfully loaded {model_name} via direct loading")
+            return True, f"✅ {model_name} loaded successfully (direct)"
+            
+        except Exception as e2:
+            error_msg = f"❌ Failed to load {model_name}: {str(e2)[:200]}"
+            model_status[model_name] = error_msg
+            print(error_msg)
+            return False, error_msg
+
+def generate_cadquery_code(image, model_choice, prompt_style, progress=gr.Progress()):
+    """Generate CADQuery code from image using selected model."""
+    
+    if image is None:
+        return "❌ Please upload an image first."
+    
+    # Model mapping
+    model_map = {
+        "GLM-4.5V": "zai-org/GLM-4.5V",
+        "GLM-4.5V-FP8": "zai-org/GLM-4.5V-FP8", 
+        "GLM-4.5V-AWQ": "QuantTrio/GLM-4.5V-AWQ"
+    }
+    
+    model_name = model_map[model_choice]
+    
+    progress(0.1, desc="Loading model...")
+    
+    # Load model if needed
+    success, message = load_model(model_name)
+    if not success:
+        return f"❌ {message}"
+    
+    progress(0.3, desc="Preparing prompt...")
+    
+    # Create prompt based on style
+    prompts = {
+        "Simple": "Generate CADQuery Python code for this 3D model:",
+        
+        "Detailed": """Analyze this 3D CAD model and generate Python CADQuery code to recreate it.
+
+Requirements:
+- Import cadquery as cq
+- Store result in 'result' variable
+- Use proper CADQuery syntax
+- Create complete runnable code
+
+Code:""",
+        
+        "Chain-of-Thought": """Analyze this 3D CAD model step by step:
+
+Step 1: Identify the basic geometry (box, cylinder, etc.)
+Step 2: Note any features (holes, fillets, chamfers, etc.)
+Step 3: Estimate dimensions and proportions
+Step 4: Generate clean CADQuery Python code
+
+Requirements:
+- Import cadquery as cq
+- Store final result in 'result' variable
+- Use realistic dimensions
+- Include proper method chaining
+
+```python
+import cadquery as cq
+
+# Generated CADQuery code:"""
+    }
+    
+    prompt = prompts[prompt_style]
+    
+    progress(0.5, desc="Generating code...")
+    
+    try:
+        start_time = time.time()
+        
+        model_data = models[model_name]
+        
+        if model_data["type"] == "pipeline":
+            # Use pipeline
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": image},
+                        {"type": "text", "text": prompt}
+                    ]
+                }
+            ]
+            
+            result = model_data["pipe"](messages, max_new_tokens=512, temperature=0.7)
+            generated_text = result[0]["generated_text"] if isinstance(result, list) else str(result)
+            
+        else:
+            # Use direct model
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": image},
+                        {"type": "text", "text": prompt}
+                    ]
+                }
+            ]
+            
+            inputs = model_data["processor"].apply_chat_template(
+                messages,
+                add_generation_prompt=True,
+                tokenize=True,
+                return_dict=True,
+                return_tensors="pt",
+            )
+            
+            if torch.cuda.is_available():
+                inputs = {k: v.to(model_data["model"].device) for k, v in inputs.items()}
+            
+            outputs = model_data["model"].generate(
+                **inputs,
+                max_new_tokens=512,
+                temperature=0.7,
+                do_sample=True,
+                top_p=0.9
+            )
+            
+            generated_text = model_data["processor"].tokenizer.decode(
+                outputs[0][inputs["input_ids"].shape[-1]:], 
+                skip_special_tokens=True
+            )
+        
+        generation_time = time.time() - start_time
+        
+        progress(0.8, desc="Processing output...")
+        
+        # Extract clean code
+        clean_code = extract_cadquery_code(generated_text)
+        
+        progress(1.0, desc="Complete!")
+        
+        # Format output
+        output = f"""## 🎯 Generated CADQuery Code
+
+```python
+{clean_code}
+```
+
+## 📊 Generation Info
+- **Model**: {model_choice}
+- **Time**: {generation_time:.2f} seconds
+- **Prompt Style**: {prompt_style}
+- **Device**: {"GPU" if torch.cuda.is_available() else "CPU"}
+
+## 🔧 Usage Instructions
+1. Copy the code above
+2. Install CADQuery: `pip install cadquery`
+3. Run the code to generate your 3D model
+4. Export to STL/STEP format if needed
+
+## ⚠️ Note
+The generated code is AI-generated and may require manual adjustments for complex geometries.
+"""
+        
+        return output
+        
+    except Exception as e:
+        error_trace = traceback.format_exc()
+        return f"""❌ **Generation Failed**
+
+**Error**: {str(e)}
+
+**Model**: {model_choice}
+
+**Traceback**:
+```
+{error_trace}
+```
+
+Try a different model variant or simpler image."""
+
+def extract_cadquery_code(generated_text: str) -> str:
+    """Extract clean CADQuery code from generated text."""
+    text = generated_text.strip()
+    
+    # Look for code blocks
+    if "```python" in text:
+        start = text.find("```python") + 9
+        end = text.find("```", start)
+        if end > start:
+            code = text[start:end].strip()
+        else:
+            code = text[start:].strip()
+    elif "```" in text:
+        start = text.find("```") + 3
+        end = text.find("```", start)
+        if end > start:
+            code = text[start:end].strip()
+        else:
+            code = text[start:].strip()
+    elif "import cadquery" in text.lower():
+        # Find the import statement and take everything after
+        lines = text.split('\n')
+        code_lines = []
+        started = False
+        
+        for line in lines:
+            if "import cadquery" in line.lower() or "import cq" in line.lower():
+                started = True
+            if started:
+                code_lines.append(line)
+        
+        code = '\n'.join(code_lines)
+    else:
+        code = text
+    
+    # Basic cleanup
+    lines = code.split('\n')
+    cleaned_lines = []
+    
+    for line in lines:
+        line = line.strip()
+        if line and not line.startswith('```') and not line.startswith('#') or line.startswith('# '):
+            cleaned_lines.append(line)
+    
+    final_code = '\n'.join(cleaned_lines)
+    
+    # Ensure basic structure
+    if "import cadquery" not in final_code and "import cq" not in final_code:
+        final_code = "import cadquery as cq\n\n" + final_code
+    
+    # Ensure result variable
+    if "result" not in final_code and "=" in final_code:
+        lines = final_code.split('\n')
+        for i, line in enumerate(lines):
+            if "=" in line and ("cq." in line or "Workplane" in line):
+                lines[i] = f"result = {line.split('=', 1)[1].strip()}"
+                break
+        final_code = '\n'.join(lines)
+    
+    return final_code
+
+def get_system_info():
+    """Get system information for debugging."""
+    info = check_environment()
+    
+    info_text = "## 🖥️ System Information\n\n"
+    for key, value in info.items():
+        info_text += f"- **{key}**: {value}\n"
+    
+    info_text += f"\n## 📊 Model Status\n\n"
+    if model_status:
+        for model, status in model_status.items():
+            info_text += f"- **{model}**: {status}\n"
+    else:
+        info_text += "No models loaded yet.\n"
+    
+    return info_text
+
+def test_single_model(model_choice):
+    """Test loading a single model."""
+    model_map = {
+        "GLM-4.5V": "zai-org/GLM-4.5V",
+        "GLM-4.5V-FP8": "zai-org/GLM-4.5V-FP8", 
+        "GLM-4.5V-AWQ": "QuantTrio/GLM-4.5V-AWQ"
+    }
+    
+    model_name = model_map[model_choice]
+    success, message = load_model(model_name)
+    
+    return f"## Test Result for {model_choice}\n\n{message}"
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Create Gradio interface
+def create_interface():
+    """Create the Gradio interface."""
+    
+    with gr.Blocks(title="GLM-4.5V CAD Generator", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("""
+        # 🔧 GLM-4.5V CAD Code Generator
+        
+        Upload a 3D CAD model image and generate CADQuery Python code using state-of-the-art vision-language models!
+        
+        **Available Models:**
+        - **GLM-4.5V**: Full precision (best quality, 106B parameters)
+        - **GLM-4.5V-FP8**: 8-bit quantized (balanced performance/memory)
+        - **GLM-4.5V-AWQ**: AWQ quantized (fastest inference, lowest memory)
+        """)
+        
+        with gr.Tab("🚀 CAD Generation"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    image_input = gr.Image(
+                        type="pil", 
+                        label="Upload CAD Model Image",
+                        height=400
+                    )
+                    
+                    model_choice = gr.Dropdown(
+                        choices=["GLM-4.5V", "GLM-4.5V-FP8", "GLM-4.5V-AWQ"],
+                        value="GLM-4.5V-FP8",  # Default to balanced option
+                        label="Select GLM Model Variant"
+                    )
+                    
+                    prompt_style = gr.Dropdown(
+                        choices=["Simple", "Detailed", "Chain-of-Thought"],
+                        value="Chain-of-Thought",
+                        label="Prompt Strategy"
+                    )
+                    
+                    generate_btn = gr.Button("🚀 Generate CADQuery Code", variant="primary", size="lg")
+                
+                with gr.Column(scale=2):
+                    output_text = gr.Markdown(
+                        label="Generated Code",
+                        value="Upload an image and click 'Generate' to start!"
+                    )
+            
+            generate_btn.click(
+                fn=generate_cadquery_code,
+                inputs=[image_input, model_choice, prompt_style],
+                outputs=output_text
+            )
+        
+        with gr.Tab("🔍 Model Testing"):
+            with gr.Row():
+                with gr.Column():
+                    test_model_choice = gr.Dropdown(
+                        choices=["GLM-4.5V", "GLM-4.5V-FP8", "GLM-4.5V-AWQ"],
+                        value="GLM-4.5V-FP8",
+                        label="Model to Test"
+                    )
+                    test_btn = gr.Button("🧪 Test Model Loading", variant="secondary")
+                
+                with gr.Column():
+                    test_output = gr.Markdown(value="Click 'Test Model Loading' to check if models work.")
+            
+            test_btn.click(
+                fn=test_single_model,
+                inputs=test_model_choice,
+                outputs=test_output
+            )
+        
+        with gr.Tab("⚙️ System Info"):
+            info_output = gr.Markdown()
+            refresh_btn = gr.Button("🔄 Refresh System Info")
+            
+            # Load initial system info
+            demo.load(fn=get_system_info, outputs=info_output)
+            refresh_btn.click(fn=get_system_info, outputs=info_output)
+        
+        with gr.Tab("📖 Help & Examples"):
+            gr.Markdown("""
+            ## 🎯 How to Use
+            
+            1. **Upload Image**: Use clear, well-lit 3D CAD model images (PNG, JPG, JPEG)
+            2. **Select Model**: 
+               - GLM-4.5V: Best quality, slower (if you have good GPU)
+               - GLM-4.5V-FP8: Balanced option (recommended)
+               - GLM-4.5V-AWQ: Fastest, uses least memory
+            3. **Choose Prompt Style**:
+               - Simple: Basic prompt
+               - Detailed: More specific requirements
+               - Chain-of-Thought: Step-by-step analysis (best results)
+            4. **Generate**: Click the button and wait for results
+            
+            ## 💡 Tips for Best Results
+            
+            - Use clear, uncluttered CAD images
+            - Simple geometric shapes work better than complex assemblies
+            - Good lighting and contrast help model recognition
+            - Try different prompt styles if first attempt isn't satisfactory
+            
+            ## 📝 Example Use Cases
+            
+            - Mechanical parts (brackets, housings, gears)
+            - Architectural elements (columns, beams, panels)
+            - Product design components
+            - Educational CAD modeling
+            
+            ## ⚙️ Generated Code Usage
+            
+            ```python
+            # Install CADQuery
+            pip install cadquery
+            
+            # Run generated code
+            import cadquery as cq
+            # ... your generated code ...
+            
+            # Export result
+            cq.exporters.export(result, "model.stl")
+            ```
+            
+            ## 🐛 Troubleshooting
+            
+            - **Model loading fails**: Try a different variant (FP8 or AWQ)
+            - **Generation is slow**: Use GLM-4.5V-AWQ for faster results
+            - **Code has errors**: Try Chain-of-Thought prompt style
+            - **Poor results**: Ensure image is clear and well-lit
+            """)
+    
+    return demo
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+if __name__ == "__main__":
+    print("🚀 Starting GLM-4.5V CAD Generator...")
+    print("Environment check:")
+    info = check_environment()
+    for key, value in info.items():
+        print(f"  {key}: {value}")
+    
+    # Create and launch the interface
+    demo = create_interface()
+    demo.launch(
+        share=False,  # HF Spaces automatically provides public URL
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True
+    )