Spaces:
Sleeping
Sleeping
| import spaces | |
| import gradio as gr | |
| import torch | |
| from transformers import pipeline | |
| from PIL import Image | |
| import time | |
| import traceback | |
| # Global model storage | |
| models = {} | |
| def load_glm_model(model_choice): | |
| """Load GLM model on GPU.""" | |
| 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] | |
| if model_name in models: | |
| return True, f"✅ {model_choice} already loaded" | |
| try: | |
| 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 | |
| return True, f"✅ {model_choice} loaded successfully" | |
| except Exception as e: | |
| error_msg = f"❌ Failed to load {model_choice}: {str(e)[:200]}" | |
| return False, error_msg | |
| def generate_cadquery_code(image, model_choice, prompt_style): | |
| """Generate CADQuery code from image.""" | |
| if image is None: | |
| return "❌ Please upload an image first." | |
| try: | |
| # Create prompt | |
| prompts = { | |
| "Simple": "Generate CADQuery Python code for this 3D model:", | |
| "Detailed": """Analyze this 3D CAD model and generate Python CADQuery code. | |
| Requirements: | |
| - Import cadquery as cq | |
| - Store result in 'result' variable | |
| - Use proper CADQuery syntax | |
| 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, etc.) | |
| Step 3: Generate clean CADQuery Python code | |
| ```python | |
| import cadquery as cq | |
| # Generated code:""" | |
| } | |
| prompt = prompts[prompt_style] | |
| # Load model if needed | |
| 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] | |
| # Load model if not already loaded | |
| if model_name not in models: | |
| pipe = pipeline( | |
| "image-text-to-text", | |
| model=model_name, | |
| device_map="auto", | |
| torch_dtype=torch.float16, | |
| trust_remote_code=True | |
| ) | |
| models[model_name] = pipe | |
| else: | |
| pipe = models[model_name] | |
| # Generate | |
| start_time = time.time() | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "image": image}, | |
| {"type": "text", "text": prompt} | |
| ] | |
| } | |
| ] | |
| result = pipe(messages, max_new_tokens=512, temperature=0.7, do_sample=True) | |
| if isinstance(result, list) and len(result) > 0: | |
| generated_text = result[0].get("generated_text", str(result)) | |
| else: | |
| generated_text = str(result) | |
| generation_time = time.time() - start_time | |
| clean_code = extract_cadquery_code(generated_text) | |
| output = f"""## 🎯 Generated CADQuery Code | |
| ```python | |
| {clean_code} | |
| ``` | |
| ## 📊 Generation Info | |
| - **Model**: {model_choice} | |
| - **Time**: {generation_time:.2f} seconds | |
| - **Prompt**: {prompt_style} | |
| - **Device**: {"GPU" if torch.cuda.is_available() else "CPU"} | |
| ## 🔧 Usage | |
| ```bash | |
| pip install cadquery | |
| python your_script.py | |
| ``` | |
| ## ⚠️ Note | |
| Generated code may need manual adjustments for complex geometries. | |
| """ | |
| return output | |
| except Exception as e: | |
| error_trace = traceback.format_exc() | |
| return f"""❌ **Generation Failed** | |
| **Error**: {str(e)} | |
| **Traceback**: | |
| ``` | |
| {error_trace[:1000]}... | |
| ``` | |
| Try a different model variant or check your image.""" | |
| def extract_cadquery_code(generated_text: str) -> str: | |
| """Extract clean CADQuery code from generated text.""" | |
| text = generated_text.strip() | |
| 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 "import cadquery" in text.lower(): | |
| lines = text.split('\n') | |
| code_lines = [] | |
| started = False | |
| for line in lines: | |
| if "import cadquery" in line.lower(): | |
| started = True | |
| if started: | |
| code_lines.append(line) | |
| code = '\n'.join(code_lines) | |
| else: | |
| code = text | |
| lines = code.split('\n') | |
| cleaned_lines = [] | |
| for line in lines: | |
| line = line.strip() | |
| if line and not line.startswith('```'): | |
| cleaned_lines.append(line) | |
| final_code = '\n'.join(cleaned_lines) | |
| if "import cadquery" not in final_code: | |
| final_code = "import cadquery as cq\n\n" + final_code | |
| 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 test_model_loading(model_choice): | |
| """Test loading a specific model.""" | |
| success, message = load_glm_model(model_choice) | |
| return f"## Test Result\n\n{message}" | |
| def get_system_info(): | |
| """Get system information.""" | |
| info = { | |
| "CUDA Available": torch.cuda.is_available(), | |
| "CUDA Device Count": torch.cuda.device_count() if torch.cuda.is_available() else 0, | |
| "PyTorch Version": torch.__version__, | |
| "Device": "GPU" if torch.cuda.is_available() else "CPU" | |
| } | |
| info_text = "## 🖥️ System Information\n\n" | |
| for key, value in info.items(): | |
| info_text += f"- **{key}**: {value}\n" | |
| return info_text | |
| 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 Generator | |
| Upload a 3D CAD model image and generate CADQuery Python code using GLM-4.5V models! | |
| **Available Models:** | |
| - **GLM-4.5V-AWQ**: AWQ quantized (fastest startup) | |
| - **GLM-4.5V-FP8**: 8-bit quantized (balanced) | |
| - **GLM-4.5V**: Full precision (best quality) | |
| """) | |
| with gr.Tab("🚀 Generate"): | |
| 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-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", 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("🧪 Test"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| test_model_choice = gr.Dropdown( | |
| choices=["GLM-4.5V-AWQ", "GLM-4.5V-FP8", "GLM-4.5V"], | |
| value="GLM-4.5V-AWQ", | |
| 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_model_loading, | |
| inputs=test_model_choice, | |
| outputs=test_output | |
| ) | |
| with gr.Tab("⚙️ System"): | |
| info_output = gr.Markdown() | |
| refresh_btn = gr.Button("🔄 Refresh 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"): | |
| gr.Markdown(""" | |
| ## 🎯 How to Use | |
| 1. **Upload Image**: Clear 3D CAD model images work best | |
| 2. **Select Model**: GLM-4.5V-AWQ is fastest for testing | |
| 3. **Choose Prompt**: Chain-of-Thought usually gives best results | |
| 4. **Generate**: Click the button and wait for results | |
| ## 💡 Tips for Best Results | |
| - Use clear, well-lit CAD images | |
| - Simple geometric shapes work better than complex assemblies | |
| - Try different prompt styles if first attempt isn't satisfactory | |
| ## 🔧 Using Generated Code | |
| ```bash | |
| # Install CADQuery | |
| pip install cadquery | |
| # Run your generated code | |
| python your_cad_script.py | |
| # Export to STL | |
| cq.exporters.export(result, "model.stl") | |
| ``` | |
| ## 🖥️ Hardware Requirements | |
| - This app runs on GPU-enabled Hugging Face Spaces | |
| - First model load takes 5-10 minutes | |
| - Generation takes 15-45 seconds per image | |
| """) | |
| return demo | |
| if __name__ == "__main__": | |
| print("🚀 Starting GLM-4.5V CAD Generator...") | |
| print(f"CUDA available: {torch.cuda.is_available()}") | |
| print(f"PyTorch version: {torch.__version__}") | |
| demo = create_interface() | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| show_error=True | |
| ) |