Apple commited on
Commit
c02daf6
·
1 Parent(s): 5bc3c80

Add CADFusion Gradio demo

Browse files
Files changed (4) hide show
  1. .app.py.swp +0 -0
  2. app.py +61 -0
  3. requirements.txt +9 -0
  4. spaces-config.json +7 -0
.app.py.swp ADDED
Binary file (12.3 kB). View file
 
app.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app.py
2
+ import gradio as gr
3
+ import torch
4
+ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
5
+
6
+ MODEL_ID = "microsoft/CADFusion"
7
+
8
+ def load_model():
9
+ print("Loading tokenizer...")
10
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
11
+ print("Trying to load model in 4-bit (bitsandbytes)...")
12
+ try:
13
+ bnb_config = BitsAndBytesConfig(
14
+ load_in_4bit=True,
15
+ bnb_4bit_quant_type="nf4",
16
+ bnb_4bit_use_double_quant=True,
17
+ bnb_4bit_compute_dtype=torch.float16,
18
+ )
19
+ model = AutoModelForCausalLM.from_pretrained(
20
+ MODEL_ID,
21
+ quantization_config=bnb_config,
22
+ device_map="auto",
23
+ trust_remote_code=True,
24
+ )
25
+ print("Loaded in 4-bit")
26
+ except Exception as e:
27
+ print("4-bit load failed:", e)
28
+ print("Falling back to fp16 (may require larger GPU RAM)...")
29
+ model = AutoModelForCausalLM.from_pretrained(
30
+ MODEL_ID,
31
+ device_map="auto",
32
+ torch_dtype=torch.float16,
33
+ trust_remote_code=True,
34
+ )
35
+
36
+ model.eval()
37
+ return tokenizer, model
38
+
39
+ tokenizer, model = load_model()
40
+
41
+ def generate(prompt, max_new_tokens=256):
42
+ if prompt is None or prompt.strip() == "":
43
+ return "Please provide a text description of the CAD model."
44
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
45
+ with torch.no_grad():
46
+ out = model.generate(**inputs, max_new_tokens=int(max_new_tokens), do_sample=False)
47
+ text = tokenizer.decode(out[0], skip_special_tokens=True)
48
+ return text
49
+
50
+ with gr.Blocks() as demo:
51
+ gr.Markdown("# CADFusion demo (microsoft/CADFusion)\nEnter a design description and hit Generate.")
52
+ with gr.Row():
53
+ prompt = gr.Textbox(lines=5, placeholder="e.g. 'a coffee mug with cylindrical body and curved handle'")
54
+ tokens = gr.Slider(64, 1024, value=256, label="max_new_tokens")
55
+ out = gr.Textbox(lines=20)
56
+ btn = gr.Button("Generate")
57
+ btn.click(fn=generate, inputs=[prompt, tokens], outputs=out)
58
+
59
+ if __name__ == "__main__":
60
+ demo.launch()
61
+
requirements.txt ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ gradio>=3.30
2
+ transformers>=4.32
3
+ bitsandbytes
4
+ torch
5
+ accelerate
6
+ safetensors
7
+ sentencepiece
8
+ huggingface-hub
9
+
spaces-config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "sdk": "gradio",
3
+ "app_file": "app.py",
4
+ "suggested_hardware": "gpu",
5
+ "suggested_storage": "large"
6
+ }
7
+