Update app.py

app.py

@@ -6,19 +6,21 @@ from PIL import Image
 import time
 import traceback
 
-# Global model storage
+# Global model storage for Zero GPU compatibility
 models = {}
 
 @spaces.GPU(duration=300)
-def load_glm_model(model_choice):
-    """Load GLM model on GPU."""
+def load_model_on_gpu(model_choice):
+    """Load GLM model on GPU - separated for clarity."""
     model_map = {
         "GLM-4.5V-AWQ": "QuantTrio/GLM-4.5V-AWQ",
-        "GLM-4.5V-FP8": "zai-org/GLM-4.5V-FP8",
+        "GLM-4.5V-FP8": "zai-org/GLM-4.5V-FP8", 
         "GLM-4.5V": "zai-org/GLM-4.5V"
     }
     
-    model_name = model_map[model_choice]
+    model_name = model_map.get(model_choice)
+    if not model_name:
+        return False, f"Unknown model: {model_choice}"
     
     if model_name in models:
         return True, f"✅ {model_choice} already loaded"
@@ -28,50 +30,28 @@ def load_glm_model(model_choice):
             "image-text-to-text",
             model=model_name,
             device_map="auto",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            torch_dtype=torch.float16,
             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
+        return False, f"❌ Failed to load {model_choice}: {str(e)[:200]}"
 
 @spaces.GPU(duration=120)
-def generate_cadquery_code(image, model_choice, prompt_style):
-    """Generate CADQuery code from image."""
-    
+def generate_code(image, model_choice, prompt_style):
+    """Generate CADQuery code - main GPU function."""
     if image is None:
         return "❌ Please upload an image first."
     
+    # Create prompts
+    prompts = {
+        "Simple": "Generate CADQuery Python code for this 3D model:",
+        "Detailed": "Analyze this 3D CAD model and generate Python CADQuery code.\n\nRequirements:\n- Import cadquery as cq\n- Store result in 'result' variable\n- Use proper CADQuery syntax\n\nCode:",
+        "Chain-of-Thought": "Analyze this 3D CAD model step by step:\n\nStep 1: Identify the basic geometry\nStep 2: Note any features\nStep 3: Generate clean CADQuery Python code\n\n```python\nimport cadquery as cq\n\n# Generated code:"
+    }
+    
     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",
@@ -81,12 +61,11 @@ import cadquery as cq
         
         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",
+                device_map="auto", 
                 torch_dtype=torch.float16,
                 trust_remote_code=True
             )
@@ -95,37 +74,40 @@ import cadquery as cq
             pipe = models[model_name]
         
         # Generate
-        start_time = time.time()
-        
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image", "image": image},
-                    {"type": "text", "text": prompt}
-                ]
-            }
-        ]
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image},
+                {"type": "text", "text": prompts[prompt_style]}
+            ]
+        }]
         
-        result = pipe(messages, max_new_tokens=512, temperature=0.7, do_sample=True)
+        result = pipe(messages, max_new_tokens=512, temperature=0.7)
         
         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)
+        # Simple code extraction
+        code = generated_text.strip()
+        if "```python" in code:
+            start = code.find("```python") + 9
+            end = code.find("```", start)
+            if end > start:
+                code = code[start:end].strip()
         
-        output = f"""## 🎯 Generated CADQuery Code
+        if "import cadquery" not in code:
+            code = "import cadquery as cq\n\n" + code
+        
+        return f"""## 🎯 Generated CADQuery Code
 
 ```python
-{clean_code}
+{code}
 ```
 
-## 📊 Generation Info
+## 📊 Info
 - **Model**: {model_choice}
-- **Time**: {generation_time:.2f} seconds
 - **Prompt**: {prompt_style}
 - **Device**: {"GPU" if torch.cuda.is_available() else "CPU"}
 
@@ -134,216 +116,85 @@ import cadquery as cq
 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
+        return f"❌ **Generation Failed**: {str(e)[:500]}"
 
-def test_model_loading(model_choice):
-    """Test loading a specific model."""
-    success, message = load_glm_model(model_choice)
+def test_model(model_choice):
+    """Test model loading."""
+    success, message = load_model_on_gpu(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 system_info():
+    """Get system info."""
+    info = f"""## 🖥️ System Information
 
-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.")
+- **CUDA Available**: {torch.cuda.is_available()}
+- **CUDA Devices**: {torch.cuda.device_count() if torch.cuda.is_available() else 0}
+- **PyTorch Version**: {torch.__version__}
+- **Device**: {"GPU" if torch.cuda.is_available() else "CPU"}
+"""
+    return info
+
+# Create interface
+with gr.Blocks(title="GLM-4.5V CAD Generator", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🔧 GLM-4.5V CAD Generator
+    
+    Generate CADQuery Python code from 3D CAD model images using GLM-4.5V models!
+    
+    **Models**: GLM-4.5V-AWQ (fastest) | GLM-4.5V-FP8 (balanced) | GLM-4.5V (best quality)
+    """)
+    
+    with gr.Tab("🚀 Generate"):
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(type="pil", label="Upload CAD Model Image")
+                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")
             
-            test_btn.click(
-                fn=test_model_loading,
-                inputs=test_model_choice,
-                outputs=test_output
-            )
+            with gr.Column():
+                output = gr.Markdown("Upload an image and click Generate!")
         
-        with gr.Tab("⚙️ System"):
-            info_output = gr.Markdown()
-            refresh_btn = gr.Button("🔄 Refresh System Info")
+        generate_btn.click(
+            fn=generate_code,
+            inputs=[image_input, model_choice, prompt_style],
+            outputs=output
+        )
+    
+    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")
             
-            demo.load(fn=get_system_info, outputs=info_output)
-            refresh_btn.click(fn=get_system_info, outputs=info_output)
+            with gr.Column():
+                test_output = gr.Markdown("Click Test Model to check loading.")
         
-        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
-            """)
+        test_btn.click(fn=test_model, inputs=test_model_choice, outputs=test_output)
     
-    return demo
+    with gr.Tab("⚙️ System"):
+        info_display = gr.Markdown()
+        refresh_btn = gr.Button("🔄 Refresh")
+        
+        demo.load(fn=system_info, outputs=info_display)
+        refresh_btn.click(fn=system_info, outputs=info_display)
 
 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
-    )
+    demo.launch(share=True, show_error=True)