Update app.py

app.py

@@ -1,16 +1,20 @@
 import sys
 from pathlib import Path
+import gradio as gr
+from huggingface_hub import hf_hub_download
 
-# Add CADFusion repo to Python path
-CADFUSION_DIR = Path("CADFusion")
+# -------------------------------
+# Step 1: Add CADFusion repo to sys.path
+# -------------------------------
+CADFUSION_DIR = Path("CADFusion")  # make sure this folder exists in your Space
 sys.path.append(str(CADFUSION_DIR))
 
 from cadfusion.inference import InferenceRunner
 from cadfusion.utils.config import load_config
 
-# Use HF Hub to download adapter config & checkpoint
-from huggingface_hub import hf_hub_download
-
+# -------------------------------
+# Step 2: Download config & checkpoint from Hugging Face Hub
+# -------------------------------
 CONFIG_PATH = hf_hub_download(
     repo_id="microsoft/CADFusion",
     filename="adapter_config.json",
@@ -23,7 +27,36 @@ CHECKPOINT_PATH = hf_hub_download(
     revision="v1_1"
 )
 
-# Initialize runner
+# -------------------------------
+# Step 3: Initialize CADFusion runner
+# -------------------------------
 config = load_config(CONFIG_PATH)
 runner = InferenceRunner(config=config)
 runner.load_checkpoint(CHECKPOINT_PATH)
+
+# -------------------------------
+# Step 4: Define CAD generation function
+# -------------------------------
+def generate_cad(prompt: str):
+    try:
+        outputs = runner.infer_from_text(prompt)
+        out_file = "output.stl"
+        mesh = outputs.to_trimesh()
+        mesh.export(out_file)
+        return out_file
+    except Exception as e:
+        return f"Error during inference: {str(e)}"
+
+# -------------------------------
+# Step 5: Gradio interface
+# -------------------------------
+demo = gr.Interface(
+    fn=generate_cad,
+    inputs=gr.Textbox(label="Enter CAD description"),
+    outputs=gr.File(label="Generated STL Model"),
+    title="CADFusion - Text to CAD",
+    description="Enter a natural language description and generate a 3D CAD mesh (STL)."
+)
+
+if __name__ == "__main__":
+    demo.launch()