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

# Функция для взаимодействия с моделью на Hugging Face
def process_image(image, prompt):
    # Преобразование объекта Image в байты
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    image_data = buffered.getvalue()
    
    # Отправка запроса на API Hugging Face
    headers = {
        "Authorization": f"Bearer {os.getenv('HF_TOKEN')}"
    }
    files = {
        "file": ("image.jpg", image_data, "image/jpeg"),
        "data": (None, '{"inputs": {"prompt": ' + f'"{prompt}"' + '}}', "application/json"),
    }
    response = requests.post("https://api-inference.huggingface.co/models/CrucibleAI/ControlNetMediaPipeFace", headers=headers, files=files)
    
    # Обработка ответа
    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()
"""

import gradio as gr
from transformers import StableDiffusionPipeline
from PIL import Image
import torch

# Загрузка предварительно обученной модели
pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0")
pipe = pipe.to("cuda")

def image_to_image(image, prompt):
    # Преобразование загруженного изображения в формат, совместимый с моделью
    image = Image.fromarray(image.astype('uint8'), 'RGB')
    
    # Генерация изображения на основе введенного промпта и загруженного изображения
    with torch.autocast("cuda"):
        image = pipe(prompt=prompt, init_image=image, strength=0.8)["sample"][0]  # strength может быть настроен
    
    # Преобразование результата в массив numpy для совместимости с Gradio
    return image

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil")
            prompt_input = gr.Textbox(placeholder="Enter prompt here...")
        with gr.Column():
            output_image = gr.Image(type="pil")

    with gr.Row():
        submit_button = gr.Button("Generate")

    submit_button.click(
        fn=image_to_image,
        inputs=[image_input, prompt_input],
        outputs=output_image
    )

demo.launch()