Update app.py

app.py

@@ -1,17 +1,18 @@
-import spaces  # MUST be imported FIRST before torch
+import spaces
 import gradio as gr
 import torch
 from transformers import pipeline
 from PIL import Image
-import traceback
 import time
 
 # Global model storage
 models = {}
 
 @spaces.GPU
-def load_glm_model(model_choice):
-    """Load GLM model with Zero GPU allocation."""
+def generate_cadquery_with_zero_gpu(image_data, model_choice, prompt_style):
+    """Single function that handles everything on Zero GPU."""
+    
+    # Model mapping
     model_map = {
         "GLM-4.5V-AWQ": "QuantTrio/GLM-4.5V-AWQ",
         "GLM-4.5V-FP8": "zai-org/GLM-4.5V-FP8",
@@ -20,81 +21,39 @@ def load_glm_model(model_choice):
     
     model_name = model_map[model_choice]
     
-    if model_name in models:
-        return True, f"✅ {model_choice} already loaded"
-    
     try:
-        print(f"🔄 Loading {model_name} on Zero GPU...")
-        
-        pipe = pipeline(
-            "image-text-to-text",
-            model=model_name,
-            device_map="auto",
-            torch_dtype=torch.float16,
-            trust_remote_code=True
-        )
-        
-        models[model_name] = pipe
-        print(f"✅ Successfully loaded {model_name}")
-        return True, f"✅ {model_choice} loaded successfully"
+        # 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,
+                device_map="auto",
+                torch_dtype=torch.float16,
+                trust_remote_code=True
+            )
+            models[model_name] = pipe
+            print(f"✅ Loaded {model_name}")
         
-    except Exception as e:
-        error_msg = f"❌ Failed to load {model_choice}: {str(e)[:200]}"
-        print(error_msg)
-        return False, error_msg
-
-@spaces.GPU(duration=120)
-def generate_cadquery_code_gpu(image, model_choice, prompt_style):
-    """Generate CADQuery code using Zero GPU allocation."""
-    
-    if image is None:
-        return "❌ Please upload an image first."
-    
-    success, message = load_glm_model(model_choice)
-    if not success:
-        return f"❌ {message}"
-    
-    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]
-    
-    try:
-        start_time = time.time()
+        pipe = models[model_name]
         
-        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"
+        # 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:"
         }
         
-        model_name = model_map[model_choice]
-        pipe = models[model_name]
+        prompt = prompts[prompt_style]
+        
+        # Generate
+        start_time = time.time()
         
         messages = [
             {
-                "role": "user",
+                "role": "user", 
                 "content": [
-                    {"type": "image", "image": image},
+                    {"type": "image", "image": image_data},
                     {"type": "text", "text": prompt}
                 ]
             }
@@ -102,74 +61,66 @@ import cadquery as cq
         
         result = pipe(messages, max_new_tokens=512, temperature=0.7, do_sample=True)
         
-        if isinstance(result, list) and len(result) > 0:
+        if isinstance(result, list):
             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}
 ```
 
-## 📊 Generation Info
-- **Model**: {model_choice}
-- **Time**: {generation_time:.2f} seconds
-- **Prompt**: {prompt_style}
-- **Compute**: Zero GPU (A100)
+## 📊 Info
+- **Model**: {model_choice}  
+- **Time**: {generation_time:.2f}s
+- **Style**: {prompt_style}
 
 ## 🔧 Usage
 ```bash
 pip install cadquery
-python your_script.py
+python script.py
 ```
 """
         
         return output
         
     except Exception as e:
-        return f"❌ **Generation Failed**: {str(e)[:500]}"
+        return f"❌ **Error**: {str(e)[:300]}"
 
-def extract_cadquery_code(generated_text: str) -> str:
-    """Extract clean CADQuery code from generated text."""
-    text = generated_text.strip()
+def extract_code(text):
+    """Extract CADQuery code from generated text."""
+    text = 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()
+        code = text[start:end].strip() if end > start else 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)
+    # Clean up
+    lines = [line.strip() for line in code.split('\n') if line.strip() and not line.strip().startswith('```')]
+    final_code = '\n'.join(lines)
     
+    # Ensure proper structure
     if "import cadquery" not in final_code:
         final_code = "import cadquery as cq\n\n" + final_code
     
@@ -184,92 +135,93 @@ def extract_cadquery_code(generated_text: str) -> str:
     return final_code
 
 @spaces.GPU
-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 create_interface():
-    """Create the Gradio interface."""
+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"
+    }
     
-    with gr.Blocks(title="GLM-4.5V CAD Generator", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("""
-        # 🔧 GLM-4.5V CAD Generator (Zero GPU)
-        
-        Upload a 3D CAD model image and generate CADQuery Python code!
-        
-        **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(scale=1):
-                    image_input = gr.Image(type="pil", label="Upload CAD 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="Model"
-                    )
-                    
-                    prompt_style = gr.Dropdown(
-                        choices=["Simple", "Detailed", "Chain-of-Thought"],
-                        value="Chain-of-Thought",
-                        label="Prompt Style"
-                    )
-                    
-                    generate_btn = gr.Button("🚀 Generate", variant="primary")
-                
-                with gr.Column(scale=2):
-                    output_text = gr.Markdown("Upload image and click Generate!")
+    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]}"
+
+# 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")
             
-            generate_btn.click(
-                fn=generate_cadquery_code_gpu,  # Direct call - no wrapper
-                inputs=[image_input, model_choice, prompt_style],
-                outputs=output_text
-            )
+            with gr.Column():
+                output_area = gr.Markdown("Upload image and generate!")
         
-        with gr.Tab("🧪 Test"):
-            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")
-            test_output = gr.Markdown()
-            
-            test_btn.click(
-                fn=test_model_loading,
-                inputs=test_model_choice,
-                outputs=test_output
-            )
+        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("ℹ️ Info"):
-            gr.Markdown("""
-            ## Zero GPU Benefits
-            - **A100 GPU** allocated on-demand
-            - **Pay per use** - no idle costs
-            - **Automatic scaling**
-            
-            ## Usage Tips
-            - Clear CAD images work best
-            - GLM-4.5V-AWQ is fastest for testing
-            - Chain-of-Thought prompts give best results
-            
-            ## Generated Code
-            Install CADQuery: `pip install cadquery`
-            
-            Run your generated script to create 3D models!
-            """)
+        test_btn.click(test_wrapper, test_select, test_out)
     
-    return demo
+    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
+        """)
 
 if __name__ == "__main__":
-    print("🚀 GLM-4.5V CAD Generator with Zero GPU")
-    
-    demo = create_interface()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860)