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

# =========================================================
# MODEL CONFIGURATION
# =========================================================
MAX_SEED = np.iinfo(np.int32).max
# Turbo model ဖြစ်၍ CPU ပေါ်တွင် 1-4 steps သာ သုံးရန် အကြံပြုပါသည် (မြန်ဆန်စေရန်)
DEFAULT_STEPS = 4 

# =========================================================
# LOAD PIPELINE (CPU Optimized)
# =========================================================
print("Loading Z-Image-Turbo pipeline to CPU...")

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

# Memory ချွေတာရန် (CPU အတွက် အရေးကြီးပါသည်)
pipe.enable_attention_slicing()
pipe.to("cpu")

# =========================================================
# PROMPT EXAMPLES (User ပေးထားသော list ထဲမှ အချို့ကို နမူနာယူထားသည်)
# =========================================================
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
# =========================================================
def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, num_images):
    if not prompt:
        raise gr.Error("Please enter a prompt.")

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

    # CPU ပေါ်တွင် RAM မပြည့်စေရန် ပုံအရေအတွက်ကို ကန့်သတ်ခြင်း
    num_images = min(max(1, int(num_images)), 2)

    generator = torch.Generator("cpu").manual_seed(int(seed))

    # CPU inference ဖြစ်၍ အချိန်ကြာနိုင်ကြောင်း သတိပြုပါ
    result = pipe(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(num_inference_steps),
        guidance_scale=0.0,
        generator=generator,
        max_sequence_length=512, # CPU အတွက် length လျှော့ထားခြင်းက ပိုမြန်စေသည်
        num_images_per_prompt=num_images,
        output_type="pil",
    )

    return result.images, seed

# ============================================
# 🎨 UI Design (Original CSS and Layout)
# ============================================
css = """
/* User ပေးထားသော 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 (CPU Version)", elem_classes="header-text")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="✏️ Prompt", lines=3)
            random_button = gr.Button("🎲 RANDOM PROMPT")
            
            with gr.Row():
                height_input = gr.Slider(256, 1024, 512, step=64, label="Height")
                width_input = gr.Slider(256, 1024, 512, step=64, label="Width")
            
            num_images_input = gr.Slider(1, 2, 1, step=1, label="Images Count")
            
            with gr.Accordion("⚙️ Settings", open=False):
                steps_slider = gr.Slider(1, 10, DEFAULT_STEPS, step=1, label="Steps (Keep low for CPU)")
                seed_input = gr.Number(value=42, label="Seed")
                randomize_seed_checkbox = gr.Checkbox(label="Randomize Seed", value=True)

            generate_button = gr.Button("✨ GENERATE", variant="primary")
            
        with gr.Column():
            output_gallery = gr.Gallery(label="Output", columns=1)
            used_seed_output = gr.Number(label="Seed Used")

    random_button.click(fn=get_random_prompt, outputs=[prompt_input])
    generate_button.click(
        fn=generate_image,
        inputs=[prompt_input, height_input, width_input, steps_slider, seed_input, randomize_seed_checkbox, num_images_input],
        outputs=[output_gallery, used_seed_output]
    )

if __name__ == "__main__":
    # show_api ကို ဖယ်ရှားလိုက်ပါပြီ
    demo.queue(max_size=10).launch(
        debug=False,
        share=False  # Hugging Face Space မှာ run ရင် share=True လုပ်စရာမလိုပါ
    )