import os, io, base64, requests, gradio as gr
from PIL import Image

API_BASE = os.getenv("API_BASE", "")
API_KEY = os.getenv("API_KEY", "")

def shrink_if_needed(pil_img):
    if pil_img is None: return None
    w, h = pil_img.size
    if w <= 1080 and h <= 1080: return pil_img
    scale = min(1080 / w, 1080 / h)
    new_w, new_h = max(1, int(w * scale)), max(1, int(h * scale))
    return pil_img.resize((new_w, new_h), Image.LANCZOS)

def img_to_b64(pil_img):
    if pil_img is None: return None
    buf = io.BytesIO()
    pil_img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")

def b64_to_img(b64_str):
    return Image.open(io.BytesIO(base64.b64decode(b64_str))).convert("RGB")

def call_staticaliza(payload):
    endpoint = "image-generation"
    url = f"{API_BASE}/{endpoint}"
    headers = {"Content-Type": "application/json"}
    if API_KEY: headers["Authorization"] = f"Bearer {API_KEY}"
    r = requests.post(url, headers=headers, json=payload, timeout=60)
    r.raise_for_status()
    j = r.json()
    out = j.get("data", {}).get("output")
    if not out: raise RuntimeError("Empty output from API")
    return b64_to_img(out)

def generate(prompt, negative_prompt, steps, width, height, reference_image, profile: gr.OAuthProfile | None = None, oauth_token: gr.OAuthToken | None = None):
    if profile is None: raise gr.Error("Please sign in with Hugging Face to use this Space.")
    prompt = (prompt or "").strip()[:512]
    negative_prompt = (negative_prompt or "").strip()[:512]
    ref_img = shrink_if_needed(reference_image) if reference_image is not None else None
    ref_b64 = img_to_b64(ref_img) if ref_img is not None else None
    print("---")
    print("Prompt: " + prompt)
    print("Negative Prompt: " + negative_prompt)
    print("Width, Height, Steps: " + str(width) + ", " + str(height) + ", " + str(steps))
    print("---")
    payload = {
        "input": prompt,
        "negative_input": negative_prompt,
        "reference_image": ref_b64,
        "resolution": None if ref_b64 else [int(width), int(height)],
        "steps": int(steps),
        "guidance": 1,
        "lossy": True
    }
    return call_staticaliza(payload)
with gr.Blocks() as demo:
    gr.Markdown("# 👁️‍🗨️ Qwen Image Lightning Dedicated")
    gr.Markdown("- Use this, heck yeah!")
    gr.Markdown("- Warning, your requests may be privately logged for security and safety purposes only within the Spaces runtime log, do not insert private information, please confirm that Staticaliza API is safe, reliable, and trustable!")
    gr.Markdown("- Warning, this model may produce sensitive content, avoid using bad prompts in accordance to Hugging Face Terms and Service (TOS)!")
    
    login = gr.LoginButton()
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="Prompt")
            negative_prompt = gr.Textbox(label="Negative Prompt")
            steps = gr.Slider(1, 8, step=1, value=8, label="Steps")
            width = gr.Slider(4, 1080, step=1, value=1080, label="Width")
            height = gr.Slider(4, 1080, step=1, value=1080, label="Height")
            reference_image = gr.Image(label="Optional Reference Image", type="pil")
            run = gr.Button("Generate / Edit")
        with gr.Column(scale=1):
            out = gr.Image(label="Output", format="png")
    run.click(generate, [prompt, negative_prompt, steps, width, height, reference_image], out)

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