import gradio as gr
from PIL import Image
import requests
from io import BytesIO
import os

# Функция для взаимодействия с моделью на Hugging Face
def process_image(image, prompt):
    # Загрузка изображения
    image_data = image.read()
    
    # Отправка запроса на API Hugging Face
    headers = {
        "Authorization": f"Bearer {os.getenv('HF_TOKEN')}"  # Использование API ключа из переменной окружения
    }
    data = {
        "inputs": {
            "image": image_data,
            "prompt": prompt
        }
    }
    response = requests.post("https://api-inference.huggingface.co/models/CrucibleAI/ControlNetMediaPipeFace", headers=headers, files=data)
    
    # Обработка ответа
    if response.status_code == 200:
        # Преобразование ответа в изображение
        image = Image.open(BytesIO(response.content))
        return image
    else:
        # В случае ошибки возвращаем информацию об ошибке
        return f"Error: {response.text}"

# Создание Gradio Blocks интерфейса
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload Image")
            prompt_input = gr.Textbox(label="Enter Prompt")
        with gr.Column():
            output_image = gr.Image(type="pil", label="Output Image")
    
    submit_button = gr.Button("Submit")
    submit_button.click(fn=process_image, inputs=[image_input, prompt_input], outputs=output_image)

# Запуск интерфейса
demo.launch()