Update app.py

app.py

@@ -9,12 +9,12 @@ import time
 # Global model storage
 models = {}
 
-@spaces.GPU  # Zero GPU decorator - allocates GPU only when this function runs
+@spaces.GPU
 def load_glm_model(model_choice):
     """Load GLM model with Zero GPU allocation."""
     model_map = {
-        "GLM-4.5V-FP8": "zai-org/GLM-4.5V-FP8",
         "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"
     }
     
@@ -26,7 +26,6 @@ def load_glm_model(model_choice):
     try:
         print(f"🔄 Loading {model_name} on Zero GPU...")
         
-        # Load with Zero GPU optimization
         pipe = pipeline(
             "image-text-to-text",
             model=model_name,
@@ -37,55 +36,44 @@ def load_glm_model(model_choice):
         
         models[model_name] = pipe
         print(f"✅ Successfully loaded {model_name}")
-        return True, f"✅ {model_choice} loaded successfully on Zero GPU"
+        return True, f"✅ {model_choice} loaded successfully"
         
     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)  # Allocate GPU for up to 2 minutes for generation
+@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."
     
-    # Load model if needed
     success, message = load_glm_model(model_choice)
     if not success:
         return f"❌ {message}"
     
-    # Create prompts
     prompts = {
         "Simple": "Generate CADQuery Python code for this 3D model:",
-        
-        "Detailed": """Analyze this 3D CAD model and generate Python CADQuery code to recreate it.
+        "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
-- 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
+Step 2: Note any features (holes, fillets, etc.)
+Step 3: Generate clean CADQuery Python code
 
 ```python
 import cadquery as cq
 
-# Generated CADQuery code:"""
+# Generated code:"""
     }
     
     prompt = prompts[prompt_style]
@@ -94,15 +82,14 @@ import cadquery as cq
         start_time = time.time()
         
         model_map = {
+            "GLM-4.5V-AWQ": "QuantTrio/GLM-4.5V-AWQ",
             "GLM-4.5V-FP8": "zai-org/GLM-4.5V-FP8",
-            "GLM-4.5V-AWQ": "QuantTrio/GLM-4.5V-AWQ", 
             "GLM-4.5V": "zai-org/GLM-4.5V"
         }
         
         model_name = model_map[model_choice]
         pipe = models[model_name]
         
-        # Create the input for GLM
         messages = [
             {
                 "role": "user",
@@ -113,21 +100,16 @@ import cadquery as cq
             }
         ]
         
-        # Generate with the pipeline
         result = pipe(messages, max_new_tokens=512, temperature=0.7, do_sample=True)
         
-        # Extract the generated text
         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
-        
-        # Extract clean code
         clean_code = extract_cadquery_code(generated_text)
         
-        # Format output
         output = f"""## 🎯 Generated CADQuery Code
 
 ```python
@@ -137,50 +119,25 @@ import cadquery as cq
 ## 📊 Generation Info
 - **Model**: {model_choice}
 - **Time**: {generation_time:.2f} seconds
-- **Prompt Style**: {prompt_style}
+- **Prompt**: {prompt_style}
 - **Compute**: Zero GPU (A100)
 
-## 🔧 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
-Generated using Zero GPU - GPU allocated only during generation for cost efficiency.
+## 🔧 Usage
+```bash
+pip install cadquery
+python your_script.py
+```
 """
         
         return output
         
     except Exception as e:
-        error_trace = traceback.format_exc()
-        return f"""❌ **Generation Failed**
-
-**Error**: {str(e)}
-
-**Model**: {model_choice}
-
-**Traceback**:
-```
-{error_trace[:1000]}...
-```
-
-Try a different model variant or simpler image."""
-
-def generate_cadquery_code(image, model_choice, prompt_style):
-    """Wrapper function WITHOUT progress tracking to avoid Zero GPU conflicts."""
-    
-    # Simple status update instead of progress
-    print(f"🚀 Starting generation with {model_choice}...")
-    result = generate_cadquery_code_gpu(image, model_choice, prompt_style)
-    print("✅ Generation complete!")
-    return result
+        return f"❌ **Generation Failed**: {str(e)[:500]}"
 
 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)
@@ -188,21 +145,13 @@ def extract_cadquery_code(generated_text: str) -> str:
             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() or "import cq" in text.lower():
-        # Find import and take everything after
+    elif "import cadquery" in text.lower():
         lines = text.split('\n')
         code_lines = []
         started = False
         
         for line in lines:
