import gradio as gr
from huggingface_hub import InferenceClient
import spaces
from openai import OpenAI
import io
import requests
from PIL import Image

@spaces.GPU
def generate_image(prompt, hf_token, progress=gr.Progress(track_tqdm=True)):
    if not hf_token:
        raise gr.Error("Please enter your Hugging Face API Token!")
    if not prompt:
        raise gr.Error("Please enter a prompt!")
        
    client = InferenceClient(token=hf_token.strip())
    
    image = client.text_to_image(
        prompt, 
        model="black-forest-labs/FLUX.1-schnell"
    )
    
    return image

def edit_image(image, prompt, omni_url, omni_key, omni_model, progress=gr.Progress(track_tqdm=True)):
    if not omni_url or not omni_key:
        raise gr.Error("Please provide both your OmniRoute URL and API Key!")
    if not image:
        raise gr.Error("Please upload an image to edit!")
    if not prompt:
        raise gr.Error("Please enter a prompt for editing!")
        
    client = OpenAI(
        base_url=omni_url.strip(),
        api_key=omni_key.strip()
    )
    
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    img_byte_arr = img_byte_arr.getvalue()
    
    try:
        response = client.images.edit(
            image=img_byte_arr,
            prompt=prompt,
            model=omni_model.strip()
        )
        
        if response.data and len(response.data) > 0:
            url = response.data[0].url
            if url:
                img_response = requests.get(url)
                edited = Image.open(io.BytesIO(img_response.content))
                return edited
            elif response.data[0].b64_json:
                import base64
                img_bytes = base64.b64decode(response.data[0].b64_json)
                edited = Image.open(io.BytesIO(img_bytes))
                return edited
                
        raise gr.Error("OmniRoute returned an empty response.")
    except Exception as e:
        raise gr.Error(f"OmniRoute API Error: {str(e)}")

custom_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
    font=gr.themes.GoogleFont("Inter"),
)

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # ⚡ FLUX.1 Image Studio (Grok Quality)
        Generate ultra-fast images from text using the real-time FLUX.1-schnell model via Hugging Face Serverless API. *No local GPU Required!*
        """
    )
    
    with gr.Tabs():
        with gr.Tab("✨ Generate New Image"):
            with gr.Row():
                with gr.Column(scale=1):
                    hf_token = gr.Textbox(
                        label="🔑 Hugging Face Access Token",
                        placeholder="hf_...",
                        type="password",
                        info="Paste your Hugging Face Token here"
                    )
                    prompt = gr.Textbox(
                        label="✨ Prompt", 
                        lines=3, 
                        placeholder="e.g. A futuristic cyberpunk city at night...", 
                        autofocus=True
                    )
                    
                    generate_btn = gr.Button("🎨 Generate Image", variant="primary", size="lg")
                    
                with gr.Column(scale=1):
                    output_image = gr.Image(label="Result", type="pil", interactive=False)

            generate_btn.click(
                fn=generate_image,
                inputs=[prompt, hf_token],
                outputs=[output_image]
            )
            prompt.submit(
                fn=generate_image,
                inputs=[prompt, hf_token],
                outputs=[output_image]
            )
            
        with gr.Tab("🖌️ Edit Existing Image (OmniRoute)"):
            with gr.Row():
                with gr.Column(scale=1):
                    with gr.Accordion("⚙️ OmniRoute Configuration", open=True):
                        omni_url = gr.Textbox(
                            label="OmniRoute Base URL",
                            value="http://localhost:20128/v1",
                            info="The endpoint URL of your OmniRoute gateway"
                        )
                        omni_key = gr.Textbox(
                            label="🔑 OmniRoute API Key",
                            placeholder="Enter your key...",
                            type="password"
                        )
                        omni_model = gr.Textbox(
                            label="Model ID",
                            value="chatgpt-web",
                            info="The underlying model to route to (e.g. chatgpt-web)"
                        )
                        
                    input_image = gr.Image(label="Image to Edit", type="pil")
                    edit_prompt = gr.Textbox(
                        label="✨ Edit Prompt", 
                        lines=2, 
                        placeholder="e.g. change the sky to red..."
                    )
                    
                    edit_btn = gr.Button("🖌️ Edit Image", variant="primary", size="lg")
                    
                with gr.Column(scale=1):
                    edit_output = gr.Image(label="Edited Result", type="pil", interactive=False)
            
            edit_btn.click(
                fn=edit_image,
                inputs=[input_image, edit_prompt, omni_url, omni_key, omni_model],
                outputs=[edit_output]
            )

if __name__ == "__main__":
    demo.launch(theme=custom_theme)