import gradio as gr
from huggingface_hub import InferenceClient
import os  # Hugging Face par ye zaroori hai

# 1. Secret key uthana (Hugging Face ke environment se)
HF_TOKEN = os.getenv("Image") 

# 2. Model Setup
client = InferenceClient(api_key=HF_TOKEN)

# 3. Image Generation Function
def generate_ai_image(prompt, style):
    styles = {
        "Realistic": "high-quality photorealistic, 8k, cinematic lighting",
        "Anime": "vibrant anime style, detailed digital art, aesthetic",
        "3D Disney": "cute 3d character render, Pixar style, soft shadows",
        "Oil Painting": "rich oil painting, artistic brush strokes, masterpiece"
    }
    
    final_prompt = f"{prompt}, {styles[style]}"
    
    try:
        image = client.text_to_image(final_prompt, model="black-forest-labs/FLUX.1-schnell")
        return image
    except Exception as e:
        return None

# 4. Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎨 FLUX.1-Schnell Image Studio")
    gr.Markdown("Created for University Project - Powered by Hugging Face API")
    
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(label="Enter Prompt", placeholder="e.g. A futuristic robot in Faisalabad streets")
            style_dropdown = gr.Dropdown(choices=["Realistic", "Anime", "3D Disney", "Oil Painting"], 
                                         label="Select Format", value="Realistic")
            generate_btn = gr.Button("Generate Magic ✨", variant="primary")
        
        with gr.Column():
            image_output = gr.Image(label="Your AI Image")

    generate_btn.click(fn=generate_ai_image, inputs=[text_input, style_dropdown], outputs=image_output)

# 5. Launch (Hugging Face par .launch() kafi hai)
demo.launch()