| """ |
| Inference script for parametric floorplan generation. |
| Generates a JSON floorplan from parametric constraints using a fine-tuned model. |
| """ |
| import json |
| import argparse |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| def build_prompt(params: dict) -> str: |
| """Build natural-language prompt from ProjectCreate-like parameters.""" |
| lines = [ |
| f"Generate a floor plan for project '{params.get('name', 'Project')}.'", |
| f"Plot dimensions: {params['plot_length']}m x {params['plot_width']}m, shape: {params.get('plot_shape', 'rectangular')}.", |
| f"Setbacks: front={params['setback_front']}m, rear={params['setback_rear']}m, left={params['setback_left']}m, right={params['setback_right']}m.", |
| f"Road side: {params['road_side']}, North direction: {params.get('north_direction', 'N')}.", |
| f"Requirements: {params['num_bedrooms']} bedrooms, {params['toilets']} toilets.", |
| ] |
| if params.get("parking"): |
| lines.append("Parking is required.") |
| if params.get("has_pooja"): |
| lines.append("Include a Pooja room.") |
| if params.get("has_study"): |
| lines.append("Include a Study room.") |
| if params.get("has_balcony"): |
| lines.append("Include a Balcony.") |
| if params.get("has_stilt"): |
| lines.append("Stilt parking required.") |
| if params.get("has_basement"): |
| lines.append("Include a basement.") |
| lines.append(f"Number of floors: {params.get('num_floors', 1)} (1=G, 2=G+1, 3=G+2).") |
| if params.get("vastu_enabled"): |
| lines.append("Vastu compliance is enabled.") |
| city = params.get("city", "other") |
| municipality = params.get("municipality") |
| lines.append(f"City: {city}, Municipality: {municipality or 'N/A'}.") |
| return "\n".join(lines) |
|
|
| def generate_floorplan(model_id: str, prompt: str, max_new_tokens: int = 2048, |
| temperature: float = 0.7, top_p: float = 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 |
|
|
| system_msg = ( |
| "You are a parametric floorplan generator for Indian residential construction. " |
| "Given plot dimensions, setbacks, road direction, number of bedrooms/toilets, " |
| "and optional rooms (pooja, study, balcony, parking, basement, stilt), " |
| "output a valid JSON floorplan with plot boundary, buildable boundary, rooms as polygons " |
| "with dimensions and positions, doors, windows, and area summaries." |
| ) |
|
|
| messages = [ |
| {"role": "system", "content": system_msg}, |
| {"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=4096).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, |
| ) |
|
|
| generated = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True) |
| return generated |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Generate a floorplan from parametric input") |
| parser.add_argument("--model_id", type=str, default="Karthik8nitt/parametric-floorplan-generator") |
| parser.add_argument("--name", type=str, default="MyHouse") |
| parser.add_argument("--plot_length", type=float, default=15.0) |
| parser.add_argument("--plot_width", type=float, default=12.0) |
| parser.add_argument("--setback_front", type=float, default=1.5) |
| parser.add_argument("--setback_rear", type=float, default=1.0) |
| parser.add_argument("--setback_left", type=float, default=1.0) |
| parser.add_argument("--setback_right", type=float, default=1.0) |
| parser.add_argument("--road_side", type=str, default="N", choices=["N","S","E","W"]) |
| parser.add_argument("--north_direction", type=str, default="N", choices=["N","S","E","W"]) |
| parser.add_argument("--num_bedrooms", type=int, default=3) |
| parser.add_argument("--toilets", type=int, default=3) |
| parser.add_argument("--parking", action="store_true") |
| parser.add_argument("--has_pooja", action="store_true") |
| parser.add_argument("--has_study", action="store_true") |
| parser.add_argument("--has_balcony", action="store_true") |
| parser.add_argument("--has_stilt", action="store_true") |
| parser.add_argument("--has_basement", action="store_true") |
| parser.add_argument("--num_floors", type=int, default=1, choices=[1,2,3]) |
| parser.add_argument("--vastu_enabled", action="store_true") |
| parser.add_argument("--city", type=str, default="Delhi") |
| parser.add_argument("--municipality", type=str, default=None) |
| parser.add_argument("--max_new_tokens", type=int, default=2048) |
| parser.add_argument("--temperature", type=float, default=0.7) |
| parser.add_argument("--top_p", type=float, default=0.9) |
| args = parser.parse_args() |
|
|
| params = vars(args) |
| prompt = build_prompt(params) |
| 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: |
| data = json.loads(result) |
| print("\n--- Parsed JSON (summary) ---") |
| print(f"Project: {data['project_name']}") |
| print(f"Plot shape: {data['plot']['shape']}") |
| print(f"Rooms: {len(data['rooms'])}") |
| print(f"Doors: {len(data['doors'])}") |
| print(f"Windows: {len(data['windows'])}") |
| print(f"Total built-up area: {data['dimensions']['total_built_up_area_sqm']} m²") |
| print(f"Total carpet area: {data['dimensions']['total_carpet_area_sqm']} m²") |
| except Exception as e: |
| print(f"\nWarning: could not parse as JSON: {e}") |
|
|
| if __name__ == "__main__": |
| main() |
|
|
