import gradio as gr
from transformers import pipeline
from PIL import Image

# -----------------------------
# 1. Модельді жүктеу: Style Transfer (Фотосуретті сурет стиліне өзгерту)
# -----------------------------
# Бұл модель GPU-сыз, CPU-да жұмыс істеуге бейімделген
try:
    style_transfer_model = pipeline(
        "image-to-image", 
        model="akhaliq/style-transfer-transformer-pytorch",
        # Trust_remote_code=True керек болуы мүмкін
    )
except Exception as e:
    # Егер модель жүктелмесе, қатені көрсету
    print(f"Стиль ауыстыру моделін жүктеу қатесі: {e}")
    style_transfer_model = None


# -----------------------------
# 2. Анализ функциясы
# -----------------------------
def apply_style_transfer(image, style_prompt):
    """
    Берілген суретке таңдалған стильді қолданады.
    """
    if style_transfer_model is None:
        return None, "Қате: Модель жүктелмеді. Space журналдарын тексеріңіз."

    # Егер сурет жүктелмесе
    if image is None:
        return None, "Қайта өңдеу үшін сурет жүктеңіз."

    try:
        # Модельді іске қосу
        result = style_transfer_model(
            image,
            prompt=style_prompt,
            num_inference_steps=20  # CPU үшін қадамдарды азайту
        )
        
        # Нәтиже суреті
        styled_image = result[0]['generated_image']
        
        # Мәтіндік нәтиже
        output_text = f"✨ Сіздің суретіңіз **'{style_prompt}'** стилінде сәтті қайта өңделді."
        
        return styled_image, output_text
        
    except Exception as e:
        return None, f"Өңдеу кезінде қате шықты: {e}"

# -----------------------------
# 3. Интерфейс
# -----------------------------
with gr.Blocks(title="StyleTransfer AI | Фотосуреттен Өнерге") as demo:
    gr.Markdown(
        "<h1 style='text-align:center; color:purple;'>🖼️ Фотосуреттен Өнерге: Әлемдегі Ғажайып Стиль</h1>"
    )
    gr.Markdown(
        "<p style='text-align:center;'>Сурет жүктеңіз және оны қай өнер стилінде көргіңіз келетінін жазыңыз (мысалы: 'Vincent van Gogh', 'Impressionism', 'Cyberpunk').</p>"
    )
    
    with gr.Row():
        # Кіріс: Сурет және Стильдік сөз
        with gr.Column():
            image_input = gr.Image(label="1. Өңделетін фотосурет", type="pil", sources=["upload", "webcam"])
            prompt_input = gr.Textbox(
                label="2. Қажетті стильді жазыңыз", 
                value="Impressionist painting of a Kazakh steppe",
                placeholder="Мысалы: A sketch drawing, Van Gogh style, Cyberpunk art"
            )
            process_btn = gr.Button("🎨 Стильді қолдану", variant="primary")
            
        # Шығыс: Өңделген сурет
        with gr.Column():
            styled_image_output = gr.Image(label="✨ Өңделген нәтиже", type="pil", interactive=False)
            output_text = gr.Markdown("Нәтиже осы жерде пайда болады.")
            
    # Функцияны батырмаға байланыстыру
    process_btn.click(
        fn=apply_style_transfer, 
        inputs=[image_input, prompt_input], 
        outputs=[styled_image_output, output_text]
    )

demo.launch()