import streamlit as st
import os
from io import BytesIO
from PIL import Image
from huggingface_hub import InferenceClient

HF_TOKEN = os.getenv("HF_TOKEN")

st.set_page_config(
    page_title="KING2-IMAGE",
    page_icon="👑",
    layout="centered",
)

st.markdown("# 👑 KING2-IMAGE — SDXL LoRA")
st.markdown("A curated Stable Diffusion XL LoRA fine-tuned for majestic, high-quality text-to-image generation.")
st.markdown("[📦 Model Card](https://huggingface.co/RASHID778/king2-image)")

prompt = st.text_area(
    "Prompt",
    value="a futuristic royal palace at sunset, highly detailed, 8k, golden hour, epic composition",
    height=100,
    placeholder="Describe the image...",
)

col1, col2 = st.columns(2)
with col1:
    guidance = st.slider("Guidance Scale", 1.0, 15.0, 7.5, 0.5)
with col2:
    steps = st.slider("Inference Steps", 10, 50, 30, 1)

examples = [
    "a futuristic royal palace at sunset, highly detailed, 8k, golden hour, epic composition",
    "majestic arabian knight on horseback, desert landscape, cinematic lighting, photorealistic",
    "ancient library with floating books, mystical atmosphere, volumetric lighting, detailed",
    "cosmic king on a throne of stars, nebula background, majestic, epic fantasy, concept art",
]

selected = st.selectbox("Or try an example:", [""] + examples)
if selected:
    prompt = selected

if st.button("🎨 Generate", type="primary", use_container_width=True):
    if not prompt.strip():
        st.error("Prompt is required")
    elif not HF_TOKEN:
        st.error("HF_TOKEN not configured in Space secrets")
    else:
        with st.spinner("Generating... (this may take 30-60 seconds)"):
            try:
                client = InferenceClient(token=HF_TOKEN)
                img = client.text_to_image(
                    prompt,
                    model="stabilityai/stable-diffusion-xl-base-1.0",
                    negative_prompt="blurry, low quality, distorted, ugly, bad anatomy, watermark, text",
                    guidance_scale=guidance,
                    num_inference_steps=steps,
                    extra_body={"adapter_id": "RASHID778/king2-image"},
                )
                st.image(img, caption="Generated Image", use_column_width=True)

                buf = BytesIO()
                img.save(buf, format="PNG")
                st.download_button(
                    "⬇️ Download PNG",
                    data=buf.getvalue(),
                    file_name="king2_image.png",
                    mime="image/png",
                )
            except Exception as e:
                error_msg = str(e)
                if "402" in error_msg or "Payment Required" in error_msg or "credits" in error_msg:
                    st.error("💳 **Inference API credits depleted.**")
                    st.markdown("""
                    **To generate images, you have two options:**

                    **Option 1 — Add credits (instant)**
                    - Go to [HF Billing Settings](https://huggingface.co/settings/billing)
                    - Add pre-paid credits (min $5)
                    - Come back and try again

                    **Option 2 — Run locally (free)**
                    ```bash
                    pip install diffusers transformers accelerate torch safetensors
                    ```
                    ```python
                    from diffusers import DiffusionPipeline
                    import torch

                    pipe = DiffusionPipeline.from_pretrained(
                        'stabilityai/stable-diffusion-xl-base-1.0',
                        torch_dtype=torch.float16
                    ).to('cuda')
                    pipe.load_lora_weights('RASHID778/king2-image')

                    image = pipe('your prompt').images[0]
                    image.save('output.png')
                    ```
                    """)
                else:
                    st.error(f"Error: {error_msg}")
                    st.info("Run locally: `pip install diffusers transformers accelerate torch`")

st.markdown("---")
st.markdown("👑 Made by RASHID778 — Part of the KING2 AI Series")