Update app.py

app.py

@@ -4,15 +4,14 @@ import torch
 from transformers import pipeline
 from PIL import Image
 import time
+import traceback
 
 # Global model storage
 models = {}
 
-@spaces.GPU
-def generate_cadquery_with_zero_gpu(image_data, model_choice, prompt_style):
-    """Single function that handles everything on Zero GPU."""
-    
-    # Model mapping
+@spaces.GPU(duration=300)
+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",
@@ -21,10 +20,69 @@ def generate_cadquery_with_zero_gpu(image_data, model_choice, prompt_style):
     
     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
+
+@spaces.GPU(duration=120)
+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:
-            print(f"🔄 Loading {model_name}...")
             pipe = pipeline(
                 "image-text-to-text",
                 model=model_name,
@@ -33,27 +91,17 @@ def generate_cadquery_with_zero_gpu(image_data, model_choice, prompt_style):
                 trust_remote_code=True
             )
             models[model_name] = pipe
-            print(f"✅ Loaded {model_name}")
-        
-        pipe = models[model_name]
-        
-        # 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 syntax.",
-            "Chain-of-Thought": "Analyze this 3D CAD model step by step: 1) Identify geometry 2) Note features 3) Generate CADQuery code. ```python\nimport cadquery as cq\n# Generated code:"
-        }
-        
-        prompt = prompts[prompt_style]
+        else:
+            pipe = models[model_name]
         
         # Generate
         start_time = time.time()
         
         messages = [
             {
-                "role": "user", 
+                "role": "user",
                 "content": [
-                    {"type": "image", "image": image_data},
+                    {"type": "image", "image": image},
                     {"type": "text", "text": prompt}
                 ]
             }
@@ -61,66 +109,87 @@ def generate_cadquery_with_zero_gpu(image_data, model_choice, prompt_style):
         
         result = pipe(messages, max_new_tokens=512, temperature=0.7, do_sample=True)
         
-        if isinstance(result, list):
+        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)
         
-        # Extract code
-        clean_code = extract_code(generated_text)
-        
-        # Format output
         output = f"""## 🎯 Generated CADQuery Code
 
 ```python
 {clean_code}
 ```
 
-## 📊 Info
-- **Model**: {model_choice}  
-- **Time**: {generation_time:.2f}s
-- **Style**: {prompt_style}
+## 📊 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 script.py
+python your_script.py
 ```
+
+## ⚠️ Note
+Generated code may need manual adjustments for complex geometries.
 """
         
         return output
         
     except Exception as e:
-        return f"❌ **Error**: {str(e)[:300]}"
+        error_trace = traceback.format_exc()
+        return f"""❌ **Generation Failed**
+
+**Error**: {str(e)}
+
+**Traceback**:
+```
+{error_trace[:1000]}...
+```
 
-def extract_code(text):
-    """Extract CADQuery code from generated text."""
-    text = text.strip()
+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)
-        code = text[start:end].strip() if end > start else text[start:].strip()
+        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
     
-    # Clean up
-    lines = [line.strip() for line in code.split('\n') if line.strip() and not line.strip().startswith('```')]
-    final_code = '\n'.join(lines)
+    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)
     
-    # Ensure proper structure
     if "import cadquery" not in final_code:
         final_code = "import cadquery as cq\n\n" + final_code
     
@@ -134,94 +203,147 @@ def extract_code(text):
     
     return final_code
 
-@spaces.GPU
-def test_model(model_choice):
-    """Test model loading."""
-    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"
+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"
     }
     
-    try:
-        model_name = model_map[model_choice]
-        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
-        return f"✅ **{model_choice}** loaded successfully!"
-    except Exception as e:
-        return f"❌ **{model_choice}** failed: {str(e)[:200]}"
+    info_text = "## 🖥️ System Information\n\n"
+    for key, value in info.items():
+        info_text += f"- **{key}**: {value}\n"
+    
+    return info_text
 
-# Simple wrapper functions that don't use Gradio context
-def generate_wrapper(image, model, style):
-    if image is None:
-        return "❌ Please upload an image first."
-    return generate_cadquery_with_zero_gpu(image, model, style)
-
-def test_wrapper(model):
-    return test_model(model)
-
-# Create interface
-with gr.Blocks(title="GLM CAD Generator", theme=gr.themes.Soft()) as demo:
-    
-    gr.Markdown("""
-    # 🔧 GLM-4.5V CAD Generator
-    
-    Generate CADQuery code from 3D model images using Zero GPU!
-    
-    **Models:** AWQ (fastest) • FP8 (balanced) • Full (best quality)
-    """)
-    
-    with gr.Tab("🚀 Generate"):
-        with gr.Row():
-            with gr.Column():
-                image_input = gr.Image(type="pil", label="CAD Image")
-                model_select = gr.Dropdown(
-                    ["GLM-4.5V-AWQ", "GLM-4.5V-FP8", "GLM-4.5V"],
-                    value="GLM-4.5V-AWQ",
-                    label="Model"
-                )
-                style_select = gr.Dropdown(
-                    ["Simple", "Detailed", "Chain-of-Thought"],
-                    value="Chain-of-Thought", 
-                    label="Style"
-                )
-                gen_btn = gr.Button("🚀 Generate", variant="primary")
+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!"
+                    )
             
-            with gr.Column():
-                output_area = gr.Markdown("Upload image and generate!")
+            generate_btn.click(
+                fn=generate_cadquery_code,
+                inputs=[image_input, model_choice, prompt_style],
+                outputs=output_text
+            )
         
-        gen_btn.click(generate_wrapper, [image_input, model_select, style_select], output_area)
-    
-    with gr.Tab("🧪 Test"):
-        test_select = gr.Dropdown(
-            ["GLM-4.5V-AWQ", "GLM-4.5V-FP8", "GLM-4.5V"],
-            value="GLM-4.5V-AWQ",
-            label="Test Model"
-        )
-        test_btn = gr.Button("Test")
-        test_out = gr.Markdown()
+        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)
         
-        test_btn.click(test_wrapper, test_select, test_out)
+        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
+            """)
     
-    with gr.Tab("ℹ️ Help"):
-        gr.Markdown("""
-        ## How to Use
-        1. Upload clear CAD model image
-        2. Select GLM model variant  
-        3. Choose prompt style
-        4. Click Generate
-        
-        ## Zero GPU
-        - A100 allocated automatically
-        - Pay only when generating
-        - No idle costs
-        
-        ## Tips
-        - AWQ model is fastest
-        - Chain-of-Thought works best
-        - Clear images get better results
-        """)
+    return demo
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    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
+    )