import torch
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from controlnet_aux import CannyDetector
from PIL import Image
import gradio as gr
import os

# Define the Hugging Face repository ID for your model
HF_REPO_ID = "saad206121/ai_interior_design_model"

# Load the entire pipeline directly from the Hugging Face Hub
# The LoRA weights are already integrated into this saved pipeline
try:
    pipe = StableDiffusionControlNetPipeline.from_pretrained(HF_REPO_ID, torch_dtype=torch.float16)
    print(f"✅ Model loaded successfully from Hugging Face Hub: {HF_REPO_ID}")
except Exception as e:
    print(f"⚠️ Could not load model from Hugging Face Hub ({HF_REPO_ID}): {e}. Please check the repo ID and permissions.")
    # Fallback to loading base model and controlnet if custom model fails
    controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
    pipe = StableDiffusionControlNetPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
    )
    print("Using base Stable Diffusion v1.5 with ControlNet.")

pipe.to("cuda")
canny = CannyDetector()

def design_room(image, user_prompt, negative_prompt, num_steps, guidance_scale):
    if image is None:
        return None

    # Preprocess: Get edges
    image = Image.fromarray(image)
    canny_image = canny(image, detect_resolution=512, image_resolution=512)

    # Combine user prompt with trigger word
    full_prompt = f"sks bedroom interior, {user_prompt}"

    result = pipe(
        full_prompt,
        image=canny_image,
        negative_prompt=negative_prompt,
        num_inference_steps=num_steps,
        guidance_scale=guidance_scale
    ).images[0]

    return result

# --- GRADIO UI ---
iface = gr.Interface(
    fn=design_room,
    inputs=[
        gr.Image(label="Upload Room Photo"),
        gr.Textbox(label="Prompt (e.g. 'modern luxury, blue furniture')"),
        gr.Textbox(label="Negative Prompt", value="low quality, blurry, distorted, messy"),
        gr.Slider(10, 50, value=50, label="Steps"),
        gr.Slider(1, 15, value=15, label="Guidance Scale")
    ],
    outputs=gr.Image(label="Designed Room"),
    title="AI Interior Designer (Custom Trained)",
    description="Upload a photo of a room. The AI will redesign it using your custom trained style while keeping the furniture in the same place."
)

iface.launch(share=False) # share=False for Hugging Face Spaces