import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Load base Stable Diffusion model
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    revision="fp16",
    safety_checker=None
)
pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# Load the FloorPlan LoRA weights
pipe.load_lora_weights("ejazhabibdar/sd-FloorPlan-model", weight_name="pytorch_lora_weights.safetensors")

# Prompt builder
def generate_floorplan(bedrooms, bathrooms, floors, plot_size):
    prompt = f"Floor plan of a {plot_size.lower()} house with {bedrooms} bedrooms, {bathrooms} bathrooms, {floors} floors."
    image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
    return image

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🏠 AI Floor Plan Generator")
    with gr.Row():
        bedrooms = gr.Slider(label="Bedrooms", minimum=1, maximum=10, value=3)
        bathrooms = gr.Slider(label="Bathrooms", minimum=1, maximum=5, value=2)
    with gr.Row():
        floors = gr.Slider(label="Floors", minimum=1, maximum=3, value=1)
        plot_size = gr.Dropdown(label="Plot Size", choices=["Small", "Medium", "Large"], value="Medium")
    output_image = gr.Image(label="Generated Floor Plan", type="pil")
    generate_btn = gr.Button("Generate Floor Plan")

    generate_btn.click(
        fn=generate_floorplan,
        inputs=[bedrooms, bathrooms, floors, plot_size],
        outputs=output_image
    )

demo.launch()