Initial CADCoder Space with Gradio

app.py

@@ -1,39 +1,122 @@
+# app.py
+import os
+import subprocess
+import traceback
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-
-# Load CAD-Coder model from Hugging Face
-MODEL_NAME = "CADCODER/CAD-Coder"
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    trust_remote_code=True,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
-
-def generate_code(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9
-        )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+MODEL_ID = "CADCODER/CAD-Coder"   # HF model id
+REPO_GIT = "https://github.com/CADCODER/CAD-Coder.git"
+REPO_DIR = "CAD-Coder"
+
+# 1) git-clone the repo if missing (your preference)
+if not os.path.isdir(REPO_DIR):
+    try:
+        print("Cloning CAD-Coder repo...")
+        subprocess.run(["git", "clone", REPO_GIT, REPO_DIR], check=True)
+    except Exception as e:
+        print("Could not clone repository:", e)
+
+# 2) Prepare model loader with graceful fallback to HF Inference API
+hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HF_HUB_API_TOKEN")
+local_generate = None
+api_generate = None
+
+# Try to load model locally (8-bit if possible)
+try:
+    import torch
+    from transformers import AutoTokenizer, AutoModelForCausalLM
+    try:
+        import bitsandbytes as bnb  # optional; enables 8-bit loading
+        has_bnb = True
+    except Exception:
+        has_bnb = False
+
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_auth_token=hf_token, trust_remote_code=True)
+
+    load_kwargs = {"device_map": "auto", "trust_remote_code": True}
+    if has_bnb:
+        print("bitsandbytes available — will attempt 8-bit load (saves memory).")
+        load_kwargs.update({"load_in_8bit": True, "torch_dtype": torch.float16})
+    else:
+        # attempt fp16 auto if GPU present
+        if torch.cuda.is_available():
+            load_kwargs["torch_dtype"] = torch.float16
+
+    print("Loading model (this can take a while)...")
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, use_auth_token=hf_token, **load_kwargs)
+
+    if hasattr(model, "to"):
+        # make sure model moved to devices by device_map
+        pass
+
+    device = next(model.parameters()).device
+    print("Model loaded on device:", device)
+
+    def local_generate_fn(prompt, max_new_tokens=512):
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+        gen = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
+        return tokenizer.decode(gen[0], skip_special_tokens=True)
+
+    local_generate = local_generate_fn
+
+except Exception as e:
+    print("Local model load failed or not feasible in this environment.")
+    traceback.print_exc()
+
+# Fallback: Hugging Face Inference API (works without loading weights locally)
+if local_generate is None:
+    try:
+        from huggingface_hub import InferenceApi
+        print("Setting up HF Inference API client as fallback...")
+        api = InferenceApi(repo_id=MODEL_ID, token=hf_token)
+
+        def api_generate_fn(prompt, max_new_tokens=512):
+            # call the hosted inference endpoint
+            out = api(inputs=prompt, params={"max_new_tokens": max_new_tokens})
+            # Response can be a dict or list depending on pipeline; extract defensively
+            if isinstance(out, list):
+                first = out[0]
+                if isinstance(first, dict):
+                    return first.get("generated_text") or str(first)
+                return str(first)
+            elif isinstance(out, dict):
+                return out.get("generated_text") or str(out)
+            else:
+                return str(out)
+
+        api_generate = api_generate_fn
+        print("Inference API fallback ready.")
+    except Exception as e:
+        print("HF Inference API not available:", e)
+        traceback.print_exc()
+
+# Final generate function: prefer local, otherwise API fallback, otherwise error
+def generate(prompt, max_new_tokens=512):
+    if local_generate:
+        return local_generate(prompt, max_new_tokens=max_new_tokens)
+    elif api_generate:
+        return api_generate(prompt, max_new_tokens=max_new_tokens)
+    else:
+        return "ERROR: No model loaded and no API fallback available. Check HF_TOKEN and Space hardware."
 
 # Gradio UI
-demo = gr.Interface(
-    fn=generate_code,
-    inputs=gr.Textbox(lines=5, placeholder="Enter your CAD design prompt..."),
-    outputs="text",
-    title="CAD-Coder Inference",
-    description="Generate CAD code from natural language using CAD-Coder."
-)
+def run_prompt(prompt, max_tokens=512):
+    if not prompt or prompt.strip() == "":
+        return "Enter a prompt describing the CAD sketch you want (e.g., 'rectangle width 10 height 5 with hole radius 1')."
+    try:
+        return generate(prompt, max_new_tokens=int(max_tokens))
+    except Exception as e:
+        traceback.print_exc()
+        return f"Generation error: {e}"
 
-if __name__ == "__main__":
-    demo.launch()
+with gr.Blocks() as demo:
+    gr.Markdown("# CAD-Coder (Text → CadQuery code)")
+    prompt = gr.Textbox(label="Natural language prompt", lines=4, placeholder="e.g. 'create a rectangular plate 100x50 with a centered 10mm hole'...")
+    max_tokens = gr.Slider(minimum=64, maximum=2048, step=64, value=512, label="Max new tokens")
+    out = gr.Textbox(label="Generated CadQuery code", lines=18)
+    btn = gr.Button("Generate")
+    btn.click(run_prompt, inputs=[prompt, max_tokens], outputs=out)
 
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))

requirements.txt

@@ -1,5 +1,7 @@
-torch
-transformers
+transformers>=4.30.0
+accelerate
+huggingface-hub
 gradio
-git+https://github.com/huggingface/accelerate.git
-
+torch
+bitsandbytes
+gitpython