import gradio as gr
from huggingface_hub import InferenceClient
import os

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

def generate_image(prompt):
    try:
        # Generate image using Stable Diffusion 2.1 (free tier compatible)
        image = client.text_to_image(
            prompt,
            model="stabilityai/stable-diffusion-2-1",
            negative_prompt="blurry, low quality",  # Optional quality improvement
            guidance_scale=7.5,  # Controls creativity vs prompt adherence
            height=512,  # Standard size
            width=512
        )
        return image
    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}. Please try again later.")

with gr.Blocks(title="Free AI Image Generator") as demo:
    gr.Markdown("## 🖼️ Free AI Image Generator (Powered by Hugging Face)")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Describe your image",
                placeholder="A astronaut riding a horse on Mars",
                lines=2
            )
            generate_btn = gr.Button("Generate Image", variant="primary")
            
        with gr.Column():
            output_image = gr.Image(
                label="Generated Image",
                height=512,
                width=512
            )
    
    # Additional controls
    with gr.Accordion("Advanced Options", open=False):
        negative_prompt = gr.Textbox(
            label="What to exclude from image",
            placeholder="blurry, distorted, low quality"
        )
        guidance = gr.Slider(3, 20, value=7.5, label="Creativity vs Accuracy")
        steps = gr.Slider(10, 50, value=25, label="Generation Steps")
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, guidance, steps],
        outputs=output_image
    )

demo.launch(share=True)  # share=True creates public link