import gradio as gr
from huggingface_hub import InferenceClient
import os
from PIL import Image

# Initialize the HF Inference Client
client = InferenceClient(token=os.getenv("HF_TOKEN"))

MODELS = {
    "Flux.1 [Schnell]": "black-forest-labs/FLUX.1-schnell",
    "Pony XL (Animagine)": "cagliostrolab/animagine-xl-3.1"
}

def generate_image(prompt, model_choice):
    if not prompt or prompt.strip() == "":
        return None
        
    model_id = MODELS.get(model_choice, MODELS["Flux.1 [Schnell]"])
    
    try:
        image = client.text_to_image(prompt, model=model_id)
        return image
    except Exception as e:
        raise gr.Error(f"Neural Core Error: {str(e)}")

# UI Styling
custom_css = """
footer {visibility: hidden}
body { background-color: #030303 !important; }
.gradio-container { border: none !important; }
#gen-button { background: #ffffff !important; color: #000000 !important; font-weight: 900 !important; }
"""

# Fixed: Removed gr.Div (unsupported) and moved theme/css to launch() for Gradio 6 compatibility
with gr.Blocks() as demo:
    gr.Markdown("# CODEX_VISION", elem_id="title")
    gr.Markdown("### NEURAL PROJECTION INTERFACE")
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt", 
                placeholder="Describe what you want to see...", 
                lines=3
            )
            model = gr.Radio(
                choices=list(MODELS.keys()), 
                value="Flux.1 [Schnell]", 
                label="Engine"
            )
            btn = gr.Button("Initialize Visualization", variant="primary", elem_id="gen-button")
            
        with gr.Column(scale=1):
            output = gr.Image(label="Projection Result", type="pil")
            
    btn.click(
        fn=generate_image,
        inputs=[prompt, model],
        outputs=output
    )

if __name__ == "__main__":
    # Gradio 6.0 compatible launch
    demo.queue().launch(
        css=custom_css, 
        theme=gr.themes.Default(primary_hue="zinc", secondary_hue="zinc")
    )