generate.py

"""
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()