import torch
import gradio as gr
from diffusers import DiffusionPipeline
import numpy as np
import random

# =========================================================
# MODEL CONFIGURATION (အမြန်ဆုံးနှုန်းအတွက် ပြင်ဆင်ချက်)
# =========================================================
MAX_SEED = np.iinfo(np.int32).max
# Turbo မော်ဒယ်ဖြစ်သောကြောင့် အလွန်မြန်ဆန်စေရန် (၂) ဆင့်သာ သုံးပါမည်
DEFAULT_STEPS = 2 

# =========================================================
# LOAD PIPELINE (Ultra CPU Optimized for Speed)
# =========================================================
print("Loading Z-Image-Turbo pipeline to CPU for MAXIMUM SPEED...")

# CPU ပေါ်တွင် အမြန်ဆုံး အလုပ်လုပ်ရန် float32 သုံးသည်
pipe = DiffusionPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True
)

# ⚠️ အရေးကြီးသော ပြင်ဆင်ချက်: 
# pipe.enable_attention_slicing() ကို အမြန်နှုန်းအတွက် တမင်ဖယ်ရှားထားပါသည်။ 
# ၎င်းသည် Memory ကို ချွေတာပေးသော်လည်း အချိန်ပိုကြာစေသောကြောင့် ဖြစ်သည်။
pipe.to("cpu")

# =========================================================
# PROMPT EXAMPLES
# =========================================================
prompt_examples = [
    "Moody mature anime scene of two lovers kissing under neon rain, sensual atmosphere",
    "A woman in a blue hanbok sits on a wooden floor, gazing out of a window.",
    "A traditional Japanese onsen, with steam rising, a young woman in a colorful kimono."
]

def get_random_prompt():
    return random.choice(prompt_examples)

# =========================================================
# IMAGE GENERATOR (Speed Focused)
# =========================================================
def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed):
    if not prompt:
        raise gr.Error("Please enter a prompt.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # အမြန်ဆုံးဖြစ်စေရန် ပုံ (၁) ပုံတည်းကိုသာ အတင်းအကျပ် ဖန်တီးခိုင်းပါမည်
    generator = torch.Generator("cpu").manual_seed(int(seed))

    result = pipe(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(num_inference_steps),
        guidance_scale=0.0,
        generator=generator,
        num_images_per_prompt=1, # အမြန်နှုန်းအတွက် ၁ ပုံသာ
        output_type="pil",
    )
    
    return result.images, seed

# ============================================
# 🎨 UI Design 
# ============================================
css = """
@import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@400;700&display=swap');
.gradio-container { background-color: #FEF9C3 !important; font-family: 'Comic Neue', cursive !important; }
.header-text h1 { font-family: 'Bangers', cursive !important; text-align: center; font-size: 3rem; }
.warning-box { background: #FEF3C7; border: 3px solid #F59E0B; padding: 10px; text-align: center; }
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("# ⚡ AI Image Generator (Ultra Fast CPU)", elem_classes="header-text")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="✏️ The Vision (Prompt)", lines=3)
            random_button = gr.Button("🎲 RANDOM PROMPT")
            
            with gr.Row():
                # ပုံသေးလေ ပိုမြန်လေဖြစ်၍ Default ကို 512 အစား 384 သို့ လျှော့ချထားပါသည်
                height_input = gr.Slider(256, 1024, 384, step=64, label="Height")
                width_input = gr.Slider(256, 1024, 384, step=64, label="Width")
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                # Turbo မော်ဒယ်အတွက် အမြင့်ဆုံး ၁၀ ဆင့်ထိသာ ပေးထားသည်
                steps_slider = gr.Slider(1, 5, DEFAULT_STEPS, step=1, label="Steps (Lower = Faster)")
                seed_input = gr.Number(value=42, label="Seed")
                randomize_seed_checkbox = gr.Checkbox(label="Randomize Seed", value=True)
                
            generate_button = gr.Button("✨ အလျင်အမြန် ဖန်တီးမည် (GENERATE FAST)", variant="primary")
            
        with gr.Column():
            output_gallery = gr.Gallery(label="The Masterpiece", columns=1)
            used_seed_output = gr.Number(label="Seed Used")

    random_button.click(fn=get_random_prompt, outputs=[prompt_input])
    
    # UI တွင် Image Count ကို ဖယ်ရှားပြီး Code ထဲတွင် ၁ ပုံတည်းဟု Fix လုပ်ထားသည်
    generate_button.click(
        fn=generate_image,
        inputs=[prompt_input, height_input, width_input, steps_slider, seed_input, randomize_seed_checkbox],
        outputs=[output_gallery, used_seed_output]
    )

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