import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import DiffusionPipeline, FlowMatchEulerDiscreteScheduler
from PIL import Image
import io
import os
import subprocess
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

huggingface_token = os.getenv("HUGGINGFACE_TOKEN")

Title = "Text-2-Image Generation With Transformer Rectifier | Nirajan Dhakal"

pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, token = huggingface_token).to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

@spaces.GPU(duration=190)
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=5.0, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    image = pipe(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=guidance_scale
    ).images[0]
    return image, seed

def download_image(image, file_format):
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format=file_format)
    img_byte_arr = img_byte_arr.getvalue()
    return img_byte_arr

examples = [
    "a galaxy swirling with vibrant blue and purple hues",
    "a futuristic cityscape under a dark sky",
    "a serene forest with a magical glowing tree",
    "a futuristic cityscape with sleek skyscrapers and flying cars",
    "a portrait of a smiling woman with a colorful floral crown",
    "a fantastical creature with the body of a dragon and the wings of a butterfly",
]

css = """
body {
    background-color: #f4faff;
    color: #005662;
    font-family: Arial, sans-serif;
}
#col-container {
    margin: 0 auto;
    max-width: 100%;
    padding: 20px;
}
.gr-button {
    background-color: #0288d1;
    color: white;
    border-radius: 8px;
    transition: background-color 0.3s ease;
}
.gr-button:hover {
    background-color: #0277bd;
}
.gr-examples-card {
    background-color: #ffffff;
    border: 1px solid #0288d1;
    border-radius: 12px;
    padding: 16px;
    margin-bottom: 12px;
}
.gr-examples-card:hover {
    background-color: #e0f7fa;
    border-color: #0277bd;
}
.gr-progress-bar, .gr-progress-bar-fill {
    background-color: #0288d1 !important;
}
.gr-slider, .gr-slider-track {
    background-color: #0288d1 !important;
}
.gr-slider-thumb {
    background-color: #005662 !important;
}
.gr-text-input, .gr-image {
    width: 100%;
    box-sizing: border-box;
    margin-bottom: 10px;
}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="blue", secondary_hue="gray")) as demo:
    gr.HTML(title=Title)
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# FLUX.1 [dev]
12B param rectified flow transformer guidance-distilled from FLUX.1 [pro]
        
<a href="https://huggingface.co/black-forest-labs/FLUX.1-dev" style="text-decoration:none;">
<div class="gr-examples-card">
    <h3>View Model Details</h3>
    <p>Explore more about this model on Hugging Face.</p>
</div>
</a>
        """)

            
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=15,
                    step=0.1,
                    value=3.5,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )
        
        download_format = gr.Radio(
            label="Download Format",
            choices=["PNG", "JPEG", "SVG", "WEBP"],
            value="PNG",
            type="value",
        )

        download_button = gr.Button("Download Image")

        download_button.click(
            fn=download_image,
            inputs=[result, download_format],
            outputs=gr.File(label="Download"),
        )

        gr.Examples(
            examples=examples,
            fn=infer,
            inputs=[prompt],
            outputs=[result, seed],
            cache_examples="lazy"
        )

        share_buttons = gr.Row()
        with share_buttons:
            twitter_button = gr.Button("Share on Twitter")
            facebook_button = gr.Button("Share on Facebook")

        twitter_button.click(
            fn=lambda img: f"https://twitter.com/intent/tweet?url={img}",
            inputs=[result],
            outputs=None,
            _js="(img) => window.open(img, '_blank')"
        )

        facebook_button.click(
            fn=lambda img: f"https://www.facebook.com/sharer/sharer.php?u={img}",
            inputs=[result],
            outputs=None,
            _js="(img) => window.open(img, '_blank')"
        )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=[result, seed]
    )

    demo.load(
        fn=lambda: None,
        inputs=None,
        outputs=None,
        _js="""
        function() {
            document.addEventListener('keydown', function(event) {
                if (event.key === 'Enter') {
                    document.querySelector('button').click();
                }
            });
        }
        """
    )


demo.launch()