import spaces
import gradio as gr
import torch
from PIL import Image
from diffusers import DiffusionPipeline
import random

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.backends.cuda.matmul.allow_tf32 = True

base_model = "black-forest-labs/FLUX.1-dev"
pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)

lora_repo = "strangerzonehf/Flux-Pixel-Background-LoRA"
trigger_word = ""
pipe.load_lora_weights(lora_repo)

pipe.to("cuda")

MAX_SEED = 2**32-1

@spaces.GPU()
def run_lora(prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device="cuda").manual_seed(seed)

    progress(0, "Starting image generation...")

    for i in range(1, steps + 1):
        if i % (steps // 10) == 0:
            progress(i / steps * 100, f"Processing step {i} of {steps}...")

    image = pipe(
        prompt=f"{prompt} {trigger_word}",
        num_inference_steps=steps,
        guidance_scale=cfg_scale,
        width=width,
        height=height,
        generator=generator,
        joint_attention_kwargs={"scale": lora_scale},
    ).images[0]

    progress(100, "Completed!")
    return image, seed

example_image_path = "example0.webp"
example_prompt = """Pixel Background, a silhouette of a surfer is seen riding a wave on a red surfboard. The surfers shadow is cast on the left side of the image, adding a touch of depth to the composition. The background is a vibrant orange, pink, and blue, with a sun setting in the upper right corner of the frame. The silhouette of the surfer, a palm tree casts a shadow onto the wave, adding depth and contrast to the scene."""
example_cfg_scale = 3.2
example_steps = 32
example_width = 1152
example_height = 896
example_seed = 3981632454
example_lora_scale = 0.85

def load_example():
    example_image = Image.open(example_image_path)
    return example_prompt, example_cfg_scale, example_steps, True, example_seed, example_width, example_height, example_lora_scale, example_image

css = """
.container {max-width: 1200px; margin: auto; padding: 20px;}
.header {text-align: center; margin-bottom: 30px;}
.generate-btn {background-color: #2ecc71 !important; color: white !important;}
.generate-btn:hover {background-color: #27ae60 !important;}
.parameter-box {background-color: #f5f6fa; padding: 20px; border-radius: 10px; margin: 10px 0;}
.result-box {background-color: #f5f6fa; padding: 20px; border-radius: 10px;}
"""

with gr.Blocks(css=css) as app:
    with gr.Column(elem_classes="container"):
        gr.Markdown("# 🎨 Flux ART Image Generator", elem_classes="header")
        
        with gr.Row(equal_height=True):
            with gr.Column(scale=3):
                with gr.Group(elem_classes="parameter-box"):
                    prompt = gr.TextArea(
                        label="✍️ Your Prompt",
                        placeholder="Describe the image you want to generate...",
                        lines=5
                    )
                    
                with gr.Group(elem_classes="parameter-box"):
                    gr.Markdown("### 🎛️ Generation Parameters")
                    with gr.Row():
                        with gr.Column():
                            cfg_scale = gr.Slider(
                                label="CFG Scale",
                                minimum=1,
                                maximum=20,
                                step=0.5,
                                value=example_cfg_scale
                            )
                            steps = gr.Slider(
                                label="Steps",
                                minimum=1,
                                maximum=100,
                                step=1,
                                value=example_steps
                            )
                            lora_scale = gr.Slider(
                                label="LoRA Scale",
                                minimum=0,
                                maximum=1,
                                step=0.01,
                                value=example_lora_scale
                            )
                
                with gr.Group(elem_classes="parameter-box"):
                    gr.Markdown("### 📐 Image Dimensions")
                    with gr.Row():
                        width = gr.Slider(
                            label="Width",
                            minimum=256,
                            maximum=1536,
                            step=64,
                            value=example_width
                        )
                        height = gr.Slider(
                            label="Height",
                            minimum=256,
                            maximum=1536,
                            step=64,
                            value=example_height
                        )
                
                with gr.Group(elem_classes="parameter-box"):
                    gr.Markdown("### 🎲 Seed Settings")
                    with gr.Row():
                        randomize_seed = gr.Checkbox(
                            True,
                            label="Randomize seed"
                        )
                        seed = gr.Slider(
                            label="Seed",
                            minimum=0,
                            maximum=MAX_SEED,
                            step=1,
                            value=example_seed
                        )
                
                generate_button = gr.Button(
                    "🚀 Generate Image",
                    elem_classes="generate-btn"
                )
            
            with gr.Column(scale=2):
                with gr.Group(elem_classes="result-box"):
                    gr.Markdown("### 🖼️ Generated Image")
                    result = gr.Image(label="Result")

    app.load(
        load_example,
        inputs=[],
        outputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, result]
    )
    
    generate_button.click(
        run_lora,
        inputs=[prompt, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale],
        outputs=[result, seed]
    )

app.queue()
app.launch()