import torch
from diffusers import StableDiffusionInpaintPipeline
from PIL import Image
import gradio as gr

# Инициализация модели для CPU
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-inpainting",
    torch_dtype=torch.float32,  # Использование стандартной точности на CPU
)

def process_image(image, mask_image, prompt):
    # Конвертируем изображение и маску в формат PIL и подгоняем под нужный размер
    image = Image.fromarray(image).resize((512, 512))
    mask_image = Image.fromarray(mask_image).convert("L").resize((512, 512))
    mask_image = mask_image.point(lambda x: 255 if x > 128 else 0)  # Бинаризация маски
    
    # Генерация нового изображения
    result = pipe(prompt=prompt, image=image, mask_image=mask_image).images[0]
    
    return result

# Создание интерфейса с Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 🖼️ Удаление объектов с изображений")
    
    with gr.Row():
        input_image = gr.Image(
            label="Загрузите фото", 
            type="numpy",
            height=300
        )
        mask_image = gr.Image(
            label="Нарисуйте маску (черным)", 
            interactive=True,  # Режим рисования маски
            type="numpy",
            height=300
        )
        output_image = gr.Image(label="Результат", height=300)
    
    prompt = gr.Textbox(
        label="Описание фона",
        value="реалистичный фон",
        placeholder="Что должно быть вместо удаленного объекта..."
    )
    
    btn = gr.Button("Обработать")
    btn.click(
        fn=process_image,
        inputs=[input_image, mask_image, prompt],
        outputs=output_image,
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)