import gradio as gr
import torch
import spaces
from diffusers import Flux2Pipeline

# Model configuration
repo_id = "black-forest-labs/FLUX.2-dev"
torch_dtype = torch.bfloat16

# Download and initialize model at startup
pipe = Flux2Pipeline.from_pretrained(
    repo_id, torch_dtype=torch_dtype
)
pipe.enable_model_cpu_offload()

@spaces.GPU(duration=120)
def generate_image(prompt, seed, steps, guidance_scale, progress=gr.Progress()):
    if not prompt:
        gr.Warning("Please enter a prompt")
        return None
    
    progress(0.2, desc="Preparing generation...")
    
    generator = torch.Generator(device="cpu").manual_seed(int(seed))
    
    progress(0.4, desc="Generating image...")
    
    image = pipe(
        prompt=prompt,
        generator=generator,
        num_inference_steps=int(steps),
        guidance_scale=guidance_scale,
    ).images[0]
    
    progress(1.0, desc="Done!")
    return image

# Minimal Apple-inspired CSS
apple_css = """
@import url('https://fonts.googleapis.com/css2?family=SF+Pro+Display:wght@300;400;500;600&display=swap');

* {
    font-family: -apple-system, BlinkMacSystemFont, 'SF Pro Display', sans-serif !important;
}

.gradio-container {
    background: #ffffff !important;
    max-width: 900px !important;
    margin: 0 auto !important;
}

.main-container {
    padding: 60px 40px !important;
}

.header-text {
    text-align: center;
    margin-bottom: 48px;
}

.header-text h1 {
    font-size: 56px;
    font-weight: 600;
    color: #1d1d1f;
    letter-spacing: -0.02em;
    margin: 0 0 12px 0;
}

.header-text p {
    font-size: 21px;
    color: #86868b;
    font-weight: 400;
    margin: 0;
}

.input-section textarea {
    background: #f5f5f7 !important;
    border: none !important;
    border-radius: 16px !important;
    padding: 20px !important;
    font-size: 17px !important;
    color: #1d1d1f !important;
    min-height: 60px !important;
    resize: none !important;
}

.input-section textarea:focus {
    background: #f5f5f7 !important;
    outline: none !important;
    box-shadow: 0 0 0 4px rgba(0, 125, 250, 0.2) !important;
}

.input-section textarea::placeholder {
    color: #86868b !important;
}

.generate-btn {
    background: #0071e3 !important;
    border: none !important;
    border-radius: 980px !important;
    padding: 16px 32px !important;
    font-size: 17px !important;
    font-weight: 400 !important;
    color: white !important;
    cursor: pointer !important;
    transition: all 0.2s ease !important;
    margin-top: 16px !important;
}

.generate-btn:hover {
    background: #0077ed !important;
}

.output-section {
    margin-top: 40px;
}

.output-section img {
    border-radius: 20px !important;
    box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1) !important;
}

.footer-text {
    text-align: center;
    margin-top: 48px;
    padding-top: 24px;
    border-top: 1px solid #f5f5f7;
}

.footer-text a {
    color: #0071e3;
    text-decoration: none;
    font-size: 14px;
}

.footer-text a:hover {
    text-decoration: underline;
}

/* Hide labels for cleaner look */
.input-section label {
    display: none !important;
}

/* Remove borders and shadows from containers */
.gr-group, .gr-box {
    border: none !important;
    box-shadow: none !important;
    background: transparent !important;
}

.gr-padded {
    padding: 0 !important;
}

.settings-section {
    background: #f5f5f7 !important;
    border-radius: 16px !important;
    padding: 20px !important;
    margin-top: 16px !important;
}

.settings-section label {
    color: #1d1d1f !important;
    font-size: 14px !important;
    font-weight: 500 !important;
}
"""

with gr.Blocks(title="FLUX.2 Image Studio") as demo:
    
    gr.HTML("""
        <div class="header-text">
            <h1>FLUX.2 Studio</h1>
            <p>Create stunning images with FLUX.2-dev</p>
        </div>
    """)
    
    with gr.Column(elem_classes="input-section"):
        prompt = gr.Textbox(
            placeholder="A peaceful mountain lake at sunrise...",
            lines=3,
            show_label=False,
            container=False
        )
        
        with gr.Accordion("Advanced Settings", open=False):
            with gr.Row():
                seed = gr.Number(
                    label="Seed",
                    value=42,
                    precision=0
                )
                steps = gr.Slider(
                    label="Inference Steps",
                    minimum=10,
                    maximum=100,
                    value=50,
                    step=1
                )
            guidance_scale = gr.Slider(
                label="Guidance Scale",
                minimum=1.0,
                maximum=10.0,
                value=4.0,
                step=0.5
            )
        
        generate_btn = gr.Button(
            "Create",
            elem_classes="generate-btn",
            variant="primary"
        )
    
    with gr.Column(elem_classes="output-section"):
        output_image = gr.Image(
            type="pil",
            show_label=False,
            container=False
        )
    
    gr.HTML("""
        <div class="footer-text">
            <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">Built with anycoder</a>
        </div>
    """)
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, seed, steps, guidance_scale],
        outputs=[output_image],
        api_visibility="public"
    )
    
    prompt.submit(
        fn=generate_image,
        inputs=[prompt, seed, steps, guidance_scale],
        outputs=[output_image],
        api_visibility="private"
    )

demo.launch(
    css=apple_css,
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}
    ]
)