# Import Gradio dan client API
import gradio as gr
from PIL import Image
import numpy as np
import tempfile
from gradio_client import Client, handle_file
from themes import IndonesiaTheme  # Impor tema custom
import os

# Siapkan URL untuk permintaan API Virtual Try-On
url_api = os.environ['url_api']

# Fungsi untuk memanggil API /generate_image
def generate_image(prompt, id_image, start_step, guidance, seed, true_cfg, width, height, num_steps, id_weight, neg_prompt, timestep_to_start_cfg, max_sequence_length):
    client = Client(url_api)
    
    try:
        # Jika id_image adalah numpy array, konversikan ke gambar menggunakan PIL
        if isinstance(id_image, np.ndarray):
            id_image_pil = Image.fromarray(id_image.astype('uint8'))  # Konversi ke format gambar dari numpy array

            # Simpan gambar sebagai file sementara
            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
                id_image_pil.save(temp_file.name)  # Simpan gambar PIL ke file sementara
                
                # Panggil API dengan file gambar
                result = client.predict(
                    prompt=prompt,
                    id_image=handle_file(temp_file.name),  # Kirim file gambar ke API
                    start_step=start_step,
                    guidance=guidance,
                    seed=seed,
                    true_cfg=true_cfg,
                    width=width,
                    height=height,
                    num_steps=num_steps,
                    id_weight=id_weight,
                    neg_prompt=neg_prompt,
                    timestep_to_start_cfg=timestep_to_start_cfg,
                    max_sequence_length=max_sequence_length,
                    api_name="/generate_image"
                )
        
        else:
            raise ValueError("Gambar yang diunggah tidak valid.")
        
        # Proses keluaran dari API
        generated_image_path = result[0]  # Ambil path gambar yang dihasilkan
        used_seed = result[1]  # Ambil seed yang digunakan
        output_gallery_items = []  # List untuk menampung item galeri

        # Ekstrak daftar gambar dari hasil API (jika ada lebih dari satu gambar)
        if isinstance(result[2], list):
            for item in result[2]:
                if 'image' in item:
                    output_gallery_items.append((item['image'], item.get('caption', '')))  # Tambahkan gambar dan keterangan ke galeri

        # Jika tidak ada error, kembalikan hasil gambar dan seed
        return generated_image_path, used_seed, "Success: Image generated successfully."
    
    except Exception as e:
        # Jika terjadi error, kembalikan pesan error ke "system_result"
        error_message = f"Error: {str(e)}"
        return None, None, error_message

# CSS untuk styling antarmuka
css = """
#col-left, #col-mid {
    margin: 0 auto;
    max-width: 400px;
    padding: 10px;
    border-radius: 15px;
    background-color: #f9f9f9;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}
#col-right {
    margin: 0 auto;
    max-width: 400px;
    padding: 10px;
    border-radius: 15px;
    background: linear-gradient(180deg, #B6BBC4, #EEEEEE);
    color: white;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}
#col-bott {
    margin: 0 auto;
    padding: 10px;
    border-radius: 15px;
    background-color: #f9f9f9;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}
#banner {
    width: 100%;
    text-align: center;
    margin-bottom: 20px;
}
#run-button {
    background-color: #ff4b5c;
    color: white;
    font-weight: bold;
    padding: 30px;
    border-radius: 10px;
    cursor: pointer;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2);
}
#footer {
    text-align: center;
    margin-top: 20px;
    color: silver;
}
#markdown-silver {
    color: silver; /* Mengatur warna font Markdown menjadi silver */
}
"""


