import os

# ===== CPU / Thread 最佳化（關鍵）=====
CPU_THREADS = str(os.cpu_count() or 2)

os.environ["OMP_NUM_THREADS"] = CPU_THREADS
os.environ["OPENVINO_NUM_THREADS"] = CPU_THREADS
os.environ["MKL_NUM_THREADS"] = CPU_THREADS
os.environ["NUMEXPR_NUM_THREADS"] = CPU_THREADS

# 避免 thread oversubscription
os.environ["OMP_WAIT_POLICY"] = "PASSIVE"
os.environ["KMP_BLOCKTIME"] = "0"

import torch
import gradio as gr
from optimum.intel import OVZImagePipeline

# ===== Global objects（避免重複初始化）=====
pipe = None
generator = torch.Generator("cpu")


# ===== Load model（lazy + only once）=====
def load_model():
    global pipe

    if pipe is None:
        pipe = OVZImagePipeline.from_pretrained(
            "hsuwill000/Z-Image-Turbo-ov",
            device="cpu"
        )

        # ===== Warmup（避免第一次慢）=====
        pipe(
            prompt="warmup",
            height=512,
            width=512,
            num_inference_steps=1,
            guidance_scale=0.0,
        )

    return pipe


# ===== Inference =====
def generate_image(prompt, height, width, steps, seed):
    pipe = load_model()

    if seed != -1:
        generator.manual_seed(int(seed))

    image = pipe(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(steps),
        guidance_scale=0.0,  # Turbo 必須 0
        generator=generator,
    ).images[0]

    return image


# ===== Default prompt =====
default_prompt = """Cinematic portrait of a Japanese girl as Shinobu Kocho, realistic facial features, gradient purple compound eyes, intricate butterfly hair ornament, wearing silk haori with butterfly wing patterns, soft moonlight, hyper-realistic, depth of field, purple butterfly particles, 8k resolution, ethereal lighting"""


# ===== UI =====
with gr.Blocks() as demo:
    gr.Markdown("## 🚀 Z-Image Turbo (OpenVINO CPU Optimized)")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                value=default_prompt,
                lines=4
            )

            height = gr.Slider(256, 1024, value=512, step=64, label="Height")
            width = gr.Slider(256, 1024, value=512, step=64, label="Width")

            steps = gr.Slider(1, 20, value=9, step=1, label="Steps")
            seed = gr.Number(value=-1, label="Seed (-1 = random)")

            run_btn = gr.Button("Generate")

        with gr.Column():
            output = gr.Image(label="Result")

    run_btn.click(
        fn=generate_image,
        inputs=[prompt, height, width, steps, seed],
        outputs=output,
    )


# ===== Queue（最小化延遲）=====
demo.queue(concurrency_count=1, max_size=2)

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