Update app.py

app.py

@@ -1,200 +1,152 @@
 import spaces
-import gradio as gr
-import torch
-from transformers import pipeline
-from PIL import Image
-import time
-import traceback
-
-# Global model storage for Zero GPU compatibility
-models = {}
-
-@spaces.GPU(duration=300)
-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": "zai-org/GLM-4.5V"
-    }
-    
-    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"
-    
-    try:
-        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 True, f"✅ {model_choice} loaded successfully"
-    except Exception as e:
-        return False, f"❌ Failed to load {model_choice}: {str(e)[:200]}"
-
-@spaces.GPU(duration=120)
-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:
-        # 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]
-        
-        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
-        messages = [{
-            "role": "user",
-            "content": [
-                {"type": "image", "image": image},
-                {"type": "text", "text": prompts[prompt_style]}
-            ]
-        }]
-        
-        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)
-        
-        # 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()
-        
-        if "import cadquery" not in code:
-            code = "import cadquery as cq\n\n" + code
-        
-        return f"""## 🎯 Generated CADQuery Code
-
-```python
-{code}
-```
-
-## 📊 Info
-- **Model**: {model_choice}
-- **Prompt**: {prompt_style}
-- **Device**: {"GPU" if torch.cuda.is_available() else "CPU"}
-
-## 🔧 Usage
-```bash
-pip install cadquery
-python your_script.py
-```
-"""
-        
-    except Exception as e:
-        return f"❌ **Generation Failed**: {str(e)[:500]}"
-
-def test_model(model_choice):
-    """Test model loading."""
-    success, message = load_model_on_gpu(model_choice)
-    return f"## Test Result\n\n{message}"
-
-def system_info():
-    """Get system info."""
-    info = f"""## 🖥️ System Information
-
-- **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"}
-"""
+  import gradio as gr
+  import torch
+  from transformers import pipeline
+  from PIL import Image
+
+  # Global model storage
+  models = {}
+
+  @spaces.GPU(duration=120)
+  def generate_code(image, model_choice, prompt_style):
+      """Generate CADQuery code - single 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:
+          # Model mapping
+          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 or get cached model
+          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
+          messages = [{
+              "role": "user",
+              "content": [
+                  {"type": "image", "image": image},
+                  {"type": "text", "text": prompts[prompt_style]}
+              ]
+          }]
+
+          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)
+
+          # Extract code
+          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()
+
+          if "import cadquery" not in code:
+              code = "import cadquery as cq\n\n" + code
+
+          return f"""## 🎯 Generated CADQuery Code
+
+  ```python
+  {code}
+
+  📊 Info
+
+  - Model: {model_choice}
+  - Prompt: {prompt_style}
+  - Device: GPU
+
+  🔧 Usage
+
+  pip install cadquery
+  python your_script.py
+  """
+
+  except Exception as e:
+      return f"❌ **Generation Failed**: {str(e)[:500]}"
+  def system_info():
+      """Get system info - no GPU needed."""
+      info = f"""## 🖥️ System Information
+
+  - 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")
-            
-            with gr.Column():
-                output = gr.Markdown("Upload an image and click Generate!")
-        
-        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")
-            
-            with gr.Column():
-                test_output = gr.Markdown("Click Test Model to check loading.")
-        
-        test_btn.click(fn=test_model, inputs=test_model_choice, outputs=test_output)
-    
-    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()}")
-    demo.launch(share=True, show_error=True)
+  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")
+
+          with gr.Column():
+              output = gr.Markdown("Upload an image and click Generate!")
+
+      generate_btn.click(
+          fn=generate_code,
+          inputs=[image_input, model_choice, prompt_style],
+          outputs=output
+      )
+
+  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()}")
+      demo.launch()