-            if "import cadquery" in line.lower() or "import cq" in line.lower():
+            if "import cadquery" in line.lower():
                 started = True
             if started:
                 code_lines.append(line)
@@ -211,7 +160,6 @@ def extract_cadquery_code(generated_text: str) -> str:
     else:
         code = text
     
-    # Basic cleanup
     lines = code.split('\n')
     cleaned_lines = []
     
@@ -222,11 +170,9 @@ def extract_cadquery_code(generated_text: str) -> str:
     
     final_code = '\n'.join(cleaned_lines)
     
-    # Ensure basic structure
-    if "import cadquery" not in final_code and "import cq" not in final_code:
+    if "import cadquery" not in final_code:
         final_code = "import cadquery as cq\n\n" + final_code
     
-    # Ensure result variable exists
     if "result" not in final_code and "=" in final_code:
         lines = final_code.split('\n')
         for i, line in enumerate(lines):
@@ -237,135 +183,60 @@ def extract_cadquery_code(generated_text: str) -> str:
     
     return final_code
 
-@spaces.GPU  # Test model loading with Zero GPU
+@spaces.GPU
 def test_model_loading(model_choice):
-    """Test loading a specific model on Zero GPU."""
+    """Test loading a specific model."""
     success, message = load_glm_model(model_choice)
-    return f"## Zero GPU Test Result for {model_choice}\n\n{message}\n\n**Note**: GPU allocated only during this test."
-
-def get_system_info():
-    """Get system information."""
-    info_text = """## 🖥️ Zero GPU Information
-
-- **Compute**: Serverless GPU allocation
-- **GPU Type**: A100 (allocated on demand)
-- **Billing**: Pay per second of GPU usage
-- **Benefits**: No idle costs, automatic scaling
-
-## 💰 Cost Efficiency
-- GPU allocated only during generation
-- ~$0.50-2.00 per generation (estimated)
-- No cost when idle
-- Much cheaper than dedicated GPU instances
-
-## ⚡ Performance
-- Cold start: 30-60 seconds (model loading)
-- Warm start: 10-30 seconds (model cached)
-- Generation: 15-45 seconds per image
-
-## 📋 Current Status
-- Space running on CPU (free)
-- Zero GPU allocated automatically when needed
-- No ongoing GPU costs
-"""
-    
-    return info_text
-
-# Simple status tracking without Gradio progress
-def show_status(message):
-    """Simple status display."""
-    return f"**Status**: {message}"
+    return f"## Test Result\n\n{message}"
 
 def create_interface():
-    """Create the Zero GPU optimized interface."""
+    """Create the Gradio interface."""
     
-    with gr.Blocks(title="GLM-4.5V CAD Generator (Zero GPU)", theme=gr.themes.Soft()) as demo:
+    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 using GLM-4.5V models on **Zero GPU**!
+        Upload a 3D CAD model image and generate CADQuery Python code!
         
