from gradio_client import Client, handle_file
import gradio as gr
import concurrent.futures
import tempfile
import os
from PIL import Image

# Impor tema custom dari themes.py
from themes import IndonesiaTheme

# Siapkan URL dan header untuk permintaan API
url_api1 = os.environ['url_api1']
url_api2 = os.environ['url_api2']

# Fungsi untuk FLUX Std
def infer_image(prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress()):
    client = Client(url_api1)
    try:
        progress(0, "Memulai proses...")
        result = client.predict(
            prompt=prompt,
            seed=seed,
            randomize_seed=randomize_seed,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            api_name="/infer"
        )
        progress(1, "Proses selesai.")
        return result[0], seed if not randomize_seed else result[1]
    except concurrent.futures.CancelledError:
        return None, "Request was cancelled. Please try again."

# Fungsi untuk FLUX Inpainting
def inpainting_process(input_image_editor, input_text, seed_slicer, randomize_seed_checkbox, strength_slider, num_inference_steps_slider, progress=gr.Progress()):
    client = Client(url_api2)
    try:
        progress(0, "Memulai proses inpainting...")
        
        # Simpan gambar sementara ke file
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp:
            input_image_editor["composite"].save(temp.name)
            temp_image_path = temp.name
        
        # Tutup file setelah penulisan selesai
        temp.close()

        # Gunakan handle_file untuk memproses gambar dari jalur file sementara
        input_image_editor = {"background": handle_file(temp_image_path)}

        result = client.predict(
            input_image_editor=input_image_editor,
            input_text=input_text,
            seed_slicer=seed_slicer,
            randomize_seed_checkbox=randomize_seed_checkbox,
            strength_slider=strength_slider,
            num_inference_steps_slider=num_inference_steps_slider,
            api_name="/process"
        )
        progress(1, "Proses inpainting selesai.")
        return result[0], result[1]
    except concurrent.futures.CancelledError:
        return None, "Request was cancelled. Please try again."
    finally:
        # Hapus file sementara setelah selesai digunakan
        import os
        os.remove(temp_image_path)

# CSS untuk styling antarmuka
css = """
#col-left, #col-mid, #col-right {
    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);
}
#banner {
    width: 100%;
    text-align: center;
    margin-bottom: 20px;
}
#run-button {
    background-color: #ff4b5c;
    color: white;
    font-weight: bold;
    padding: 10px;
    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;
}
"""

# Interface Gradio
with gr.Blocks(css=css, theme=IndonesiaTheme()) as app:
    # Tambahkan banner
    gr.HTML("""
        <div style='text-align: center;'>
            <img src='https://i.postimg.cc/ZYFNKrjK/banner.jpg' alt='Banner' style='width: 100%; height: auto;'/>
        </div>
    """)

    # Ganti judul
    gr.Markdown("<h2 style='text-align: center;'>FLUX.1 Untuk membuat gambar terbaik di dunia.</h2>")
    gr.HTML("""
        <h4 style='text-align: center;'>
            <a href="https://huggingface.co/sayakpaul/FLUX.1-merged">FLUX.1 [merged]</a> Merge by 
            <a href="https://huggingface.co/sayakpaul">Sayak Paul</a> of 2 of the 12B param rectified flow transformers 
            <a href="https://huggingface.co/black-forest-labs/FLUX.1-dev">FLUX.1 [dev]</a> and 
            <a href="https://huggingface.co/black-forest-labs/FLUX.1-schnell">FLUX.1 [schnell]</a> by 
            <a href="https://blackforestlabs.ai/">Black Forest Labs</a>, created by 
            <a href="https://github.com/drat">Deddy</a>
        </h4>
    """)

    with gr.Tabs():
        with gr.TabItem("FLUX Std"):
            with gr.Row():
                with gr.Column():
                    prompt_input = gr.Textbox(label="Prompt", placeholder="Ceritakan tentang image apa yang ingin dibuat...")
                    seed_slider = gr.Slider(label="Seed", minimum=0, maximum=10000, step=1, value=0)
                    randomize_seed_checkbox = gr.Checkbox(label="Randomize Seed", value=True)
                    width_slider = gr.Slider(label="Lebar", minimum=512, maximum=2048, step=64, value=1024)
                    height_slider = gr.Slider(label="Tinggi", minimum=512, maximum=2048, step=64, value=1024)
                    guidance_scale_slider = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.5, value=3.5)
                    inference_steps_slider = gr.Slider(label="Jumlah Inference Steps", minimum=1, maximum=50, step=1, value=8)
                    submit_button = gr.Button("Mulai Membuat Gambar", elem_id="run-button")

                with gr.Column():
                    result_image = gr.Image(label="Result")
                    seed_output = gr.Number(label="Seed Digunakan")

            submit_button.click(
                infer_image,
                inputs=[prompt_input, seed_slider, randomize_seed_checkbox, width_slider, height_slider, guidance_scale_slider, inference_steps_slider],
                outputs=[result_image, seed_output],
                show_progress=True  # Menampilkan progress bar
            )

        with gr.TabItem("FLUX Inpainting"):
            with gr.Row():
                with gr.Column():
                    input_image_editor = gr.ImageEditor(
                        label='Image',
                        type='pil',
                        sources=["upload", "webcam"],
                        image_mode='RGB',
                        layers=False,
                        brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"))

                    input_text = gr.Textbox(label="Prompt", placeholder="Deskripsi untuk inpainting...")
                    seed_slicer = gr.Slider(label="Seed", minimum=0, maximum=10000, step=1, value=42)
                    randomize_seed_checkbox = gr.Checkbox(label="Randomize Seed", value=True)
                    strength_slider = gr.Slider(label="Strength", minimum=0.0, maximum=1.0, step=0.01, value=0.85)
                    num_inference_steps_slider = gr.Slider(label="Number of Inference Steps", minimum=1, maximum=50, step=1, value=20)
                    inpainting_button = gr.Button("Mulai Proses Inpainting", elem_id="run-button")

                with gr.Column():
                    generated_image = gr.Image(label="Generated Image")
                    input_mask = gr.Image(label="Input Mask")

            inpainting_button.click(
                inpainting_process,
                inputs=[input_image_editor, input_text, seed_slicer, randomize_seed_checkbox, strength_slider, num_inference_steps_slider],
                outputs=[generated_image, input_mask],
                show_progress=True  # Menampilkan progress bar
            )

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

# Meluncurkan aplikasi
if __name__ == "__main__":
    app.launch()