Add inference script

generate.py

@@ -0,0 +1,74 @@
+"""
+Inference script for parametric floorplan generation.
+Given parametric constraints, generates a JSON floorplan using a fine-tuned model.
+Usage:
+    python generate.py --room_count 4 --total_area 100 --room_types Bedroom Bathroom Kitchen LivingRoom
+"""
+import json
+import argparse
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+def build_prompt(room_count, total_area, room_types, room_details=None, edges=None):
+    lines = [
+        f"Generate a floor plan with {room_count} rooms and a total area of {total_area} square meters.",
+        f"The room types are: {', '.join(room_types)}."
+    ]
+    if room_details:
+        lines.append("Room details:")
+        for i, rd in enumerate(room_details):
+            lines.append(f"  - Room {i+1} ({rd.get('room_type','unknown')}): area ~{rd.get('area','unspecified')} m², width ~{rd.get('width','unspecified')} m, height ~{rd.get('height','unspecified')} m")
+    if edges:
+        lines.append(f"Adjacency requirements (room indices): {edges}")
+    return "\n".join(lines)
+
+def generate_floorplan(model_id, prompt, max_new_tokens=1024, temperature=0.7, top_p=0.9):
+    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True,
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    messages = [
+        {"role": "system", "content": "You are a parametric floorplan generator. Given constraints about room count, area, room types, and adjacencies, output a valid JSON floorplan with room polygons, areas, and adjacency edges."},
+        {"role": "user", "content": prompt},
+    ]
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+    return tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_id", type=str, default="Karthik8nitt/parametric-floorplan-generator")
+    parser.add_argument("--room_count", type=int, default=4)
+    parser.add_argument("--total_area", type=float, default=100.0)
+    parser.add_argument("--room_types", nargs="+", default=["Bedroom", "Bathroom", "Kitchen", "LivingRoom"])
+    parser.add_argument("--max_new_tokens", type=int, default=1024)
+    parser.add_argument("--temperature", type=float, default=0.7)
+    parser.add_argument("--top_p", type=float, default=0.9)
+    args = parser.parse_args()
+
+    prompt = build_prompt(args.room_count, args.total_area, args.room_types)
+    print("Prompt:\n", prompt)
+    print("\n--- Generating floorplan ---\n")
+    result = generate_floorplan(args.model_id, prompt, args.max_new_tokens, args.temperature, args.top_p)
+    print(result)
+    try:
+        print("\n--- Parsed JSON ---")
+        print(json.dumps(json.loads(result), indent=2))
+    except Exception as e:
+        print(f"\nWarning: could not parse as JSON: {e}")
+
+if __name__ == "__main__":
+    main()