-        **Zero GPU Benefits:**
-        - 🚀 **A100 GPU** allocated on-demand
-        - 💰 **Pay per use** - no idle costs
-        - ⚡ **Automatic scaling** - no setup required
-        
-        **Available Models:**
-        - **GLM-4.5V-FP8**: 8-bit quantized (recommended)
-        - **GLM-4.5V-AWQ**: AWQ quantized (fastest startup)
-        - **GLM-4.5V**: Full precision (best quality)
+        **Models:** GLM-4.5V-AWQ (fastest) • GLM-4.5V-FP8 (balanced) • GLM-4.5V (best quality)
         """)
         
-        with gr.Tab("🚀 CAD Generation"):
+        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
-                    )
+                    image_input = gr.Image(type="pil", label="Upload CAD Image", height=400)
                     
                     model_choice = gr.Dropdown(
-                        choices=["GLM-4.5V-FP8", "GLM-4.5V-AWQ", "GLM-4.5V"],
-                        value="GLM-4.5V-AWQ",  # AWQ for fastest startup
-                        label="Select GLM Model Variant"
+                        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 Strategy"
+                        label="Prompt Style"
                     )
                     
-                    generate_btn = gr.Button("🚀 Generate with Zero GPU", variant="primary", size="lg")
+                    generate_btn = gr.Button("🚀 Generate", variant="primary")
                 
                 with gr.Column(scale=2):
-                    # Status display
-                    status_display = gr.Markdown(value="**Status**: Ready to generate")
-                    
-                    output_text = gr.Markdown(
-                        label="Generated Code",
-                        value="Upload an image and click 'Generate' to allocate Zero GPU and start generation!"
-                    )
-            
-            # Note about Zero GPU
-            gr.Markdown("""
-            **⚡ Zero GPU Note**: GPU will be allocated automatically when you click generate. 
-            First generation may take 30-60 seconds due to model loading. Subsequent generations will be faster.
-            """)
-            
-            # Update status on click
-            def update_status_and_generate(image, model_choice, prompt_style):
-                if image is None:
-                    return "**Status**: Please upload an image first", "❌ Please upload an image first."
-                
-                # Show generating status
-                yield "**Status**: 🔄 Allocating Zero GPU and generating...", "🔄 Generating CADQuery code..."
-                
-                # Call the actual generation
-                result = generate_cadquery_code(image, model_choice, prompt_style)
-                
-                # Show completion status
-                yield "**Status**: ✅ Generation complete!", result
+                    output_text = gr.Markdown("Upload image and click Generate!")
             
             generate_btn.click(
-                fn=update_status_and_generate,
+                fn=generate_cadquery_code_gpu,  # Direct call - no wrapper
                 inputs=[image_input, model_choice, prompt_style],
-                outputs=[status_display, output_text]
+                outputs=output_text
             )
         
-        with gr.Tab("🔍 Model Testing"):
-            with gr.Row():
-                with gr.Column():
-                    test_model_choice = gr.Dropdown(
-                        choices=["GLM-4.5V-FP8", "GLM-4.5V-AWQ", "GLM-4.5V"],
-                        value="GLM-4.5V-AWQ",
-                        label="Model to Test"
-                    )
-                    test_btn = gr.Button("🧪 Test on Zero GPU", variant="secondary")
-                
-                with gr.Column():
-                    test_output = gr.Markdown(value="Click 'Test on Zero GPU' to verify model loading works.")
+        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,
@@ -373,61 +244,28 @@ def create_interface():
                 outputs=test_output
             )
         
-        with gr.Tab("⚙️ Zero GPU Info"):
-            info_output = gr.Markdown()
-            
-            demo.load(fn=get_system_info, outputs=info_output)
-        
-        with gr.Tab("📖 Help"):
+        with gr.Tab("ℹ️ Info"):
             gr.Markdown("""
-            ## 🎯 How Zero GPU Works
-            
-            1. **Upload Image**: Standard image upload
-            2. **Click Generate**: Zero GPU automatically allocates A100 GPU
-            3. **Model Loading**: First time takes 30-60 seconds
-            4. **Generation**: 15-45 seconds for code generation
-            5. **GPU Release**: GPU automatically released after generation
-            
-            ## 💰 Cost Information
-            
-            - **No idle costs**: Pay only when generating
-            - **Estimated cost**: $0.50-2.00 per generation
-            - **GPU Type**: NVIDIA A100 (40GB VRAM)
-            - **Billing**: Per second of actual GPU usage
-            
-            ## 🚀 Performance Tips
-            
-            - **First generation**: Slower due to model download
-            - **Subsequent generations**: Much faster (model cached)
-            - **GLM-4.5V-AWQ**: Fastest startup time
-            - **Simple images**: Process faster than complex ones
-            
-            ## 🔧 Generated Code Usage
-            
-            ```bash
-            # Install CADQuery
-            pip install cadquery
-            
-            # Run your generated code
-            python your_cad_script.py
+            ## Zero GPU Benefits
+            - **A100 GPU** allocated on-demand
+            - **Pay per use** - no idle costs
+            - **Automatic scaling**
             
-            # Export to STL
-            cq.exporters.export(result, "model.stl")
-            ```
+            ## Usage Tips
+            - Clear CAD images work best
+            - GLM-4.5V-AWQ is fastest for testing
+            - Chain-of-Thought prompts give best results
             
-            ## 💡 Best Practices
+            ## Generated Code
+            Install CADQuery: `pip install cadquery`
             
-            - Start with **GLM-4.5V-AWQ** for testing
-            - Use **Chain-of-Thought** prompts for best results
-            - Upload clear, well-lit CAD images
-            - Try simpler geometries first
+            Run your generated script to create 3D models!
             """)
     
     return demo
 
 if __name__ == "__main__":
-    print("🚀 Starting GLM-4.5V CAD Generator with Zero GPU...")
-    print("Zero GPU will be allocated automatically during generation.")
+    print("🚀 GLM-4.5V CAD Generator with Zero GPU")
     
     demo = create_interface()
     demo.launch(