import gradio as gr
import torch
from PIL import Image
from unsloth import FastVisionModel


MODEL_NAME = "sabaridsnfuji/FloorPlanVisionAIAdaptor"


print("Loading model...")

model, tokenizer = FastVisionModel.from_pretrained(
    MODEL_NAME,
    load_in_4bit=True,
    use_gradient_checkpointing="unsloth"
)

FastVisionModel.for_inference(model)

print("Model loaded successfully!")
print("CUDA available:", torch.cuda.is_available())

if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))


def analyze_floorplan(image, instruction):

    if image is None:
        return "Please upload a floor plan image."

    if not instruction:
        instruction = """
        You are an expert in architecture and interior design.
        Analyze the floor plan image carefully.

        Describe:
        1. Number of rooms
        2. Room names
        3. Approximate layout
        4. Connections between rooms
        5. Doors and windows if visible
        6. Important architectural features
        7. Any other useful observations
        """

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": instruction}
            ]
        }
    ]

    input_text = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True
    )

    inputs = tokenizer(
        image,
        input_text,
        add_special_tokens=False,
        return_tensors="pt"
    ).to("cuda")

    outputs = model.generate(
        **inputs,
        max_new_tokens=1024,
        use_cache=True
    )

    generated_text = tokenizer.decode(
        outputs[0],
        skip_special_tokens=True
    )

    return generated_text


demo = gr.Interface(
    fn=analyze_floorplan,
    inputs=[
        gr.Image(
            type="pil",
            label="Upload Floor Plan"
        ),
        gr.Textbox(
            label="Instruction",
            placeholder="Ask something about the floor plan...",
            value="Analyze this floor plan in detail."
        )
    ],
    outputs=gr.Textbox(
        label="Analysis"
    ),
    title="Floor Plan Vision AI",
    description="AI-powered floor plan analysis using FloorPlanVisionAIAdaptor."
)


if __name__ == "__main__":
    demo.launch()