import torch
import numpy as np
import cv2
import gradio as gr
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
from tqdm import tqdm
from diffusers import StableDiffusionInpaintPipeline, DDIMScheduler
from PIL import Image
from huggingface_hub import hf_hub_download

# 🛠 Kiểm tra thiết bị
device = "cuda" if torch.cuda.is_available() else "cpu"

# 🏗 Tải mô hình SAM từ Hugging Face
sam_checkpoint = hf_hub_download(repo_id="facebook/sam", filename="sam_vit_h_4b8939.pth")
sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint)
mask_generator = SamAutomaticMaskGenerator(sam)

# 🏗 Tải mô hình Stable Diffusion từ Hugging Face
scheduler = DDIMScheduler.from_pretrained("runwayml/stable-diffusion-inpainting", subfolder="scheduler")
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-inpainting",
    scheduler=scheduler,
    torch_dtype=torch.float32,
    cache_dir="./models",
    low_cpu_mem_usage=True
).to(device)

pipe.enable_attention_slicing()

# Hàm xử lý ảnh
def inpaint(image, mask, prompt):
    # Convert to PIL
    original_image = Image.fromarray(image).convert("RGB")
    mask_image = Image.fromarray(mask).convert("L")

    # Resize images to 512x512
    original_image = original_image.resize((512, 512))
    mask_image = mask_image.resize((512, 512))

    # Inpainting AI
    output = pipe(prompt=prompt, image=original_image, mask_image=mask_image, num_inference_steps=25).images[0]
    return np.array(output)

# UI với Gradio
interface = gr.Interface(
    fn=inpaint,
    inputs=[
        gr.Image(type="numpy", label="Upload Image"),
        gr.Image(type="numpy", label="Upload Mask"),
        gr.Textbox(label="Prompt (Describe what to add)")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="AI Furniture Inpainting",
    description="Upload an image of a room and a mask where furniture should be added."
)

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