Initial CADFusion Space with Gradio

app.py

@@ -1,34 +1,39 @@
+import os, sys, subprocess
+
+# Clone CADFusion repo if not already cloned
+if not os.path.exists("CADFusion"):
+    subprocess.check_call(["git", "clone", "https://github.com/microsoft/CADFusion.git"])
+
+# Add CADFusion repo to Python path
+sys.path.append("CADFusion")
+
+# Now import from repo
+from models import CADFusionModel  # repo folder has `models.py`
+
 import gradio as gr
 import torch
-from cadfusion.models import CADFusionModel
-
-# Load model (from HF weights + GitHub code)
-print("Loading CADFusion model...")
-device = "cuda" if torch.cuda.is_available() else "cpu"
+from transformers import AutoTokenizer
 
-model = CADFusionModel.from_pretrained("microsoft/CADFusion")
-model = model.to(device)
-model.eval()
+# Load HF checkpoint
+checkpoint = "microsoft/CADFusion"
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+model = CADFusionModel.from_pretrained(checkpoint)
 
-def generate(prompt):
-    """Run CADFusion inference on user prompt"""
+# Define inference function
+def run_cadfusion(prompt: str):
+    inputs = tokenizer(prompt, return_tensors="pt")
     with torch.no_grad():
-        output = model.generate(
-            prompt,
-            max_new_tokens=256,
-            temperature=0.7,
-            top_p=0.9,
-        )
-    return output
+        output = model.generate(**inputs, max_new_tokens=128)
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
 # Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("## 🏗️ CADFusion Demo\nEnter a CAD prompt below:")
-    inp = gr.Textbox(label="Your CAD prompt")
-    out = gr.Textbox(label="Model Output")
-
-    btn = gr.Button("Generate")
-    btn.click(fn=generate, inputs=inp, outputs=out)
+demo = gr.Interface(
+    fn=run_cadfusion,
+    inputs="text",
+    outputs="text",
+    title="CADFusion Demo",
+    description="Run Microsoft's CADFusion model"
+)
 
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -1,4 +1,4 @@
-git+https://github.com/microsoft/CADFusion.git
 gradio
 torch
-transformers
+transformers
+gitpython