import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

print("Model load ho raha hai...")

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float32,
    safety_checker=None
)

lora_path = hf_hub_download(
    repo_id="maria26/Floor_Plan_LoRA",
    filename="pytorch_lora_weights.safetensors"
)
pipe.load_lora_weights(load_file(lora_path))
pipe = pipe.to("cpu")
print("Model ready!")

def generate(bedrooms, bathrooms, marla, additional):
    prompt = (
        f"{int(marla)} marla house floor plan, "
        f"{int(bedrooms)} bedrooms, {int(bathrooms)} bathrooms, "
        f"{additional}, architectural blueprint, "
        f"top view, black and white, 2D plan, detailed"
    )
    image = pipe(
        prompt,
        num_inference_steps=20,
        guidance_scale=7.5
    ).images[0]
    return image

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Number(label="Bedrooms", value=3, precision=0),
        gr.Number(label="Bathrooms", value=2, precision=0),
        gr.Number(label="Marla", value=5, precision=0),
        gr.Textbox(
            label="Additional Features",
            value="living area, kitchen, stairs, garage"
        )
    ],
    outputs=gr.Image(label="Generated Floor Plan"),
    title="Floor Plan Generator",
    description="Apna floor plan AI se banayein"
)

demo.launch(server_name="0.0.0.0", server_port=7860)