# Antarmuka Gradio
with gr.Blocks(css=css, theme=IndonesiaTheme()) as PuLidFluxApp:
        # Tambahkan banner
    gr.HTML("""
        <div style='text-align: center;'>
            <img src='https://i.ibb.co.com/ynqKvrr/banner-pulid.jpg' alt='Banner' style='width: 100%; height: auto;'/>
        </div>
    """)
    gr.HTML("<h2 style='text-align: center;'>Aplikasi Pembuatan Gambar Menggunakan PuLID-FLUX</h2>")
    
    with gr.Row():
        with gr.Column(elem_id="col-left"):
            gr.Markdown("### ➡️ Deskripsi Gambar")
            prompt = gr.Textbox(label="Prompt", value="portrait, color, cinematic")
            gr.Markdown("### ➡️ Sumber Wajah")
            id_image = gr.Image(label="Foto Wajah")
            neg_prompt = gr.Textbox(label="Negative Prompt", value="bad quality, worst quality, text, signature, watermark, extra limbs")
        
        with gr.Column(elem_id="col-mid"):
            gr.Markdown("### ➡️ Advanced Parameters")
            start_step = gr.Slider(label="Timestep to Start", minimum=0, maximum=10, value=0)
            guidance = gr.Slider(label="Guidance", minimum=0, maximum=10, value=4)
            seed = gr.Textbox(label="Seed (-1 for random)", value="-1")
            true_cfg = gr.Slider(label="True CFG Scale", minimum=0, maximum=10, value=1)
            width = gr.Slider(label="Width", minimum=100, maximum=2000, value=896)
            height = gr.Slider(label="Height", minimum=100, maximum=2000, value=1152)
            num_steps = gr.Slider(label="Number of Steps", minimum=1, maximum=100, value=20)
            id_weight = gr.Slider(label="ID Weight", minimum=0, maximum=1, value=1)
            timestep_to_start_cfg = gr.Slider(label="Timestep to Start CFG", minimum=0, maximum=10, value=1)
            max_sequence_length = gr.Slider(label="Max Sequence Length", minimum=1, maximum=512, value=128)
        
        with gr.Column(elem_id="col-right"):
            gr.Markdown("""
                    ### Paper: PuLID: [Pure and Lightning ID Customization via Contrastive Alignment](https://arxiv.org/abs/2404.16022) | [Codes: GitHub](https://github.com/ToTheBeginning/PuLID)

                    ### 💡Tips:

                    - **Timestep to start inserting ID**: 
                      - Semakin kecil nilainya, semakin tinggi kesetiaannya, namun semakin rendah kemampuan untuk mengedit; semakin tinggi nilainya, semakin rendah kesetiaannya, namun semakin tinggi kemampuan untuk mengedit.
                      - Rentang yang disarankan untuk nilai ini adalah antara 0 hingga 4. Untuk adegan fotorealistik, kami merekomendasikan menggunakan 4; untuk adegan bergaya, kami merekomendasikan menggunakan 0-1.
                      - Jika Anda tidak puas dengan kemiripannya, Anda dapat menurunkan nilai ini; sebaliknya, jika Anda tidak puas dengan kemampuan pengeditan, Anda dapat menaikkan nilai ini.

                    - **True CFG scale**:
                      - Dalam sebagian besar skenario, disarankan menggunakan fake CFG, yaitu dengan mengatur true CFG scale ke 1, dan hanya menyesuaikan guidance scale. Hal ini juga lebih efisien.
                      - Namun, dalam beberapa kasus, penggunaan true CFG dapat menghasilkan hasil yang lebih baik. Untuk informasi lebih detail, silakan merujuk ke dokumentasi.

                    - **Pelajari lebih lanjut tentang model**:
                      - Silakan merujuk ke dokumentasi GitHub untuk detail lebih lanjut dan informasi tentang model. Kami menyediakan penjelasan detail tentang kedua parameter di atas dalam dokumen.

                    - **Contoh**: 
                      - Kami menyediakan beberapa contoh (sudah di-cache, jadi cukup klik untuk melihat apa yang dapat dilakukan model) di bagian bawah. Anda bisa mencoba prompt contoh tersebut terlebih dahulu.
                    """, elem_id="markdown-silver")


    # Tombol untuk memulai proses
    with gr.Row():
        with gr.Column(elem_id="col-bott"):
            run_button = gr.Button("⭐ Mulai Generate Image ⭐", elem_id="run-button")
            gr.Markdown("### ✅ Hasil Generated PuLID")
            generated_image = gr.Image(label="Generated Image")
            used_seed = gr.Textbox(label="Used Seed")
            system_result = gr.Textbox(label="System Result")  # Output tambahan untuk menampilkan hasil/error

    # Menghubungkan tombol dengan fungsi pemanggilan API
    run_button.click(
        fn=generate_image,
        inputs=[prompt, id_image, start_step, guidance, seed, true_cfg, width, height, num_steps, id_weight, neg_prompt, timestep_to_start_cfg, max_sequence_length],
        outputs=[generated_image, used_seed, system_result]  # Tambahkan system_result sebagai output
    )


    # Tambahkan footer di bagian bawah
    gr.HTML("""
    <footer id="footer">
        Transfer Energi Semesta Digital © 2024 | 🇮🇩 Untuk Indonesia Jaya!
    </footer>
    """)

# Menjalankan aplikasi
if __name__ == "__main__":
    PuLidFluxApp.queue(api_open=False).launch(show_api=False)